Initial commit: GovAI 政务AI平台

ppt-worker/llm_client.py

@@ -0,0 +1,80 @@
+"""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()