agent.py

"""
CodeMentor Automata - Agent & Orchestration (LangGraph Upgraded)
Implements a state-of-the-art LangGraph ReAct agent bound to DuckDuckGo 
to search unfamiliar libraries and return verified documentation.
"""
import re
from ddgs import DDGS
from langgraph.prebuilt import create_react_agent
from langchain_community.tools import DuckDuckGoSearchRun
from langchain_groq import ChatGroq

from config import GROQ_API_KEY1, GROQ_MODEL
from prompts import (
    LEVEL_PROMPTS, JSON_OUTPUT_INSTRUCTIONS, PLAN_PROMPT,
    CLEANUP_PROMPT, OPTIMIZATION_PROMPT, RESOURCE_QUERY_HINT,
    TOPIC_EXTRACTION_PROMPT,
)
from llm import call_groq_json

_search_tool = DuckDuckGoSearchRun()
_search_tool.name = "duckduckgo_search"  # Required for LangGraph routing

# Common stdlib modules we don't bother looking up
_KNOWN_STDLIB = {
    "os", "sys", "json", "re", "math", "time", "random", "collections",
    "itertools", "functools", "typing", "datetime", "string", "copy",
}


def _get_llm(temperature=0.2):
    if not GROQ_API_KEY1:
        raise RuntimeError("GROQ_API_KEY not set.")
    return ChatGroq(model=GROQ_MODEL, temperature=temperature, api_key=GROQ_API_KEY1)


def _build_doc_agent():
    """Builds a modern LangGraph ReAct agent."""
    llm = _get_llm()
    tools = [_search_tool]

    # LangGraph handles the system prompting and ReAct loop automatically
    return create_react_agent(llm, tools)


def _detect_candidate_libraries(code: str) -> list:
    """Heuristic: pull import statements as agent search candidates."""
    imports = re.findall(r"^\s*(?:import|from)\s+([\w\.]+)", code, re.MULTILINE)
    top_level = {imp.split(".")[0] for imp in imports}
    return sorted(top_level - _KNOWN_STDLIB)


def _youtube_thumbnail(url: str):
    match = re.search(r"(?:v=|youtu\.be/)([\w-]{11})", url)
    return f"https://img.youtube.com/vi/{match.group(1)}/hqdefault.jpg" if match else None


def _structured_citations(query: str, max_results: int = 3) -> list:
    """Direct DDGS search — real titles + links, so citations are always structured."""
    try:
        with DDGS() as ddgs:
            raw = list(ddgs.text(query, max_results=max_results))
    except Exception:
        return []
    return [
        {
            "title": r.get("title", "Untitled"),
            "url": r.get("href", ""),
            "snippet": (r.get("body", "") or "")[:160],
            "thumbnail": _youtube_thumbnail(r.get("href", "")),
        }
        for r in raw
    ]


def search_documentation(code: str) -> list:
    candidates = _detect_candidate_libraries(code)
    if not candidates:
        return []

    try:
        agent = _build_doc_agent()
    except RuntimeError as e:
        return [{"library": lib, "summary": f"(doc search unavailable: {e})", "citations": []} for lib in candidates]

    findings = []
    for lib in candidates:
        try:
            query = (
                f"Search for official documentation on the '{lib}' Python library/module "
                f"and summarize in 2-3 sentences what it's for and the most relevant "
                f"function/API that would show up in a typical code snippet."
            )
            response = agent.invoke({"messages": [("user", query)]})
            summary = response["messages"][-1].content
        except Exception as e:
            summary = f"(doc search failed: {e})"

        citations = _structured_citations(f"{lib} python library documentation", max_results=3)
        findings.append({"library": lib, "summary": summary, "citations": citations})

    return findings


def extract_code_topic(code: str, language: str) -> str:
    """
    Identifies the core algorithm/data-structure/concept the code implements
    (e.g. "Dijkstra's Algorithm", "Binary Search"), so resource search can
    target that exact topic instead of just the language name or raw import
    names — this is what makes 'Further Resources' show topic-specific
    videos/articles rather than generic language tutorials.
    """
    try:
        result = call_groq_json(TOPIC_EXTRACTION_PROMPT, f"Language: {language}\n\nCode:\n{code}")
        topic = (result or {}).get("topic", "").strip()
        return topic or f"{language} programming"
    except Exception:
        return f"{language} programming"


def find_resources(language: str, topic: str) -> dict:
    """
    Finds learning resources for the CORE TOPIC extracted from the code
    (e.g. "Dijkstra's Algorithm"), not just the language name. A dedicated
    video-biased query runs first — a walkthrough video of a named
    algorithm/concept is usually the most useful resource — then a broader
    query fills in with articles/docs, skipping anything already found.
    """
    video_query = f"{topic} {language} explained tutorial site:youtube.com"
    article_query = f"{topic} {language} tutorial {RESOURCE_QUERY_HINT}"

    video_citations = _structured_citations(video_query, max_results=3)
    seen_urls = {c["url"] for c in video_citations}
    article_citations = [
        c for c in _structured_citations(article_query, max_results=4)
        if c["url"] not in seen_urls
    ]

    citations = video_citations + article_citations
    return {
        "query": video_query,
        "topic": topic,
        "citations": citations,
        "ok": bool(citations),
    }


def generate_breakdown(code: str, language: str, level: str) -> dict:
    system_prompt = LEVEL_PROMPTS[level] + "\n\n" + JSON_OUTPUT_INSTRUCTIONS
    user_content = f"Language: {language}\n\nCode:\n{code}"
    return call_groq_json(system_prompt, user_content)


def generate_plan(code: str, language: str) -> dict:
    user_content = f"Language: {language}\n\nCode:\n{code}"
    return call_groq_json(PLAN_PROMPT, user_content)


def cleanup_code(code: str, language: str) -> dict:
    user_content = f"Language: {language}\n\nCode:\n{code}"
    return call_groq_json(CLEANUP_PROMPT, user_content)


def suggest_optimal_solution(code: str, language: str) -> dict:
    user_content = f"Language: {language}\n\nCode:\n{code}"
    return call_groq_json(OPTIMIZATION_PROMPT, user_content)