haiku.rag/haiku_rag_slim/haiku/rag/graph/deep_qa/state.py
Yiorgis Gozadinos 0c8d2838b2
Better typing
2025-11-13 13:22:52 +02:00

56 lines
1.8 KiB
Python

import asyncio
from dataclasses import dataclass
from typing import TYPE_CHECKING
from pydantic import BaseModel, Field
from haiku.rag.client import HaikuRAG
from haiku.rag.graph.deep_qa.dependencies import DeepQAContext
from haiku.rag.graph.deep_qa.models import DeepQAAnswer
if TYPE_CHECKING:
from haiku.rag.config.models import AppConfig
from haiku.rag.graph.agui.emitter import AGUIEmitter
@dataclass
class DeepQADeps:
client: HaikuRAG
agui_emitter: "AGUIEmitter[DeepQAState, DeepQAAnswer] | None" = None
semaphore: asyncio.Semaphore | None = None
class DeepQAState(BaseModel):
"""Deep QA state for multi-agent question answering."""
model_config = {"arbitrary_types_allowed": True}
context: DeepQAContext = Field(description="Shared QA context")
max_sub_questions: int = Field(
default=3, description="Maximum number of sub-questions"
)
max_iterations: int = Field(
default=2, description="Maximum number of QA iterations"
)
max_concurrency: int = Field(
default=1, description="Maximum parallel sub-question searches"
)
iterations: int = Field(default=0, description="Current iteration number")
@classmethod
def from_config(cls, context: DeepQAContext, config: "AppConfig") -> "DeepQAState":
"""Create a DeepQAState from an AppConfig.
Args:
context: The DeepQAContext containing the question and settings
config: The AppConfig object (uses config.qa for state parameters)
Returns:
A configured DeepQAState instance
"""
return cls(
context=context,
max_sub_questions=config.qa.max_sub_questions,
max_iterations=config.qa.max_iterations,
max_concurrency=config.qa.max_concurrency,
)