from __future__ import annotations import uuid from datetime import UTC, datetime from typing import Any from sqlalchemy.orm import Session from app.core.agent_enums import AgentPermissionLevel, AgentRunStatus from app.core.logging import get_logger from app.models.agent_run import AgentRun, AgentToolCall, SemanticParseLog from app.repositories.agent_run import AgentRunRepository from app.schemas.agent_run import AgentRunRead, AgentToolCallRead, SemanticParseRead from app.services.agent_foundation import AgentFoundationService logger = get_logger("app.services.agent_runs") class AgentRunService: def __init__(self, db: Session) -> None: self.db = db self.repository = AgentRunRepository(db) def list_runs( self, *, agent: str | None = None, status: str | None = None, source: str | None = None, limit: int = 20, ) -> list[AgentRunRead]: self._ensure_ready() runs = self.repository.list(agent=agent, status=status, source=source, limit=limit) return [self._serialize_run(item) for item in runs] def get_run(self, run_id: str) -> AgentRunRead | None: self._ensure_ready() run = self.repository.get_by_run_id(run_id) if run is None: return None return self._serialize_run(run) def create_run( self, *, agent: str, source: str, user_id: str | None = None, task_id: str | None = None, ontology_json: dict[str, Any] | None = None, route_json: dict[str, Any] | None = None, permission_level: str = AgentPermissionLevel.READ.value, status: str = AgentRunStatus.RUNNING.value, result_summary: str | None = None, error_message: str | None = None, started_at: datetime | None = None, finished_at: datetime | None = None, ) -> AgentRunRead: self._ensure_ready() run = AgentRun( run_id=f"run_{uuid.uuid4().hex[:16]}", agent=agent, source=source, user_id=user_id, task_id=task_id, ontology_json=ontology_json or {}, route_json=route_json or {}, permission_level=permission_level, status=status, result_summary=result_summary, error_message=error_message, started_at=started_at or datetime.now(UTC), finished_at=finished_at, ) created = self.repository.create_run(run) logger.info("Created agent run id=%s run_id=%s", created.id, created.run_id) return self._serialize_run(created) def update_run( self, run_id: str, *, agent: str | None = None, ontology_json: dict[str, Any] | None = None, route_json: dict[str, Any] | None = None, permission_level: str | None = None, status: str | None = None, result_summary: str | None = None, error_message: str | None = None, finished_at: datetime | None = None, ) -> AgentRunRead: self._ensure_ready() run = self.repository.get_by_run_id(run_id) if run is None: raise LookupError("Run not found") if agent is not None: run.agent = agent if ontology_json is not None: run.ontology_json = ontology_json if route_json is not None: run.route_json = route_json if permission_level is not None: run.permission_level = permission_level if status is not None: run.status = status if result_summary is not None: run.result_summary = result_summary if error_message is not None: run.error_message = error_message if finished_at is not None: run.finished_at = finished_at updated = self.repository.save_run(run) logger.info("Updated agent run run_id=%s status=%s", updated.run_id, updated.status) return self._serialize_run(updated) def merge_route_json( self, run_id: str, route_patch: dict[str, Any], *, status: str | None = None, result_summary: str | None = None, error_message: str | None = None, finished_at: datetime | None = None, ) -> AgentRunRead: self._ensure_ready() run = self.repository.get_by_run_id(run_id) if run is None: raise LookupError("Run not found") route_json = dict(run.route_json or {}) route_json.update(route_patch or {}) run.route_json = route_json if status is not None: run.status = status if result_summary is not None: run.result_summary = result_summary if error_message is not None: run.error_message = error_message if finished_at is not None: run.finished_at = finished_at updated = self.repository.save_run(run) logger.info("Merged route_json for agent run run_id=%s status=%s", updated.run_id, updated.status) return self._serialize_run(updated) def record_tool_call( self, *, run_id: str, tool_type: str, tool_name: str, request_json: dict[str, Any] | None = None, response_json: dict[str, Any] | None = None, status: str, duration_ms: int = 0, error_message: str | None = None, ) -> AgentToolCallRead: self._ensure_ready() tool_call = AgentToolCall( run_id=run_id, tool_type=tool_type, tool_name=tool_name, request_json=request_json or {}, response_json=response_json or {}, status=status, duration_ms=duration_ms, error_message=error_message, ) created = self.repository.create_tool_call(tool_call) logger.info("Recorded tool call run_id=%s tool=%s", run_id, tool_name) return AgentToolCallRead.model_validate(created) def record_semantic_parse( self, *, run_id: str, user_id: str | None, raw_query: str, scenario: str, intent: str, entities_json: list[Any] | None = None, time_range_json: dict[str, Any] | None = None, metrics_json: list[Any] | None = None, constraints_json: list[Any] | None = None, risk_flags_json: list[Any] | None = None, permission_json: dict[str, Any] | None = None, confidence: float = 0.0, ) -> SemanticParseRead: self._ensure_ready() semantic_parse = SemanticParseLog( run_id=run_id, user_id=user_id, raw_query=raw_query, scenario=scenario, intent=intent, entities_json=entities_json or [], time_range_json=time_range_json or {}, metrics_json=metrics_json or [], constraints_json=constraints_json or [], risk_flags_json=risk_flags_json or [], permission_json=permission_json or {}, confidence=confidence, ) created = self.repository.create_semantic_parse(semantic_parse) logger.info( "Recorded semantic parse run_id=%s scenario=%s intent=%s", run_id, scenario, intent ) return SemanticParseRead.model_validate(created) def _ensure_ready(self) -> None: AgentFoundationService(self.db).ensure_foundation_ready() @staticmethod def _serialize_run(run: AgentRun) -> AgentRunRead: semantic_parse = run.semantic_parse_logs[0] if run.semantic_parse_logs else None return AgentRunRead( id=run.id, run_id=run.run_id, agent=run.agent, source=run.source, user_id=run.user_id, task_id=run.task_id, ontology_json=run.ontology_json, route_json=run.route_json, permission_level=run.permission_level, status=run.status, result_summary=run.result_summary, error_message=run.error_message, started_at=run.started_at, finished_at=run.finished_at, tool_calls=[AgentToolCallRead.model_validate(item) for item in run.tool_calls], semantic_parse=SemanticParseRead.model_validate(semantic_parse) if semantic_parse else None, )