glm.py
3.92 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
"""
GLM-4.7 (智谱AI) 集成
"""
import logging
from langchain_core.language_models import BaseChatModel
from langchain_core.messages import BaseMessage, AIMessage
from zhipuai import ZhipuAI
from .base import BaseLLMClient, LLMResponse, LLMUsage
from ..config import get_settings
logger = logging.getLogger(__name__)
class ChatZhipuAI(BaseChatModel):
"""智谱AI 聊天模型 LangChain 包装器"""
client: ZhipuAI = None
model: str = "glm-4"
temperature: float = 0.7
max_tokens: int = 4096
def __init__(self, api_key: str, model: str = "glm-4", **kwargs):
super().__init__(**kwargs)
self.client = ZhipuAI(api_key=api_key)
self.model = model
@property
def _llm_type(self) -> str:
return "zhipuai"
def _generate(self, messages, stop=None, run_manager=None, **kwargs):
from langchain_core.outputs import ChatGeneration, ChatResult
formatted_messages = []
for msg in messages:
if hasattr(msg, 'type'):
role = "user" if msg.type == "human" else "assistant" if msg.type == "ai" else "system"
else:
role = "user"
formatted_messages.append({"role": role, "content": msg.content})
response = self.client.chat.completions.create(
model=self.model,
messages=formatted_messages,
temperature=self.temperature,
max_tokens=self.max_tokens,
)
content = response.choices[0].message.content
generation = ChatGeneration(message=AIMessage(content=content))
return ChatResult(
generations=[generation],
llm_output={
"token_usage": {
"prompt_tokens": response.usage.prompt_tokens,
"completion_tokens": response.usage.completion_tokens,
"total_tokens": response.usage.total_tokens,
}
}
)
class GLMClient(BaseLLMClient):
"""GLM-4.7 客户端"""
def __init__(self, api_key: str = None, model: str = None):
settings = get_settings()
self._api_key = api_key or settings.glm_api_key
self._model = model or settings.glm_model
self._client = ZhipuAI(api_key=self._api_key)
self._chat_model = None
@property
def provider(self) -> str:
return "glm"
@property
def model_name(self) -> str:
return self._model
def get_chat_model(self) -> BaseChatModel:
"""获取 LangChain 聊天模型"""
if self._chat_model is None:
self._chat_model = ChatZhipuAI(api_key=self._api_key, model=self._model)
return self._chat_model
def invoke(self, messages: list[BaseMessage]) -> LLMResponse:
"""调用 GLM"""
try:
formatted_messages = []
for msg in messages:
if hasattr(msg, 'type'):
role = "user" if msg.type == "human" else "assistant" if msg.type == "ai" else "system"
else:
role = "user"
formatted_messages.append({"role": role, "content": msg.content})
response = self._client.chat.completions.create(
model=self._model,
messages=formatted_messages,
temperature=0.7,
max_tokens=4096,
)
content = response.choices[0].message.content
usage = self.create_usage(
prompt_tokens=response.usage.prompt_tokens,
completion_tokens=response.usage.completion_tokens,
)
logger.info(
f"GLM 调用完成: model={self._model}, "
f"tokens={usage.total_tokens}"
)
return LLMResponse(content=content, usage=usage, raw_response=response)
except Exception as e:
logger.error(f"GLM 调用失败: {e}")
raise