> ## Documentation Index
> Fetch the complete documentation index at: https://genai.codewithsiva.dev/llms.txt
> Use this file to discover all available pages before exploring further.

# Multi Agent

> Build a Retrieval-Augmented Generation agent using Chroma and LangGraph

In this section, we will see how to build a RAG (Retrieval-Augmented Generation) agent that can load a PDF document, store embeddings in a vector database, and retrieve info using custom tools.

## Objectives

1. Split document pages and create a Chroma database.
2. Wrap database query retrieval in a custom `@tool`.
3. Set up a StateGraph that invokes LLM reasoning and retrieval tools sequentially.

## Agent IV: RAG Agent

#### Goal

Build a RAG agent querying details from the `Stock_Market_Performance_2024.pdf` document.

#### Sample Input

```python theme={null}
{"messages": [("user", "How did the tech sector perform in 2024?")]}
```

#### Sample Output

Outputs summarizing tech sector performance: `In 2024, the tech sector gained 25% driven by AI breakthroughs.` (cited from the document).

#### Plan

1. Load and chunk the PDF document (`Stock_Market_Performance_2024.pdf`).
2. Embed chunks using `OpenAIEmbeddings` and store them in a local Chroma vector database.
3. Define a retriever tool function (`retriever_tool`) and bind it to the LLM.
4. Set up a StateGraph containing an LLM call node, a retriever action node, and conditional routing edges.

### Code Implementation

#### 1. Load and Chunk PDF

We split document text to feed to embeddings model:

```python theme={null}
import os
from langchain_community.document_loaders import PyPDFLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter

pdf_path = "Stock_Market_Performance_2024.pdf"
loader = PyPDFLoader(pdf_path)
pages = loader.load()

text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
pages_split = text_splitter.split_documents(pages)
```

#### 2. Setup ChromaDB Vector Store

We embed the split documents and store them in a Chroma DB:

```python theme={null}
from langchain_openai import OpenAIEmbeddings
from langchain_chroma import Chroma

embeddings = OpenAIEmbeddings(model="text-embedding-3-small")
vectorstore = Chroma.from_documents(
    documents=pages_split,
    embedding=embeddings,
    collection_name="stock_market"
)
retriever = vectorstore.as_retriever(search_type="similarity", search_kwargs={"k": 5})
```

#### 3. Define Retriever Tool

We wrap retrieval search in a custom tool and bind it to the model:

```python theme={null}
from langchain_core.tools import tool
from langchain_openai import ChatOpenAI

@tool
def retriever_tool(query: str) -> str:
    """This tool searches and returns the information from the Stock Market Performance 2024 document."""
    docs = retriever.invoke(query)
    results = [f"Document {i+1}:\n{doc.page_content}" for i, doc in enumerate(docs)]
    return "

".join(results)

tools = [retriever_tool]
tools_dict = {t.name: t for t in tools}
llm = ChatOpenAI(model="gpt-4o", temperature=0).bind_tools(tools)
```

#### 4. Define Nodes and Build Graph

We define the graph node calls and setup the graph structure:

```python theme={null}
from typing import Annotated, Sequence, TypedDict
from langchain_core.messages import BaseMessage, SystemMessage, ToolMessage
from langgraph.graph.message import add_messages
from langgraph.graph import StateGraph, END

class AgentState(TypedDict):
    messages: Annotated[Sequence[BaseMessage], add_messages]

def call_llm(state: AgentState) -> dict:
    sys = SystemMessage(content="You are an AI assistant answering questions about Stock Market Performance.")
    message = llm.invoke([sys] + list(state["messages"]))
    return {"messages": [message]}

def take_action(state: AgentState) -> dict:
    tool_calls = state["messages"][-1].tool_calls
    results = []
    for t in tool_calls:
        res = tools_dict[t["name"]].invoke(t["args"].get("query", ""))
        results.append(ToolMessage(tool_call_id=t["id"], name=t["name"], content=str(res)))
    return {"messages": results}

def should_continue(state: AgentState):
    last_msg = state["messages"][-1]
    if hasattr(last_msg, "tool_calls") and last_msg.tool_calls:
        return "retriever_agent"
    return "end"

graph = StateGraph(AgentState)
graph.add_node("llm", call_llm)
graph.add_node("retriever_agent", take_action)

graph.set_entry_point("llm")
graph.add_conditional_edges(
    "llm",
    should_continue,
    {
        "retriever_agent": "retriever_agent",
        "end": END
    }
)
graph.add_edge("retriever_agent", "llm")
app = graph.compile()
```

#### 5. Invoke the Agent

```python theme={null}
result = app.invoke({"messages": [HumanMessage(content="How did the tech sector perform in 2024?")]})
print(result["messages"][-1].content)
```

## Exercise: Multi-Document RAG with Document Type Router 📂

#### Goal

Add a second Corporate Email records search tool to the RAG Agent, allowing the model to choose the correct resource (Stock Performance DB or Email database) dynamically based on user query.

#### Sample Input

```python theme={null}
{"messages": [("user", "Check the meeting email: why was the user unable to attend?")]}
```

#### Sample Output

Outputs citing the email document (e.g., `"The user was unable to attend due to a conflicting doctor's appointment."`).

#### Plan

1. Create a retriever tool `@tool` named `email_retriever` returning content from corporate email records.
2. Bind both `[retriever_tool, email_retriever]` to the LLM.
3. Configure `take_action` to map and invoke both tools correctly.
4. Compile the graph, invoke with the sample query, and print the results.

<Accordion title="Solution">
  ```python theme={null}
  from typing import Annotated, Sequence, TypedDict
  from langchain_core.messages import BaseMessage, SystemMessage, ToolMessage
  from langgraph.graph.message import add_messages
  from langchain_openai import ChatOpenAI
  from langchain_core.tools import tool
  from langgraph.graph import StateGraph, START, END

  # 1. State
  class AgentState(TypedDict):
      messages: Annotated[Sequence[BaseMessage], add_messages]

  # 2. Tools
  @tool
  def retriever_tool(query: str) -> str:
      """This tool searches the Stock Market Performance 2024 database."""
      return "In 2024, the tech sector gained 25% driven by AI breakthroughs."

  @tool
  def email_retriever(query: str) -> str:
      """This tool searches corporate email records."""
      return "Subject: Unable to Attend Meeting\n\nHi, I cannot make it to the meeting today because I have a doctor's appointment scheduled at that time."

  tools = [retriever_tool, email_retriever]
  tools_dict = {t.name: t for t in tools}
  llm = ChatOpenAI(model="gpt-4o", temperature=0).bind_tools(tools)

  # 3. Nodes
  def call_llm(state: AgentState) -> dict:
      sys = SystemMessage(content="Use retriever_tool for stock performance, and email_retriever for email records.")
      msg = llm.invoke([sys] + list(state["messages"]))
      return {"messages": [msg]}

  def take_action(state: AgentState) -> dict:
      tool_calls = state["messages"][-1].tool_calls
      results = []
      for t in tool_calls:
          res_content = tools_dict[t["name"]].invoke(t["args"].get("query", ""))
          results.append(ToolMessage(tool_call_id=t["id"], name=t["name"], content=str(res_content)))
      return {"messages": results}

  def should_continue(state: AgentState):
      last_msg = state["messages"][-1]
      return "retriever_agent" if (hasattr(last_msg, "tool_calls") and last_msg.tool_calls) else "end"

  # 4. Build
  graph = StateGraph(AgentState)
  graph.add_node("llm", call_llm)
  graph.add_node("retriever_agent", take_action)
  graph.set_entry_point("llm")
  graph.add_conditional_edges(
      "llm",
      should_continue,
      {
          "retriever_agent": "retriever_agent",
          "end": END
      }
  )
  graph.add_edge("retriever_agent", "llm")
  app = graph.compile()

  # 5. Test
  res = app.invoke({"messages": [("user", "Check the meeting email: why was the user unable to attend?")]})
  print(res["messages"][-1].content)
  ```
</Accordion>

## Practice & Exercises

To reinforce what you've learned in this section, practice with the interactive notebook:

<CardGroup cols={1}>
  <Card title="Practice & Exercises" icon="laptop-code">
    Practice loading source documents, chunking text, setting up a local vector database retriever, and routing questions dynamically.

    [💻 VS Code](vscode://file/Users/sivaprasad/Downloads/GenAI%20With%20Python/public/notebooks/langgraph/Multi_Agent.ipynb) | [🚀 Colab](https://colab.research.google.com/github/prasad230776/genai-course/blob/master/public/notebooks/langgraph/Multi_Agent.ipynb) | <a href="/public/notebooks/langgraph/Multi_Agent.ipynb" download>📥 Download</a>
  </Card>
</CardGroup>