add moonshot tool call support

generate/chat_completion/models/moonshot.py

@@ -1,13 +1,12 @@
 from __future__ import annotations
 
-from typing import AsyncIterator, ClassVar, Iterator, List, Optional
-
+from typing import AsyncIterator, ClassVar, Iterator, List, Optional, Union, Literal
 from pydantic import PositiveInt
 from typing_extensions import Unpack, override
 
 from generate.chat_completion.message import Prompt
 from generate.chat_completion.model_output import ChatCompletionOutput, ChatCompletionStreamOutput
-from generate.chat_completion.models.openai_like import OpenAILikeChat
+from generate.chat_completion.models.openai_like import OpenAILikeChat, OpenAIToolChoice, OpenAITool, SupportOpenAIToolCall
 from generate.http import HttpClient
 from generate.model import ModelParameters, RemoteModelParametersDict
 from generate.platforms import MoonshotSettings
@@ -18,15 +17,18 @@ class MoonshotChatParameters(ModelParameters):
     temperature: Optional[Temperature] = None
     top_p: Optional[Probability] = None
     max_tokens: Optional[PositiveInt] = None
-
+    tools: Optional[List[OpenAITool]] = None
+    tool_choice: Union[Literal['auto', 'none'], OpenAIToolChoice, None] = None
 
 class MoonshotChatParametersDict(RemoteModelParametersDict, total=False):
     temperature: Temperature
     top_p: Probability
     max_tokens: PositiveInt
+    tools: Optional[List[OpenAITool]]
+    tool_choice: Union[Literal['auto'], OpenAIToolChoice, None]
 
 
-class MoonshotChat(OpenAILikeChat):
+class MoonshotChat(OpenAILikeChat, SupportOpenAIToolCall):
     model_type: ClassVar[str] = 'moonshot'
     available_models: ClassVar[List[str]] = ['moonshot-v1-8k', 'moonshot-v1-32k', 'moonshot-v1-128k']
 

tests/test_function.py

@@ -2,7 +2,7 @@
 
 from pydantic import BaseModel
 
-from generate import OpenAIChat, tool
+from generate import OpenAIChat, tool, MoonshotChat
 from generate.chat_completion.tool import get_json_schema
 
 

tests/test_moonshot.py

@@ -0,0 +1,25 @@
+from typing import Literal
+
+from pydantic import BaseModel
+
+from generate import OpenAIChat, tool, MoonshotChat
+from generate.chat_completion.tool import get_json_schema
+from dotenv import load_dotenv
+
+load_dotenv()  # 加载.env文件
+
+@tool
+def get_weather(loc: str) -> str:
+    """
+    获取指定地区的天气信息
+
+    Parameters:
+        loc: 地区，比如北京，上海等
+    """
+    return f'{loc}，晴朗，27度'
+
+def test_moonshot() -> None:
+    moonshot = MoonshotChat()
+    resp = moonshot.generate('今天北京天气怎么样？', tools=[{"type":"function", "function": get_weather.json_schema}], tool_choice='auto')
+    assert resp.message.tool_calls[0].function.name == "get_weather" # type: ignore
+