高级搜索工具 Web search 业务] (#41)

* 高级搜索工具 Web search 业务

* 高级搜索工具 Web search 业务集成测试

* websearch

Author: glide-the

tests/integration_tests/test_tools.py

@@ -0,0 +1,57 @@
+from zhipuai import ZhipuAI
+import zhipuai
+
+import logging
+import logging.config
+
+
+def test_tools(logging_conf):
+    logging.config.dictConfig(logging_conf)  # type: ignore
+    client = ZhipuAI()  # 填写您自己的APIKey
+    try:
+        response = client.tools.web_search(
+            model="web-search-pro",
+            messages=[
+                {
+                    "content": "你好",
+                    "role": "user",
+                }
+            ],
+            stream=False
+        )
+        print(response)
+
+
+
+    except zhipuai.core._errors.APIRequestFailedError as err:
+        print(err)
+    except zhipuai.core._errors.APIInternalError as err:
+        print(err)
+    except zhipuai.core._errors.APIStatusError as err:
+        print(err)
+
+
+def test_tools_stream(logging_conf):
+    logging.config.dictConfig(logging_conf)  # type: ignore
+    client = ZhipuAI()  # 填写您自己的APIKey
+    try:
+        response = client.tools.web_search(
+            model="web-search-pro",
+            messages=[
+                {
+                    "content": "你好",
+                    "role": "user",
+                }
+            ],
+            stream=True
+        )
+        for item in response:
+            print(item)
+
+
+    except zhipuai.core._errors.APIRequestFailedError as err:
+        print(err)
+    except zhipuai.core._errors.APIInternalError as err:
+        print(err)
+    except zhipuai.core._errors.APIStatusError as err:
+        print(err)

zhipuai/_client.py

@@ -59,6 +59,7 @@ def __init__(
         self.fine_tuning = api_resource.fine_tuning.FineTuning(self)
         self.batches = api_resource.Batches(self)
         self.knowledge = api_resource.Knowledge(self)
+        self.tools = api_resource.Tools(self)
 
     @property
     @override

zhipuai/api_resource/__init__.py

@@ -24,7 +24,9 @@
 from .knowledge import (
     Knowledge
 )
-
+from .tools import (
+    Tools
+)
 __all__ = [
     'AsyncCompletions',
     'Chat',
@@ -36,5 +38,6 @@
     'FineTuning',
     'Batches',
     'Knowledge',
+    'Tools',
 
 ]

zhipuai/api_resource/tools/__init__.py

@@ -0,0 +1,3 @@
+from .tools import Tools
+
+__all__ = ['Tools']

zhipuai/api_resource/tools/tools.py

@@ -0,0 +1,64 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, List, Union, Dict, Optional
+from typing_extensions import Literal
+
+from ...core import NOT_GIVEN, Body, Headers, NotGiven, BaseAPI, maybe_transform, StreamResponse, deepcopy_minimal
+
+import httpx
+
+from ...core import (
+    make_request_options,
+)
+import logging
+
+from ...types.tools import tools_web_search_params, WebSearch, WebSearchChunk
+
+logger = logging.getLogger(__name__)
+
+if TYPE_CHECKING:
+    from ..._client import ZhipuAI
+
+__all__ = ["Tools"]
+
+
+class Tools(BaseAPI):
+
+    def __init__(self, client: "ZhipuAI") -> None:
+        super().__init__(client)
+
+    def web_search(
+            self,
+            *,
+            model: str,
+            request_id: Optional[str] | NotGiven = NOT_GIVEN,
+            stream: Optional[Literal[False]] | Literal[True] | NotGiven = NOT_GIVEN,
+            messages: Union[str, List[str], List[int], object, None],
+            scope: Optional[str] | NotGiven = NOT_GIVEN,
+            location: Optional[str] | NotGiven = NOT_GIVEN,
+            recent_days: Optional[int] | NotGiven = NOT_GIVEN,
+            extra_headers: Headers | None = None,
+            extra_body: Body | None = None,
+            timeout: float | httpx.Timeout | None | NotGiven = NOT_GIVEN,
+    ) -> WebSearch | StreamResponse[WebSearchChunk]:
+
+        body = deepcopy_minimal(
+            {
+                "model": model,
+                "request_id": request_id,
+                "messages": messages,
+                "stream": stream,
+                "scope": scope,
+                "location": location,
+                "recent_days": recent_days,
+            })
+        return self._post(
+            "/tools",
+            body= maybe_transform(body, tools_web_search_params.WebSearchParams),
+            options=make_request_options(
+                extra_headers=extra_headers, extra_body=extra_body, timeout=timeout
+            ),
+            cast_type=WebSearch,
+            stream=stream or False,
+            stream_cls=StreamResponse[WebSearchChunk],
+        )

zhipuai/types/tools/__init__.py

@@ -0,0 +1,18 @@
+from .web_search import (
+    WebSearch,
+    SearchIntent,
+    SearchResult,
+    SearchRecommend,
+)
+
+from .web_search_chunk import (
+    WebSearchChunk
+)
+
+__all__ = [
+    'WebSearch',
+    'SearchIntent',
+    'SearchResult',
+    'SearchRecommend',
+    'WebSearchChunk'
+]

zhipuai/types/tools/tools_web_search_params.py

@@ -0,0 +1,34 @@
+
+from __future__ import annotations
+
+from typing import Union, Optional, List
+from typing_extensions import Literal, Required, TypedDict
+
+__all__ = ["WebSearchParams"]
+
+
+class WebSearchParams(TypedDict):
+    """
+    工具名：web-search-pro参数类型定义
+
+    Attributes:
+        :param model: str, 模型名称
+        :param request_id: Optional[str], 请求ID
+        :param stream: Optional[bool], 是否流式
+        :param messages: Union[str, List[str], List[int], object, None],
+                        包含历史对话上下文的内容，按照 {"role": "user", "content": "你好"} 的json 数组形式进行传参
+                        当前版本仅支持 User Message 单轮对话，工具会理解User Message并进行搜索，
+                        请尽可能传入不带指令格式的用户原始提问，以提高搜索准确率。
+        :param scope: Optional[str], 指定搜索范围，全网、学术等，默认全网
+        :param location: Optional[str], 指定搜索用户地区 location 提高相关性
+        :param recent_days: Optional[int],支持指定返回 N 天（1-30）更新的搜索结果
+
+
+    """
+    model: str
+    request_id: Optional[str]
+    stream: Optional[bool]
+    messages: Union[str, List[str], List[int], object, None]
+    scope: Optional[str] = None
+    location: Optional[str] = None
+    recent_days: Optional[int] = None

zhipuai/types/tools/web_search.py

@@ -0,0 +1,72 @@
+from typing import List, Optional
+
+from ..chat.chat_completion import Function
+from ...core import BaseModel
+
+__all__ = [
+    "WebSearch",
+    "SearchIntent",
+    "SearchResult",
+    "SearchRecommend",
+]
+
+
+class SearchIntent(BaseModel):
+    index: int
+    # 搜索轮次，默认为 0
+    query: str
+    # 搜索优化 query
+    intent: str
+    # 判断的意图类型
+    keywords: str
+    # 搜索关键词
+
+
+class SearchResult(BaseModel):
+    index: int
+    # 搜索轮次，默认为 0
+    title: str
+    # 标题
+    link: str
+    # 链接
+    content: str
+    # 内容
+    icon: str
+    # 图标
+    media: str
+    # 来源媒体
+    refer: str
+    # 角标序号 [ref_1]
+
+
+class SearchRecommend(BaseModel):
+    index: int
+    # 搜索轮次，默认为 0
+    query: str
+    # 推荐query
+
+
+class WebSearchMessageToolCall(BaseModel):
+    id: str
+    search_intent: Optional[SearchIntent]
+    search_result: Optional[SearchResult]
+    search_recommend: Optional[SearchRecommend]
+    type: str
+
+
+class WebSearchMessage(BaseModel):
+    role: str
+    tool_calls: Optional[List[WebSearchMessageToolCall]] = None
+
+
+class WebSearchChoice(BaseModel):
+    index: int
+    finish_reason: str
+    message: WebSearchMessage
+
+
+class WebSearch(BaseModel):
+    created: Optional[int] = None
+    choices: List[WebSearchChoice]
+    request_id: Optional[str] = None
+    id: Optional[str] = None

zhipuai/types/tools/web_search_chunk.py

@@ -0,0 +1,35 @@
+from typing import List, Optional, Dict, Any
+
+from .web_search import SearchIntent, SearchResult, SearchRecommend
+from ...core import BaseModel
+
+__all__ = [
+    "WebSearchChunk"
+]
+
+
+class ChoiceDeltaToolCall(BaseModel):
+    index: int
+    id: Optional[str] = None
+
+    search_intent: Optional[SearchIntent] = None
+    search_result: Optional[SearchResult] = None
+    search_recommend: Optional[SearchRecommend] = None
+    type: Optional[str] = None
+
+
+class ChoiceDelta(BaseModel):
+    role: Optional[str] = None
+    tool_calls: Optional[List[ChoiceDeltaToolCall]] = None
+
+
+class Choice(BaseModel):
+    delta: ChoiceDelta
+    finish_reason: Optional[str] = None
+    index: int
+
+
+class WebSearchChunk(BaseModel):
+    id: Optional[str] = None
+    choices: List[Choice]
+    created: Optional[int] = None