init commit
This commit is contained in:
0
app/services/__init__.py
Normal file
0
app/services/__init__.py
Normal file
124
app/services/llm.py
Normal file
124
app/services/llm.py
Normal file
@@ -0,0 +1,124 @@
|
||||
"""LLM client wrapper for Ollama / OpenAI-compatible endpoints."""
|
||||
|
||||
import json
|
||||
from collections.abc import AsyncIterator
|
||||
from contextlib import asynccontextmanager
|
||||
from typing import Self
|
||||
|
||||
import httpx
|
||||
|
||||
from app.models.config import LLMConfig
|
||||
from app.models.message import Message
|
||||
from app.utils.logging import get_logger
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
|
||||
# --- Exception hierarchy ---
|
||||
|
||||
|
||||
class LLMError(Exception):
|
||||
"""Base exception for LLM client errors."""
|
||||
|
||||
|
||||
class LLMConnectionError(LLMError):
|
||||
"""Connection or timeout failure when reaching the LLM endpoint."""
|
||||
|
||||
|
||||
class LLMResponseError(LLMError):
|
||||
"""Non-2xx HTTP response from the LLM endpoint."""
|
||||
|
||||
def __init__(self, message: str, status_code: int | None = None) -> None:
|
||||
super().__init__(message)
|
||||
self.status_code = status_code
|
||||
|
||||
|
||||
class LLMStreamError(LLMError):
|
||||
"""Malformed SSE data during streaming."""
|
||||
|
||||
|
||||
# --- Client ---
|
||||
|
||||
|
||||
class LLMClient:
|
||||
"""Async streaming client for OpenAI-compatible chat completions.
|
||||
|
||||
Designed for Ollama but works with any endpoint implementing the
|
||||
OpenAI /v1/chat/completions SSE streaming protocol.
|
||||
"""
|
||||
|
||||
def __init__(self, config: LLMConfig) -> None:
|
||||
"""Initialize the LLM client.
|
||||
|
||||
Args:
|
||||
config: LLM configuration (model, endpoint, timeout, etc.).
|
||||
"""
|
||||
self._config = config
|
||||
self._client = httpx.AsyncClient(
|
||||
base_url=config.endpoint,
|
||||
timeout=httpx.Timeout(config.timeout, connect=10.0),
|
||||
)
|
||||
|
||||
async def stream_chat(self, messages: list[Message]) -> AsyncIterator[dict]:
|
||||
"""Stream a chat completion request, yielding parsed SSE chunks.
|
||||
|
||||
Args:
|
||||
messages: Conversation history to send to the model.
|
||||
|
||||
Yields:
|
||||
Parsed JSON dicts from each SSE data line.
|
||||
|
||||
Raises:
|
||||
LLMConnectionError: On connection or timeout failures.
|
||||
LLMResponseError: On non-2xx HTTP status.
|
||||
LLMStreamError: On malformed SSE data (only if every line fails).
|
||||
"""
|
||||
payload = {
|
||||
"model": self._config.model,
|
||||
"messages": [m.to_api_dict() for m in messages],
|
||||
"stream": True,
|
||||
"temperature": self._config.temperature,
|
||||
"max_tokens": self._config.max_tokens,
|
||||
}
|
||||
|
||||
try:
|
||||
async with self._client.stream(
|
||||
"POST", self._config.api_path, json=payload
|
||||
) as response:
|
||||
if response.status_code != 200:
|
||||
body = await response.aread()
|
||||
raise LLMResponseError(
|
||||
f"LLM returned {response.status_code}: {body.decode(errors='replace')}",
|
||||
status_code=response.status_code,
|
||||
)
|
||||
|
||||
async for line in response.aiter_lines():
|
||||
if not line.startswith("data: "):
|
||||
continue
|
||||
|
||||
data = line[6:] # strip "data: " prefix
|
||||
|
||||
if data.strip() == "[DONE]":
|
||||
return
|
||||
|
||||
try:
|
||||
yield json.loads(data)
|
||||
except json.JSONDecodeError:
|
||||
logger.warning("malformed_sse_chunk", data=data[:200])
|
||||
|
||||
except httpx.ConnectError as e:
|
||||
raise LLMConnectionError(f"Cannot connect to LLM endpoint: {e}") from e
|
||||
except httpx.TimeoutException as e:
|
||||
raise LLMConnectionError(f"LLM request timed out: {e}") from e
|
||||
except httpx.HTTPError as e:
|
||||
raise LLMError(f"HTTP error communicating with LLM: {e}") from e
|
||||
|
||||
async def close(self) -> None:
|
||||
"""Close the underlying HTTP client."""
|
||||
await self._client.aclose()
|
||||
|
||||
async def __aenter__(self) -> Self:
|
||||
return self
|
||||
|
||||
async def __aexit__(self, *exc: object) -> None:
|
||||
await self.close()
|
||||
150
app/services/streaming.py
Normal file
150
app/services/streaming.py
Normal file
@@ -0,0 +1,150 @@
|
||||
"""Streaming response handler — accumulates SSE chunks into a complete Message."""
|
||||
|
||||
from collections.abc import AsyncIterator
|
||||
|
||||
from rich.live import Live
|
||||
from rich.markdown import Markdown
|
||||
|
||||
from app.models.config import DisplayConfig
|
||||
from app.models.message import Message
|
||||
from app.models.tool_call import ToolCall, ToolCallFunction
|
||||
from app.utils.display import print_assistant_message
|
||||
from app.utils.logging import console, get_logger
|
||||
from app.utils.token_counter import TokenUsage
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
|
||||
class StreamHandler:
|
||||
"""Processes an SSE chunk stream into a Rich live display and final Message.
|
||||
|
||||
Accumulates content deltas and tool call fragments, renders a live Markdown
|
||||
panel during streaming, and produces a complete assistant Message on finish.
|
||||
"""
|
||||
|
||||
def __init__(self, display_config: DisplayConfig) -> None:
|
||||
"""Initialize the stream handler.
|
||||
|
||||
Args:
|
||||
display_config: Display preferences (streaming toggle, etc.).
|
||||
"""
|
||||
self._display_config = display_config
|
||||
self._accumulated_content: str = ""
|
||||
self._accumulated_reasoning: str = ""
|
||||
self._tool_calls: dict[int, dict[str, str]] = {}
|
||||
self._usage: TokenUsage | None = None
|
||||
|
||||
async def process_stream(self, chunk_iter: AsyncIterator[dict]) -> Message:
|
||||
"""Consume a chunk iterator, rendering live output and returning the final Message.
|
||||
|
||||
Args:
|
||||
chunk_iter: Async iterator of parsed SSE chunk dicts.
|
||||
|
||||
Returns:
|
||||
Complete assistant Message with accumulated content and tool calls.
|
||||
"""
|
||||
with Live(console=console, refresh_per_second=8) as live:
|
||||
async for chunk in chunk_iter:
|
||||
self._process_chunk(chunk)
|
||||
|
||||
# Show reasoning while waiting for content
|
||||
display_text = self._accumulated_content
|
||||
if not display_text and self._accumulated_reasoning:
|
||||
display_text = f"*thinking...*"
|
||||
|
||||
if display_text and self._display_config.stream_output:
|
||||
live.update(Markdown(display_text))
|
||||
|
||||
# Final static render
|
||||
if self._accumulated_content:
|
||||
print_assistant_message(self._accumulated_content)
|
||||
|
||||
tool_calls = self._build_tool_calls() or None
|
||||
return Message(
|
||||
role="assistant",
|
||||
content=self._accumulated_content or None,
|
||||
tool_calls=tool_calls,
|
||||
)
|
||||
|
||||
def _process_chunk(self, chunk: dict) -> None:
|
||||
"""Extract content, tool calls, and usage from a single SSE chunk.
|
||||
|
||||
Args:
|
||||
chunk: Parsed JSON dict from one SSE data line.
|
||||
"""
|
||||
# Content delta
|
||||
choices = chunk.get("choices", [])
|
||||
if choices:
|
||||
delta = choices[0].get("delta", {})
|
||||
|
||||
content_piece = delta.get("content")
|
||||
if content_piece:
|
||||
self._accumulated_content += content_piece
|
||||
|
||||
# Reasoning tokens (e.g. qwen3.5 thinking mode)
|
||||
reasoning_piece = delta.get("reasoning")
|
||||
if reasoning_piece:
|
||||
self._accumulated_reasoning += reasoning_piece
|
||||
|
||||
# Tool call deltas (accumulated by index)
|
||||
for tc_delta in delta.get("tool_calls", []):
|
||||
idx = tc_delta.get("index", 0)
|
||||
if idx not in self._tool_calls:
|
||||
self._tool_calls[idx] = {
|
||||
"id": tc_delta.get("id", ""),
|
||||
"name": "",
|
||||
"arguments": "",
|
||||
}
|
||||
entry = self._tool_calls[idx]
|
||||
if tc_delta.get("id"):
|
||||
entry["id"] = tc_delta["id"]
|
||||
func = tc_delta.get("function", {})
|
||||
if func.get("name"):
|
||||
entry["name"] += func["name"]
|
||||
if func.get("arguments"):
|
||||
entry["arguments"] += func["arguments"]
|
||||
|
||||
# Token usage (typically in the final chunk)
|
||||
usage_data = chunk.get("usage")
|
||||
if usage_data:
|
||||
self._usage = TokenUsage(
|
||||
prompt_tokens=usage_data.get("prompt_tokens", 0),
|
||||
completion_tokens=usage_data.get("completion_tokens", 0),
|
||||
total_tokens=usage_data.get("total_tokens", 0),
|
||||
)
|
||||
|
||||
def _build_tool_calls(self) -> list[ToolCall]:
|
||||
"""Convert accumulated tool call fragments into sorted ToolCall list.
|
||||
|
||||
Returns:
|
||||
List of ToolCall objects sorted by stream index.
|
||||
"""
|
||||
if not self._tool_calls:
|
||||
return []
|
||||
|
||||
result: list[ToolCall] = []
|
||||
for idx in sorted(self._tool_calls):
|
||||
entry = self._tool_calls[idx]
|
||||
result.append(
|
||||
ToolCall(
|
||||
id=entry["id"],
|
||||
type="function",
|
||||
function=ToolCallFunction(
|
||||
name=entry["name"],
|
||||
arguments=entry["arguments"],
|
||||
),
|
||||
)
|
||||
)
|
||||
return result
|
||||
|
||||
@property
|
||||
def usage(self) -> TokenUsage | None:
|
||||
"""Token usage reported by the API, if available."""
|
||||
return self._usage
|
||||
|
||||
def reset(self) -> None:
|
||||
"""Clear all accumulators for the next turn."""
|
||||
self._accumulated_content = ""
|
||||
self._accumulated_reasoning = ""
|
||||
self._tool_calls.clear()
|
||||
self._usage = None
|
||||
Reference in New Issue
Block a user