RAG powered tool calling

Introduction

RAG (Retrieval Augmented Generation) Powered Tool Calling is an advanced approach where Large Language Models (LLMs) leverage external knowledge retrieved through RAG to make more informed and effective decisions about invoking tools (e.g., APIs, functions). Instead of relying solely on its pre-trained knowledge or a fixed set of tool descriptions, the LLM first consults a dynamic knowledge base to understand the context, available tools, or specific parameters required, leading to more accurate and contextually appropriate tool usage.

Why Combine RAG with Tool Calling?

Combining RAG with tool calling offers several significant advantages:

Enhanced Contextual Understanding: RAG provides the LLM with relevant, up-to-date information from external knowledge sources. This allows the LLM to better understand the user’s query in the context of available data, leading to more appropriate tool selection.
Dynamic Tool Discovery & Parameterization: Instead of hardcoding tool descriptions or relying on the LLM’s potentially outdated training data, RAG can retrieve information about available tools, their functionalities, and required parameters. This is especially useful in environments where tools change frequently or are numerous.
Reduced Hallucination in Tool Usage: LLMs can sometimes hallucinate tool names or parameters. By grounding the LLM with factual information retrieved by RAG, the likelihood of such errors is reduced.
Improved Accuracy for Complex Queries: For queries that require specific data to be passed to a tool (e.g., an ID, a specific format), RAG can fetch this data, ensuring the tool is called correctly.
Handling Ambiguity: If a user’s request is ambiguous regarding which tool to use or what parameters are needed, RAG can retrieve clarifying information or options, helping the LLM (or even the user through the LLM) to disambiguate.
Access to Proprietary or Real-time Information for Tools: Tools often need to operate on data that is not part of the LLM’s training. RAG can fetch this proprietary or real-time data, which can then be used to inform the tool call or provide necessary inputs.

When to Use RAG Powered Tool Calling?

This approach is particularly beneficial in scenarios such as:

Complex Enterprise Systems: Where an LLM needs to interact with numerous internal APIs, databases, or services, and the specifics of these tools are documented in a knowledge base.
Dynamic Environments: When the set of available tools or their specifications (e.g., API endpoints, parameters) changes over time.
Knowledge-Intensive Tooling: For tools that require specific, factual data as input that the LLM might not inherently know (e.g., looking up a product ID before calling an inventory check tool).
Automated Customer Support: Where agents need to use tools to perform actions (e.g., check order status, process refunds) based on information retrieved from FAQs, order databases, or policy documents.
Data Analysis and Reporting: When an LLM needs to select and use data analysis tools based on the structure or type of data retrieved from a database or document store.
Personalized Assistants: Where tools need to be called based on user-specific data or preferences retrieved via RAG.

How Does it Work? (Conceptual Architecture)

A typical workflow for RAG powered tool calling might look like this:

User Query: The user submits a request.
Initial Processing & RAG Trigger: The LLM (or a preceding logic layer) analyzes the query. If it determines that external knowledge might be needed to select or use a tool, it triggers the RAG process.
Retrieval: The RAG system searches its configured knowledge base (e.g., API documentation, FAQs, database schemas) for information relevant to the user’s query and potential tool usage.
Context Augmentation: The retrieved information (context) is provided to the LLM along with the original query.
Informed Tool Decision: The LLM, now equipped with the augmented context, decides:
- Whether a tool needs to be called.
- Which specific tool to call (if multiple are candidates).
- What parameters to use for the tool, potentially extracting or inferring them from the query and the RAG context.
Tool Execution: The selected tool is executed with the determined parameters.
Response Generation: The LLM receives the output from the tool. It then synthesizes this output, the original query, and potentially the RAG context to generate a final, coherent response for the user.

This can be an iterative process, where multiple RAG lookups or tool calls might be needed for complex queries.

Python Code Example: RAG for API Call Decision

Let’s imagine a scenario where an LLM needs to call an internal company API. The company has many APIs, and their details (endpoint, purpose) are documented. RAG will help the LLM find the correct API and its endpoint based on the user’s natural language request.

import os
from typing import List, TypedDict, Annotated, Sequence
from langchain_core.documents import Document
from langchain_openai import OpenAIEmbeddings, ChatOpenAI
from langchain_community.vectorstores import Chroma
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser, JsonOutputParser
from langchain_core.tools import tool
from langchain_core.messages import BaseMessage, HumanMessage, AIMessage, ToolMessage
from langgraph.graph import StateGraph, END
from pydantic import BaseModel, Field
 
# --- 0. Environment Setup (Ensure OPENAI_API_KEY is set) ---
# os.environ["OPENAI_API_KEY"] = "your_openai_api_key"
 
# --- 1. Knowledge Base for RAG: API Documentation ---
api_docs_content = [
    Document(
        page_content="User Info API: Use this API to get details about a user. Endpoint: /api/users/{user_id}. Requires user_id.",
        metadata={"api_name": "UserInfoAPI", "source": "api_docs.txt"}
    ),
    Document(
        page_content="Order Status API: Use this API to check the status of an order. Endpoint: /api/orders/{order_id}. Requires order_id.",
        metadata={"api_name": "OrderStatusAPI", "source": "api_docs.txt"}
    ),
    Document(
        page_content="Product Search API: Use this API to search for products. Endpoint: /api/products/search?query={search_term}. Requires search_term.",
        metadata={"api_name": "ProductSearchAPI", "source": "api_docs.txt"}
    ),
    Document(
        page_content="General Help API: Use this API for general help topics or FAQs. Endpoint: /api/help?topic={topic_name}. Use for non-specific queries.",
        metadata={"api_name": "GeneralHelpAPI", "source": "api_docs.txt"}
    )
]
 
# --- 2. RAG Setup: Embeddings and Vector Store ---
embedding_function = OpenAIEmbeddings()
db = Chroma.from_documents(api_docs_content, embedding_function)
retriever = db.as_retriever(search_kwargs={"k": 2}) # Retrieve top 2 relevant docs
 
# --- 3. Define a Mock Tool: Simulate API Call ---
@tool
def call_internal_api(api_name: str, endpoint: str, method: str = "GET", params: dict = None) -> dict:
    """Simulates calling an internal API.
    Args:
        api_name: The name of the API being called (e.g., UserInfoAPI).
        endpoint: The specific endpoint to hit (e.g., /api/users/123).
        method: HTTP method (default: GET).
        params: Dictionary of parameters for the API call.
    Returns:
        A dictionary with the mock API response.
    """
    print(f"\n[Tool Call] Calling API: {api_name}, Endpoint: {endpoint}, Method: {method}, Params: {params}")
    if api_name == "OrderStatusAPI" and "/api/orders/" in endpoint:
        order_id = endpoint.split("/")[-1]
        return {"status": "success", "order_id": order_id, "details": f"Order {order_id} is shipped."}
    elif api_name == "UserInfoAPI" and "/api/users/" in endpoint:
        user_id = endpoint.split("/")[-1]
        return {"status": "success", "user_id": user_id, "name": f"User {user_id} Name", "email": f"user{user_id}@example.com"}
    elif api_name == "ProductSearchAPI" and params and "query" in params:
        return {"status": "success", "query": params["query"], "results": [f"Product A for {params['query']}", f"Product B for {params['query']}"]}
    return {"status": "error", "message": "Mock API endpoint not recognized or missing params."}
 
tools = [call_internal_api]
 
# --- 4. LLM and Agent State Definition ---
llm = ChatOpenAI(model="gpt-4o")
llm_with_tools = llm.bind_tools(tools)
 
class AgentState(TypedDict):
    messages: Sequence[BaseMessage]
    retrieved_docs: List[Document] # To store documents from RAG
 
# --- 5. Graph Nodes ---
 
# Node 1: Retrieve relevant API documentation using RAG
def retrieve_api_docs_node(state: AgentState):
    print("---NODE: RETRIEVE API DOCS---")
    last_message = state["messages"][-1]
    query = last_message.content
    docs = retriever.invoke(query)
    print(f"Retrieved {len(docs)} documents for query: '{query}'")
    # for doc in docs: print(f"  - {doc.page_content}") # Optional: print retrieved docs
    return {"retrieved_docs": docs, "messages": state["messages"]}
 
# Node 2: Decide whether to call a tool, and which one, using RAG context
def decide_tool_or_respond_node(state: AgentState):
    print("---NODE: DECIDE TOOL OR RESPOND---")
    messages = list(state["messages"])
    retrieved_docs = state["retrieved_docs"]
 
    # Augment the prompt with RAG context
    context_str = "\n\n".join([doc.page_content for doc in retrieved_docs])
    augmented_prompt = (
        f"Based on the user query and the following API documentation, decide if a tool should be called. "
        f"If so, prepare the exact tool call. If not, or if more information is needed, respond directly.\n\n"
        f"API Documentation Context:\n{context_str}\n\n"
        f"User Query: {messages[-1].content}"
    )
    
    print(f"Augmented prompt for tool decision: {augmented_prompt[:300]}...") # Print snippet
    
    # Add this augmented information as a system message for the LLM's tool decision process
    # We create a new list of messages for this specific LLM call to avoid polluting the main chat history yet
    decision_messages = [
        HumanMessage(content=augmented_prompt)
    ]
 
    ai_response = llm_with_tools.invoke(decision_messages)
    
    if not ai_response.tool_calls:
        print("LLM decided to respond directly or needs clarification.")
        # If no tool call, we might want to add the AI's reasoning or direct response to messages
        # For simplicity, we'll just pass the original messages and let a later node handle final response if no tool.
        # Or, we can add its direct response here.
        # For this example, if no tool call, we assume it's a direct answer based on RAG context.
        final_response_prompt = (
            f"Based on the user query and the following API documentation, provide a helpful answer.\n\n"
            f"API Documentation Context:\n{context_str}\n\n"
            f"User Query: {messages[-1].content}"
        )
        direct_answer = llm.invoke([HumanMessage(content=final_response_prompt)])
        messages.append(direct_answer) # Add direct answer to messages
    else:
        print(f"LLM decided to call tool(s): {ai_response.tool_calls}")
        messages.append(ai_response) # Add AI message with tool_calls to history
 
    return {"messages": messages, "retrieved_docs": retrieved_docs} 
 
# Node 3: Execute the tool
def execute_tool_node(state: AgentState):
    print("---NODE: EXECUTE TOOL---")
    last_message = state["messages"][-1]
    if not isinstance(last_message, AIMessage) or not last_message.tool_calls:
        # Should not happen if logic is correct, but good to have a fallback.
        print("No tool call found in the last AI message.")
        return state
 
    tool_call = last_message.tool_calls[0] # Assuming one tool call for simplicity
    tool_name = tool_call["name"]
    tool_args = tool_call["args"]
 
    # Find the actual tool function from the `tools` list
    selected_tool_func = next((t for t in tools if t.name == tool_name), None)
    if not selected_tool_func:
        error_message = f"Tool '{tool_name}' not found."
        print(error_message)
        tool_response_message = ToolMessage(content=error_message, tool_call_id=tool_call["id"])
    else:
        try:
            tool_output = selected_tool_func.invoke(tool_args)
            tool_response_message = ToolMessage(content=str(tool_output), tool_call_id=tool_call["id"])
        except Exception as e:
            print(f"Error executing tool {tool_name}: {e}")
            tool_response_message = ToolMessage(content=f"Error: {e}", tool_call_id=tool_call["id"])
    
    updated_messages = list(state["messages"]) + [tool_response_message]
    return {"messages": updated_messages, "retrieved_docs": state["retrieved_docs"]}
 
# Node 4: Generate final response after tool call (or if no tool was called initially)
def generate_final_response_node(state: AgentState):
    print("---NODE: GENERATE FINAL RESPONSE---")
    # The context for the final response includes the original query, RAG docs, and any tool interactions.
    # For simplicity, we'll just use the current message history which includes tool calls and responses.
    response = llm.invoke(state["messages"]) 
    updated_messages = list(state["messages"]) + [response]
    return {"messages": updated_messages, "retrieved_docs": state["retrieved_docs"]}
 
# --- 6. Define Graph Edges and Conditional Logic ---
def should_call_tool_edge(state: AgentState) -> str:
    print("---EDGE: SHOULD CALL TOOL?---")
    last_message = state["messages"][-1]
    if isinstance(last_message, AIMessage) and last_message.tool_calls:
        print("Decision: Call tool.")
        return "call_tool"
    print("Decision: No tool call / Respond directly.")
    return "respond_directly"
 
# --- 7. Construct the Graph ---
workflow = StateGraph(AgentState)
 
workflow.add_node("retrieve_docs", retrieve_api_docs_node)
workflow.add_node("decide_tool_or_respond", decide_tool_or_respond_node)
workflow.add_node("execute_tool", execute_tool_node)
workflow.add_node("generate_final_response", generate_final_response_node)
 
workflow.set_entry_point("retrieve_docs")
workflow.add_edge("retrieve_docs", "decide_tool_or_respond")
 
workflow.add_conditional_edges(
    "decide_tool_or_respond",
    should_call_tool_edge,
    {
        "call_tool": "execute_tool",
        "respond_directly": "generate_final_response" # If no tool, generate response based on RAG context
    }
)
 
workflow.add_edge("execute_tool", "generate_final_response") # After tool execution, generate final response
workflow.add_edge("generate_final_response", END)
 
app = workflow.compile()
 
# --- 8. Run the Agent ---
if __name__ == "__main__":
    print("\n--- Running RAG Powered Tool Calling Agent ---")
    
    queries = [
        "What is the status of my order 7890?",
        "Can you tell me about user 456?",
        "I need help with billing.", # Should use GeneralHelpAPI or respond directly
        "Search for blue running shoes."
    ]
 
    for query_text in queries:
        print(f"\n\n--- USER QUERY: {query_text} ---")
        inputs = {"messages": [HumanMessage(content=query_text)]}
        final_state = app.invoke(inputs)
        print("\n--- FINAL RESPONSE ---")
        # The final AI response is usually the last message in the list
        final_ai_message = final_state["messages"][-1]
        if isinstance(final_ai_message, AIMessage):
            print(f"AI: {final_ai_message.content}")
        else:
            print(f"Final state: {final_state['messages'][-1]}") # Print last if not AIMessage
        print("------------------------")

Explanation of the Code

Knowledge Base (API Docs): We define a list of Document objects, each describing an API, its endpoint, and purpose. This serves as our external knowledge for RAG.
RAG Setup: OpenAIEmbeddings and Chroma vector store are used to create a retriever for these API documents.
Mock Tool (call_internal_api): A simple Python function decorated with @tool simulates making an API call. It returns mock data based on the api_name and endpoint.
Agent State (AgentState): A TypedDict that holds the conversation messages and retrieved_docs (from RAG) as the graph progresses.
Graph Nodes:
- retrieve_api_docs_node: Takes the user query, retrieves relevant API documents using the RAG retriever, and adds them to the state.
- decide_tool_or_respond_node: This is the core RAG-powered decision node. It constructs an augmented prompt containing the original user query AND the retrieved API documentation. It then asks the LLM (bound with tools) to decide if a tool call is necessary based on this augmented context. If the LLM decides to call a tool, its AIMessage with tool_calls is added to the history. If not, it attempts to generate a direct response based on the RAG context.
- execute_tool_node: If the previous node decided on a tool call, this node extracts the tool name and arguments, executes the mock tool, and adds the ToolMessage (with the tool’s output) to the message history.
- generate_final_response_node: After any tool interaction (or if no tool was called), this node generates the final response to the user, considering the entire conversation history (which includes RAG insights via the augmented prompt and tool results).
Conditional Edge (should_call_tool_edge): This function checks the last message in the state. If it’s an AIMessage containing tool_calls, it directs the graph to execute_tool_node. Otherwise, it directs to generate_final_response_node (implying a direct answer is needed).
Graph Construction: The nodes and edges are assembled into a StateGraph.
Running the Agent: The if __name__ == "__main__": block shows how to invoke the compiled graph with different user queries. The agent will first retrieve API docs, then use that context to decide on and execute tool calls, and finally generate a response.

Benefits of this Approach

Accuracy: The LLM makes tool-related decisions based on specific, retrieved documentation rather than just its general knowledge, leading to more accurate tool selection and parameterization.
Flexibility: Easily update or expand the available tools by simply updating the API documentation (the RAG knowledge base) without retraining the LLM or extensively modifying the agent’s core logic.
Clarity: The process is more transparent as the retrieved context used for decision-making can be inspected.
Reduced Errors: Less chance of the LLM attempting to call non-existent tools or using incorrect parameters if the documentation is accurate.

Considerations

Quality of RAG Source: The effectiveness heavily depends on the quality, accuracy, and comprehensiveness of the documentation provided to the RAG system.
Retrieval Relevance: The retriever must be effective in finding the most relevant documents for the LLM’s decision-making process.
Prompt Engineering: Crafting the augmented prompt for the decide_tool_or_respond_node is crucial. It needs to clearly instruct the LLM on how to use the RAG context for tool decisions.
Complexity: This architecture is more complex than simple tool calling, involving an additional retrieval step.
Latency: The RAG step adds latency to the response time.
Cost: More LLM calls might be involved (e.g., one for RAG-augmented decision, another for final response generation).

By carefully designing the RAG source, retrieval strategy, and agent logic, RAG powered tool calling can significantly enhance the capabilities and reliability of LLM-based agents.

Kaushik's Notes

Explorer