| | """LangGraph nodes for RAG workflow + ReAct Agent inside generate_content"""
|
| |
|
| | from typing import List, Optional
|
| | from src.state.rag_state import RAGState
|
| |
|
| | from langchain_core.documents import Document
|
| | from langchain_core.tools import Tool
|
| | from langchain_core.messages import HumanMessage
|
| | from langgraph.prebuilt import create_react_agent
|
| |
|
| |
|
| | from langchain_community.utilities import WikipediaAPIWrapper
|
| | from langchain_community.tools.wikipedia.tool import WikipediaQueryRun
|
| |
|
| |
|
| | class RAGNodes:
|
| | """Contains node functions for RAG workflow"""
|
| |
|
| | def __init__(self, retriever, llm):
|
| | self.retriever = retriever
|
| | self.llm = llm
|
| | self._agent = None
|
| |
|
| | def retrieve_docs(self, state: RAGState) -> RAGState:
|
| | """Classic retriever node"""
|
| | docs = self.retriever.invoke(state.question)
|
| | return RAGState(
|
| | question=state.question,
|
| | retrieved_docs=docs
|
| | )
|
| |
|
| | def _build_tools(self) -> List[Tool]:
|
| | """Build retriever + wikipedia tools"""
|
| |
|
| | def retriever_tool_fn(query: str) -> str:
|
| | docs: List[Document] = self.retriever.invoke(query)
|
| | if not docs:
|
| | return "No documents found."
|
| | merged = []
|
| | for i, d in enumerate(docs[:8], start=1):
|
| | meta = d.metadata if hasattr(d, "metadata") else {}
|
| | title = meta.get("title") or meta.get("source") or f"doc_{i}"
|
| | merged.append(f"[{i}] {title}\n{d.page_content}")
|
| | return "\n\n".join(merged)
|
| |
|
| | retriever_tool = Tool(
|
| | name="retriever",
|
| | description="Fetch passages from indexed corpus.",
|
| | func=retriever_tool_fn,
|
| | )
|
| |
|
| | wiki = WikipediaQueryRun(
|
| | api_wrapper=WikipediaAPIWrapper(top_k_results=3, lang="en")
|
| | )
|
| | wikipedia_tool = Tool(
|
| | name="wikipedia",
|
| | description="Search Wikipedia for general knowledge.",
|
| | func=wiki.run,
|
| | )
|
| |
|
| | return [retriever_tool, wikipedia_tool]
|
| |
|
| | def _build_agent(self):
|
| | """ReAct agent with tools"""
|
| | tools = self._build_tools()
|
| | system_prompt = (
|
| | "You are a helpful RAG agent. "
|
| | "Prefer 'retriever' for user-provided docs; use 'wikipedia' for general knowledge. "
|
| | "Return only the final useful answer."
|
| | )
|
| | self._agent = create_react_agent(self.llm, tools=tools,prompt=system_prompt)
|
| |
|
| | def generate_answer(self, state: RAGState) -> RAGState:
|
| | """
|
| | Generate answer using ReAct agent with retriever + wikipedia.
|
| | """
|
| | if self._agent is None:
|
| | self._build_agent()
|
| |
|
| | result = self._agent.invoke({"messages": [HumanMessage(content=state.question)]})
|
| |
|
| | messages = result.get("messages", [])
|
| | answer: Optional[str] = None
|
| | if messages:
|
| | answer_msg = messages[-1]
|
| | answer = getattr(answer_msg, "content", None)
|
| |
|
| | return RAGState(
|
| | question=state.question,
|
| | retrieved_docs=state.retrieved_docs,
|
| | answer=answer or "Could not generate answer."
|
| | )
|
| |
|
