from typing import Any from time import perf_counter from urllib.parse import urlparse from pydantic import BaseModel from pydantic_ai.messages import ( FunctionToolResultEvent, ModelMessage, ModelRequest, ModelResponse, PartDeltaEvent, PartStartEvent, TextPart, TextPartDelta, ThinkingPart, ThinkingPartDelta, UserPromptPart, ) from app.agent.deps import AgentDeps from app.agent.factory import build_agent from app.agent.model_trace import ( ModelTrace, reset_model_trace, start_model_trace, ) from app.config import ApiSettings from app.db import PersistenceStore from app.events import EventSender from app.models import ClientUserMessage, SearchResponse from app.services.search import SearchClient MAX_HISTORY_CHARS = 24_000 class AristotleAgentRuntime: def __init__( self, search_client: SearchClient, settings: ApiSettings, document_store: PersistenceStore | None = None, ): self.search_client = search_client self.settings = settings self.document_store = document_store async def stream_response( self, user_message: ClientUserMessage, events: EventSender ) -> str: trace, trace_token = start_model_trace(self.settings) try: agent = build_agent(self.settings) options = user_message.options if options.file_ids and self.document_store is None: raise RuntimeError("Document persistence is not configured.") prompt = await self._message_with_file_context( user_message.message, options.file_ids, ) deps = AgentDeps( search_client=self.search_client, http_client=self.search_client.http, events=events, settings=self.settings, max_search_results=options.max_search_results, web_tools_enabled=True, document_store=self.document_store, file_ids=options.file_ids, ) final_parts: list[str] = [] model_started = perf_counter() model_selection_sent = False first_token_sent = False await events.send( "agent.started", input={ "web_tools_enabled": True, "max_search_results": options.max_search_results, "file_ids": options.file_ids, "primary_model": trace.primary.model, "fallback_model": trace.fallback.model if trace.fallback else None, "history_messages": len(user_message.history), "history_chars": sum( len(message.content) for message in user_message.history ), }, ) async with agent.run_stream_events( prompt, message_history=_message_history(user_message), deps=deps, model_settings={"temperature": self.settings.agent_temperature}, conversation_id=user_message.conversation_id, ) as stream: async for event in stream: if trace.selected and not model_selection_sent: await _send_model_selection(events, trace) model_selection_sent = True if _has_stream_token(event) and not first_token_sent: if trace.selected and not model_selection_sent: await _send_model_selection(events, trace) model_selection_sent = True await events.send( "model.first_token", provider=trace.selected.provider if trace.selected else None, model=trace.selected.model if trace.selected else None, url=trace.selected.url if trace.selected else None, latency_ms=int((perf_counter() - model_started) * 1000), ) first_token_sent = True text_delta = await self._handle_event(event, events) if text_delta: final_parts.append(text_delta) if trace.selected and not model_selection_sent: await _send_model_selection(events, trace) return "".join(final_parts) finally: reset_model_trace(trace_token) async def _message_with_file_context( self, message: str, file_ids: list[str], ) -> str: if not file_ids: return message if self.document_store is None: return message files: list[dict[str, Any]] = [] for file_id in file_ids: record = await self.document_store.get_file(file_id) if record is not None: files.append(record) if not files: return message file_lines = "\n".join( f"- {file['filename']} (file_id: {file['id']})" for file in files ) return ( "Attached files for this user message:\n" f"{file_lines}\n\n" "This user message includes uploaded files. Inspect the attached files " "with DocumentTools before asking for clarification about vague " "references such as 'this', 'it', 'the file', or 'is this true'. " "Use list_files first if needed, then search_document or read_file for " "evidence. Do not answer from the filename alone. For truth or accuracy " "questions, first identify what the document itself says and whether its " "internal claims are supported by the document. If the claim depends on " "current prices, dates, public facts, or real-world availability, verify " "that externally with web search after inspecting the file. Keep document " "evidence distinct from externally verified evidence.\n\n" f"User message:\n{message}" ) async def _handle_event(self, event: Any, events: EventSender) -> str: if isinstance(event, FunctionToolResultEvent): await events.send( "tool.result", tool=event.part.tool_name, result_count=_result_count(event.part.content), result_preview=_result_preview(event.part.content), ) return "" if isinstance(event, PartStartEvent): if isinstance(event.part, ThinkingPart) and event.part.content: await events.send("reasoning.delta", text=event.part.content) if isinstance(event.part, TextPart) and event.part.content: await events.send("message.delta", text=event.part.content) return event.part.content return "" if isinstance(event, PartDeltaEvent): if isinstance(event.delta, ThinkingPartDelta) and event.delta.content_delta: await events.send("reasoning.delta", text=event.delta.content_delta) if isinstance(event.delta, TextPartDelta) and event.delta.content_delta: await events.send("message.delta", text=event.delta.content_delta) return event.delta.content_delta return "" def _result_count(content: Any) -> int | None: if isinstance(content, SearchResponse): return len(content.results) data = _as_dict(content) if data: for key in ("results", "chunks", "sources", "citations", "facts", "failures"): values = data.get(key) if isinstance(values, list): return len(values) return None async def _send_model_selection(events: EventSender, trace: ModelTrace) -> None: if trace.selected is None: return if trace.selected.provider == "fallback" and trace.fallback_reason: await events.send( "model.fallback", provider=trace.selected.provider, model=trace.selected.model, url=trace.selected.url, reason=trace.fallback_reason, latency_ms=trace.selected_latency_ms, ) await events.send( "model.selected", provider=trace.selected.provider, model=trace.selected.model, url=trace.selected.url, latency_ms=trace.selected_latency_ms, ) def _has_stream_token(event: Any) -> bool: if isinstance(event, PartStartEvent): return ( isinstance(event.part, ThinkingPart) and bool(event.part.content) or isinstance(event.part, TextPart) and bool(event.part.content) ) if isinstance(event, PartDeltaEvent): return ( isinstance(event.delta, ThinkingPartDelta) and bool(event.delta.content_delta) or isinstance(event.delta, TextPartDelta) and bool(event.delta.content_delta) ) return False def _message_history(user_message: ClientUserMessage) -> list[ModelMessage]: history: list[ModelMessage] = [] remaining_chars = MAX_HISTORY_CHARS for message in reversed(user_message.history): content = message.content.strip() if not content: continue if len(content) > remaining_chars: content = content[-remaining_chars:].strip() if not content: break if message.role == "user": history.append(ModelRequest(parts=[UserPromptPart(content=content)])) else: history.append(ModelResponse(parts=[TextPart(content=content)])) remaining_chars -= len(content) if remaining_chars <= 0: break history.reverse() return history def _result_preview(content: Any) -> list[dict[str, Any]] | None: if isinstance(content, SearchResponse): return [_source_preview(result.model_dump(), status="searched") for result in content.results[:5]] data = _as_dict(content) if not data: return None preview: list[dict[str, Any]] = [] for key, status in ( ("results", "searched"), ("chunks", "fetched"), ("sources", "ranked"), ("citations", "cited"), ("facts", "cited"), ): values = data.get(key) if not isinstance(values, list): continue for item in values: if isinstance(item, dict): source = _source_preview(item, status=status) if source: preview.append(source) failures = data.get("failures") if isinstance(failures, list): for failure in failures: if isinstance(failure, dict) and failure.get("input"): preview.append( { "title": failure.get("input"), "url": failure.get("input"), "snippet": failure.get("error"), "status": "failed", } ) return _dedupe_source_previews(preview)[:8] or None def _as_dict(content: Any) -> dict[str, Any] | None: if isinstance(content, dict): return content if isinstance(content, BaseModel): return content.model_dump() return None def _source_preview(item: dict[str, Any], *, status: str) -> dict[str, Any]: source_type = item.get("source_type") file_id = item.get("file_id") chunk_id = item.get("chunk_id") or item.get("id") if source_type == "document" or file_id or item.get("locator"): title = item.get("title") or item.get("filename") snippet = item.get("snippet") or item.get("text_preview") or item.get("quote") return { "id": chunk_id if isinstance(chunk_id, str) else item.get("id"), "title": title if isinstance(title, str) else "Document", "url": None, "domain": None, "source": "document", "source_type": "document", "file_id": file_id if isinstance(file_id, str) else None, "chunk_id": chunk_id if isinstance(chunk_id, str) else None, "locator": item.get("locator") if isinstance(item.get("locator"), str) else None, "page": item.get("page"), "section": item.get("section"), "row_start": item.get("row_start"), "row_end": item.get("row_end"), "snippet": _preview_text(snippet), "status": status, } url = item.get("url") if not isinstance(url, str) or not url: return {} title = item.get("title") snippet = item.get("snippet") or item.get("fact") or item.get("content") source = { "title": title if isinstance(title, str) else None, "url": url, "domain": _domain(url), "source": item.get("source") if isinstance(item.get("source"), str) else None, "snippet": _preview_text(snippet), "status": status, } if item.get("marker"): source["marker"] = item.get("marker") return source def _dedupe_source_previews( previews: list[dict[str, Any]], ) -> list[dict[str, Any]]: deduped: dict[str, dict[str, Any]] = {} for preview in previews: url = preview.get("url") fallback_id = preview.get("chunk_id") or preview.get("id") if isinstance(url, str) and url: key = url elif isinstance(fallback_id, str) and fallback_id: key = fallback_id else: continue current = deduped.get(key, {}) deduped[key] = { **preview, **{key: value for key, value in current.items() if value}, } return list(deduped.values()) def _domain(url: str) -> str | None: parsed = urlparse(url) if not parsed.hostname: return None return parsed.hostname.removeprefix("www.") def _preview_text(value: Any) -> str | None: if not isinstance(value, str): return None cleaned = " ".join(value.split()) if not cleaned: return None return cleaned[:280]