File size: 3,260 Bytes
8d56dc3 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 |
"""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
# Wikipedia tool
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 # lazy-init agent
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."
)
|