hangshuo652
42 lines
1.4 KiB
Python
42 lines
1.4 KiB
Python
import json, hashlib, os
|
|
from pathlib import Path
|
|
import httpx
|
|
|
|
|
|
class LLMClient:
|
|
def __init__(self, model="gpt-4o-mini", timeout=15, cache_dir=".cache/llm"):
|
|
self.model = model
|
|
self.timeout = timeout
|
|
self.dir = Path(cache_dir)
|
|
self.dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
def _key(self, msgs):
|
|
return hashlib.sha256(json.dumps(msgs, sort_keys=True).encode()).hexdigest()
|
|
|
|
def _get(self, k):
|
|
p = self.dir / f"{k}.json"
|
|
return json.loads(p.read_text())["response"] if p.exists() else None
|
|
|
|
def _set(self, k, v):
|
|
(self.dir / f"{k}.json").write_text(json.dumps({"response": v}))
|
|
|
|
def call(self, messages, retries=1):
|
|
k = self._key(messages)
|
|
c = self._get(k)
|
|
if c:
|
|
return c
|
|
key = os.environ.get("LLM_API_KEY", os.environ.get("OPENAI_API_KEY", ""))
|
|
base = os.environ.get("LLM_API_BASE", "https://api.openai.com/v1")
|
|
for a in range(retries + 1):
|
|
try:
|
|
r = httpx.post(f"{base}/chat/completions", json={"model": self.model, "messages": messages},
|
|
headers={"Authorization": f"Bearer {key}"}, timeout=self.timeout)
|
|
r.raise_for_status()
|
|
v = r.json()["choices"][0]["message"]["content"]
|
|
self._set(k, v)
|
|
return v
|
|
except Exception:
|
|
if a == retries:
|
|
raise
|
|
return ""
|