intelligence/orchestrator.py

"""
Agent orchestrator orchestration using LangGraph.

Defines the multi-agent orchestrator that:
1. Checks document relevance
2. Generates multiple answer candidates using research agent
3. Selects the best answer through verification
4. Provides feedback loop for iterative improvement
"""
from langgraph.graph import StateGraph, END, START
from langgraph.types import Send
from typing import TypedDict, List, Dict, Any, Optional, Annotated
import operator
from langchain_core.documents import Document
from langchain_core.retrievers import BaseRetriever
import logging

from .knowledge_synthesizer import ResearchAgent
from .accuracy_verifier import VerificationAgent
from .context_validator import ContextValidator
from langchain_google_genai import ChatGoogleGenerativeAI
from configuration.parameters import parameters

logger = logging.getLogger(__name__)


class SubQResult(TypedDict):
    idx: int
    question: str
    answer: str
    report: str


class AgentState(TypedDict, total=False):
    question: str
    documents: List[Document]
    draft_answer: str
    verification_report: str
    is_relevant: bool
    retriever: BaseRetriever
    feedback: Optional[str]
    research_attempts: int
    query_used: str
    candidate_answers: List[str]
    selection_reasoning: str
    is_multi_query: bool
    sub_queries: List[str]
    sub_results: Annotated[List[SubQResult], operator.add]


class AgentWorkflow:
    MAX_RESEARCH_ATTEMPTS: int = parameters.MAX_RESEARCH_ATTEMPTS
    NUM_RESEARCH_CANDIDATES: int = parameters.NUM_RESEARCH_CANDIDATES

    def __init__(self, num_candidates: int = None) -> None:
        logger.info("Initializing AgentWorkflow...")
        self.researcher = ResearchAgent()
        self.verifier = VerificationAgent()
        self.context_validator = ContextValidator()
        self.compiled_single = None
        self.compiled_main = None
        self.llm = ChatGoogleGenerativeAI(
            model=parameters.LLM_MODEL_NAME,
            google_api_key=parameters.GOOGLE_API_KEY,
            temperature=0.1,
            max_output_tokens=256
        )
        if num_candidates is not None:
            self.NUM_RESEARCH_CANDIDATES = num_candidates
        logger.info(f"AgentWorkflow initialized (candidates={self.NUM_RESEARCH_CANDIDATES})")

    def _retrieve_docs(self, state: AgentState) -> Dict[str, Any]:
        docs = state["retriever"].invoke(state["question"])
        return {
            "documents": docs,
            "draft_answer": "",
            "verification_report": "",
            "is_relevant": False,
            "feedback": None,
            "feedback_for_research": None,
            "contradictions_for_research": [],
            "unsupported_claims_for_research": [],
            "research_attempts": 0,
            "candidate_answers": [],
            "selection_reasoning": "",
            "query_used": state["question"],
        }

    def _build_single_question_graph(self):
        g = StateGraph(AgentState)
        g.add_node("retrieve_docs", self._retrieve_docs)
        g.add_node("check_relevance", self._check_relevance_step)
        g.add_node("research", self._research_step)
        g.add_node("verify", self._verification_step)
        g.add_edge(START, "retrieve_docs")
        g.add_edge("retrieve_docs", "check_relevance")
        g.add_conditional_edges(
            "check_relevance",
            self._decide_after_relevance_check,
            {"relevant": "research", "irrelevant": END},
        )
        g.add_edge("research", "verify")
        g.add_conditional_edges(
            "verify",
            self._decide_next_step,
            {"re_research": "research", "end": END},
        )
        return g.compile()

    def _assign_workers(self, state: AgentState):
        sends = []
        for i, q in enumerate(state.get("sub_queries", [])):
            sends.append(Send("subq_worker", {"question": q, "subq_idx": i, "retriever": state["retriever"]}))
        return sends

    def _subq_worker(self, state: AgentState) -> Dict[str, Any]:
        subq_idx = state["subq_idx"]
        q = state["question"]
        result_state = self.compiled_single.invoke({
            "question": q,
            "retriever": state["retriever"],
            "research_attempts": 0,
        })
        return {
            "sub_results": [{
                "idx": subq_idx,
                "question": q,
                "answer": result_state.get("draft_answer", ""),
                "report": result_state.get("verification_report", ""),
            }]
        }

    def _combine_answers(self, state: AgentState) -> Dict[str, Any]:
        sub_results = sorted(state.get("sub_results", []), key=lambda r: r["idx"])
        combined = "\n\n".join(
            f"Q{i+1}: {r['question']}\nA: {r['answer']}"
            for i, r in enumerate(sub_results)
        )
        return {
            "draft_answer": combined,
            "verification_report": "Multi-question answer combined."
        }

    def build_orchestrator(self) -> Any:
        self.compiled_single = self._build_single_question_graph()
        g = StateGraph(AgentState)
        g.add_node("detect_query_type", self._detect_query_type)
        g.add_node("subq_worker", self._subq_worker)
        g.add_node("combine_answers", self._combine_answers)
        def run_single(state: AgentState) -> Dict[str, Any]:
            out = self.compiled_single.invoke({
                "question": state["question"],
                "retriever": state["retriever"],
                "research_attempts": 0,
            })
            return {
                "draft_answer": out.get("draft_answer", ""),
                "verification_report": out.get("verification_report", ""),
            }
        g.add_node("run_single", run_single)
        g.set_entry_point("detect_query_type")
        g.add_conditional_edges(
            "detect_query_type",
            lambda s: "multi" if s.get("is_multi_query") else "single",
            {"multi": "fanout", "single": "run_single"},
        )
        g.add_node("fanout", lambda s: {})
        g.add_conditional_edges("fanout", self._assign_workers, ["subq_worker"])
        g.add_edge("subq_worker", "combine_answers")
        g.add_edge("combine_answers", END)
        g.add_edge("run_single", END)
        return g.compile()

    def _detect_query_type(self, state: AgentState) -> Dict[str, Any]:
        prompt = f"""
You are an expert assistant for document Q&A. Analyze the following question and determine:
1. Is it a single question or does it contain multiple sub-questions?
2. If it contains multiple questions, decompose it into a list of clear, standalone sub-questions (no overlap, no ambiguity).

Return your answer as a JSON object with two fields:
- is_multi_query: true or false
- sub_queries: a list of strings (the sub-questions, or a single-item list if only one)

Question: {state['question']}
"""
        try:
            response = self.llm.invoke(prompt)
            import json
            content = response.content if hasattr(response, "content") else str(response)
            start = content.find('{')
            end = content.rfind('}')
            if start != -1 and end != -1:
                json_str = content[start:end+1]
                result = json.loads(json_str)
                is_multi = bool(result.get("is_multi_query", False))
                sub_queries = result.get("sub_queries", [])
            else:
                is_multi = False
                sub_queries = [state["question"]]
        except Exception as e:
            logger.error(f"LLM decomposition failed: {e}")
            is_multi = False
            sub_queries = [state["question"]]
        if is_multi:
            logger.info(f"[LLM Decompose] Multi-question detected: {len(sub_queries)} sub-queries")
        else:
            logger.info("[LLM Decompose] Single question detected; no decomposition needed.")
        return {"is_multi_query": is_multi, "sub_queries": sub_queries}

    def _check_relevance_step(self, state: AgentState) -> Dict[str, Any]:
        logger.debug("Checking context relevance...")
        result = self.context_validator.context_validate_with_rewrite(
            question=state["question"],
            retriever=state["retriever"],
            k=parameters.RELEVANCE_CHECK_K,
            max_rewrites=parameters.MAX_QUERY_REWRITES,
        )
        classification = result.get("classification", "NO_MATCH")
        query_used = result.get("query_used", state["question"])
        logger.info(f"Relevance: {classification} (query_used={query_used[:80]})")
        if classification in ("CAN_ANSWER", "PARTIAL"):
            documents = state["retriever"].invoke(query_used)
            return {
                "is_relevant": True,
                "query_used": query_used,
                "documents": documents
            }
        return {
            "is_relevant": False,
            "query_used": query_used,
            "draft_answer": "This question isn't related to the uploaded documents. Please ask another question.",
        }

    def _decide_after_relevance_check(self, state: AgentState) -> str:
        return "relevant" if state["is_relevant"] else "irrelevant"

    def run_workflow(self, question: str, retriever: BaseRetriever) -> Dict[str, str]:
        if self.compiled_main is None:
            self.compiled_main = self.build_orchestrator()
        initial_state: AgentState = {
            "question": question,
            "retriever": retriever,
            "sub_results": [],
            "sub_queries": [],
            "is_multi_query": False,
        }
        final = self.compiled_main.invoke(initial_state)
        return {
            "draft_answer": final.get("draft_answer", ""),
            "verification_report": final.get("verification_report", ""),
        }

    def _verification_step(self, state: AgentState) -> Dict[str, Any]:
        logger.debug("Selecting best answer from candidates...")
        candidate_answers = state.get("candidate_answers", []) or [state.get("draft_answer", "")]
        selection_result = self.verifier.select_best_answer(
            candidate_answers=candidate_answers,
            documents=state["documents"],
            question=state["question"]
        )
        best_answer = selection_result["selected_answer"]
        selection_reasoning = selection_result.get("reasoning", "")
        logger.info(f"Selected candidate {selection_result['selected_index'] + 1} as best answer")
        verification_result = self.verifier.check(
            answer=best_answer,
            documents=state["documents"],
            question=state["question"]
        )
        verification_report = verification_result["verification_report"]
        verification_report = f"**Candidates Evaluated:** {len(candidate_answers)}\n" + \
                             f"**Selected Candidate:** {selection_result['selected_index'] + 1}\n" + \
                             f"**Selection Confidence:** {selection_result.get('confidence', 'N/A')}\n" + \
                             f"**Selection Reasoning:** {selection_reasoning}\n\n" + \
                             verification_report
        feedback_for_research = verification_result.get("feedback")
        return {
            "draft_answer": best_answer,
            "verification_report": verification_report,
            "feedback_for_research": feedback_for_research,
            "selection_reasoning": selection_reasoning,
            "should_retry": verification_result.get("should_retry", False),
        }

    def _decide_next_step(self, state: AgentState) -> str:
        research_attempts = state.get("research_attempts", 1)
        should_retry = bool(state.get("should_retry", False))
        if should_retry and research_attempts < self.MAX_RESEARCH_ATTEMPTS:
            return "re_research"
        return "end"

    def _research_step(self, state: AgentState) -> Dict[str, Any]:
        attempts = state.get("research_attempts", 0) + 1
        feedback_for_research = state.get("feedback_for_research")
        previous_answer = state.get("draft_answer") if feedback_for_research else None
        logger.info(f"Research step (attempt {attempts}/{self.MAX_RESEARCH_ATTEMPTS})")
        logger.info(f"Generating {self.NUM_RESEARCH_CANDIDATES} candidate answers in parallel...")
        
        # Parallel candidate generation for 2× speedup
        import concurrent.futures
        candidate_answers = []
        
        def generate_candidate(index):
            logger.info(f"Generating candidate {index + 1}/{self.NUM_RESEARCH_CANDIDATES}")
            result = self.researcher.generate(
                question=state["question"],
                documents=state["documents"],
                feedback=feedback_for_research,
                previous_answer=previous_answer
            )
            return result["draft_answer"]
        
        # Use ThreadPoolExecutor for parallel LLM API calls (I/O-bound)
        with concurrent.futures.ThreadPoolExecutor(max_workers=self.NUM_RESEARCH_CANDIDATES) as executor:
            futures = [executor.submit(generate_candidate, i) for i in range(self.NUM_RESEARCH_CANDIDATES)]
            for future in concurrent.futures.as_completed(futures):
                try:
                    candidate_answers.append(future.result())
                except Exception as e:
                    logger.error(f"Candidate generation failed: {e}")
                    # Continue with other candidates even if one fails
        
        logger.info(f"Generated {len(candidate_answers)} candidate answers in parallel")
        return {
            "candidate_answers": candidate_answers,
            "research_attempts": attempts,
            "feedback": None
        }