Example: OpenAI Assistant with Tavily Tool Calling

The OpenAI Assistants API allows you to build AI agents that maintain state (conversations) and can independently decide to call external functions ("Tools").

Tavily is designed to be the ultimate search tool for these agents. In this example, we define Tavily Search as a tool, instruct the Assistant on how to use it, and implement a CLI loop that handles the asynchronous tool execution lifecycle.

The Implementation

import os
import json
import time
from openai import OpenAI
from tavily import TavilyClient

# 1. Initialize Clients
client = OpenAI(api_key=os.environ["OPENAI_API_KEY"])
tavily_client = TavilyClient(api_key=os.environ["TAVILY_API_KEY"])

# 2. Define the Agent's Persona
assistant_prompt_instruction = """You are a finance expert. 
Your goal is to provide answers based on information from the internet. 
You must use the provided Tavily search API function to find relevant online information. 
You should never use your own knowledge to answer questions.
Please include relevant url sources in the end of your answers.
"""

# 3. Define the Python wrapper function
def tavily_search(query):
    # We use get_search_context to get a flat, token-limited string of facts
    search_result = tavily_client.get_search_context(
        query, 
        search_depth="advanced", 
        max_tokens=8000
    )
    return search_result

# 4. Create the Assistant and Register the Tool Schema
assistant = client.beta.assistants.create(
    instructions=assistant_prompt_instruction,
    model="gpt-4-1106-preview",
    tools=[{
        "type": "function",
        "function": {
            "name": "tavily_search",
            "description": "Get information on recent events from the web.",
            "parameters": {
                "type": "object",
                "properties": {
                    # The LLM will generate this 'query' parameter autonomously
                    "query": {"type": "string", "description": "The search query to use. For example: 'Latest news on Nvidia stock performance'"},
                },
                "required": ["query"]
            }
        }
    }]
)

# A Thread represents an ongoing conversation
thread = client.beta.threads.create()
print("Assistant initialized. Type 'exit' to quit.")

# 5. The Conversation Loop
while True:
    user_input = input("\nYou: ")
    if user_input.lower() == 'exit':
        break

    # Add the user's message to the thread
    client.beta.threads.messages.create(
        thread_id=thread.id,
        role="user",
        content=user_input,
    )

    # Start a Run (tell the Assistant to process the thread)
    run = client.beta.threads.runs.create(
        thread_id=thread.id,
        assistant_id=assistant.id,
    )

    # 6. Polling Loop to wait for the Assistant's decision
    while True:
        time.sleep(1)
        # Retrieve the current status of the run
        run = client.beta.threads.runs.retrieve(thread_id=thread.id, run_id=run.id)

        # If the Assistant decides it needs external data, it pauses and 'requires_action'
        if run.status == 'requires_action':
            print("[Agent is searching the web...]")

            tool_calls = run.required_action.submit_tool_outputs.tool_calls
            tool_outputs = []

            # Iterate over the tools the LLM requested to call
            for tool in tool_calls:
                if tool.function.name == "tavily_search":
                    # Parse the arguments the LLM generated
                    args = json.loads(tool.function.arguments)

                    # Execute our local Python function
                    output = tavily_search(query=args["query"])

                    # Append the result to send back to the LLM
                    tool_outputs.append({
                        "tool_call_id": tool.id, 
                        "output": output
                    })

            # Submit the search results back to unpause the LLM run
            client.beta.threads.runs.submit_tool_outputs(
                thread_id=thread.id,
                run_id=run.id,
                tool_outputs=tool_outputs
            )

        elif run.status == 'completed':
            # The LLM finished generating its final response
            break
        elif run.status == 'failed':
            print(f"Run failed: {run.last_error}")
            break

    # 7. Print the final message from the Assistant
    messages = client.beta.threads.messages.list(thread_id=thread.id)
    # The first message in the list is the most recent one
    print(f"\nAssistant: {messages.data[0].content[0].text.value}")

Under the Hood

When you ask this script "Why did Tesla stock drop yesterday?", the following happens:

1. The gpt-4 model realizes it doesn't know the answer.
2. Because we provided the JSON schema for tavily_search, the model pauses execution (requires_action) and outputs a JSON blob requesting a search for "Tesla stock news yesterday".
3. Our Python while loop catches this status, parses the JSON, and fires the tavily_client.get_search_context() method.
4. Tavily returns a dense string of recent, factual news snippets.
5. We submit this string back to OpenAI.
6. The gpt-4 model resumes execution, reads the Tavily context, synthesizes a detailed response, appends the URLs as instructed by its prompt, and marks the run as completed.