danswer/backend/onyx/agents/agent_search/models.py

from uuid import UUID

from pydantic import BaseModel
from pydantic import model_validator
from sqlalchemy.orm import Session

from onyx.agents.agent_search.shared_graph_utils.models import PersonaExpressions
from onyx.chat.prompt_builder.answer_prompt_builder import AnswerPromptBuilder
from onyx.context.search.models import SearchRequest
from onyx.file_store.utils import InMemoryChatFile
from onyx.llm.interfaces import LLM
from onyx.tools.force import ForceUseTool
from onyx.tools.tool import Tool
from onyx.tools.tool_implementations.search.search_tool import SearchTool


class AgentSearchConfig(BaseModel):
    """
    Configuration for the Agent Search feature.
    """

    # The search request that was used to generate the Agent Search
    search_request: SearchRequest

    primary_llm: LLM
    fast_llm: LLM

    # Whether to force use of a tool, or to
    # force tool args IF the tool is used
    force_use_tool: ForceUseTool

    # contains message history for the current chat session
    # has the following (at most one is non-None)
    # message_history: list[PreviousMessage] | None = None
    # single_message_history: str | None = None
    prompt_builder: AnswerPromptBuilder

    search_tool: SearchTool | None = None

    use_agentic_search: bool = False

    # For persisting agent search data
    chat_session_id: UUID | None = None

    # The message ID of the user message that triggered the Pro Search
    message_id: int | None = None

    # Whether to persist data for Agentic Search
    use_agentic_persistence: bool = True

    # The database session for Agentic Search
    db_session: Session | None = None

    # Whether to perform initial search to inform decomposition
    # perform_initial_search_path_decision: bool = True

    # Whether to perform initial search to inform decomposition
    perform_initial_search_decomposition: bool = True

    # Whether to allow creation of refinement questions (and entity extraction, etc.)
    allow_refinement: bool = True

    # Tools available for use
    tools: list[Tool] | None = None

    using_tool_calling_llm: bool = False

    files: list[InMemoryChatFile] | None = None

    structured_response_format: dict | None = None

    skip_gen_ai_answer_generation: bool = False

    @model_validator(mode="after")
    def validate_db_session(self) -> "AgentSearchConfig":
        if self.use_agentic_persistence and self.db_session is None:
            raise ValueError(
                "db_session must be provided for pro search when using persistence"
            )
        return self

    @model_validator(mode="after")
    def validate_search_tool(self) -> "AgentSearchConfig":
        if self.use_agentic_search and self.search_tool is None:
            raise ValueError("search_tool must be provided for agentic search")
        return self

    class Config:
        arbitrary_types_allowed = True


class GraphInputs(BaseModel):
    """Input data required for the graph execution"""

    search_request: SearchRequest
    prompt_builder: AnswerPromptBuilder
    files: list[InMemoryChatFile] | None = None
    structured_response_format: dict | None = None

    class Config:
        arbitrary_types_allowed = True


class GraphTooling(BaseModel):
    """Tools and LLMs available to the graph"""

    primary_llm: LLM
    fast_llm: LLM
    search_tool: SearchTool | None = None
    tools: list[Tool] | None = None
    force_use_tool: ForceUseTool
    using_tool_calling_llm: bool = False

    class Config:
        arbitrary_types_allowed = True


class GraphPersistence(BaseModel):
    """Configuration for data persistence"""

    chat_session_id: UUID | None = None
    message_id: int | None = None
    use_agentic_persistence: bool = True
    db_session: Session | None = None

    class Config:
        arbitrary_types_allowed = True

    @model_validator(mode="after")
    def validate_db_session(self) -> "GraphPersistence":
        if self.use_agentic_persistence and self.db_session is None:
            raise ValueError(
                "db_session must be provided for pro search when using persistence"
            )
        return self


class SearchBehaviorConfig(BaseModel):
    """Configuration controlling search behavior"""

    use_agentic_search: bool = False
    perform_initial_search_decomposition: bool = True
    allow_refinement: bool = True
    skip_gen_ai_answer_generation: bool = False


class GraphConfig(BaseModel):
    """
    Main configuration class that combines all config components for Langgraph execution
    """

    inputs: GraphInputs
    tooling: GraphTooling
    persistence: GraphPersistence
    behavior: SearchBehaviorConfig

    @model_validator(mode="after")
    def validate_search_tool(self) -> "GraphConfig":
        if self.behavior.use_agentic_search and self.tooling.search_tool is None:
            raise ValueError("search_tool must be provided for agentic search")
        return self

    class Config:
        arbitrary_types_allowed = True


class AgentDocumentCitations(BaseModel):
    document_id: str
    document_title: str
    link: str


class AgentPromptEnrichmentComponents(BaseModel):
    persona_prompts: PersonaExpressions
    history: str
    date_str: str