agents/analyst.py

"""
ANALYST AGENT
=============

Purpose:
This agent takes raw research output and extracts structured insights.

Why this agent exists:
The researcher returns a pile of raw text from web + vector search.
That raw text is noisy, unfiltered, and full of redundancy.
The analyst's job is to:
• identify patterns
• extract key facts
• run code if math/data analysis is needed
• return clean, structured findings

Why a separate agent (not just the writer)?
Separation of concerns — analysis ≠ writing.
An analyst focuses on WHAT is true.
A writer focuses on HOW to communicate it.
Splitting these produces far better final reports.
"""

# =========================
# Imports
# =========================
from langchain_groq import ChatGroq
from langchain_core.prompts import ChatPromptTemplate
from pydantic import BaseModel
from typing import List

# Import the sandboxed code executor.
# WHY:
# If the planner flagged requires_code=True, the analyst
# needs to run computations (e.g. averages, charts, statistics).
# We never let the LLM invent numbers — we make it run real code.
from tools.python_repl import python_repl


# =========================
# 1️⃣ Structured Output Schema
# =========================
# WHY:
# We force the LLM to return structured JSON instead of free text.
# This makes the analyst output reliably parseable by the writer.
# A free-text response would require brittle string parsing.

class AnalysisResult(BaseModel):
    # 3-5 bullet-point insights extracted from research
    key_insights: List[str]

    # A short paragraph summarizing the overall findings
    summary: str

    # Any code output if math/data was computed (empty string if not)
    code_output: str

    # How confident is the agent in these findings (low/medium/high)
    confidence: str


# =========================
# 2️⃣ Analyst Function
# =========================
def run_analyst(
    research_text: str,
    original_query: str,
    requires_code: bool = False,
    code_task: str = ""
) -> AnalysisResult:
    """
    Analyze raw research output → structured insights.

    Parameters
    ----------
    research_text : str
        Raw output from the researcher agent.

    original_query : str
        The user's original question (keeps analyst focused).

    requires_code : bool
        If True, the analyst will also run Python for computation.
        This flag comes directly from the planner's ResearchPlan.

    code_task : str
        Optional Python code to execute if requires_code is True.
        Generated by the LLM in a first pass before structured output.

    Returns
    -------
    AnalysisResult
        Structured insights, summary, code output, confidence level.
    """

    # =========================
    # Step A: Optional Code Execution
    # =========================
    # WHY:
    # If the planner flagged requires_code=True, we run any
    # computation BEFORE calling the LLM for analysis.
    #
    # This matters because:
    # • LLMs hallucinate numbers — code execution gives real results
    # • We inject the real output INTO the LLM prompt
    # • This grounds the analysis in verified data

    code_result = ""
    if requires_code and code_task:
        print("[Analyst] Running code task...")

        # python_repl is sandboxed — safe to run agent-generated code.
        # It only allows safe libraries and has a 3-second timeout.
        code_result = python_repl.run(code_task)

        print(f"[Analyst] Code output: {code_result}")


    # =========================
    # Step B: LLM Setup
    # =========================
    # WHY:
    # We use claude-sonnet-4-6 here because analysis requires
    # more reasoning depth than retrieval.
    #
    # with_structured_output() forces the LLM to return valid JSON
    # matching the AnalysisResult schema — no parsing errors.
    llm = ChatGroq(
        model="llama-3.3-70b-versatile"
    ).with_structured_output(AnalysisResult)


    # =========================
    # Step C: Prompt Construction
    # =========================
    # WHY:
    # We inject three things into the prompt:
    # 1. The original query → keeps the analysis on-topic
    # 2. The raw research text → the material to analyze
    # 3. Code output (if any) → grounds findings in real computation
    #
    # The system message instructs the LLM to act as a critical analyst,
    # not a summarizer. This produces sharper, more useful insights.

    prompt = ChatPromptTemplate.from_messages([
        (
            "system",
            "You are a critical research analyst. "
            "Your job is to extract 3–5 key insights from raw research. "
            "Be precise. Avoid repetition. "
            "If code output is provided, incorporate it into your insights. "
            "Rate your confidence (low/medium/high) based on source quality."
        ),
        (
            "human",
            "Original question: {query}\n\n"
            "Research gathered:\n{research}\n\n"
            "Code execution output (if any):\n{code_output}"
        )
    ])


    # =========================
    # Step D: Run the Chain
    # =========================
    # WHY:
    # prompt | llm → LCEL pipeline (LangChain Expression Language)
    # This is the standard, composable way to chain components.
    # invoke() executes it synchronously and returns AnalysisResult.

    result = (prompt | llm).invoke({
        "query": original_query,
        "research": research_text,
        "code_output": code_result if code_result else "No code was executed."
    })

    # Attach the real code output to the result before returning.
    # WHY:
    # with_structured_output fills code_output from the LLM's JSON.
    # We overwrite it with the REAL executed output to ensure accuracy.
    # LLMs cannot fabricate execution results this way.
    result.code_output = code_result

    return result


# =========================
# Example test
# =========================
if __name__ == "__main__":

    # Simulated researcher output for offline testing
    sample_research = """
    Title: AI Job Market 2024
    Summary: AI engineer salaries range from $120k to $300k in the US.
    Major employers include Google, OpenAI, Microsoft, and Meta.
    Source: techcrunch.com

    Title: ML Engineer Demand
    Summary: Demand for ML roles grew 35% YoY. Python and PyTorch are top skills.
    Source: linkedin.com
    """

    # Test with code execution enabled
    result = run_analyst(
        research_text=sample_research,
        original_query="Analyze AI job trends and salary ranges",
        requires_code=True,
        code_task="salaries = [120000, 180000, 220000, 300000]\nprint(f'Mean: {sum(salaries)/len(salaries)}')"
    )

    print("\n=== Analysis Result ===")
    print("Key Insights:", result.key_insights)
    print("Summary:", result.summary)
    print("Code Output:", result.code_output)
    print("Confidence:", result.confidence)