OpenAI Assistants API: Retrieval Tool

For general information on how to use the Assistants API, check out my guide on getting set up here: https://www.jimmymills.io/blog/openai-assistants-api

In this post, we'll delve into the Retrieval tool, a powerful feature of the OpenAI Assistants API that allows you to reference various files as context for your runs. The Retrieval tool enables your assistant to access information from specific files, providing valuable context for generating responses.

Setting Up Retrieval

To use the Retrieval tool, you need to enable it on your assistant and provide one or more files for reference. These files can be attached to either the assistant itself, allowing them to be used in every run, or to a specific message, making them available only within that thread.

Use Cases

Some good use cases for the Retrieval tool would be:

Referencing internal content or new information not part of the OpenAI training set
Answering questions about long content to save time reading
Incorporating information to control LLM output for sensitive subjects

Example: Text File Assistant

Here's an example of an assistant that can answer questions about a PDF file:

from openai import OpenAI
client = OpenAI()

# Make sure your file is opened with "rb" for compatibility
# since the files API expects bytes
file = client.files.create(file=open("info.pdf", "rb"), purpose="assistants")

# Create the assistant with the Retrieval tool and attached files
assistant = client.beta.assistants.create(
  name="Page Content Assistant",
  instructions="You are a helpful assistant. When asked a question, you will answer using the contents of your knowledge base.",
  model="gpt-4-1106-preview",
  tools=[{"type": "retrieval"}],
  file_ids=[file.id]
)

# Create a new thread
thread = client.beta.threads.create()

# Place your first message into your thread
client.beta.threads.messages.create(
  thread_id=thread.id,
  role="user",
  content="Summarize the PDF",
)

# Create a run
run = client.beta.threads.runs.create(
  thread_id=thread.id,
  assistant_id=assistant.id,
)

# Wait for your run to finish
while run.status != 'completed':
  time.sleep(5) # Sleep for 5 seconds between checks to limit calls
  run = client.beta.threads.retrieve(
    thread_id=thread.id,
    run_id=run.id,
  )

# Print out the response
messages = client.beta.threads.messages.list(thread_id=thread.id)
print(messages[0].content[0].text.value)

Example: Page Content Assistant

Let's create an example assistant named "Page Content Assistant" that can answer questions using the contents of a web page for a specific thread:

from io import BytesIO
import requests
from openai import OpenAI
client = OpenAI()

# Pull the URL and create a file-like object
url = "https://www.jimmymills.io/blog/openai-assistants-api/retrieval.mdx"
response = requests.get(url)
data = BytesIO(response.text.encode("utf-8"))
data.name = url + '.html'  # Give it a name for identification in the dashboard
file = client.files.create(file=data, purpose="assistants")

# Create the assistant with the Retrieval tool, files will be attached in the message below
assistant = client.beta.assistants.create(
  name="Page Content Assistant",
  instructions="You are a helpful assistant. When asked a question, you will answer using the contents of your knowledge base.",
  model="gpt-4-1106-preview",
  tools=[{"type": "retrieval"}],
)

# Create a new thread
thread = client.beta.threads.create()

# Place your first message into your thread
client.beta.threads.messages.create(
  thread_id=thread.id,
  role="user",
  content="What does the website say is a good use case for the retrieval tool?",
  file_ids=[file.id] # Here is where you attach the file for a specific thread
)

# Create a run
run = client.beta.threads.runs.create(
  thread_id=thread.id,
  assistant_id=assistant.id,
)

# Wait for your run to finish
while run.status != 'completed':
  time.sleep(5) # Sleep for 5 seconds between checks to limit calls
  run = client.beta.threads.retrieve(
    thread_id=thread.id,
    run_id=run.id,
  )

# Print out the response
messages = client.beta.threads.messages.list(thread_id=thread.id)
print(messages[0].content[0].text.value)

Conclusion

The Retrieval tool enhances the capabilities of the OpenAI Assistants API by allowing you to incorporate external information seamlessly. Experiment with this feature to provide richer context to your assistants, enabling them to generate more accurate and context-aware responses.