initial
This commit is contained in:
Mikan
243
backend/app/core/llm.py
Normal file
243
backend/app/core/llm.py
Normal file
@@ -0,0 +1,243 @@
|
||||
"""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,
|
||||
},
|
||||
}
|
||||
Reference in New Issue
Block a user