ira/sim-search-api/app/schemas/query.py

100 lines
3.2 KiB
Python

"""
Query schemas for the sim-search API.
This module defines the Pydantic models for query-related operations.
"""
from typing import Dict, List, Optional, Any
from pydantic import BaseModel
class QueryBase(BaseModel):
"""Base query schema."""
query: str
class QueryProcess(QueryBase):
"""Query processing schema."""
pass
class QueryClassify(QueryBase):
"""Query classification schema."""
pass
class SubQuestion(BaseModel):
"""Sub-question schema."""
sub_question: str
aspect: str
priority: float
class StructuredQuery(BaseModel):
"""Structured query schema."""
original_query: str
enhanced_query: Optional[str] = None
type: Optional[str] = None
intent: Optional[str] = None
domain: Optional[str] = None
confidence: Optional[float] = None
reasoning: Optional[str] = None
entities: Optional[List[str]] = None
sub_questions: Optional[List[SubQuestion]] = None
search_queries: Optional[Dict[str, str]] = None
is_academic: Optional[bool] = None
is_code: Optional[bool] = None
is_current_events: Optional[bool] = None
class ProcessedQuery(BaseModel):
"""Processed query schema."""
original_query: str
structured_query: StructuredQuery
class Config:
"""Pydantic config."""
schema_extra = {
"example": {
"original_query": "What are the latest advancements in quantum computing?",
"structured_query": {
"original_query": "What are the latest advancements in quantum computing?",
"enhanced_query": "What are the recent breakthroughs and developments in quantum computing technology, algorithms, and applications in the past 2 years?",
"type": "exploratory",
"intent": "research",
"domain": "academic",
"confidence": 0.95,
"reasoning": "This query is asking about recent developments in a scientific field, which is typical of academic research.",
"entities": ["quantum computing", "advancements"],
"sub_questions": [
{
"sub_question": "What are the latest hardware advancements in quantum computing?",
"aspect": "hardware",
"priority": 0.9
},
{
"sub_question": "What are the recent algorithmic breakthroughs in quantum computing?",
"aspect": "algorithms",
"priority": 0.8
}
],
"search_queries": {
"google": "latest advancements in quantum computing 2024",
"scholar": "recent quantum computing breakthroughs",
"arxiv": "quantum computing hardware algorithms"
},
"is_academic": True,
"is_code": False,
"is_current_events": False
}
}
}