models/data_models.py

from dataclasses import dataclass, field
from typing import List, Dict

# @dataclass
# class ContextInput:
#     processed_query: str
#     constraints: List[str]
#     search_sources: List[str] = field(default_factory=lambda: ['wikipedia'])
#     max_results: int = 5
@dataclass
class ContextInput:
    processed_query: str
    constraints: List[str]
    domain: str  # e.g., "oil_and_gas", "fine_art"
    max_results: int = 5
    min_confidence: float = 0.7

@dataclass
class UserInput:
    query: str
    constraints: str
    #top_k: int

@dataclass
class ImageAnalysisInput:
    images: List[str]  # List of image paths
    context: Dict      # Context from ContextLearnAgent
    constraints: List[str]
    top_k: int 

@dataclass
class AgentState:
    """Tracks the current state of the agent"""
    intent: str = ""
    thoughts: List[str] = field(default_factory=list)
    decisions: List[Dict] = field(default_factory=list)
    errors: List[str] = field(default_factory=list)

@dataclass
class AssemblerInput:
    user_input_results: Dict  # From UserInputAgent
    context_results: Dict     # From ContextLearnAgent
    image_results: Dict      # From ImageAnalyzerAgent
    report_format: str = "detailed"  # or "summary"