┌─────────────────────────────────────────────────────────────────────────────┐
│                    RESEARCH ASSISTANT ARCHITECTURE                          │
├─────────────────────────────────────────────────────────────────────────────┤
│                                                                             │
│  ┌─────────────────────────────────────────────────────────────────────┐    │
│  │ RESEARCH REQUEST                                                    │    │
│  │   Research Query ──────┬──────── Scope Definition                   │    │
│  └────────────────────────┴────────────────────────────────────────────┘    │
│                           │                                                 │
│                           ▼                                                 │
│  ┌─────────────────────────────────────────────────────────────────────┐    │
│  │ QUERY PLANNING                                                      │    │
│  │   Query Decomposition ──► Priority Ranking ──► Search Strategy      │    │
│  └─────────────────────────────────────────────────────────────────────┘    │
│                           │                                                 │
│                           ▼                                                 │
│  ┌────────────────────────────────────────────────────────────────────┐     │
│  │ DATA GATHERING                                                     │     │
│  │   ┌───────────┐ ┌───────────┐ ┌───────────┐ ┌───────────┐          │     │
│  │   │Web Search │ │API Sources│ │Document   │ │ News APIs │          │     │
│  │   │           │ │           │ │RAG        │ │           │          │     │
│  │   └─────┬─────┘ └─────┬─────┘ └─────┬─────┘ └─────┬─────┘          │     │
│  └─────────┴─────────────┴─────────────┴─────────────┴────────────────┘     │
│                           │                                                 │
│                           ▼                                                 │
│  ┌─────────────────────────────────────────────────────────────────────┐    │
│  │ INFORMATION PROCESSING                                              │    │
│  │   Entity Extraction ──► Deduplication ──► Cross-Validation          │    │
│  └─────────────────────────────────────────────────────────────────────┘    │
│                           │                                                 │
│                           ▼                                                 │
│  ┌─────────────────────────────────────────────────────────────────────┐    │
│  │ REPORT SYNTHESIS                                                    │    │
│  │   Report Outline ──► Draft Generation ──► Citation Assembly         │    │
│  └─────────────────────────────────────────────────────────────────────┘    │
│                           │                                                 │
│                           ▼                                                 │
│  ┌─────────────────────────────────────────────────────────────────────┐    │
│  │ DELIVERABLES                                                        │    │
│  │   ┌────────────────┐ ┌────────────────┐ ┌────────────────┐          │    │
│  │   │Research Report │ │Executive       │ │Raw Data Export │          │    │
│  │   │                │ │Summary         │ │                │          │    │
│  │   └────────────────┘ └────────────────┘ └────────────────┘          │    │
│  └─────────────────────────────────────────────────────────────────────┘    │
│                                                                             │
└─────────────────────────────────────────────────────────────────────────────┘

research-assistant/
├── src/
│   ├── __init__.py
│   ├── config.py
│   ├── planning/
│   │   ├── __init__.py
│   │   ├── decomposer.py      # Query decomposition
│   │   ├── strategist.py      # Search strategy
│   │   └── prioritizer.py     # Sub-query ranking
│   ├── gathering/
│   │   ├── __init__.py
│   │   ├── web_search.py      # Web search integration
│   │   ├── api_sources.py     # External APIs
│   │   ├── document_rag.py    # Internal document search
│   │   └── news_api.py        # News aggregation
│   ├── processing/
│   │   ├── __init__.py
│   │   ├── extractor.py       # Entity extraction
│   │   ├── deduplicator.py    # Information deduplication
│   │   └── validator.py       # Cross-validation
│   ├── synthesis/
│   │   ├── __init__.py
│   │   ├── outliner.py        # Report structure
│   │   ├── writer.py          # Content generation
│   │   └── citations.py       # Citation management
│   ├── agent/
│   │   ├── __init__.py
│   │   ├── state.py           # Agent state management
│   │   ├── workflow.py        # LangGraph workflow
│   │   └── tools.py           # Agent tools
│   └── api/
│       ├── __init__.py
│       └── main.py            # FastAPI endpoints
├── tests/
├── docker-compose.yml
└── requirements.txt

# src/config.py
from pydantic_settings import BaseSettings
from typing import Optional

class Settings(BaseSettings):
    # LLM Settings
    openai_api_key: str
    openai_model: str = "gpt-4o"

    # Search Settings
    brave_api_key: str  # Brave Search API (https://brave.com/search/api)
    news_api_key: Optional[str] = None

    # Vector Store
    chroma_persist_dir: str = "./data/chroma"

    # Research Settings
    max_search_results: int = 10
    max_sub_queries: int = 5
    validation_threshold: float = 0.7

    # Report Settings
    max_report_length: int = 5000
    include_raw_sources: bool = True

    class Config:
        env_file = ".env"

settings = Settings()

# src/agent/state.py
from typing import TypedDict, List, Optional, Annotated
from pydantic import BaseModel
from enum import Enum
import operator

class ResearchPhase(str, Enum):
    PLANNING = "planning"
    GATHERING = "gathering"
    PROCESSING = "processing"
    SYNTHESIS = "synthesis"
    COMPLETE = "complete"

class SubQuery(BaseModel):
    """A decomposed sub-query for research."""
    query: str
    priority: int  # 1-5, higher = more important
    search_type: str  # web, api, docs, news
    completed: bool = False
    results: List[dict] = []

class SourceInfo(BaseModel):
    """Information from a single source."""
    content: str
    source_url: str
    source_type: str
    relevance_score: float
    extracted_entities: List[str] = []
    timestamp: Optional[str] = None

class ResearchState(TypedDict):
    """State for the research agent workflow."""
    # Input
    original_query: str
    scope: str

    # Planning
    sub_queries: List[SubQuery]
    search_strategy: str

    # Gathering
    raw_sources: Annotated[List[SourceInfo], operator.add]

    # Processing
    validated_facts: List[dict]
    entities: List[str]

    # Synthesis
    report_outline: List[str]
    draft_sections: dict

    # Output
    final_report: str
    executive_summary: str
    citations: List[dict]

    # Control
    current_phase: ResearchPhase
    iteration: int
    errors: List[str]

# src/planning/decomposer.py
from typing import List
from langchain_openai import ChatOpenAI
from langchain_core.prompts import ChatPromptTemplate
from pydantic import BaseModel, Field
from ..config import settings

class DecomposedQueries(BaseModel):
    """Structured output for query decomposition."""
    sub_queries: List[str] = Field(
        description="List of focused sub-queries"
    )
    search_types: List[str] = Field(
        description="Recommended search type for each sub-query"
    )
    priorities: List[int] = Field(
        description="Priority 1-5 for each sub-query"
    )
    reasoning: str = Field(
        description="Explanation of decomposition strategy"
    )

class QueryDecomposer:
    """Decomposes complex research queries into searchable sub-queries."""

    def __init__(self):
        self.llm = ChatOpenAI(
            model=settings.openai_model,
            api_key=settings.openai_api_key,
            temperature=0.2
        ).with_structured_output(DecomposedQueries)

        self.prompt = ChatPromptTemplate.from_messages([
            ("system", """You are a research planning expert. Your job is to
decompose complex research queries into focused, searchable sub-queries.

Guidelines:
- Create 3-5 sub-queries that together cover the main query
- Each sub-query should be specific and searchable
- Assign search types: 'web' for general info, 'news' for recent events,
  'api' for structured data, 'docs' for internal documents
- Prioritize sub-queries (5=critical, 1=nice-to-have)
- Consider temporal aspects (historical vs current data)"""),
            ("human", """Research Query: {query}

Scope/Context: {scope}

Decompose this into focused sub-queries for comprehensive research.""")
        ])

    async def decompose(
        self,
        query: str,
        scope: str = "general"
    ) -> DecomposedQueries:
        """Decompose a research query into sub-queries."""
        chain = self.prompt | self.llm
        result = await chain.ainvoke({
            "query": query,
            "scope": scope
        })
        return result


class SearchStrategist:
    """Determines optimal search strategy based on query analysis."""

    def __init__(self):
        self.llm = ChatOpenAI(
            model=settings.openai_model,
            api_key=settings.openai_api_key,
            temperature=0
        )

    async def plan_strategy(
        self,
        query: str,
        sub_queries: List[str]
    ) -> str:
        """Create a search strategy document."""
        prompt = ChatPromptTemplate.from_messages([
            ("system", """Create a brief search strategy for this research task.
Include: key sources to prioritize, potential challenges, validation approach."""),
            ("human", "Main query: {query}\nSub-queries: {sub_queries}")
        ])

        chain = prompt | self.llm
        result = await chain.ainvoke({
            "query": query,
            "sub_queries": "\n".join(sub_queries)
        })
        return result.content

┌─────────────────────────────────────────────────────────────┐
│ QUERY DECOMPOSITION EXAMPLE                                 │
├─────────────────────────────────────────────────────────────┤
│                                                             │
│  Original: "What's the competitive landscape for            │
│             AI-powered CRM solutions in 2024?"               │
│        │                                                     │
│        ▼                                                     │
│  ┌─────────────────────────────────────────────────────┐    │
│  │ QueryDecomposer (GPT-4o, structured output)         │    │
│  │                                                     │    │
│  │ Sub-queries:                                        │    │
│  │ 1. "Top AI CRM vendors market share 2024"           │    │
│  │    search_type: web, priority: 5                    │    │
│  │                                                     │    │
│  │ 2. "Salesforce Einstein AI features pricing"        │    │
│  │    search_type: web, priority: 4                    │    │
│  │                                                     │    │
│  │ 3. "HubSpot AI tools comparison"                    │    │
│  │    search_type: web, priority: 4                    │    │
│  │                                                     │    │
│  │ 4. "AI CRM startup funding 2024"                    │    │
│  │    search_type: news, priority: 3                   │    │
│  │                                                     │    │
│  │ 5. "Enterprise AI CRM adoption trends"              │    │
│  │    search_type: docs, priority: 3                   │    │
│  └─────────────────────────────────────────────────────┘    │
│                                                             │
└─────────────────────────────────────────────────────────────┘

# src/gathering/web_search.py
from typing import List, Dict
import httpx
from langchain_core.tools import tool
from ..config import settings

class WebSearcher:
    """Web search integration using Brave Search API.

    Brave Search provides privacy-focused search with an independent
    index, making it ideal for research applications where diverse
    sources and privacy are important.
    """

    def __init__(self):
        self.api_key = settings.brave_api_key
        self.base_url = "https://api.search.brave.com/res/v1"

    async def search(
        self,
        query: str,
        max_results: int = 5
    ) -> List[Dict]:
        """Search the web for information using Brave Search."""
        try:
            async with httpx.AsyncClient() as client:
                response = await client.get(
                    f"{self.base_url}/web/search",
                    headers={
                        "X-Subscription-Token": self.api_key,
                        "Accept": "application/json"
                    },
                    params={
                        "q": query,
                        "count": max_results,
                        "text_decorations": False,
                        "search_lang": "en"
                    },
                    timeout=30.0
                )
                response.raise_for_status()
                data = response.json()

            results = []
            for i, item in enumerate(data.get("web", {}).get("results", [])):
                results.append({
                    "content": item.get("description", ""),
                    "url": item.get("url", ""),
                    "title": item.get("title", ""),
                    # Position-based relevance (Brave doesn't return scores)
                    "score": 1.0 - (i * 0.1)
                })

            return results
        except Exception as e:
            return [{"error": str(e)}]

@tool
async def web_search(query: str) -> str:
    """Search the web for information about a topic using Brave Search.

    Args:
        query: The search query

    Returns:
        Search results as formatted text
    """
    searcher = WebSearcher()
    results = await searcher.search(query)

    formatted = []
    for i, r in enumerate(results, 1):
        formatted.append(f"[{i}] {r.get('title', 'No title')}")
        formatted.append(f"    URL: {r.get('url', '')}")
        formatted.append(f"    {r.get('content', '')[:500]}...")
        formatted.append("")

    return "\n".join(formatted)

# src/gathering/document_rag.py
from typing import List, Dict
from langchain_openai import OpenAIEmbeddings
from langchain_chroma import Chroma
from langchain_core.tools import tool
from ..config import settings

class DocumentRetriever:
    """RAG-based retrieval from internal documents."""

    def __init__(self):
        self.embeddings = OpenAIEmbeddings(
            api_key=settings.openai_api_key
        )
        self.vectorstore = Chroma(
            persist_directory=settings.chroma_persist_dir,
            embedding_function=self.embeddings,
            collection_name="research_docs"
        )

    async def search(
        self,
        query: str,
        k: int = 5
    ) -> List[Dict]:
        """Search internal documents."""
        docs = self.vectorstore.similarity_search_with_score(
            query,
            k=k
        )

        results = []
        for doc, score in docs:
            results.append({
                "content": doc.page_content,
                "source": doc.metadata.get("source", "unknown"),
                "score": float(score),
                "metadata": doc.metadata
            })

        return results

@tool
async def search_documents(query: str) -> str:
    """Search internal company documents for relevant information.

    Args:
        query: The search query

    Returns:
        Relevant document excerpts
    """
    retriever = DocumentRetriever()
    results = await retriever.search(query)

    formatted = []
    for i, r in enumerate(results, 1):
        formatted.append(f"[Doc {i}] Source: {r['source']}")
        formatted.append(f"    Relevance: {r['score']:.2f}")
        formatted.append(f"    {r['content'][:500]}...")
        formatted.append("")

    return "\n".join(formatted)

# src/processing/validator.py
from typing import List, Dict
from langchain_openai import ChatOpenAI
from langchain_core.prompts import ChatPromptTemplate
from pydantic import BaseModel, Field
from ..config import settings

class ValidationResult(BaseModel):
    """Result of cross-validation."""
    claim: str
    is_valid: bool
    confidence: float = Field(ge=0, le=1)
    supporting_sources: List[str]
    conflicting_sources: List[str]
    notes: str

class InformationValidator:
    """Cross-validates information across multiple sources."""

    def __init__(self):
        self.llm = ChatOpenAI(
            model=settings.openai_model,
            api_key=settings.openai_api_key,
            temperature=0
        ).with_structured_output(ValidationResult)

        self.prompt = ChatPromptTemplate.from_messages([
            ("system", """You are a fact-checking expert. Analyze the claim
and supporting sources to determine validity.

Consider:
- Source credibility and recency
- Consistency across sources
- Potential biases or conflicts
- Data vs opinion

Be conservative - only mark as valid if well-supported."""),
            ("human", """Claim to validate: {claim}

Sources:
{sources}

Validate this claim.""")
        ])

    async def validate(
        self,
        claim: str,
        sources: List[Dict]
    ) -> ValidationResult:
        """Validate a claim against multiple sources."""
        sources_text = "\n\n".join([
            f"Source {i+1} ({s.get('source_type', 'unknown')}):\n{s.get('content', '')[:1000]}"
            for i, s in enumerate(sources)
        ])

        chain = self.prompt | self.llm
        result = await chain.ainvoke({
            "claim": claim,
            "sources": sources_text
        })
        return result


class EntityExtractor:
    """Extracts key entities from research content."""

    def __init__(self):
        self.llm = ChatOpenAI(
            model=settings.openai_model,
            api_key=settings.openai_api_key,
            temperature=0
        )

    async def extract(self, content: str) -> Dict[str, List[str]]:
        """Extract entities from content."""
        prompt = ChatPromptTemplate.from_messages([
            ("system", """Extract key entities from the content.
Return as JSON with categories: companies, people, products,
technologies, locations, dates, metrics."""),
            ("human", "{content}")
        ])

        chain = prompt | self.llm
        result = await chain.ainvoke({"content": content})

        # Parse result (simplified)
        import json
        try:
            return json.loads(result.content)
        except:
            return {"entities": []}

┌─────────────────────────────────────────────────────────────┐
│ INFORMATION VALIDATION                                      │
├─────────────────────────────────────────────────────────────┤
│                                                             │
│  Claim: "Salesforce has 23% of the AI CRM market"           │
│        │                                                     │
│        ▼                                                     │
│  ┌─────────────────────────────────────────────────────┐    │
│  │ InformationValidator                                │    │
│  │                                                     │    │
│  │ Sources checked:                                    │    │
│  │ • Gartner report: "Salesforce leads with ~20-25%"   │    │
│  │ • TechCrunch article: "Salesforce dominates..."     │    │
│  │ • Company press release: "23.4% market share"       │    │
│  │                                                     │    │
│  │ Validation result:                                  │    │
│  │ ┌─────────────────────────────────────────────────┐ │    │
│  │ │ is_valid: true                                  │ │    │
│  │ │ confidence: 0.85                                │ │    │
│  │ │ supporting_sources: [Gartner, press release]    │ │    │
│  │ │ conflicting_sources: []                         │ │    │
│  │ │ notes: "Consistent across sources"              │ │    │
│  │ └─────────────────────────────────────────────────┘ │    │
│  └─────────────────────────────────────────────────────┘    │
│                                                             │
└─────────────────────────────────────────────────────────────┘

# src/synthesis/writer.py
from typing import List, Dict
from langchain_openai import ChatOpenAI
from langchain_core.prompts import ChatPromptTemplate
from ..config import settings

class ReportWriter:
    """Generates research reports from validated findings."""

    def __init__(self):
        self.llm = ChatOpenAI(
            model=settings.openai_model,
            api_key=settings.openai_api_key,
            temperature=0.3
        )

    async def create_outline(
        self,
        query: str,
        findings: List[Dict]
    ) -> List[str]:
        """Create report outline based on findings."""
        prompt = ChatPromptTemplate.from_messages([
            ("system", """Create a logical report outline for the research findings.
Include: Executive Summary, Background, Key Findings (3-5 sections),
Analysis, Recommendations, Appendix."""),
            ("human", "Research query: {query}\n\nKey findings summary: {findings}")
        ])

        findings_summary = "\n".join([
            f"- {f.get('claim', f.get('content', ''))[:200]}"
            for f in findings[:10]
        ])

        chain = prompt | self.llm
        result = await chain.ainvoke({
            "query": query,
            "findings": findings_summary
        })

        # Parse outline into sections
        sections = [
            line.strip()
            for line in result.content.split("\n")
            if line.strip()
        ]
        return sections

    async def write_section(
        self,
        section_title: str,
        findings: List[Dict],
        context: str
    ) -> str:
        """Write a single report section."""
        prompt = ChatPromptTemplate.from_messages([
            ("system", """Write a professional research report section.
- Use clear, objective language
- Include specific data and examples
- Cite sources using [Source N] format
- Be concise but comprehensive"""),
            ("human", """Section: {section}

Relevant findings:
{findings}

Context: {context}

Write this section.""")
        ])

        findings_text = "\n".join([
            f"[Source {i+1}]: {f.get('content', '')[:500]}"
            for i, f in enumerate(findings[:5])
        ])

        chain = prompt | self.llm
        result = await chain.ainvoke({
            "section": section_title,
            "findings": findings_text,
            "context": context
        })
        return result.content

    async def write_executive_summary(
        self,
        query: str,
        full_report: str
    ) -> str:
        """Generate executive summary from full report."""
        prompt = ChatPromptTemplate.from_messages([
            ("system", """Write a concise executive summary (3-4 paragraphs).
Include: key question, main findings, critical insights, recommendations.
Write for busy executives who need quick understanding."""),
            ("human", "Research question: {query}\n\nFull report:\n{report}")
        ])

        chain = prompt | self.llm
        result = await chain.ainvoke({
            "query": query,
            "report": full_report[:8000]  # Limit for context
        })
        return result.content

# src/agent/workflow.py
from typing import Literal
from langgraph.graph import StateGraph, END
from langchain_openai import ChatOpenAI
from langchain_core.messages import HumanMessage, AIMessage

from .state import ResearchState, ResearchPhase, SubQuery, SourceInfo
from ..planning.decomposer import QueryDecomposer, SearchStrategist
from ..gathering.web_search import WebSearcher
from ..gathering.document_rag import DocumentRetriever
from ..processing.validator import InformationValidator, EntityExtractor
from ..synthesis.writer import ReportWriter
from ..config import settings

# Initialize components
decomposer = QueryDecomposer()
strategist = SearchStrategist()
web_searcher = WebSearcher()
doc_retriever = DocumentRetriever()
validator = InformationValidator()
extractor = EntityExtractor()
writer = ReportWriter()

async def planning_node(state: ResearchState) -> ResearchState:
    """Decompose query and plan search strategy."""
    # Decompose query
    decomposed = await decomposer.decompose(
        state["original_query"],
        state.get("scope", "general")
    )

    # Create sub-queries
    sub_queries = [
        SubQuery(
            query=q,
            priority=p,
            search_type=t
        )
        for q, p, t in zip(
            decomposed.sub_queries,
            decomposed.priorities,
            decomposed.search_types
        )
    ]

    # Create strategy
    strategy = await strategist.plan_strategy(
        state["original_query"],
        decomposed.sub_queries
    )

    return {
        **state,
        "sub_queries": sub_queries,
        "search_strategy": strategy,
        "current_phase": ResearchPhase.GATHERING
    }

async def gathering_node(state: ResearchState) -> ResearchState:
    """Execute searches for each sub-query."""
    raw_sources = []

    for sq in state["sub_queries"]:
        if sq.completed:
            continue

        if sq.search_type == "web":
            results = await web_searcher.search(sq.query)
            for r in results:
                raw_sources.append(SourceInfo(
                    content=r.get("content", ""),
                    source_url=r.get("url", ""),
                    source_type="web",
                    relevance_score=r.get("score", 0.5)
                ))

        elif sq.search_type == "docs":
            results = await doc_retriever.search(sq.query)
            for r in results:
                raw_sources.append(SourceInfo(
                    content=r.get("content", ""),
                    source_url=r.get("source", "internal"),
                    source_type="document",
                    relevance_score=r.get("score", 0.5)
                ))

        sq.completed = True

    return {
        **state,
        "raw_sources": raw_sources,
        "current_phase": ResearchPhase.PROCESSING
    }

async def processing_node(state: ResearchState) -> ResearchState:
    """Process and validate gathered information."""
    validated_facts = []
    all_entities = set()

    # Group sources by topic for validation
    sources = state["raw_sources"]

    # Extract entities from all sources
    for source in sources[:20]:  # Limit for efficiency
        entities = await extractor.extract(source.content)
        for category, items in entities.items():
            all_entities.update(items if isinstance(items, list) else [])

    # Validate key claims (simplified)
    # In production, would extract claims and cross-validate
    for source in sources[:10]:
        validated_facts.append({
            "content": source.content,
            "source": source.source_url,
            "validated": True,
            "confidence": source.relevance_score
        })

    return {
        **state,
        "validated_facts": validated_facts,
        "entities": list(all_entities),
        "current_phase": ResearchPhase.SYNTHESIS
    }

async def synthesis_node(state: ResearchState) -> ResearchState:
    """Generate the final research report."""
    # Create outline
    outline = await writer.create_outline(
        state["original_query"],
        state["validated_facts"]
    )

    # Write each section
    sections = {}
    for section in outline[:7]:  # Limit sections
        section_content = await writer.write_section(
            section,
            state["validated_facts"],
            state["original_query"]
        )
        sections[section] = section_content

    # Assemble full report
    full_report = "\n\n".join([
        f"## {title}\n\n{content}"
        for title, content in sections.items()
    ])

    # Generate executive summary
    exec_summary = await writer.write_executive_summary(
        state["original_query"],
        full_report
    )

    # Compile citations
    citations = [
        {
            "id": i + 1,
            "source": f["source"],
            "content": f["content"][:200]
        }
        for i, f in enumerate(state["validated_facts"])
    ]

    return {
        **state,
        "report_outline": outline,
        "draft_sections": sections,
        "final_report": full_report,
        "executive_summary": exec_summary,
        "citations": citations,
        "current_phase": ResearchPhase.COMPLETE
    }

def should_continue(state: ResearchState) -> Literal["continue", "end"]:
    """Determine if workflow should continue."""
    if state["current_phase"] == ResearchPhase.COMPLETE:
        return "end"
    if state.get("iteration", 0) > 5:
        return "end"
    return "continue"

def route_by_phase(
    state: ResearchState
) -> Literal["planning", "gathering", "processing", "synthesis"]:
    """Route to appropriate node based on phase."""
    phase = state["current_phase"]
    if phase == ResearchPhase.PLANNING:
        return "planning"
    elif phase == ResearchPhase.GATHERING:
        return "gathering"
    elif phase == ResearchPhase.PROCESSING:
        return "processing"
    else:
        return "synthesis"

# Build the graph
def create_research_workflow() -> StateGraph:
    """Create the research agent workflow."""
    workflow = StateGraph(ResearchState)

    # Add nodes
    workflow.add_node("planning", planning_node)
    workflow.add_node("gathering", gathering_node)
    workflow.add_node("processing", processing_node)
    workflow.add_node("synthesis", synthesis_node)

    # Set entry point
    workflow.set_entry_point("planning")

    # Add edges
    workflow.add_edge("planning", "gathering")
    workflow.add_edge("gathering", "processing")
    workflow.add_edge("processing", "synthesis")
    workflow.add_edge("synthesis", END)

    return workflow.compile()

# Create the agent
research_agent = create_research_workflow()

┌─────────────────────────────────────────────────────────────┐
│ RESEARCH AGENT STATE MACHINE                                │
├─────────────────────────────────────────────────────────────┤
│                                                             │
│  Entry Point                                                │
│       │                                                     │
│       ▼                                                     │
│  ┌─────────────────┐                                        │
│  │    PLANNING     │  • Decompose query into sub-queries    │
│  │                 │  • Assign priorities and search types  │
│  │                 │  • Create search strategy              │
│  └────────┬────────┘                                        │
│           │                                                 │
│           ▼                                                 │
│  ┌─────────────────┐                                        │
│  │   GATHERING     │  • Execute web searches                │
│  │                 │  • Query internal documents            │
│  │                 │  • Call external APIs                  │
│  └────────┬────────┘                                        │
│           │                                                 │
│           ▼                                                 │
│  ┌─────────────────┐                                        │
│  │   PROCESSING    │  • Extract entities                    │
│  │                 │  • Cross-validate claims               │
│  │                 │  • Deduplicate information             │
│  └────────┬────────┘                                        │
│           │                                                 │
│           ▼                                                 │
│  ┌─────────────────┐                                        │
│  │   SYNTHESIS     │  • Create report outline               │
│  │                 │  • Write sections with citations       │
│  │                 │  • Generate executive summary          │
│  └────────┬────────┘                                        │
│           │                                                 │
│           ▼                                                 │
│         END                                                 │
│                                                             │
└─────────────────────────────────────────────────────────────┘

# src/api/main.py
from fastapi import FastAPI, HTTPException, BackgroundTasks
from fastapi.responses import StreamingResponse
from pydantic import BaseModel
from typing import Optional
import asyncio
import json

from ..agent.workflow import research_agent
from ..agent.state import ResearchState, ResearchPhase

app = FastAPI(
    title="Research Assistant API",
    description="AI-powered autonomous research agent"
)

class ResearchRequest(BaseModel):
    query: str
    scope: Optional[str] = "general"

class ResearchResponse(BaseModel):
    executive_summary: str
    full_report: str
    citations: list
    entities: list

# Store for async research jobs
research_jobs = {}

@app.post("/research")
async def start_research(request: ResearchRequest):
    """Start a new research task."""
    initial_state: ResearchState = {
        "original_query": request.query,
        "scope": request.scope,
        "sub_queries": [],
        "search_strategy": "",
        "raw_sources": [],
        "validated_facts": [],
        "entities": [],
        "report_outline": [],
        "draft_sections": {},
        "final_report": "",
        "executive_summary": "",
        "citations": [],
        "current_phase": ResearchPhase.PLANNING,
        "iteration": 0,
        "errors": []
    }

    # Run research
    result = await research_agent.ainvoke(initial_state)

    return ResearchResponse(
        executive_summary=result["executive_summary"],
        full_report=result["final_report"],
        citations=result["citations"],
        entities=result["entities"]
    )

@app.post("/research/stream")
async def stream_research(request: ResearchRequest):
    """Stream research progress."""
    async def generate():
        initial_state: ResearchState = {
            "original_query": request.query,
            "scope": request.scope,
            "sub_queries": [],
            "search_strategy": "",
            "raw_sources": [],
            "validated_facts": [],
            "entities": [],
            "report_outline": [],
            "draft_sections": {},
            "final_report": "",
            "executive_summary": "",
            "citations": [],
            "current_phase": ResearchPhase.PLANNING,
            "iteration": 0,
            "errors": []
        }

        async for event in research_agent.astream(initial_state):
            for node_name, state in event.items():
                yield f"data: {json.dumps({'phase': node_name, 'status': 'processing'})}\n\n"

        yield f"data: {json.dumps({'phase': 'complete', 'status': 'done'})}\n\n"

    return StreamingResponse(
        generate(),
        media_type="text/event-stream"
    )

@app.get("/health")
async def health():
    return {"status": "healthy"}

# docker-compose.yml
version: '3.8'

services:
  research-api:
    build: .
    ports:
      - "8000:8000"
    environment:
      - OPENAI_API_KEY=${OPENAI_API_KEY}
      - BRAVE_API_KEY=${BRAVE_API_KEY}
    volumes:
      - ./data:/app/data
    healthcheck:
      test: ["CMD", "curl", "-f", "http://localhost:8000/health"]
      interval: 30s
      timeout: 10s
      retries: 3

  chroma:
    image: chromadb/chroma:latest
    ports:
      - "8001:8000"
    volumes:
      - chroma_data:/chroma/chroma

volumes:
  chroma_data:

FROM python:3.11-slim

WORKDIR /app

COPY requirements.txt .
RUN pip install --no-cache-dir -r requirements.txt

COPY src/ ./src/

EXPOSE 8000

CMD ["uvicorn", "src.api.main:app", "--host", "0.0.0.0", "--port", "8000"]