File size: 1,798 Bytes
f35d7bb 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 |
"""LangGraph nodes for RAG workflow"""
from src.state.rag_state import RAGState
class RAGNodes:
"""Contains node functions for RAG workflow"""
def __init__(self, retriever, llm):
"""
Initialize RAG nodes
Args:
retriever: Document retriever instance
llm: Language model instance
"""
self.retriever = retriever
self.llm = llm
def retrieve_docs(self, state: RAGState) -> RAGState:
"""
Retrieve relevant documents node
Args:
state: Current RAG state
Returns:
Updated RAG state with retrieved documents
"""
docs = self.retriever.invoke(state.question)
return RAGState(
question=state.question,
retrieved_docs=docs
)
def generate_answer(self, state: RAGState) -> RAGState:
"""
Generate answer from retrieved documents node
Args:
state: Current RAG state with retrieved documents
Returns:
Updated RAG state with generated answer
"""
# Combine retrieved documents into context
context = "\n\n".join([doc.page_content for doc in state.retrieved_docs])
# Create prompt
prompt = f"""You are a professional Project Analyst.
Answer strictly using the context.
If unknown, say you don't know.
Context:
{context}
Question: {state.question}"""
# Generate response
response = self.llm.invoke(prompt)
return RAGState(
question=state.question,
retrieved_docs=state.retrieved_docs,
answer=response.content
) |