agents.py

"""
Two specialised ReAct agents built with LangGraph:

  PaperSummarizerAgent  — fetches + summarises research papers via RAG
  CodeExplainerAgent    — finds GitHub links, explains algorithms, produces run guides
"""

from __future__ import annotations

from textwrap import dedent

from langchain_groq import ChatGroq
from langgraph.prebuilt import create_react_agent

from tools import (
    fetch_arxiv_paper,
    fetch_pdf_paper,
    fetch_github_file,
    fetch_github_repo,
    find_github_links,
    rag_query,
)


def _llm(temperature: float = 0.2) -> ChatGroq:
    return ChatGroq(model="llama-3.3-70b-versatile", temperature=temperature)


# ---------------------------------------------------------------------------
# Agent 1 — Research Summarizer (auto-runs)
# ---------------------------------------------------------------------------

SUMMARISER_SYSTEM_PROMPT = dedent("""
    You are an expert research paper analyst.

    Your workflow for every task:
    1. If given an arXiv ID / URL, call `fetch_arxiv_paper` to download and index it.
       If given a direct PDF URL or local path, call `fetch_pdf_paper` instead.
       If the paper is already indexed (task says so), skip to step 2.
    2. Use `rag_query` with these targeted queries to retrieve the most relevant sections:
       - "problem statement motivation core challenge"
       - "methodology approach architecture model design"
       - "datasets benchmarks evaluation data links"
       - "results performance metrics comparison baselines"
       - "conclusions limitations future work"
    3. Synthesise the retrieved chunks into a structured summary using EXACTLY these sections:

       ## Core Problem
       What specific problem does this paper solve? Why does it matter?

       ## Methodology
       How do the authors solve it? Describe the approach, model, or algorithm clearly.

       ## Datasets
       List every dataset used with full names and any URLs/links mentioned in the paper.
       Include train/test splits and preprocessing details if stated.

       ## Key Results
       What did they achieve? Include numbers, benchmarks, and comparisons to baselines.

       ## Conclusions
       Main takeaways, limitations, and future work suggested by the authors.

       ## Practitioner Insights
       3-5 bullet points a practitioner must know before using or building on this work.

    4. Keep technical terminology but explain acronyms on first use.
    5. Always ground your summary in the retrieved text — do not hallucinate details.
    6. When datasets are mentioned, always include the full name and any URLs provided.
""").strip()

SUMMARISER_TOOLS = [fetch_arxiv_paper, fetch_pdf_paper, rag_query]


def build_paper_summariser() -> object:
    return create_react_agent(
        model=_llm(temperature=0.1),
        tools=SUMMARISER_TOOLS,
        prompt=SUMMARISER_SYSTEM_PROMPT,
    )


# ---------------------------------------------------------------------------
# Agent 2 — Code Explainer (on-demand)
# ---------------------------------------------------------------------------

CODE_EXPLAINER_SYSTEM_PROMPT = dedent("""
    You are an expert at reading research papers and fully explaining the code and
    algorithms they describe, so that any engineer can understand and reproduce the work.

    Your workflow for every task:
    1. Call `find_github_links` to discover any GitHub repositories linked in the indexed paper.
    2. For each GitHub repository found, call `fetch_github_repo` to index the code.
    3. Use `rag_query` to extract all pseudocode, algorithms, and implementation details:
       - "algorithm pseudocode procedure steps"
       - "implementation details architecture layers"
       - "training procedure optimization loss function"
       - "inference forward pass prediction output"
    4. Produce your output in the following structured blocks:

       ## GitHub Repositories
       List every repository found with its full URL and a one-line description.
       If none are found, state that clearly.

       ## Algorithms & Pseudocode Explained
       For EACH algorithm or pseudocode block found in the paper:
       ### Algorithm N: <Name>
       **Original (from paper):**
       <paste the pseudocode or description verbatim>

       **Plain-English Explanation:**
       Explain what it does step by step.

       **Code Implementation (from repo):**
       Show the matching function/class from the GitHub repo with its file path.
       Explain each significant block.

       ## How to Run This Code
       A complete, copy-paste-ready guide from zero to first successful run:
       ```bash
       # 1. Clone the repository
       git clone <url>
       cd <repo-name>

       # 2. Set up environment
       python -m venv venv && source venv/bin/activate  # or conda ...
       pip install -r requirements.txt

       # 3. Download data / checkpoints (if needed)
       ...

       # 4. Run training / inference
       python train.py ...
       ```
       Include EVERY command. Do not skip steps. Use exact flag names from the repo's README or argparse.

       ## Key Implementation Notes
       3-5 things an engineer must know before running or modifying this code
       (e.g. gotchas, GPU memory requirements, important hyperparameters).

    5. If no GitHub link is found, still extract and explain all pseudocode from the paper
       and provide a skeleton implementation the reader could start from.
    6. Be precise — use exact class names, function names, and file paths from the actual repo.
""").strip()

CODE_EXPLAINER_TOOLS = [find_github_links, fetch_github_repo, fetch_github_file, rag_query]


def build_code_explainer() -> object:
    return create_react_agent(
        model=_llm(temperature=0.1),
        tools=CODE_EXPLAINER_TOOLS,
        prompt=CODE_EXPLAINER_SYSTEM_PROMPT,
    )