evyzacq
146 lines
4.7 KiB
Python
146 lines
4.7 KiB
Python
# OpenAI-compatible LLM client with retry, token tracking, and vision support.
|
||
|
||
import base64
|
||
import logging
|
||
import os
|
||
import time
|
||
from typing import Optional
|
||
|
||
from openai import OpenAI
|
||
|
||
logger = logging.getLogger("testflow")
|
||
|
||
|
||
class LLMClient:
|
||
"""Generic OpenAI-compatible LLM client.
|
||
|
||
Usage::
|
||
|
||
client = LLMClient(api_key="sk-xxx", base_url="https://api.deepseek.com/v1")
|
||
text = client.chat("deepseek-chat", [{"role": "user", "content": "Hello"}])
|
||
print(client.usage)
|
||
"""
|
||
|
||
TIMEOUT = 120
|
||
MAX_RETRIES = 3
|
||
|
||
def __init__(
|
||
self,
|
||
*,
|
||
api_key: str = "",
|
||
base_url: str = "",
|
||
model: str = "",
|
||
timeout: int | None = None,
|
||
):
|
||
if not api_key:
|
||
raise ValueError(f"API key is required for LLMClient")
|
||
self._client = OpenAI(api_key=api_key, base_url=base_url)
|
||
self._timeout = timeout or self.TIMEOUT
|
||
self._model = model
|
||
self._prompt_tokens = 0
|
||
self._completion_tokens = 0
|
||
|
||
@property
|
||
def usage(self) -> dict:
|
||
return {
|
||
"prompt_tokens": self._prompt_tokens,
|
||
"completion_tokens": self._completion_tokens,
|
||
"total_tokens": self._prompt_tokens + self._completion_tokens,
|
||
}
|
||
|
||
@staticmethod
|
||
def estimate_tokens(text: str) -> int:
|
||
cjk = sum(1 for c in text if '一' <= c <= '鿿' or ' ' <= c <= '〿')
|
||
other = len(text) - cjk
|
||
return max(1, int(cjk / 1.7 + other / 3.0))
|
||
|
||
@staticmethod
|
||
def estimate_image_tokens() -> int:
|
||
return 500
|
||
|
||
def chat(
|
||
self,
|
||
model: str,
|
||
messages: list[dict],
|
||
*,
|
||
timeout: int | None = None,
|
||
response_format: dict | None = None,
|
||
temperature: float = 0.0,
|
||
) -> str:
|
||
# Estimate prompt size
|
||
prompt_chars = sum(len(str(m.get("content", ""))) for m in messages)
|
||
label = f"chat({model})"
|
||
|
||
def _call():
|
||
t0 = time.time()
|
||
kwargs = dict(
|
||
model=model,
|
||
messages=messages,
|
||
timeout=timeout or self._timeout,
|
||
temperature=temperature,
|
||
)
|
||
if response_format is not None:
|
||
kwargs["response_format"] = response_format
|
||
logger.info("[LLM] → %s (prompt ~%d chars / ~%d tokens)",
|
||
model, prompt_chars, prompt_chars // 3)
|
||
resp = self._client.chat.completions.create(**kwargs)
|
||
content = resp.choices[0].message.content
|
||
usg = resp.usage
|
||
if usg:
|
||
self._prompt_tokens += usg.prompt_tokens
|
||
self._completion_tokens += usg.completion_tokens
|
||
elapsed = time.time() - t0
|
||
logger.info("[LLM] ← %s: %d chars, p:%d c:%d tokens, %.1fs",
|
||
model, len(content) if content else 0,
|
||
usg.prompt_tokens if usg else 0,
|
||
usg.completion_tokens if usg else 0,
|
||
elapsed)
|
||
if not content:
|
||
raise RuntimeError("Empty response from LLM")
|
||
return content
|
||
|
||
return self._retry(_call, label)
|
||
|
||
def chat_with_image(
|
||
self,
|
||
model: str,
|
||
image_path: str,
|
||
prompt: str,
|
||
*,
|
||
timeout: int | None = None,
|
||
) -> str:
|
||
"""Send a prompt with an image attachment (for vision models)."""
|
||
with open(image_path, "rb") as f:
|
||
img_b64 = base64.b64encode(f.read()).decode()
|
||
mime = self._mime_type(image_path)
|
||
|
||
messages = [{
|
||
"role": "user",
|
||
"content": [
|
||
{"type": "image_url", "image_url": {"url": f"data:{mime};base64,{img_b64}"}},
|
||
{"type": "text", "text": prompt},
|
||
],
|
||
}]
|
||
return self.chat(model, messages, timeout=timeout)
|
||
|
||
@staticmethod
|
||
def _mime_type(image_path: str) -> str:
|
||
ext = os.path.splitext(image_path)[1].lstrip(".").lower()
|
||
return {
|
||
"png": "image/png", "jpg": "image/jpeg", "jpeg": "image/jpeg",
|
||
"gif": "image/gif", "bmp": "image/bmp",
|
||
"webp": "image/webp", "svg": "image/svg+xml", "tiff": "image/tiff",
|
||
}.get(ext, "image/png")
|
||
|
||
def _retry(self, fn, label: str) -> str:
|
||
last_error: Optional[Exception] = None
|
||
for attempt in range(self.MAX_RETRIES):
|
||
try:
|
||
return fn()
|
||
except Exception as e:
|
||
last_error = e
|
||
logger.warning("%s error (attempt %d/%d): %s", label, attempt + 1, self.MAX_RETRIES, e)
|
||
if attempt < self.MAX_RETRIES - 1:
|
||
time.sleep(2 ** attempt)
|
||
raise RuntimeError(f"{label}: all retries exhausted") from last_error
|