62 lines
2.2 KiB
Python
62 lines
2.2 KiB
Python
import os
|
|
import json
|
|
import time
|
|
import uuid
|
|
from typing import Dict, Any, Optional
|
|
import requests
|
|
|
|
class LLMClient:
|
|
def __init__(self, cfg: Dict[str, Any], logs_dir: str):
|
|
self.cfg = cfg
|
|
self.logs_dir = logs_dir
|
|
os.makedirs(self.logs_dir, exist_ok=True)
|
|
|
|
def _log(self, role: str, content: str, uid: str):
|
|
path = os.path.join(self.logs_dir, f"{time.strftime('%Y%m%d-%H%M%S')}-{uid}-{role}.md")
|
|
with open(path, "w", encoding="utf-8") as f:
|
|
f.write(content)
|
|
|
|
def chat(self, system: str, user: str) -> str:
|
|
uid = str(uuid.uuid4())[:8]
|
|
backend = self.cfg.get("backend", "ollama")
|
|
if backend == "ollama":
|
|
host = self.cfg["ollama"]["host"]
|
|
model = self.cfg["ollama"]["model"]
|
|
payload = {
|
|
"model": model,
|
|
"messages": [
|
|
{"role": "system", "content": system},
|
|
{"role": "user", "content": user},
|
|
],
|
|
"stream": False,
|
|
}
|
|
self._log("system", system, uid)
|
|
self._log("user", user, uid)
|
|
r = requests.post(f"{host}/v1/chat/completions", json=payload, timeout=600)
|
|
r.raise_for_status()
|
|
data = r.json()
|
|
content = data["choices"][0]["message"]["content"]
|
|
self._log("assistant", content, uid)
|
|
return content
|
|
|
|
elif backend == "openai":
|
|
import openai # requires openai>=1.0
|
|
openai.api_key = os.environ.get("OPENAI_API_KEY")
|
|
if not openai.api_key:
|
|
raise RuntimeError("Missing OPENAI_API_KEY")
|
|
model = self.cfg["openai"]["model"]
|
|
messages = [
|
|
{"role": "system", "content": system},
|
|
{"role": "user", "content": user},
|
|
]
|
|
self._log("system", system, uid)
|
|
self._log("user", user, uid)
|
|
client = openai.OpenAI()
|
|
resp = client.chat.completions.create(model=model, messages=messages)
|
|
content = resp.choices[0].message.content
|
|
self._log("assistant", content, uid)
|
|
return content
|
|
|
|
else:
|
|
raise ValueError(f"Unknown backend: {backend}")
|