rag_pipeline.py

"""
rag_pipeline.py — Milestone 2 RAG Pipeline with Tool Calling
Hybrid retrieval (BM25 + Semantic) + Groq LLM + Tavily web search.
Uses langgraph.prebuilt.create_react_agent (modern LangChain agent pattern).
"""

import os
from pathlib import Path

from langchain_groq import ChatGroq
from langchain_core.tools.retriever import create_retriever_tool
from langchain_core.messages import HumanMessage, AIMessage, ToolMessage
from langgraph.prebuilt import create_react_agent

from hybrid import create_hybrid_retriever
from semantic import load_semantic
from utils import load_documents

ROOT = Path(__file__).resolve().parent.parent
PROC = ROOT / "data" / "processed"

SYSTEM_PROMPT = """You are a helpful Movies & TV assistant powered by Amazon review data.

You have two tools available:
- search_movie_reviews: Search Amazon Movies & TV reviews. Use this first for questions about movies, TV shows, ratings, or user opinions.
- web_search: Search the web for current information (release dates, streaming availability, recent news) not covered by reviews.

Always try search_movie_reviews first. Use web_search as a supplement when reviews lack the information.
Base your final answer on tool results. Be concise (3-5 sentences).
If neither tool returns useful results, say so honestly."""


def load_rag_agent(
    faiss_path: str,
    docs_path: str,
    groq_api_key: str | None = None,
    tavily_api_key: str | None = None,
    model: str = "llama-3.3-70b-versatile",
    k: int = 20,
    semantic_weight: float = 0.6,
    temp: int = 0.2,
    max_tokens: int = 512,
):
    """Build and return a LangGraph react agent with hybrid retriever + web search tools."""

    groq_key = groq_api_key or os.getenv("GROQ_API_KEY")
    if not groq_key:
        raise ValueError("GROQ_API_KEY not set.")

    # Build hybrid retriever
    vector_store, _ = load_semantic(faiss_path)
    documents = load_documents(docs_path)

    hybrid_retriever = create_hybrid_retriever(
        vectorstore=vector_store,
        documents=documents,
        semantic_weight=semantic_weight,
        bm25_weight=1.0 - semantic_weight,
    )
    hybrid_retriever.retrievers[0].search_kwargs = {"k": k}
    hybrid_retriever.retrievers[1].k = k

    # Retriever tool
    retriever_tool = create_retriever_tool(
        hybrid_retriever,
        name="search_movie_reviews",
        description=(
            "Search Amazon Movies & TV reviews for product recommendations, ratings, "
            "and user opinions. Use for questions about specific movies, TV shows, "
            "actors, genres, or entertainment products."
        ),
    )

    tools = [retriever_tool]

    # Web search tool
    tavily_key = tavily_api_key or os.getenv("TAVILY_API_KEY")
    if tavily_key:
        os.environ["TAVILY_API_KEY"] = tavily_key
        from web_search_tool import web_search
        tools.append(web_search)

    llm = ChatGroq(model=model, 
                   api_key=groq_key, 
                   temperature=temp, 
                   max_tokens=max_tokens)

    agent = create_react_agent(
        model=llm,
        tools=tools,
        prompt=SYSTEM_PROMPT,
    )

    return agent


def run_rag_agent(agent, query: str) -> dict:
    """
    Invoke the agent and return:
      - answer:      final response string
      - tools_used:  ordered list of tool names called
      - sources:     retrieved passages (from search_movie_reviews tool)
    """
    result = agent.invoke({"messages": [HumanMessage(content=query)]})

    tools_used = []
    sources = []

    for msg in result["messages"]:
        if isinstance(msg, AIMessage) and msg.tool_calls:
            for tc in msg.tool_calls:
                tools_used.append(tc["name"])

        if isinstance(msg, ToolMessage) and msg.name == "search_movie_reviews":
            for block in msg.content.split("\n\n"):
                if block.strip():
                    sources.append({"text": block.strip()[:300]})

    # Last AIMessage without tool calls is the final answer
    answer = ""
    for msg in reversed(result["messages"]):
        if isinstance(msg, AIMessage) and not msg.tool_calls and msg.content:
            answer = msg.content
            break

    return {
        "answer": answer,
        "tools_used": tools_used,
        "sources": sources,
    }