mirror of
https://github.com/danswer-ai/danswer.git
synced 2025-05-08 10:50:26 +02:00
113 lines
2.7 KiB
Python
113 lines
2.7 KiB
Python
from typing import Literal
|
|
|
|
from pydantic import BaseModel
|
|
|
|
from onyx.agents.agent_search.deep_search_a.main.models import AgentAdditionalMetrics
|
|
from onyx.agents.agent_search.deep_search_a.main.models import AgentBaseMetrics
|
|
from onyx.agents.agent_search.deep_search_a.main.models import AgentRefinedMetrics
|
|
from onyx.agents.agent_search.deep_search_a.main.models import AgentTimings
|
|
from onyx.context.search.models import InferenceSection
|
|
from onyx.tools.models import SearchQueryInfo
|
|
|
|
|
|
# Pydantic models for structured outputs
|
|
class RewrittenQueries(BaseModel):
|
|
rewritten_queries: list[str]
|
|
|
|
|
|
class BinaryDecision(BaseModel):
|
|
decision: Literal["yes", "no"]
|
|
|
|
|
|
class BinaryDecisionWithReasoning(BaseModel):
|
|
reasoning: str
|
|
decision: Literal["yes", "no"]
|
|
|
|
|
|
class RetrievalFitScoreMetrics(BaseModel):
|
|
scores: dict[str, float]
|
|
chunk_ids: list[str]
|
|
|
|
|
|
class RetrievalFitStats(BaseModel):
|
|
fit_score_lift: float
|
|
rerank_effect: float
|
|
fit_scores: dict[str, RetrievalFitScoreMetrics]
|
|
|
|
|
|
class AgentChunkScores(BaseModel):
|
|
scores: dict[str, dict[str, list[int | float]]]
|
|
|
|
|
|
class AgentChunkStats(BaseModel):
|
|
verified_count: int | None
|
|
verified_avg_scores: float | None
|
|
rejected_count: int | None
|
|
rejected_avg_scores: float | None
|
|
verified_doc_chunk_ids: list[str]
|
|
dismissed_doc_chunk_ids: list[str]
|
|
|
|
|
|
class InitialAgentResultStats(BaseModel):
|
|
sub_questions: dict[str, float | int | None]
|
|
original_question: dict[str, float | int | None]
|
|
agent_effectiveness: dict[str, float | int | None]
|
|
|
|
|
|
class RefinedAgentStats(BaseModel):
|
|
revision_doc_efficiency: float | None
|
|
revision_question_efficiency: float | None
|
|
|
|
|
|
class Term(BaseModel):
|
|
term_name: str
|
|
term_type: str
|
|
term_similar_to: list[str]
|
|
|
|
|
|
### Models ###
|
|
|
|
|
|
class Entity(BaseModel):
|
|
entity_name: str
|
|
entity_type: str
|
|
|
|
|
|
class Relationship(BaseModel):
|
|
relationship_name: str
|
|
relationship_type: str
|
|
relationship_entities: list[str]
|
|
|
|
|
|
class EntityRelationshipTermExtraction(BaseModel):
|
|
entities: list[Entity]
|
|
relationships: list[Relationship]
|
|
terms: list[Term]
|
|
|
|
|
|
### Models ###
|
|
|
|
|
|
class QueryResult(BaseModel):
|
|
query: str
|
|
search_results: list[InferenceSection]
|
|
stats: RetrievalFitStats | None
|
|
query_info: SearchQueryInfo | None
|
|
|
|
|
|
class QuestionAnswerResults(BaseModel):
|
|
question: str
|
|
question_id: str
|
|
answer: str
|
|
quality: str
|
|
expanded_retrieval_results: list[QueryResult]
|
|
documents: list[InferenceSection]
|
|
sub_question_retrieval_stats: AgentChunkStats
|
|
|
|
|
|
class CombinedAgentMetrics(BaseModel):
|
|
timings: AgentTimings
|
|
base_metrics: AgentBaseMetrics | None
|
|
refined_metrics: AgentRefinedMetrics
|
|
additional_metrics: AgentAdditionalMetrics
|