├── .env.example ├── .gitignore ├── README.md ├── demo.ipynb ├── langgraph.json ├── pandas_rag_langgraph └── agent.py ├── poetry.lock ├── pyproject.toml └── static └── pandas-rag-agent.png /.env.example: -------------------------------------------------------------------------------- 1 | OPENAI_API_KEY=placeholder 2 | TAVILY_API_KEY=placeholder 3 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | __pycache__ 2 | .langgraph-data 3 | .env 4 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # Pandas RAG with LangGraph 2 | 3 | This is a demo app for a RAG system that answers questions about [Pandas](https://pandas.pydata.org/) documentation powered by [LangGraph](https://github.com/langchain-ai/langgraph) and [LangChain](https://github.com/langchain-ai/langchain). 4 | 5 | ## Data 6 | 7 | This demo uses a **very** small subset of [Pandas](https://pandas.pydata.org/) documentation for question answering, namely the following 3 pages from user guides: 8 | 9 | * [Indexing](https://pandas.pydata.org/docs/user_guide/indexing.html) 10 | * [GroupBy](https://pandas.pydata.org/docs/user_guide/groupby.html) 11 | * [Merging](https://pandas.pydata.org/docs/user_guide/merging.html) 12 | 13 | Such a small subset is chosen on purpose to demonstrate how the agent architecture below can flexibly handle model hallucinations. Specifically, there is a lot of knowledge about popular open-source libraries already present in the model weights. Therefore, if the model fails to look up data from vectorstore, it might still attempt to answer the question. In certain cases this behavior might be not desireable, and this demo demonstrates how you can handle those cases. 14 | 15 | ## Components 16 | 17 | The demo uses the following components: 18 | 19 | - LLM: Anthropic's Claude 3.5 Sonnet (`claude-3-5-sonnet-20240620`) via LangChain's [`ChatAnthropic`](https://python.langchain.com/v0.2/docs/integrations/chat/anthropic/). Specifically, the LLM is used for three different tasks: 20 | - candidate answer generation 21 | - grading answer hallucinations 22 | - grading answer relevance 23 | - Embeddings: OpenAI Embeddings (`text-embedding-ada-002`) via LangChain's [`OpenAIEmbeddings`](https://python.langchain.com/v0.2/docs/integrations/text_embedding/openai/) 24 | - Vectorstore: [Chroma DB](https://www.trychroma.com/) (via LangChain's [`Chroma`](https://python.langchain.com/v0.2/docs/integrations/vectorstores/chroma/) 25 | - **NOTE**: for production use cases you would likely need to deploy your own vector database instance or use a hosted solution 26 | - Web search: [Tavily Search](https://tavily.com/) (via LangChain's [`TavilySearchResults`](https://python.langchain.com/v0.2/docs/integrations/tools/tavily_search/)) 27 | 28 | ## Architecture 29 | 30 | This demo implements a custom RAG architecture that combines ideas from [Self-RAG](https://arxiv.org/abs/2310.11511) and [Corrective RAG](https://arxiv.org/abs/2401.15884). For simplicity, it omits the document relevance grading step, but you can find full implementation of those papers in LangGraph [here](https://langchain-ai.github.io/langgraph/tutorials/rag/langgraph_self_rag/) and [here](https://langchain-ai.github.io/langgraph/tutorials/rag/langgraph_crag/). 31 | 32 | The flow is as follows: 33 | 34 | 1. Search vectorstore for documents relevant to the user question 35 | 2. Generate a candidate answer based on those documents 36 | 3. Grade the candidate answer for hallucinations: is it grounded in the documents or did the model hallucinate? 37 | - if grounded, proceed to the next step (grading answer relevance) 38 | - if hallucination, re-generate the answer (return to step 2). Attempt re-generation at most N times (user-controlled) 39 | 4. Grade the candidate answer for relevance: did it address user's question? 40 | - if yes, return the candidate answer to the user 41 | - if no, rewrite the query and attempt the search (return to step 1) again. Attempt re-writing at most N times (user-controlled) 42 | 5. (Optional) If the answer is not grounded in the documents and/or not relevant after N tries, pass the user question to web search 43 | 6. (Optional) Generate an answer based on the web search results and return it to the user 44 | 45 | See this flow chart for a visual representation. 46 | 47 | ![rag-agent](/static/pandas-rag-agent.png) 48 | 49 | ## Interacting with the graph 50 | 51 | First, make sure that your environment variables are set in `.env` file. See `.env.example`. 52 | 53 | ```python 54 | from dotenv import load_dotenv 55 | 56 | load_dotenv() 57 | 58 | from pandas_rag_langgraph.agent import graph 59 | 60 | inputs = {"messages": [("human", "how do i calculate sum by groups")]} 61 | for output in graph.stream(inputs): 62 | print(output) 63 | ``` 64 | 65 | ## Deploying LangGraph applications 66 | 67 | If you'd like to deploy a LangGraph application like the one in this demo, you can use [LangGraph Cloud](https://langchain-ai.github.io/langgraph/cloud/) 68 | 69 | ## Further work 70 | 71 | This demo can be improved in a variety of ways: 72 | - add chat history to ask follow up questions (see example of that in [Chat LangChain](https://github.com/langchain-ai/chat-langchain)) 73 | - add better document chunking 74 | - add document relevance grader ([here](https://langchain-ai.github.io/langgraph/tutorials/rag/langgraph_self_rag/) and [here](https://langchain-ai.github.io/langgraph/tutorials/rag/langgraph_crag/)) 75 | - add a mixed keyword search ([BM25](https://en.wikipedia.org/wiki/Okapi_BM25)) & vector search solution for document retrieval -------------------------------------------------------------------------------- /langgraph.json: -------------------------------------------------------------------------------- 1 | { 2 | "dependencies": ["."], 3 | "graphs": { 4 | "agent": "./pandas_rag_langgraph/agent.py:graph" 5 | }, 6 | "env": ".env" 7 | } 8 | -------------------------------------------------------------------------------- /pandas_rag_langgraph/agent.py: -------------------------------------------------------------------------------- 1 | import re 2 | from typing import Annotated, Iterator, Literal, TypedDict 3 | 4 | from langchain import hub 5 | from langchain_community.document_loaders import web_base 6 | from langchain_community.vectorstores import Chroma 7 | from langchain_community.tools.tavily_search import TavilySearchResults 8 | from langchain_core.documents import Document 9 | from langchain_core.output_parsers import StrOutputParser 10 | from langchain_core.messages import BaseMessage, AIMessage, convert_to_messages 11 | from langchain_core.prompts import ChatPromptTemplate 12 | from langchain_core.pydantic_v1 import BaseModel, Field 13 | from langchain_core.retrievers import BaseRetriever 14 | from langchain_text_splitters import RecursiveCharacterTextSplitter 15 | from langchain_anthropic import ChatAnthropic 16 | from langchain_openai import OpenAIEmbeddings 17 | from langgraph.graph import END, StateGraph, add_messages 18 | 19 | MAX_RETRIES = 3 20 | 21 | # Index 3 pages from Pandas user guides 22 | SOURCE_URLS = [ 23 | 'https://pandas.pydata.org/docs/user_guide/indexing.html', 24 | 'https://pandas.pydata.org/docs/user_guide/groupby.html', 25 | 'https://pandas.pydata.org/docs/user_guide/merging.html' 26 | ] 27 | 28 | NEWLINE_RE = re.compile("\n+") 29 | 30 | class PandasDocsLoader(web_base.WebBaseLoader): 31 | def lazy_load(self) -> Iterator[Document]: 32 | """Lazy load text from the url(s) in web_path.""" 33 | for path in self.web_paths: 34 | soup = self._scrape(path, bs_kwargs=self.bs_kwargs) 35 | text = soup.get_text(**self.bs_get_text_kwargs) 36 | text = NEWLINE_RE.sub("\n", text) 37 | metadata = web_base._build_metadata(soup, path) 38 | yield Document(page_content=text, metadata=metadata) 39 | 40 | 41 | def prepare_documents(urls: list[str]) -> list[Document]: 42 | text_splitter = RecursiveCharacterTextSplitter( 43 | separators=[ 44 | r"In \[[0-9]+\]", 45 | r"\n+", 46 | r"\s+" 47 | ], 48 | is_separator_regex=True, 49 | chunk_size=1000 50 | ) 51 | docs = [PandasDocsLoader(url).load() for url in urls] 52 | docs_list = [item for sublist in docs for item in sublist] 53 | return text_splitter.split_documents(docs_list) 54 | 55 | 56 | def get_retriever() -> BaseRetriever: 57 | documents = prepare_documents(SOURCE_URLS) 58 | vectorstore = Chroma.from_documents( 59 | documents=documents, 60 | collection_name="pandas-rag-chroma", 61 | embedding=OpenAIEmbeddings(), 62 | ) 63 | retriever = vectorstore.as_retriever() 64 | return retriever 65 | 66 | 67 | # LLM / Retriever / Tools 68 | llm = ChatAnthropic(model="claude-3-5-sonnet-20240620", temperature=0) 69 | retriever = get_retriever() 70 | tavily_search_tool = TavilySearchResults(max_results=3) 71 | 72 | # Prompts / data models 73 | 74 | RAG_PROMPT: ChatPromptTemplate = hub.pull("rlm/rag-prompt") 75 | 76 | 77 | class GradeHallucinations(BaseModel): 78 | """Binary score for hallucination present in generation answer.""" 79 | 80 | binary_score: str = Field( 81 | description="Answer is grounded in the facts, 'yes' or 'no'" 82 | ) 83 | 84 | 85 | HALLUCINATION_GRADER_SYSTEM = ( 86 | """ 87 | You are a grader assessing whether an LLM generation is grounded in / supported by a set of retrieved facts. 88 | Give a binary score 'yes' or 'no', where 'yes' means that the answer is grounded in / supported by the set of facts. 89 | 90 | IF the generation includes code examples, make sure those examples are FULLY present in the set of facts, otherwise always return score 'no'. 91 | """ 92 | ) 93 | HALLUCINATION_GRADER_PROMPT = ChatPromptTemplate.from_messages( 94 | [ 95 | ("system", HALLUCINATION_GRADER_SYSTEM), 96 | ("human", "Set of facts: \n\n {documents} \n\n LLM generation: {generation}"), 97 | ] 98 | ) 99 | 100 | 101 | class GradeAnswer(BaseModel): 102 | """Binary score to assess answer addresses question.""" 103 | 104 | binary_score: str = Field( 105 | description="Answer addresses the question, 'yes' or 'no'" 106 | ) 107 | 108 | 109 | ANSWER_GRADER_SYSTEM = ( 110 | """ 111 | You are a grader assessing whether an answer addresses / resolves a question. 112 | Give a binary score 'yes' or 'no', where 'yes' means that the answer resolves the question. 113 | """ 114 | ) 115 | ANSWER_GRADER_PROMPT = ChatPromptTemplate.from_messages( 116 | [ 117 | ("system", ANSWER_GRADER_SYSTEM), 118 | ("human", "User question: \n\n {question} \n\n LLM generation: {generation}"), 119 | ] 120 | ) 121 | 122 | 123 | QUERY_REWRITER_SYSTEM = ( 124 | """ 125 | You a question re-writer that converts an input question to a better version that is optimized for vectorstore retrieval. 126 | Look at the input and try to reason about the underlying semantic intent / meaning. 127 | """ 128 | ) 129 | QUERY_REWRITER_PROMPT = ChatPromptTemplate.from_messages( 130 | [ 131 | ("system", QUERY_REWRITER_SYSTEM), 132 | ( 133 | "human", 134 | "Here is the initial question: \n\n {question} \n Formulate an improved question.", 135 | ), 136 | ] 137 | ) 138 | 139 | 140 | class GraphState(TypedDict): 141 | messages: Annotated[list[BaseMessage], add_messages] 142 | question: str 143 | documents: list[Document] 144 | candidate_answer: str 145 | retries: int 146 | web_fallback: bool 147 | 148 | 149 | class GraphConfig(TypedDict): 150 | max_retries: int 151 | 152 | 153 | def document_search(state: GraphState): 154 | """ 155 | Retrieve documents 156 | 157 | Args: 158 | state (dict): The current graph state 159 | 160 | Returns: 161 | state (dict): New key added to state, documents, that contains retrieved documents 162 | """ 163 | print("---RETRIEVE---") 164 | question = convert_to_messages(state["messages"])[-1].content 165 | 166 | # Retrieval 167 | documents = retriever.invoke(question) 168 | return {"documents": documents, "question": question, "web_fallback": True} 169 | 170 | 171 | def generate(state: GraphState): 172 | """ 173 | Generate answer 174 | 175 | Args: 176 | state (dict): The current graph state 177 | 178 | Returns: 179 | state (dict): New key added to state, generation, that contains LLM generation 180 | """ 181 | print("---GENERATE---") 182 | question = state["question"] 183 | documents = state["documents"] 184 | retries = state["retries"] if state.get("retries") is not None else -1 185 | 186 | rag_chain = RAG_PROMPT | llm | StrOutputParser() 187 | generation = rag_chain.invoke({"context": documents, "question": question}) 188 | return {"retries": retries + 1, "candidate_answer": generation} 189 | 190 | 191 | def transform_query(state: GraphState): 192 | """ 193 | Transform the query to produce a better question. 194 | 195 | Args: 196 | state (dict): The current graph state 197 | 198 | Returns: 199 | state (dict): Updates question key with a re-phrased question 200 | """ 201 | print("---TRANSFORM QUERY---") 202 | question = state["question"] 203 | 204 | # Re-write question 205 | query_rewriter = QUERY_REWRITER_PROMPT | llm | StrOutputParser() 206 | better_question = query_rewriter.invoke({"question": question}) 207 | return {"question": better_question} 208 | 209 | 210 | def web_search(state: GraphState): 211 | print("---RUNNING WEB SEARCH---") 212 | question = state["question"] 213 | documents = state["documents"] 214 | search_results = tavily_search_tool.invoke(question) 215 | search_content = "\n".join([d["content"] for d in search_results]) 216 | documents.append(Document(page_content=search_content, metadata={"source": "websearch"})) 217 | return {"documents": documents, "web_fallback": False} 218 | 219 | 220 | ### Edges 221 | 222 | 223 | def grade_generation_v_documents_and_question(state: GraphState, config) -> Literal["generate", "transform_query", "web_search", "finalize_response"]: 224 | """ 225 | Determines whether the generation is grounded in the document and answers question. 226 | 227 | Args: 228 | state (dict): The current graph state 229 | 230 | Returns: 231 | str: Decision for next node to call 232 | """ 233 | question = state["question"] 234 | documents = state["documents"] 235 | generation = state["candidate_answer"] 236 | web_fallback = state["web_fallback"] 237 | retries = state["retries"] if state.get("retries") is not None else -1 238 | max_retries = config.get("configurable", {}).get("max_retries", MAX_RETRIES) 239 | 240 | # this means we've already gone through web fallback and can return to the user 241 | if not web_fallback: 242 | return "finalize_response" 243 | 244 | print("---CHECK HALLUCINATIONS---") 245 | hallucination_grader = HALLUCINATION_GRADER_PROMPT | llm.with_structured_output(GradeHallucinations) 246 | hallucination_grade: GradeHallucinations = hallucination_grader.invoke( 247 | {"documents": documents, "generation": generation} 248 | ) 249 | 250 | # Check hallucination 251 | if hallucination_grade.binary_score == "no": 252 | return "generate" if retries < max_retries else "web_search" 253 | 254 | print("---DECISION: GENERATION IS GROUNDED IN DOCUMENTS---") 255 | 256 | # Check question-answering 257 | print("---GRADE GENERATION vs QUESTION---") 258 | 259 | answer_grader = ANSWER_GRADER_PROMPT | llm.with_structured_output(GradeAnswer) 260 | answer_grade: GradeAnswer = answer_grader.invoke({"question": question, "generation": generation}) 261 | if answer_grade.binary_score == "yes": 262 | print("---DECISION: GENERATION ADDRESSES QUESTION---") 263 | return "finalize_response" 264 | else: 265 | print("---DECISION: GENERATION DOES NOT ADDRESS QUESTION---") 266 | return "transform_query" if retries < max_retries else "web_search" 267 | 268 | 269 | def finalize_response(state: GraphState): 270 | print("---FINALIZING THE RESPONSE---") 271 | return {"messages": [AIMessage(content=state["candidate_answer"])]} 272 | 273 | 274 | # Define graph 275 | 276 | workflow = StateGraph(GraphState, config_schema=GraphConfig) 277 | 278 | # Define the nodes 279 | workflow.add_node("document_search", document_search) 280 | workflow.add_node("generate", generate) 281 | workflow.add_node("transform_query", transform_query) 282 | workflow.add_node("web_search", web_search) 283 | workflow.add_node("finalize_response", finalize_response) 284 | 285 | # Build graph 286 | workflow.set_entry_point("document_search") 287 | workflow.add_edge("document_search", "generate") 288 | workflow.add_edge("transform_query", "document_search") 289 | workflow.add_edge("web_search", "generate") 290 | workflow.add_edge("finalize_response", END) 291 | 292 | workflow.add_conditional_edges( 293 | "generate", 294 | grade_generation_v_documents_and_question 295 | ) 296 | 297 | # Compile 298 | graph = workflow.compile() -------------------------------------------------------------------------------- /pyproject.toml: -------------------------------------------------------------------------------- 1 | [tool.poetry] 2 | name = "pandas-rag-langgraph" 3 | version = "0.1.0" 4 | description = "" 5 | authors = ["vbarda "] 6 | readme = "README.md" 7 | 8 | 9 | [tool.poetry.dependencies] 10 | python = "^3.11" 11 | beautifulsoup4 = "^4.12.3" 12 | chromadb = "^0.5.0" 13 | tiktoken = "^0.7.0" 14 | langchain-core = "^0.2.5" 15 | langchain-community = "^0.2.4" 16 | langchain-text-splitters = "^0.2.1" 17 | langgraph = "^0.1.0" 18 | langchain-anthropic = "^0.1.15" 19 | langchain-openai = "^0.1.9" 20 | langchainhub = "^0.1.20" 21 | 22 | 23 | [tool.poetry.group.dev.dependencies] 24 | langgraph-cli = "^0.1.35" 25 | langgraph-sdk = "^0.1.18" 26 | 27 | [build-system] 28 | requires = ["poetry-core"] 29 | build-backend = "poetry.core.masonry.api" 30 | -------------------------------------------------------------------------------- /static/pandas-rag-agent.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/vbarda/pandas-rag-langgraph/553e98b467b02e3a5da4ef84e98c4c3f12d13c8a/static/pandas-rag-agent.png --------------------------------------------------------------------------------