freedakgmail
81 lines
2.3 KiB
Python
81 lines
2.3 KiB
Python
"""LLM 客户端 — 用于 PPT Master 管线中的策略师和执行器阶段"""
|
|
|
|
import httpx
|
|
import json
|
|
from typing import Generator
|
|
from config import config
|
|
|
|
|
|
class LLMClient:
|
|
"""OpenAI 兼容的 LLM 客户端"""
|
|
|
|
def __init__(self):
|
|
self.api_key = config.OPENAI_API_KEY
|
|
self.base_url = config.OPENAI_BASE_URL.rstrip("/")
|
|
self.model = config.OPENAI_MODEL
|
|
self.client = httpx.Client(timeout=300.0)
|
|
|
|
def chat(
|
|
self,
|
|
messages: list[dict],
|
|
temperature: float = 0.7,
|
|
max_tokens: int = 16384,
|
|
stream: bool = False,
|
|
) -> str:
|
|
"""同步聊天完成"""
|
|
headers = {
|
|
"Authorization": f"Bearer {self.api_key}",
|
|
"Content-Type": "application/json",
|
|
}
|
|
payload = {
|
|
"model": self.model,
|
|
"messages": messages,
|
|
"temperature": temperature,
|
|
"max_tokens": max_tokens,
|
|
"stream": stream,
|
|
}
|
|
|
|
if stream:
|
|
return self._stream_chat(headers, payload)
|
|
|
|
resp = self.client.post(
|
|
f"{self.base_url}/chat/completions",
|
|
headers=headers,
|
|
json=payload,
|
|
)
|
|
resp.raise_for_status()
|
|
data = resp.json()
|
|
return data["choices"][0]["message"]["content"]
|
|
|
|
def _stream_chat(self, headers: dict, payload: dict) -> str:
|
|
"""流式聊天，收集完整响应"""
|
|
full_content = ""
|
|
with self.client.stream(
|
|
"POST",
|
|
f"{self.base_url}/chat/completions",
|
|
headers=headers,
|
|
json=payload,
|
|
) as resp:
|
|
resp.raise_for_status()
|
|
for line in resp.iter_lines():
|
|
if not line.startswith("data: "):
|
|
continue
|
|
data_str = line[6:]
|
|
if data_str == "[DONE]":
|
|
break
|
|
try:
|
|
chunk = json.loads(data_str)
|
|
delta = chunk["choices"][0].get("delta", {})
|
|
content = delta.get("content", "")
|
|
if content:
|
|
full_content += content
|
|
except (json.JSONDecodeError, KeyError, IndexError):
|
|
continue
|
|
return full_content
|
|
|
|
def close(self):
|
|
self.client.close()
|
|
|
|
|
|
llm_client = LLMClient()
|