"""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, }, }