Skip to content

Distiller API

This documentation provides an overview of our Distiller API. This API allows you to leverage our agentic framework to create intelligent, autonomous agents capable of complex reasoning and decision-making.

To learn more about the Distiller component of AI Refinery, visit the distiller page. Check the Tutorial pages on how to use the Distiller API to create highly customizable multi-agent projects.

Example Usage

import os
import asyncio

from openai import AsyncOpenAI

from air import login, DistillerClient
from air import utils


auth = login(
    account=str(os.getenv("ACCOUNT")),
    api_key=str(os.getenv("API_KEY")),
)
base_url = os.getenv("AIREFINERY_ADDRESS", "")

async def simple_agent(query: str):
    global auth

    prompt = f"""Your task is to generate some synthetic data so that it will be useful to answer the user question. Do not mention this is synthetic data in your answer.\n\n{query}"""
    # This is for local serving
    client = AsyncOpenAI(**auth.openai(base_url=base_url))

    # This is for production endpoint
    # client = AsyncOpenAI(**auth.openai())

    response = await client.chat.completions.create(
        messages=[{"role": "user", "content": prompt}],
        model="meta-llama/Llama-3.1-70B-Instruct",
    )

    return response.choices[0].message.content

def interactive():
    global base_url
    distiller_client = DistillerClient(
        base_url=base_url,
    )

    # upload your config file to register a new distiller project
    distiller_client.create_project(config_path="example.yaml", project="example")

    # Define a mapping between your custom agent to Callable.
    # When the custom agent is summoned by the super agent / orchestrator,
    # distiller-sdk will run the custom agent and send its response back to the
    # multi-agent system.
    custom_agent_gallery = {"Data Scientist Agent": simple_agent}

    distiller_client.interactive(
        project="example",
        uuid="test_user",
        custom_agent_gallery=custom_agent_gallery,
    )

if __name__ == "__main__":
    print("Interactive")
    interactive()

Class Overview

DistillerClient

The DistillerClient class provides an interface for interacting with the AI Refinery's distiller service, allowing users to create projects, download configurations and run distiller sessions. 

class DistillerClient:
     """
    Distiller SDK for AI Refinery.

    This class provides an interface for interacting with the AI Refinery's
    distiller service, allowing users to create projects, download configurations,
    and run distiller sessions.
    """

Methods

__init__

Initializes the DistillerClient instance with optional base_url parameters.

def __init__(self, *, base_url: str = "") -> None:
    ...
Parameters:

create_project

Creates a project based on the configuration file specified by the config path. (REST API)

def create_project(self, *, config_path: str, project: str) -> bool:
    ...
Parameters:
  • config_path (str): YAML configuration file path.
  • project (str): Project name of your choice.
Returns:
  • bool: True if project creation is successful.

download_project

Downloads the configuration from the server for a given project.

def download_project(self, project: str) -> dict:
    ...
Parameters:
  • project (str): Name of the project to download the configuration for.
Returns:
  • dict: The downloaded configuration as a JSON object.

connect

Connects to the account/project/uuid-specific URL. 

async def connect(  
    self,  
    project: str,  
    uuid: str,  
    custom_agent_gallery: Optional[dict[str, Callable]] = None,  
) -> None:  
    ...
Parameters:
  • project (str): Name of the project.
  • uuid (str): Unique identifier for the session.
  • custom_agent_gallery (Optional[dict[str, Callable]], optional): Custom agent handlers. Defaults to None.
Returns:
  • None: This function does not return any value.

close

Closes the websocket connection if it exists. 

async def close(self) -> None:  
    ...
Returns:
  • None: This function does not return any value.

send

Sends a payload over the established websocket connection. 

async def send(self, payload: dict[str, Any]) -> None:  
    ...
Parameters:
  • payload (dict[str, Any]): The payload to send.
Raises:
  • ConnectionError: If the websocket connection is not established.
Returns:
  • None: This function does not return any value.

recv

Receives a message from the websocket connection. 

async def recv(self) -> dict[str, Any]:  
    ...
Returns:
  • dict[str, Any]: The received message.
Raises:
  • ConnectionError: If the websocket connection is not established.
  • ValueError: If the received message is not JSON.

request

Submits a request to the websocket. 

async def request(self, request_type: str, request_args: dict, **kwargs):  
    ...
Parameters:
  • request_type (str): Type of the request.
  • request_args (dict): Arguments for the request.
  • **kwargs: Additional keyword arguments.
Returns:
  • None: This function does not return any value.

query

Sends a query request to the websocket. 

async def query(self, query: str, image: Optional[str] = None, **kwargs):  
    ...
Parameters:
  • query (str): The query string.
  • image (Optional[str], optional): Image to include in the query. Defaults to None.
  • **kwargs: Additional keyword arguments.
Returns:
  • Coroutine: The request coroutine.

retrieve_memory

Retrieves memory from the websocket. 

async def retrieve_memory(self, **kwargs):  
    ...
Parameters:
  • **kwargs: Keyword arguments.
Returns:
  • Coroutine: The request coroutine.

add_memory

Adds memory to the websocket. 

async def add_memory(self, **kwargs):  
    ...
Parameters:
  • **kwargs: Keyword arguments.
Returns:
  • Coroutine: The request coroutine.

reset_memory

Resets memory in the websocket. 

async def reset_memory(self, **kwargs):  
    ...
Parameters:
  • **kwargs: Keyword arguments.
Returns:
  • Coroutine: The request coroutine.

retrieve_history

Retrieves chat history from the database. 

async def retrieve_history(  
    self,  
    db_client: DatabaseClient,  
    project: str,  
    uuid: str,  
    n_messages: int,  
    as_string=False,  
) -> str | list[dict]:  
    ...
Parameters:
  • db_client (DatabaseClient): Database client.
  • project (str): Name of the project.
  • uuid (str): Unique identifier for the session.
  • n_messages (int): Number of messages to retrieve.
  • as_string (bool, optional): Whether to return the history as a string. Defaults to False.
Returns:
  • str | list[dict]: Chat history as a string or list of dictionaries.

interactive

Enters interactive mode, allowing the user to interact with the agents through the terminal. 

def interactive(  
    self, project: str, uuid: str, custom_agent_gallery: dict[str, Callable] = {}  
):  
    ...
Parameters:
  • project (str): Name of the project.
  • uuid (str): Unique identifier for the session.
  • custom_agent_gallery (dict[str, Callable], optional): Custom agent handlers. Defaults to {}.

Conclusion

The Distiller and its components provide a robust framework for managing and orchestrating various tasks within your projects. By leveraging the power of the Orchestrator, Super Agents, and Utility Agents, you can streamline your workflows and achieve higher efficiency and productivity.