mirror of
https://github.com/danswer-ai/danswer.git
synced 2025-05-03 08:20:40 +02:00
129 lines
3.3 KiB
Python
129 lines
3.3 KiB
Python
from pydantic import BaseModel
|
|
|
|
from onyx.agents.agent_search.deep_search.main.models import (
|
|
AgentAdditionalMetrics,
|
|
)
|
|
from onyx.agents.agent_search.deep_search.main.models import AgentBaseMetrics
|
|
from onyx.agents.agent_search.deep_search.main.models import (
|
|
AgentRefinedMetrics,
|
|
)
|
|
from onyx.agents.agent_search.deep_search.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 AgentChunkRetrievalStats(BaseModel):
|
|
verified_count: int | None = None
|
|
verified_avg_scores: float | None = None
|
|
rejected_count: int | None = None
|
|
rejected_avg_scores: float | None = 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] = []
|
|
|
|
|
|
class EntityExtractionResult(BaseModel):
|
|
retrieved_entities_relationships: EntityRelationshipTermExtraction
|
|
|
|
|
|
class QueryRetrievalResult(BaseModel):
|
|
query: str
|
|
retrieved_documents: list[InferenceSection]
|
|
stats: RetrievalFitStats | None
|
|
query_info: SearchQueryInfo | None
|
|
|
|
|
|
class SubQuestionAnswerResults(BaseModel):
|
|
question: str
|
|
question_id: str
|
|
answer: str
|
|
verified_high_quality: bool
|
|
sub_query_retrieval_results: list[QueryRetrievalResult]
|
|
verified_reranked_documents: list[InferenceSection]
|
|
context_documents: list[InferenceSection]
|
|
cited_documents: list[InferenceSection]
|
|
sub_question_retrieval_stats: AgentChunkRetrievalStats
|
|
|
|
|
|
class CombinedAgentMetrics(BaseModel):
|
|
timings: AgentTimings
|
|
base_metrics: AgentBaseMetrics | None
|
|
refined_metrics: AgentRefinedMetrics
|
|
additional_metrics: AgentAdditionalMetrics
|
|
|
|
|
|
class PersonaPromptExpressions(BaseModel):
|
|
contextualized_prompt: str
|
|
base_prompt: str
|
|
|
|
|
|
class AgentPromptEnrichmentComponents(BaseModel):
|
|
persona_prompts: PersonaPromptExpressions
|
|
history: str
|
|
date_str: str
|