hangshuo652
48 lines
1.7 KiB
Python
48 lines
1.7 KiB
Python
import json, hashlib, os
|
|
from pathlib import Path
|
|
from typing import Optional
|
|
import httpx
|
|
|
|
|
|
class LLMClient:
|
|
def __init__(self, model: str = "gpt-4o-mini", timeout: int = 15,
|
|
cache_dir: str = ".cache/llm"):
|
|
self.model = model
|
|
self.timeout = timeout
|
|
self.cache_dir = Path(cache_dir)
|
|
self.cache_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
def _cache_key(self, messages: list) -> str:
|
|
return hashlib.sha256(json.dumps(messages, sort_keys=True).encode()).hexdigest()
|
|
|
|
def _cache_get(self, key: str) -> Optional[str]:
|
|
path = self.cache_dir / f"{key}.json"
|
|
if path.exists():
|
|
return json.loads(path.read_text()).get("response")
|
|
return None
|
|
|
|
def _cache_set(self, key: str, response: str):
|
|
(self.cache_dir / f"{key}.json").write_text(json.dumps({"response": response}))
|
|
|
|
def call(self, messages: list, retries: int = 1) -> str:
|
|
key = self._cache_key(messages)
|
|
cached = self._cache_get(key)
|
|
if cached:
|
|
return cached
|
|
api_key = os.environ.get("OPENAI_API_KEY", "")
|
|
for attempt in range(retries + 1):
|
|
try:
|
|
resp = httpx.post(
|
|
"https://api.openai.com/v1/chat/completions",
|
|
json={"model": self.model, "messages": messages},
|
|
headers={"Authorization": f"Bearer {api_key}"},
|
|
timeout=self.timeout)
|
|
resp.raise_for_status()
|
|
result = resp.json()["choices"][0]["message"]["content"]
|
|
self._cache_set(key, result)
|
|
return result
|
|
except Exception:
|
|
if attempt == retries:
|
|
raise
|
|
return ""
|