File size: 3,543 Bytes
c419ba2 661eb14 c419ba2 661eb14 c419ba2 661eb14 c419ba2 661eb14 c419ba2 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 | import base64
import json
import time
from pathlib import Path
import openai
from core.config import DEEPSEEK_API_KEY, DEEPSEEK_BASE_URL, MODEL_NAME
class LLMUnavailable(Exception):
pass
class LLM:
def __init__(self, api_key: str | None = None):
resolved = api_key if api_key is not None else DEEPSEEK_API_KEY
self._api_key = resolved
self._client = openai.OpenAI(
api_key=resolved or "no-key",
base_url=DEEPSEEK_BASE_URL,
)
def _check_key(self) -> None:
if not self._api_key:
raise LLMUnavailable("No API key configured")
def chat_json(self, system: str, user: str, max_retries: int = 2) -> dict:
self._check_key()
last_exc: Exception | None = None
for attempt in range(max_retries + 1):
try:
resp = self._client.chat.completions.create(
model=MODEL_NAME,
messages=[
{"role": "system", "content": system},
{"role": "user", "content": user},
],
temperature=0,
response_format={"type": "json_object"},
)
content = resp.choices[0].message.content or ""
try:
return json.loads(content)
except json.JSONDecodeError:
if attempt < max_retries:
system = system + " Respond ONLY with valid JSON, no prose."
continue
raise LLMUnavailable("Malformed JSON after retries")
except openai.AuthenticationError as e:
raise LLMUnavailable("Invalid API key") from e
except (openai.APIStatusError, openai.APIConnectionError) as e:
last_exc = e
if attempt < max_retries:
time.sleep(2 ** attempt)
continue
raise LLMUnavailable(f"API error after retries: {last_exc}") from last_exc
def chat_vision(
self,
system: str,
user_text: str,
image: bytes | str | Path,
max_retries: int = 2,
) -> str:
self._check_key()
if isinstance(image, (str, Path)):
raw = Path(image).read_bytes()
else:
raw = image
b64 = base64.b64encode(raw).decode()
data_uri = f"data:image/png;base64,{b64}"
last_exc: Exception | None = None
for attempt in range(max_retries + 1):
try:
resp = self._client.chat.completions.create(
model=MODEL_NAME,
messages=[
{"role": "system", "content": system},
{"role": "user", "content": [
{"type": "text", "text": user_text},
{"type": "image_url", "image_url": {"url": data_uri}},
]},
],
temperature=0,
)
return resp.choices[0].message.content or ""
except openai.AuthenticationError as e:
raise LLMUnavailable("Invalid API key") from e
except (openai.APIStatusError, openai.APIConnectionError) as e:
last_exc = e
if attempt < max_retries:
time.sleep(2 ** attempt)
continue
raise LLMUnavailable(f"API error after retries: {last_exc}") from last_exc
|