Skip to main content

Integrate Retrievals to Assistant

Integrate to Assistant when creating an Assistant​

You can create an assistant with the desired collection by using the create_assistant method.

When Integrating retrieval to an assistant, you need to specify your RetrievalConfig. The RetrievalConfig contains the following parameters:

  1. Method: When retrieval system is fetching related chunks from vector database, it needs a query as input. This field specifies how the query is generated.
    • User Message: The latest user message is used as query.
    • Function Call: The query is generated by the model, and passed to the retrieval system as a function call.
    • Memory: The query is the latest user message together with the context in the chat session.
  2. Top K: The number of chunks to be returned from the retrieval task.
  3. Max Tokens: The maximum number of tokens to be returned from the retrieval task.
  4. Score Threshold: The minimum score required for a chunk to be considered as a relevant. The score is a value between 0 and 1, where 1 is the highest score.
from taskingai.assistant import RetrievalConfig, RetrievalMethod

assistant = taskingai.assistant.create_assistant(
    # choose an available chat_completion model from your project
    model_id="YOUR_MODEL_ID",
    name="My Assistant",
    retrievals=[{"type": "collection", "id": "YOUR_COLLECTION_ID"}],
    memory={"type": "naive"},
    retrieval_configs=RetrievalConfig(top_k=3, max_tokens=4096, method=RetrievalMethod.USER_MESSAGE, score_threshold=0.5)
)

Integrate to existing Assistant​

Similarly, we can also integrate retrievals to existing assistant by using the update_assistant method.

import taskingai
from taskingai.assistant import RetrievalConfig, RetrievalMethod

assistant = taskingai.assistant.update_assistant(
    assistant_id="YOUR_ASSISTANT_ID",
    name="My Updated Assistant",
    description="This is my new assistant",
    retrievals=[{"type": "collection", "id": "YOUR_COLLECTION_ID"}],
    retrieval_configs=RetrievalConfig(top_k=4, max_tokens=8192, method=RetrievalMethod.USER_MESSAGE, score_threshold=0.6)
 
)