feat: initial release - Hunzi agent framework v0.1.0
This commit is contained in:
@@ -0,0 +1,512 @@
|
||||
"""LLM client module wrapping OpenAI-compatible APIs."""
|
||||
|
||||
import json
|
||||
import logging
|
||||
import math
|
||||
import re
|
||||
import time
|
||||
from typing import Any, AsyncGenerator, Dict, Generator, List, Optional, Union
|
||||
|
||||
import httpx
|
||||
|
||||
from app.config import LLMProviderConfig
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# ---------- helpers ----------------------------------------------------------------
|
||||
|
||||
def _encode_image(image: str) -> str:
|
||||
"""Return image string ready for the image_url field.
|
||||
|
||||
- If it already looks like a data URI or plain URL, pass through.
|
||||
- Otherwise treat the string as raw base64 and wrap it.
|
||||
"""
|
||||
if image.startswith("data:") or image.startswith("http://") or image.startswith("https://"):
|
||||
return image
|
||||
return f"data:image/png;base64,{image}"
|
||||
|
||||
|
||||
def _build_vision_content(text: str, images: List[str]) -> list:
|
||||
"""Build the OpenAI multimodal content array."""
|
||||
parts: list[dict] = [{"type": "text", "text": text}]
|
||||
for img in images:
|
||||
parts.append({"type": "image_url", "image_url": {"url": _encode_image(img)}})
|
||||
return parts
|
||||
|
||||
|
||||
# ---------- exceptions ------------------------------------------------------------
|
||||
|
||||
class LLMError(Exception):
|
||||
"""Base exception for LLM client errors."""
|
||||
|
||||
|
||||
class LLMAPIError(LLMError):
|
||||
"""Raised when the remote API returns a non-2xx response."""
|
||||
def __init__(self, status: int, detail: str, body: Any = None):
|
||||
self.status = status
|
||||
self.body = body
|
||||
super().__init__(f"API error {status}: {detail}")
|
||||
|
||||
|
||||
class LLMAvailabilityError(LLMAPIError):
|
||||
"""Transient errors that are safe to retry (5xx, rate-limit)."""
|
||||
|
||||
|
||||
class LLMPermanentError(LLMAPIError):
|
||||
"""Client-side / permanent errors (4xx excluding rate-limit)."""
|
||||
|
||||
|
||||
# ---------- retry logic -----------------------------------------------------------
|
||||
|
||||
_RETRYABLE_STATUS_CODES = {429, 500, 502, 503, 504}
|
||||
_RATELIMIT_CODES = {429}
|
||||
|
||||
|
||||
def _is_retryable(status: int) -> bool:
|
||||
return status in _RETRYABLE_STATUS_CODES
|
||||
|
||||
|
||||
def _extract_retry_after(response: httpx.Response) -> Optional[float]:
|
||||
raw = response.headers.get("Retry-After")
|
||||
if raw is None:
|
||||
return None
|
||||
try:
|
||||
return float(raw)
|
||||
except ValueError:
|
||||
return None
|
||||
|
||||
|
||||
# ---------- client ----------------------------------------------------------------
|
||||
|
||||
class LLMClient:
|
||||
"""Client for OpenAI-compatible chat completion APIs.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
config : LLMProviderConfig
|
||||
Per-provider configuration (base URL, model, API key, etc.).
|
||||
max_retries : int
|
||||
Maximum number of retry attempts for transient errors.
|
||||
timeout : float
|
||||
Request timeout in seconds.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
config: LLMProviderConfig,
|
||||
max_retries: int = 3,
|
||||
timeout: float = 120.0,
|
||||
):
|
||||
self.config = config
|
||||
self.max_retries = max_retries
|
||||
self.timeout = timeout
|
||||
self._sync_client: Optional[httpx.Client] = None
|
||||
self._async_client: Optional[httpx.AsyncClient] = None
|
||||
|
||||
# -- lazy client helpers --------------------------------------------------
|
||||
|
||||
def _sync_headers(self) -> Dict[str, str]:
|
||||
return {
|
||||
"Content-Type": "application/json",
|
||||
"Authorization": f"Bearer {self.config.api_key}",
|
||||
}
|
||||
|
||||
@property
|
||||
def _client(self) -> httpx.Client:
|
||||
if self._sync_client is None:
|
||||
self._sync_client = httpx.Client(
|
||||
timeout=httpx.Timeout(self.timeout),
|
||||
follow_redirects=True,
|
||||
)
|
||||
return self._sync_client
|
||||
|
||||
@property
|
||||
def _async_client(self) -> httpx.AsyncClient:
|
||||
if self._async_client is None:
|
||||
self._async_client = httpx.AsyncClient(
|
||||
timeout=httpx.Timeout(self.timeout),
|
||||
follow_redirects=True,
|
||||
)
|
||||
return self._async_client
|
||||
|
||||
def close(self) -> None:
|
||||
"""Close the synchronous HTTP client."""
|
||||
if self._sync_client is not None:
|
||||
self._sync_client.close()
|
||||
self._sync_client = None
|
||||
|
||||
async def aclose(self) -> None:
|
||||
"""Close the asynchronous HTTP client."""
|
||||
if self._async_client is not None:
|
||||
await self._async_client.aclose()
|
||||
self._async_client = None
|
||||
|
||||
# -- request building -------------------------------------------------------
|
||||
|
||||
def _build_request_body(
|
||||
self,
|
||||
messages: List[Dict[str, str]],
|
||||
stream: bool = False,
|
||||
**extra: Any,
|
||||
) -> Dict[str, Any]:
|
||||
body: Dict[str, Any] = {
|
||||
"model": self.config.model,
|
||||
"messages": messages,
|
||||
"stream": stream,
|
||||
}
|
||||
if self.config.temperature is not None:
|
||||
body["temperature"] = self.config.temperature
|
||||
if self.config.max_tokens is not None:
|
||||
body["max_tokens"] = self.config.max_tokens
|
||||
if self.config.top_p is not None:
|
||||
body["top_p"] = self.config.top_p
|
||||
body.update(extra)
|
||||
return body
|
||||
|
||||
# -- core sync ----------------------------------------------------------------
|
||||
|
||||
def _do_sync_request(
|
||||
self,
|
||||
body: Dict[str, Any],
|
||||
stream: bool = False,
|
||||
) -> httpx.Response:
|
||||
"""Execute a single sync request with retry logic."""
|
||||
url = f"{self.config.base_url}/chat/completions"
|
||||
headers = self._sync_headers()
|
||||
|
||||
last_exc: Optional[Exception] = None
|
||||
for attempt in range(self.max_retries + 1):
|
||||
try:
|
||||
resp = self._client.post(url, json=body, headers=headers)
|
||||
except (httpx.ConnectError, httpx.TimeoutException) as exc:
|
||||
last_exc = exc
|
||||
logger.warning(
|
||||
"Request failed (attempt %d/%d): %s",
|
||||
attempt + 1,
|
||||
self.max_retries + 1,
|
||||
exc,
|
||||
)
|
||||
if attempt < self.max_retries:
|
||||
time.sleep(min(2 ** attempt, 30))
|
||||
continue
|
||||
|
||||
if resp.status_code == 200:
|
||||
return resp
|
||||
|
||||
body_text = ""
|
||||
try:
|
||||
body_text = resp.json()
|
||||
except Exception:
|
||||
body_text = resp.text
|
||||
|
||||
if _is_retryable(resp.status_code) and attempt < self.max_retries:
|
||||
retry_after = _extract_retry_after(resp) or min(2 ** attempt, 30)
|
||||
logger.warning(
|
||||
"Retryable error %d (attempt %d/%d), waiting %.1fs: %s",
|
||||
resp.status_code,
|
||||
attempt + 1,
|
||||
self.max_retries + 1,
|
||||
retry_after,
|
||||
body_text,
|
||||
)
|
||||
time.sleep(retry_after)
|
||||
continue
|
||||
|
||||
# Non-retryable or exhausted retries
|
||||
raise LLMAvailabilityError(resp.status_code, body_text, body_text) if _is_retryable(resp.status_code) else LLMPermanentError(resp.status_code, body_text, body_text)
|
||||
|
||||
if last_exc is not None:
|
||||
raise last_exc
|
||||
raise LLMError("Unexpected control-flow exit in retry loop")
|
||||
|
||||
# -- core async -------------------------------------------------------------
|
||||
|
||||
async def _do_async_request(
|
||||
self,
|
||||
body: Dict[str, Any],
|
||||
stream: bool = False,
|
||||
) -> httpx.Response:
|
||||
"""Execute a single async request with retry logic."""
|
||||
url = f"{self.config.base_url}/chat/completions"
|
||||
headers = self._sync_headers()
|
||||
|
||||
last_exc: Optional[Exception] = None
|
||||
for attempt in range(self.max_retries + 1):
|
||||
try:
|
||||
resp = await self._async_client.post(url, json=body, headers=headers)
|
||||
except (httpx.ConnectError, httpx.TimeoutException) as exc:
|
||||
last_exc = exc
|
||||
logger.warning(
|
||||
"Async request failed (attempt %d/%d): %s",
|
||||
attempt + 1,
|
||||
self.max_retries + 1,
|
||||
exc,
|
||||
)
|
||||
if attempt < self.max_retries:
|
||||
await _async_sleep(min(2 ** attempt, 30))
|
||||
continue
|
||||
|
||||
if resp.status_code == 200:
|
||||
return resp
|
||||
|
||||
body_text = ""
|
||||
try:
|
||||
body_text = resp.json()
|
||||
except Exception:
|
||||
body_text = resp.text
|
||||
|
||||
if _is_retryable(resp.status_code) and attempt < self.max_retries:
|
||||
retry_after = _extract_retry_after(resp) or min(2 ** attempt, 30)
|
||||
logger.warning(
|
||||
"Async retryable error %d (attempt %d/%d), waiting %.1fs: %s",
|
||||
resp.status_code,
|
||||
attempt + 1,
|
||||
self.max_retries + 1,
|
||||
retry_after,
|
||||
body_text,
|
||||
)
|
||||
await _async_sleep(retry_after)
|
||||
continue
|
||||
|
||||
raise LLMAvailabilityError(resp.status_code, body_text, body_text) if _is_retryable(resp.status_code) else LLMPermanentError(resp.status_code, body_text, body_text)
|
||||
|
||||
if last_exc is not None:
|
||||
raise last_exc
|
||||
raise LLMError("Unexpected control-flow exit in async retry loop")
|
||||
|
||||
# -- public sync API --------------------------------------------------------
|
||||
|
||||
def chat(
|
||||
self,
|
||||
messages: List[Dict[str, str]],
|
||||
stream: bool = False,
|
||||
**extra: Any,
|
||||
) -> Union[Dict[str, Any], Generator[str, None, None]]:
|
||||
"""Send a chat completion request.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
messages : list[dict]
|
||||
Conversation messages in OpenAI format.
|
||||
stream : bool
|
||||
If True, returns a generator yielding content chunks.
|
||||
**extra : any
|
||||
Additional parameters forwarded to the API (e.g. ``tools``, ``response_format``).
|
||||
|
||||
Returns
|
||||
-------
|
||||
dict or Generator[str, None, None]
|
||||
The parsed JSON response body, or a generator of content strings when streaming.
|
||||
"""
|
||||
body = self._build_request_body(messages, stream=stream, **extra)
|
||||
|
||||
if stream:
|
||||
return self.chat_stream(messages, **extra)
|
||||
|
||||
resp = self._do_sync_request(body, stream=False)
|
||||
return resp.json()
|
||||
|
||||
def chat_stream(
|
||||
self,
|
||||
messages: List[Dict[str, str]],
|
||||
**extra: Any,
|
||||
) -> Generator[str, None, None]:
|
||||
"""Yield streaming content chunks from the API.
|
||||
|
||||
Each yielded string is a content delta (``delta.content``) from the
|
||||
SSE stream. Non-content chunks (role, function calls) are silently
|
||||
skipped.
|
||||
"""
|
||||
body = self._build_request_body(messages, stream=True, **extra)
|
||||
url = f"{self.config.base_url}/chat/completions"
|
||||
headers = self._sync_headers()
|
||||
|
||||
with self._client.stream(
|
||||
"POST", url, json=body, headers=headers, timeout=self.timeout
|
||||
) as resp:
|
||||
if resp.status_code != 200:
|
||||
err_body = resp.read()
|
||||
try:
|
||||
err_json = json.loads(err_body)
|
||||
detail = err_json.get("error", {}).get("message", err_body)
|
||||
except Exception:
|
||||
detail = err_body
|
||||
if _is_retryable(resp.status_code):
|
||||
raise LLMAvailabilityError(resp.status_code, detail, err_body)
|
||||
raise LLMPermanentError(resp.status_code, detail, err_body)
|
||||
|
||||
for line in resp.iter_lines():
|
||||
line = line.strip()
|
||||
if not line:
|
||||
continue
|
||||
if line.startswith("data: "):
|
||||
payload = line[6:]
|
||||
elif line.startswith("data:"):
|
||||
payload = line[5:]
|
||||
else:
|
||||
continue
|
||||
|
||||
payload = payload.strip()
|
||||
if payload == "[DONE]":
|
||||
break
|
||||
|
||||
try:
|
||||
chunk = json.loads(payload)
|
||||
except json.JSONDecodeError:
|
||||
continue
|
||||
|
||||
content = (
|
||||
chunk.get("choices", [{}])[0]
|
||||
.get("delta", {})
|
||||
.get("content")
|
||||
)
|
||||
if content is not None:
|
||||
yield content
|
||||
|
||||
def vision_chat(
|
||||
self,
|
||||
text: str,
|
||||
images: List[str],
|
||||
stream: bool = False,
|
||||
**extra: Any,
|
||||
) -> Union[Dict[str, Any], Generator[str, None, None]]:
|
||||
"""Send a multimodal (text + images) chat request.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
text : str
|
||||
The user's text prompt.
|
||||
images : list[str]
|
||||
List of image references — base64 strings or URLs.
|
||||
stream : bool
|
||||
If True, returns a generator yielding content chunks.
|
||||
**extra : any
|
||||
Additional parameters forwarded to the API.
|
||||
|
||||
Returns
|
||||
-------
|
||||
dict or Generator[str, None, None]
|
||||
"""
|
||||
content = _build_vision_content(text, images)
|
||||
messages = [{"role": "user", "content": content}]
|
||||
return self.chat(messages, stream=stream, **extra)
|
||||
|
||||
# -- public async API -------------------------------------------------------
|
||||
|
||||
async def achat(
|
||||
self,
|
||||
messages: List[Dict[str, str]],
|
||||
stream: bool = False,
|
||||
**extra: Any,
|
||||
) -> Union[Dict[str, Any], AsyncGenerator[str, None]]:
|
||||
"""Async version of :meth:`chat`."""
|
||||
body = self._build_request_body(messages, stream=stream, **extra)
|
||||
|
||||
if stream:
|
||||
return self.achat_stream(messages, **extra)
|
||||
|
||||
resp = await self._do_async_request(body, stream=False)
|
||||
return resp.json()
|
||||
|
||||
async def achat_stream(
|
||||
self,
|
||||
messages: List[Dict[str, str]],
|
||||
**extra: Any,
|
||||
) -> AsyncGenerator[str, None]:
|
||||
"""Async streaming chat — yields content delta strings."""
|
||||
body = self._build_request_body(messages, stream=True, **extra)
|
||||
url = f"{self.config.base_url}/chat/completions"
|
||||
headers = self._sync_headers()
|
||||
|
||||
async with self._async_client.stream(
|
||||
"POST", url, json=body, headers=headers
|
||||
) as resp:
|
||||
if resp.status_code != 200:
|
||||
err_body = await resp.aread()
|
||||
try:
|
||||
err_json = json.loads(err_body)
|
||||
detail = err_json.get("error", {}).get("message", err_body)
|
||||
except Exception:
|
||||
detail = err_body
|
||||
if _is_retryable(resp.status_code):
|
||||
raise LLMAvailabilityError(resp.status_code, detail, err_body)
|
||||
raise LLMPermanentError(resp.status_code, detail, err_body)
|
||||
|
||||
async for line in resp.aiter_lines():
|
||||
line = line.strip()
|
||||
if not line:
|
||||
continue
|
||||
if line.startswith("data: "):
|
||||
payload = line[6:]
|
||||
elif line.startswith("data:"):
|
||||
payload = line[5:]
|
||||
else:
|
||||
continue
|
||||
|
||||
payload = payload.strip()
|
||||
if payload == "[DONE]":
|
||||
break
|
||||
|
||||
try:
|
||||
chunk = json.loads(payload)
|
||||
except json.JSONDecodeError:
|
||||
continue
|
||||
|
||||
content = (
|
||||
chunk.get("choices", [{}])[0]
|
||||
.get("delta", {})
|
||||
.get("content")
|
||||
)
|
||||
if content is not None:
|
||||
yield content
|
||||
|
||||
async def avision_chat(
|
||||
self,
|
||||
text: str,
|
||||
images: List[str],
|
||||
stream: bool = False,
|
||||
**extra: Any,
|
||||
) -> Union[Dict[str, Any], AsyncGenerator[str, None]]:
|
||||
"""Async version of :meth:`vision_chat`."""
|
||||
content = _build_vision_content(text, images)
|
||||
messages = [{"role": "user", "content": content}]
|
||||
return await self.achat(messages, stream=stream, **extra)
|
||||
|
||||
# -- token estimation -------------------------------------------------------
|
||||
|
||||
@staticmethod
|
||||
def count_tokens(text: str) -> int:
|
||||
"""Roughly estimate the number of tokens in *text*.
|
||||
|
||||
Uses a simple heuristic based on average English word length and
|
||||
punctuation. For precise counts use a real tokenizer (e.g.
|
||||
``tiktoken``).
|
||||
"""
|
||||
if not text:
|
||||
return 0
|
||||
# Heuristic: ~4 chars per token on average for English text,
|
||||
# with a floor of 1 token per word to handle short words / tokens.
|
||||
word_count = len(text.split())
|
||||
char_estimate = math.ceil(len(text) / 4)
|
||||
return max(word_count, char_estimate)
|
||||
|
||||
# -- context manager support ------------------------------------------------
|
||||
|
||||
def __enter__(self) -> "LLMClient":
|
||||
return self
|
||||
|
||||
def __exit__(self, *args: Any) -> None:
|
||||
self.close()
|
||||
|
||||
async def __aenter__(self) -> "LLMClient":
|
||||
return self
|
||||
|
||||
async def __aexit__(self, *args: Any) -> None:
|
||||
await self.aclose()
|
||||
|
||||
|
||||
# -- async sleep helper (avoids importing asyncio at top-level) -------------------
|
||||
|
||||
async def _async_sleep(delay: float) -> None:
|
||||
import asyncio
|
||||
await asyncio.sleep(delay)
|
||||
Reference in New Issue
Block a user