244 lines
9.9 KiB
Python
244 lines
9.9 KiB
Python
|
"""OpenAI-compatible LLM client with tool calling, streaming, and logging."""
|
||
|
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import json
|
||
|
|
import time
|
||
|
|
import uuid
|
||
|
|
from typing import Any, AsyncIterator, Dict, List, Optional
|
||
|
|
|
||
|
|
import httpx
|
||
|
|
from sqlalchemy.ext.asyncio import AsyncSession
|
||
|
|
|
||
|
|
from app.core.settings_service import get_all_settings, cast_setting
|
||
|
|
from app.logging_setup import get_logger
|
||
|
|
from app.models import LlmCallLog
|
||
|
|
|
||
|
|
log = get_logger("llm")
|
||
|
|
|
||
|
|
|
||
|
|
class LlmResponse:
|
||
|
|
"""Non-streaming response wrapper."""
|
||
|
|
|
||
|
|
def __init__(self, text: str, tool_calls: List[Dict[str, Any]], usage: Optional[Dict[str, int]]):
|
||
|
|
self.text = text
|
||
|
|
self.tool_calls = tool_calls
|
||
|
|
self.usage = usage or {}
|
||
|
|
|
||
|
|
|
||
|
|
class LlmClient:
|
||
|
|
"""Lightweight OpenAI-compatible chat-completions client."""
|
||
|
|
|
||
|
|
def __init__(self, settings_map: Dict[str, Any]):
|
||
|
|
self.base_url: str = str(settings_map.get("llm.base_url", "")).rstrip("/")
|
||
|
|
self.api_key: str = str(settings_map.get("llm.api_key", "dummy"))
|
||
|
|
self.model: str = str(settings_map.get("llm.model", "local-model"))
|
||
|
|
self.temperature: float = float(cast_setting("llm.temperature", settings_map.get("llm.temperature", 0.7)))
|
||
|
|
self.max_tokens: int = int(cast_setting("llm.max_tokens", settings_map.get("llm.max_tokens", 1024)))
|
||
|
|
self.timeout: int = int(cast_setting("llm.request_timeout", settings_map.get("llm.request_timeout", 120)))
|
||
|
|
self.streaming: bool = bool(cast_setting("llm.streaming", settings_map.get("llm.streaming", True)))
|
||
|
|
|
||
|
|
@classmethod
|
||
|
|
async def from_db(cls, db: AsyncSession) -> "LlmClient":
|
||
|
|
s = await get_all_settings(db)
|
||
|
|
return cls(s)
|
||
|
|
|
||
|
|
def _headers(self) -> Dict[str, str]:
|
||
|
|
h = {"Content-Type": "application/json"}
|
||
|
|
if self.api_key and self.api_key != "dummy":
|
||
|
|
h["Authorization"] = f"Bearer {self.api_key}"
|
||
|
|
return h
|
||
|
|
|
||
|
|
async def chat(
|
||
|
|
self,
|
||
|
|
messages: List[Dict[str, Any]],
|
||
|
|
tools: Optional[List[Dict[str, Any]]] = None,
|
||
|
|
tool_choice: Any = None,
|
||
|
|
temperature: Optional[float] = None,
|
||
|
|
max_tokens: Optional[int] = None,
|
||
|
|
purpose: str = "orchestrator",
|
||
|
|
user_id: Optional[uuid.UUID] = None,
|
||
|
|
session_id: Optional[uuid.UUID] = None,
|
||
|
|
db: Optional[AsyncSession] = None,
|
||
|
|
) -> LlmResponse:
|
||
|
|
"""Non-streaming chat completion with tool support."""
|
||
|
|
url = f"{self.base_url}/chat/completions"
|
||
|
|
payload: Dict[str, Any] = {
|
||
|
|
"model": self.model,
|
||
|
|
"messages": messages,
|
||
|
|
"temperature": temperature if temperature is not None else self.temperature,
|
||
|
|
"max_tokens": max_tokens or self.max_tokens,
|
||
|
|
"stream": False,
|
||
|
|
}
|
||
|
|
if tools:
|
||
|
|
payload["tools"] = tools
|
||
|
|
if tool_choice is not None:
|
||
|
|
payload["tool_choice"] = tool_choice
|
||
|
|
started = time.monotonic()
|
||
|
|
err: Optional[str] = None
|
||
|
|
text = ""
|
||
|
|
tool_calls: List[Dict[str, Any]] = []
|
||
|
|
usage: Dict[str, int] = {}
|
||
|
|
try:
|
||
|
|
async with httpx.AsyncClient(timeout=self.timeout) as client:
|
||
|
|
resp = await client.post(url, json=payload, headers=self._headers())
|
||
|
|
resp.raise_for_status()
|
||
|
|
data = resp.json()
|
||
|
|
choice = (data.get("choices") or [{}])[0]
|
||
|
|
msg = choice.get("message", {})
|
||
|
|
text = msg.get("content") or ""
|
||
|
|
tool_calls = msg.get("tool_calls") or []
|
||
|
|
usage = data.get("usage") or {}
|
||
|
|
except Exception as e:
|
||
|
|
err = f"{type(e).__name__}: {e}"
|
||
|
|
log.error("llm_call_failed", purpose=purpose, error=err)
|
||
|
|
raise
|
||
|
|
finally:
|
||
|
|
latency_ms = int((time.monotonic() - started) * 1000)
|
||
|
|
if db is not None:
|
||
|
|
db.add(LlmCallLog(
|
||
|
|
user_id=user_id,
|
||
|
|
session_id=session_id,
|
||
|
|
purpose=purpose,
|
||
|
|
model=self.model,
|
||
|
|
base_url=self.base_url,
|
||
|
|
prompt_messages=messages,
|
||
|
|
tools=tools,
|
||
|
|
response_text=text,
|
||
|
|
tool_calls=tool_calls,
|
||
|
|
prompt_tokens=usage.get("prompt_tokens"),
|
||
|
|
completion_tokens=usage.get("completion_tokens"),
|
||
|
|
total_tokens=usage.get("total_tokens"),
|
||
|
|
latency_ms=latency_ms,
|
||
|
|
error=err,
|
||
|
|
))
|
||
|
|
try:
|
||
|
|
await db.commit()
|
||
|
|
except Exception:
|
||
|
|
await db.rollback()
|
||
|
|
return LlmResponse(text=text, tool_calls=tool_calls, usage=usage)
|
||
|
|
|
||
|
|
async def stream_chat(
|
||
|
|
self,
|
||
|
|
messages: List[Dict[str, Any]],
|
||
|
|
tools: Optional[List[Dict[str, Any]]] = None,
|
||
|
|
tool_choice: Any = None,
|
||
|
|
temperature: Optional[float] = None,
|
||
|
|
max_tokens: Optional[int] = None,
|
||
|
|
purpose: str = "orchestrator",
|
||
|
|
user_id: Optional[uuid.UUID] = None,
|
||
|
|
session_id: Optional[uuid.UUID] = None,
|
||
|
|
db: Optional[AsyncSession] = None,
|
||
|
|
) -> AsyncIterator[Dict[str, Any]]:
|
||
|
|
"""Streaming chat completion. Yields incremental deltas.
|
||
|
|
|
||
|
|
Yields dicts of the form:
|
||
|
|
{"type": "delta", "content": "..."} - text delta
|
||
|
|
{"type": "tool_calls", "tool_calls": [...]} - final tool calls (if any)
|
||
|
|
{"type": "done", "usage": {...}}
|
||
|
|
{"type": "error", "error": "..."}
|
||
|
|
"""
|
||
|
|
url = f"{self.base_url}/chat/completions"
|
||
|
|
payload: Dict[str, Any] = {
|
||
|
|
"model": self.model,
|
||
|
|
"messages": messages,
|
||
|
|
"temperature": temperature if temperature is not None else self.temperature,
|
||
|
|
"max_tokens": max_tokens or self.max_tokens,
|
||
|
|
"stream": True,
|
||
|
|
}
|
||
|
|
if tools:
|
||
|
|
payload["tools"] = tools
|
||
|
|
if tool_choice is not None:
|
||
|
|
payload["tool_choice"] = tool_choice
|
||
|
|
|
||
|
|
started = time.monotonic()
|
||
|
|
full_text_parts: List[str] = []
|
||
|
|
tool_call_accum: Dict[int, Dict[str, Any]] = {}
|
||
|
|
usage: Dict[str, int] = {}
|
||
|
|
err: Optional[str] = None
|
||
|
|
|
||
|
|
try:
|
||
|
|
async with httpx.AsyncClient(timeout=self.timeout) as client:
|
||
|
|
async with client.stream("POST", url, json=payload, headers=self._headers()) as resp:
|
||
|
|
resp.raise_for_status()
|
||
|
|
async for line in resp.aiter_lines():
|
||
|
|
if not line or not line.startswith("data:"):
|
||
|
|
continue
|
||
|
|
data_str = line[5:].strip()
|
||
|
|
if data_str == "[DONE]":
|
||
|
|
break
|
||
|
|
try:
|
||
|
|
chunk = json.loads(data_str)
|
||
|
|
except json.JSONDecodeError:
|
||
|
|
continue
|
||
|
|
choices = chunk.get("choices") or []
|
||
|
|
if not choices:
|
||
|
|
if chunk.get("usage"):
|
||
|
|
usage = chunk["usage"]
|
||
|
|
continue
|
||
|
|
delta = choices[0].get("delta", {})
|
||
|
|
if delta.get("content"):
|
||
|
|
full_text_parts.append(delta["content"])
|
||
|
|
yield {"type": "delta", "content": delta["content"]}
|
||
|
|
if delta.get("tool_calls"):
|
||
|
|
for tc in delta["tool_calls"]:
|
||
|
|
idx = tc.get("index", 0)
|
||
|
|
acc = tool_call_accum.setdefault(idx, {
|
||
|
|
"id": tc.get("id", ""),
|
||
|
|
"type": "function",
|
||
|
|
"function": {"name": "", "arguments": ""},
|
||
|
|
})
|
||
|
|
if tc.get("id"):
|
||
|
|
acc["id"] = tc["id"]
|
||
|
|
if tc.get("function", {}).get("name"):
|
||
|
|
acc["function"]["name"] += tc["function"]["name"]
|
||
|
|
if tc.get("function", {}).get("arguments"):
|
||
|
|
acc["function"]["arguments"] += tc["function"]["arguments"]
|
||
|
|
if chunk.get("usage"):
|
||
|
|
usage = chunk["usage"]
|
||
|
|
except Exception as e:
|
||
|
|
err = f"{type(e).__name__}: {e}"
|
||
|
|
log.error("llm_stream_failed", purpose=purpose, error=err)
|
||
|
|
yield {"type": "error", "error": err}
|
||
|
|
return
|
||
|
|
|
||
|
|
full_text = "".join(full_text_parts)
|
||
|
|
final_tool_calls = [tool_call_accum[i] for i in sorted(tool_call_accum.keys())]
|
||
|
|
if final_tool_calls:
|
||
|
|
yield {"type": "tool_calls", "tool_calls": final_tool_calls}
|
||
|
|
yield {"type": "done", "usage": usage, "full_text": full_text}
|
||
|
|
|
||
|
|
latency_ms = int((time.monotonic() - started) * 1000)
|
||
|
|
if db is not None:
|
||
|
|
db.add(LlmCallLog(
|
||
|
|
user_id=user_id,
|
||
|
|
session_id=session_id,
|
||
|
|
purpose=purpose,
|
||
|
|
model=self.model,
|
||
|
|
base_url=self.base_url,
|
||
|
|
prompt_messages=messages,
|
||
|
|
tools=tools,
|
||
|
|
response_text=full_text,
|
||
|
|
tool_calls=final_tool_calls,
|
||
|
|
prompt_tokens=usage.get("prompt_tokens"),
|
||
|
|
completion_tokens=usage.get("completion_tokens"),
|
||
|
|
total_tokens=usage.get("total_tokens"),
|
||
|
|
latency_ms=latency_ms,
|
||
|
|
error=err,
|
||
|
|
))
|
||
|
|
try:
|
||
|
|
await db.commit()
|
||
|
|
except Exception:
|
||
|
|
await db.rollback()
|
||
|
|
|
||
|
|
|
||
|
|
def build_tool_schema(name: str, description: str, params: Dict[str, Any]) -> Dict[str, Any]:
|
||
|
|
"""Helper to build an OpenAI-style tool schema."""
|
||
|
|
return {
|
||
|
|
"type": "function",
|
||
|
|
"function": {
|
||
|
|
"name": name,
|
||
|
|
"description": description,
|
||
|
|
"parameters": params,
|
||
|
|
},
|
||
|
|
}
|