Integrate Codex as-a-tool with smolagents

Run in Google Colab

This notebook demonstrates how to integrate Codex as a tool within an existing AI Agent using the smolagents library.

For simplicity, this example shows how to set up a basic Agent with vector store search, then showcases how adding Codex as a tool enables the Agent to handle questions beyond its knowledge base.

This tutorial presumes you have already set up an AI Agent that does RAG with smolagents. If unsure how to set up an AI Agent that does RAG with smolagents, follow our tutorial: Adding Tool Calls to RAG.

The code provided in this notebook is for an Agentic RAG application, with single-turn conversations.

If you prefer to integrate Codex without adding tool calls to your Agent, check out our other integrations for an overview.

Let’s first install packages required for this tutorial. Most of these packages are the same as those used in smolagents’ RAG tutorial.

%pip install smolagents pandas langchain langchain-community rank_bm25 --upgrade -q

%pip install litellm # Optional dependency of smolagents to use LiteLLM as a gateway to many LLMs,

%pip install cleanlab-codex

Optional: Helper methods for a smolagents agent from prior tutorial (Agentic RAG with smolagents)

from datetime import datetime
from langchain.docstore.document import Document
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_community.retrievers import BM25Retriever
from smolagents import Tool, tool

source_docs = [
    Document(
        page_content="""Simple Water Bottle - Amber (limited edition launched Jan 1st 2025)
        
        A water bottle designed with a perfect blend of functionality and aesthetics in mind. 
        Crafted from high-quality, durable plastic with a sleek honey-colored finish.

        Price: $24.99
        Dimensions: 10 inches height x 4 inches width
        """,
        metadata={"source": "bottle.txt"}
    ),
]

text_splitter = RecursiveCharacterTextSplitter(
    chunk_size=500,
    chunk_overlap=50,
    add_start_index=True,
    strip_whitespace=True,
    separators=["\n\n", "\n", ".", " ", ""],
)
docs_processed = text_splitter.split_documents(source_docs)

fallback_answer = "Based on the available information, I cannot provide a complete answer to this question."

# First, we'll define a retriever tool similar to the one in smolagents' RAG tutorial
# (https://huggingface.co/docs/smolagents/examples/rag). The key difference is that we've 
# enhanced the tool description to give the agent clear instructions about when and how to 
# use this search capability. This helps the agent make better decisions about tool usage 
# and handle cases where information isn't found.
class RetrieverTool(Tool):
    name = "retriever"
    description = f"""Uses search to retrieve relevant parts of a knowledge base to answer a query.
    Start with the original question for search. If no relevant information is found, prefer alternate tools or state "{fallback_answer}".
    Avoid making assumptions.
    """

    inputs = {
        "query": {
            "type": "string",
            "description": "The question to answer. Always use the original question, never rephrase it.",
        }
    }
    output_type = "string"

    def __init__(self, docs, **kwargs):
        super().__init__(**kwargs)
        self.retriever = BM25Retriever.from_documents(
            docs, k=10
        )

    def forward(self, query: str) -> str:
    
        assert isinstance(query, str), "Your search query must be a string"

        docs = self.retriever.invoke(
            query,
        )
        return "\nRetrieved documents:\n" + "".join(
            [
                f"\n\n===== Document {str(i)} =====\n" + doc.page_content
                for i, doc in enumerate(docs)
            ]
        )

@tool
def get_todays_date(date_format: str) -> str:
  """A tool that returns today's date in the date format requested. Use this tool when knowing today's date is necessary to answer the question, like checking if something is available now or calculating how long until a future or past date.
  
  Args:
      date_format: The date format to return today's date in. Options are: '%Y-%m-%d', '%d', '%m', '%Y'. The default is '%Y-%m-%d'.
    
  Returns:
      str: Today's date in the requested format.
  """
  datetime_str = datetime.now().strftime(date_format)
  return datetime_str

Example: Customer Service for a New Product

Let’s revisit our Agent built in the tutorial: RAG with Tool Calls in smolagents, which has the option to call a get_todays_date() tool. This example represents a customer support / e-commerce use-case where the Knowledge Base contains product listings like the following:

For simplicity, our Agent’s Knowledge Base here only contains a single document featuring this one product description.

Let’s initialize an LLM with tool-calling capabilities, as well as the retriever tool and then integrate Codex to improve the Agent’s responses.

from smolagents import LiteLLMModel
import os

os.environ["OPENAI_API_KEY"] = "<YOUR-KEY-HERE>"  # Replace with your OpenAI API key
model = LiteLLMModel("openai/gpt-4o-mini")

Create Codex Project

To use Codex, first create a Project.

Here we assume some common (question, answer) pairs about the Simple Water Bottle have already been added to a Codex Project. Learn how that was done via our tutorial: Populating Codex.

Our existing Codex Project contains the following entries:

access_key = "<YOUR-PROJECT-ACCESS-KEY>"  # Obtain from your Project's settings page: https://codex.cleanlab.ai/

Integrate Codex as an additional tool

You only need to make minimal changes to your code to include Codex as an additional tool:

Add Codex to the list of tools provided to the Agent. But before that, you should update the description of the Codex tool to include instructions for the agent on when to use it.

After that, call your Agent with these updated variables to start experiencing the benefits of Codex!

from cleanlab_codex import CodexTool
from smolagents import CodeAgent

# Instantiate the retriever tool
retriever_tool = RetrieverTool(docs_processed)

# Instantiate the Codex tool
codex_tool = CodexTool.from_access_key(access_key=access_key)

# Update the description of the Codex tool to include instructions for the agent on when to use it.
codex_tool.tool_description += "\nUse this tool once as a backup when you cannot find sufficient information to answer the question. Always use the original question as the query - do not rephrase or modify it."

# Convert the Codex tool to a Smolagents tool
codex_tool_smolagents = codex_tool.to_smolagents_tool()

agent_with_codex = CodeAgent(
    tools=[
        retriever_tool,
        codex_tool_smolagents,
        get_todays_date,  # Add other tools here
    ],
    model=model,
    verbosity_level=0,  # Suppress internal logging
)

Optional: Create another version of the Assistant without Codex (rag_without_codex)

agent_without_codex = CodeAgent(
    tools=[
        retriever_tool,
        get_todays_date,  # Add other tools here
    ],
    model=model,
    verbosity_level=0,  # Suppress internal logging
)

Note: This tutorial uses a Codex tool definition provided in smolagents format via the to_smolagents_tool() function. You can instead manually write the Codex tool definition yourself or import it in alternate provided formats.

In with agentic frameworks like smolagents, retrieval is implemented as a tool (named retrieval) that can be called by the Agent. We’re using a CodeAgent, which is more powerful than a traditional tool-calling Agent because:

1. A CodeAgent can write and execute code to orchestrate tool interactions in sophisticated ways
2. A tool-calling Agent can only make direct calls to the provided tools in a sequential manner

With this added complexity, we should avoid modifying the pre-defined system prompt unless we fully understand the implications of doing so. Instead, we recommend providing clear tool usage instructions through the tool descriptions themselves.

The smolagents framework will naturally integrate the updated description of the Codex tool into the system prompt. For smolagents, we recommend instructing the Agent to only consider the Codex tool after it has used retrieval and is still unsure how to answer.

RAG with Codex in action

Integrating Codex as-a-Tool allows your Agent to answer more questions than it was originally capable of.

Example 1

Let’s ask a question to our original Agent (before Codex was integrated).

user_question = "Can I return my simple water bottle?"
response = agent_without_codex.run(user_question)
print(response)

Based on the available information, I cannot provide a complete answer to this question.

The original Agent is unable to answer, in this case because the required information is not in its Knowledge Base.

Let’s ask the same question to our agent with Codex added as an additional tool. Note that we use the updated tool list when Codex is integrated in the Agent.

response = agent_with_codex.run(user_question)
print(response)

You can return your simple water bottle within 30 days for a full refund, no questions asked. Please contact the support team to initiate your return.

As you see, integrating Codex enables your Agent to answer questions it originally strugged with, as long as a similar question was already answered in the corresponding Codex Project.

Example 2

Let’s ask another question to our agent with Codex integrated.

user_question = "How can I order the Simple Water Bottle in bulk?"
response = agent_with_codex.run(user_question)
print(response)

Based on the available information, I cannot provide a complete answer to this question.

Our Agent is unable to answer this question because there is no relevant information in its Knowledge Base, nor has a similar question been answered in the Codex Project (see the contents of the Codex Project above).

Codex automatically recognizes this question could not be answered and logs it into the Project where it awaits an answer from a SME. Navigate to your Codex Project in the Web App where you (or a SME at your company) can enter the desired answer for this query.

As soon as an answer is provided in Codex, our Agent will be able to answer all similar questions going forward (as seen for the previous query).

Example 3

Let’s ask another query to our Agent with Codex integrated. This is a query the original Agent was able to correctly answer without Codex (since the relevant information exists in the Knowledge Base).

user_question = "How big is the water bottle?"
response = agent_with_codex.run(user_question)
print(response)

The water bottle has a height of 10 inches and a width of 4 inches.

We see that the Agent with Codex integrated is still able to correctly answer this query. Integrating Codex has no negative effect on questions your original Agent could answer.

Next Steps

Now that Codex is integrated with your Agent, you and SMEs can open the Codex Project and answer questions logged there to continuously improve your AI.

Adding Codex only improves your AI Agent. Once integrated, Codex automatically logs all user queries that your original AI Agent handles poorly. Using a simple web interface, SMEs at your company can answer the highest priority questions in the Codex Project. As soon as an answer is entered in Codex, your AI Agent will be able to properly handle all similar questions encountered in the future.

Codex is the fastest way for nontechnical SMEs to directly improve your AI application. As the Developer, you simply integrate Codex once, and from then on, SMEs can continuously improve how your AI handles common user queries without needing your help.

Need help, more capabilities, or other deployment options? Check the FAQ or email us at: support@cleanlab.ai