95 lines
3.0 KiB
Python
95 lines
3.0 KiB
Python
|
import asyncio
|
||
|
|
import json
|
||
|
|
import logging
|
||
|
|
from collections.abc import AsyncGenerator
|
||
|
|
|
||
|
|
from zhipuai import ZhipuAI
|
||
|
|
|
||
|
|
from app.engine.providers.base import AIProvider
|
||
|
|
|
||
|
|
logger = logging.getLogger(__name__)
|
||
|
|
|
||
|
|
DEFAULT_MODEL = "glm-4-flash"
|
||
|
|
|
||
|
|
|
||
|
|
class GLMProvider(AIProvider):
|
||
|
|
def __init__(self, api_key: str, model: str = "glm-4-flash", base_url: str = ""):
|
||
|
|
self.client = ZhipuAI(api_key=api_key, base_url=base_url or None)
|
||
|
|
self.model = model
|
||
|
|
|
||
|
|
async def chat(self, messages: list[dict], system_prompt: str) -> str:
|
||
|
|
full_messages = [{"role": "system", "content": system_prompt}] + messages
|
||
|
|
return await self._call(full_messages, temperature=0.3)
|
||
|
|
|
||
|
|
async def stream_chat(
|
||
|
|
self, messages: list[dict], system_prompt: str
|
||
|
|
) -> AsyncGenerator[str, None]:
|
||
|
|
full_messages = [{"role": "system", "content": system_prompt}] + messages
|
||
|
|
|
||
|
|
def _sync_stream():
|
||
|
|
resp = self.client.chat.completions.create(
|
||
|
|
model=self.model,
|
||
|
|
messages=full_messages,
|
||
|
|
stream=True,
|
||
|
|
temperature=0.3,
|
||
|
|
)
|
||
|
|
for chunk in resp:
|
||
|
|
delta = chunk.choices[0].delta
|
||
|
|
if delta and delta.content:
|
||
|
|
yield delta.content
|
||
|
|
|
||
|
|
try:
|
||
|
|
loop = asyncio.get_event_loop()
|
||
|
|
gen = _sync_stream()
|
||
|
|
|
||
|
|
def _next_token(g):
|
||
|
|
try:
|
||
|
|
return next(g)
|
||
|
|
except StopIteration:
|
||
|
|
return None
|
||
|
|
|
||
|
|
while True:
|
||
|
|
token = await loop.run_in_executor(None, _next_token, gen)
|
||
|
|
if token is None:
|
||
|
|
break
|
||
|
|
yield token
|
||
|
|
except Exception:
|
||
|
|
logger.exception("GLM streaming call failed")
|
||
|
|
raise
|
||
|
|
|
||
|
|
async def chat_with_extraction(
|
||
|
|
self, messages: list[dict], system_prompt: str
|
||
|
|
) -> tuple[str, dict]:
|
||
|
|
full_messages = [{"role": "system", "content": system_prompt}] + messages
|
||
|
|
if not messages:
|
||
|
|
full_messages.append({"role": "user", "content": "请根据以上要求输出结果。"})
|
||
|
|
|
||
|
|
raw = await self._call(full_messages, temperature=0.1)
|
||
|
|
extracted = self._parse_json(raw)
|
||
|
|
return raw, extracted
|
||
|
|
|
||
|
|
async def _call(self, messages: list[dict], temperature: float = 0.3) -> str:
|
||
|
|
def _sync_call():
|
||
|
|
resp = self.client.chat.completions.create(
|
||
|
|
model=self.model,
|
||
|
|
messages=messages,
|
||
|
|
temperature=temperature,
|
||
|
|
)
|
||
|
|
return resp.choices[0].message.content
|
||
|
|
|
||
|
|
try:
|
||
|
|
return await asyncio.to_thread(_sync_call)
|
||
|
|
except Exception:
|
||
|
|
logger.exception("GLM API call failed")
|
||
|
|
raise
|
||
|
|
|
||
|
|
@staticmethod
|
||
|
|
def _parse_json(text: str) -> dict:
|
||
|
|
try:
|
||
|
|
start = text.index("{")
|
||
|
|
end = text.rindex("}") + 1
|
||
|
|
return json.loads(text[start:end])
|
||
|
|
except (ValueError, json.JSONDecodeError):
|
||
|
|
logger.warning("Failed to parse JSON from model output")
|
||
|
|
return {}
|