X-Financial/server/src/app/services/agent_runs.py

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,
        )