Skip to main content

Retrievals

In TaskingAI, you can use retrieval to get information from external sources and use it in your assistant's responses. This guide will walk you through the steps to use retrieval in your assistant.

Create a Collection​

First, you need to create a collection to store the retrieved data. You can create a collection and insert useful records by using the create_collection and create_record method.

import taskingai
from taskingai.retrieval.collection import *

collection = taskingai.retrieval.collection.create_collection(
    embedding_model_id="YOUR_EMBEDDING_MODEL_ID",
    capacity=1000
)

taskingai.retrieval.create_record(
    collection_id=collection.collection_id,
    type="text",
    content="Machine learning is ...",
    text_splitter={"type": "token", "chunk_size": 200, "chunk_overlap": 10}
)

# create more records ...

Learn more in the documentation of the retrieval module.

Create an Assistant with a Retrieval Collection​

You can create an assistant with the collection we just created by using the create_assistant method.

from taskingai.assistant.assistant import AssistantRetrieval, AssistantRetrievalType
from taskingai.assistant.memory import AssistantZeroMemory


assistant = taskingai.assistant.create_assistant(
    # choose an available chat_completion model from your project
    model_id="YOUR_MODEL_ID",
    name="My Assistant",
    retrievals=[AssistantRetrieval(
        type=AssistantRetrievalType.COLLECTION,
        id=collection.collection_id,
    )],
    memory=AssistantZeroMemory()
)

Apply Retrieval Augment Generation​

Now you can test your assistant to generate a response.

chat = taskingai.assistant.create_chat(
    assistant_id=assistant.assistant_id,
)

taskingai.assistant.create_message(
    assistant_id=assistant.assistant_id,
    chat_id=chat.chat_id,
    text="What is machine learning?",
)
    
# generate assistant response
assistant_message = taskingai.assistant.generate_message(
    assistant_id=assistant.assistant_id,
    chat_id=chat.chat_id,
)

print(f"Assistant: {assistant_message.content.text}")

The assistant will first retrieve useful text chunks from the collection and generate a response based on them to increase the accuracy of the response.