query.py

"""Non-streaming agent helpers and shared turn-summary builders for agent runs."""

from __future__ import annotations

from enum import Enum
from typing import Optional, TypeAlias, cast

from fastapi import HTTPException
from ogx_client import APIConnectionError, APIStatusError, AsyncOgxClient
from pydantic_ai.exceptions import (
    AgentRunError,
)
from pydantic_ai.messages import ModelRequest, ModelResponse, ToolReturnPart
from pydantic_ai.run import AgentRunResult
from pydantic_ai.usage import RunUsage

from configuration import configuration
from log import get_logger
from metrics import recording
from models.api.responses.error import (
    AbstractErrorResponse,
    InternalServerErrorResponse,
    PromptTooLongResponse,
)
from models.common.agents import AgentTurnAccumulator
from models.common.moderation import ShieldModerationResult
from models.common.responses.responses_api_params import ResponsesApiParams
from models.common.responses.types import ResponseInput
from models.common.turn_summary import TurnSummary
from utils.agents.error_handler import map_agent_inference_error
from utils.agents.tool_processor import (
    process_function_tool_call,
    process_function_tool_result,
    process_native_tool_call,
    process_native_tool_result,
)
from utils.conversations import append_turn_items_to_conversation
from utils.pydantic_ai_helpers import build_agent
from utils.query import (
    extract_provider_and_model_from_model_id,
)
from utils.responses import extract_vector_store_ids_from_tools
from utils.token_counter import TokenCounter

logger = get_logger(__name__)

AgentInferenceError: TypeAlias = (
    AgentRunError | APIStatusError | APIConnectionError | RuntimeError
)


class AgentFinishReason(str, Enum):
    """Finish reason for a completed agent model response."""

    CONTENT_FILTER = "content_filter"
    CANCELLED = "cancelled"
    SUCCESS = "stop"
    LENGTH = "length"
    ERROR = "error"


def get_agent_finish_reason(response: ModelResponse) -> AgentFinishReason:
    """Get the finish reason from a completed agent model response.

    Args:
        response: Last model response from the agent run.

    Returns:
        Resolved finish reason.
    """
    raw_finish_reason = (response.provider_details or {}).get("finish_reason")
    if raw_finish_reason == "cancelled":
        return AgentFinishReason.CANCELLED
    if response.finish_reason is None:
        return AgentFinishReason.ERROR
    return AgentFinishReason(response.finish_reason)


def get_finish_reason_error(
    finish_reason: AgentFinishReason,
    model_id: str,
) -> AbstractErrorResponse:
    """Map a non-success agent finish reason to an LCS error response.

    Args:
        finish_reason: Resolved finish reason from :func:`get_agent_finish_reason`.
        model_id: Model identifier in provider/model format.

    Returns:
        Structured error response for HTTP or SSE error events.
    """
    match finish_reason:
        case AgentFinishReason.LENGTH:
            return PromptTooLongResponse(model=model_id)
        case AgentFinishReason.CONTENT_FILTER:
            return InternalServerErrorResponse.query_failed(
                "The model refused to generate a response due to content policy."
            )
        case AgentFinishReason.CANCELLED:
            return InternalServerErrorResponse.query_failed(
                "The response was cancelled before completion."
            )
        case _:
            return InternalServerErrorResponse.query_failed(
                "An unexpected error occurred while processing the request."
            )


def extract_agent_token_usage(
    usage: RunUsage,
    model: str,
    endpoint_path: str,
) -> TokenCounter:
    """Build token usage for a completed agent run and record related metrics.

    Args:
        usage: Run usage reported by the agent.
        model: Model identifier in provider/model format.
        endpoint_path: Endpoint path used for metric labeling.

    Returns:
        Aggregated token usage counter for the run.
    """
    provider_id, model_id = extract_provider_and_model_from_model_id(model)
    token_counter = TokenCounter(
        input_tokens=usage.input_tokens,
        output_tokens=usage.output_tokens,
        llm_calls=max(usage.requests, 1),
    )
    logger.debug(
        "Extracted token usage from agent run: input=%d, output=%d, requests=%d",
        token_counter.input_tokens,
        token_counter.output_tokens,
        usage.requests,
    )
    recording.record_llm_token_usage(
        provider_id,
        model_id,
        token_counter.input_tokens,
        token_counter.output_tokens,
        endpoint_path,
    )
    recording.record_llm_call(provider_id, model_id, endpoint_path)
    return token_counter


def build_turn_summary_from_agent_run(
    run_result: AgentRunResult[str],
    *,
    model_id: str,
    endpoint_path: str,
    vector_store_ids: list[str],
    rag_id_mapping: dict[str, str],
) -> TurnSummary:
    """Build a turn summary from a completed agent run.

    Args:
        run_result: Completed agent run result.
        model_id: Model identifier in provider/model format.
        endpoint_path: Endpoint path used for metric labeling.
        vector_store_ids: Vector store IDs used for source mapping.
        rag_id_mapping: Mapping from vector store IDs to user-facing source labels.

    Returns:
        Turn summary with text, tools, RAG metadata, and token usage.

    Raises:
        HTTPException: When the run failed.
    """
    finish_reason = get_agent_finish_reason(run_result.response)
    if finish_reason != AgentFinishReason.SUCCESS:
        error_response = get_finish_reason_error(finish_reason, model_id)
        raise HTTPException(**error_response.model_dump())

    state = AgentTurnAccumulator(
        vector_store_ids=vector_store_ids,
        rag_id_mapping=rag_id_mapping,
        turn_summary=TurnSummary(),
    )

    for message in run_result.new_messages():
        if isinstance(message, ModelResponse):
            if message.text:
                state.turn_summary.llm_response = message.text
            for tool_call_part in message.tool_calls:
                process_function_tool_call(state, tool_call_part)
            for call_part, return_part in message.native_tool_calls:
                process_native_tool_call(state, call_part)
                process_native_tool_result(state, return_part)
        elif isinstance(message, ModelRequest):
            for request_part in message.parts:
                if isinstance(request_part, ToolReturnPart):
                    process_function_tool_result(state, request_part)

    state.turn_summary.id = run_result.response.provider_response_id or ""
    state.turn_summary.token_usage = extract_agent_token_usage(
        run_result.usage,
        model_id,
        endpoint_path,
    )
    return state.turn_summary


async def retrieve_agent_response(
    client: AsyncOgxClient,
    responses_params: ResponsesApiParams,
    moderation_result: ShieldModerationResult,
    endpoint_path: str,
    _original_input: Optional[ResponseInput] = None,
    no_tools: bool = False,
) -> TurnSummary:
    """Retrieve a turn summary from a blocking agent run.

    Args:
        client: Llama Stack client for conversation persistence on moderation block.
        responses_params: Prepared Responses API parameters.
        moderation_result: Shield moderation outcome for the turn.
        endpoint_path: Endpoint path used for metric labeling.
        _original_input: Original user input before the explicit-input rewrite.
        no_tools: Whether to skip tool processing.
    Returns:
        Turn summary for the completed agent run.

    Raises:
        HTTPException: On moderation is not applicable; on agent or provider failure.
    """
    if moderation_result.decision == "blocked":
        await append_turn_items_to_conversation(
            client,
            responses_params.conversation,
            responses_params.input,
            [moderation_result.refusal_response],
        )
        return TurnSummary(
            id=moderation_result.moderation_id,
            llm_response=moderation_result.message,
        )
    try:
        agent = build_agent(
            client, responses_params, configuration.skills, no_tools=no_tools
        )
        logger.debug("Starting agent non-streaming response processing")
        run_result = await agent.run(cast(str, responses_params.input))
    except (AgentRunError, APIStatusError, APIConnectionError, RuntimeError) as exc:
        response = map_agent_inference_error(exc, responses_params.model)
        raise HTTPException(**response.model_dump()) from exc

    vector_store_ids = extract_vector_store_ids_from_tools(responses_params.tools)
    rag_id_mapping = configuration.rag_id_mapping
    return build_turn_summary_from_agent_run(
        run_result,
        model_id=responses_params.model,
        endpoint_path=endpoint_path,
        vector_store_ids=vector_store_ids,
        rag_id_mapping=rag_id_mapping,
    )