31 lines
1.1 KiB
Python
31 lines
1.1 KiB
Python
|
import json
|
||
|
|
from openai import OpenAI
|
||
|
|
|
||
|
|
class LLMClient:
|
||
|
|
def __init__(self, base_url, api_key: str | None = None, model: str = "gpt-4o-mini"):
|
||
|
|
self.model = model
|
||
|
|
self.client = OpenAI(base_url=base_url, api_key=api_key or "ollama")
|
||
|
|
|
||
|
|
def chat(self, system: str, user: str, json_mode: bool = False, images: list[str] | None = None) -> str | None:
|
||
|
|
images = images or []
|
||
|
|
if images:
|
||
|
|
content = [{"type": "text", "text": user}]
|
||
|
|
for b64 in images:
|
||
|
|
content.append({
|
||
|
|
"type": "image_url",
|
||
|
|
"image_url": {"url": f"data:image/png;base64,{b64}"},
|
||
|
|
})
|
||
|
|
else:
|
||
|
|
content = user
|
||
|
|
|
||
|
|
kwargs = {
|
||
|
|
"model": self.model,
|
||
|
|
"messages": [
|
||
|
|
{"role": "system", "content": system},
|
||
|
|
{"role": "user", "content": content},
|
||
|
|
],
|
||
|
|
}
|
||
|
|
if json_mode:
|
||
|
|
kwargs["response_format"] = {"type": "json_object"}
|
||
|
|
resp = self.client.chat.completions.create(**kwargs)
|
||
|
|
return resp.choices[0].message.content
|