EdgeDB AI's Python package

The edgedb.ai package is an optional binding of the AI extension in EdgeDB. To use the AI binding, you need to install edgedb-python with the ai extra dependencies:

Copy

pip install 'edgedb[ai]'

Usage

Start by importing edgedb and edgedb.ai:

Copy

import edgedb
import edgedb.ai

Blocking

The AI binding is built on top of the regular EdgeDB client objects, providing both blocking and asynchronous versions of its API. For example, a blocking AI client is initialized like this:

Copy

client = edgedb.create_client()
gpt4ai = edgedb.ai.create_ai(
    client,
    model="gpt-4-turbo-preview"
)

Add your query as context:

Copy

astronomy_ai = gpt4ai.with_context(
    query="Astronomy"
)

The default text generation prompt will ask your selected provider to limit answer to information provided in the context and will pass the queried objects’ AI index as context along with that prompt.

Call your AI client’s query_rag method, passing in a text query.

Copy

print(
    astronomy_ai.query_rag("What color is the sky on Mars?")
);

or stream back the results by using stream_rag instead:

Copy

for data in astronomy_ai.stream_rag("What color is the sky on Mars?"):
    print(data)

Async

To use an async client instead, do this:

Copy

import asyncio  # alongside the EdgeDB imports

client = edgedb.create_async_client()

async def main():
    gpt4ai = await edgedb.ai.create_async_ai(
        client,
        model="gpt-4-turbo-preview"
    )
    astronomy_ai = gpt4ai.with_context(
        query="Astronomy"
    )
    query = "What color is the sky on Mars?"
    print(
        await astronomy_ai.query_rag(query)
    );

    #or streamed
    async for data in blog_ai.stream_rag(query):
        print(data)

asyncio.run(main())

API reference

function

create_ai()

create_ai(client, ** kwargs) -> EdgeDBAI

Creates an instance of EdgeDBAI with the specified client and options.

This function ensures that the client is connected before initializing the AI with the specified options.

Parameters

client – An EdgeDB client instance.
kwargs – Keyword arguments that are passed to the AIOptions data class to configure AI-specific options. These options are:
- model: The name of the model to be used. (required)
- prompt: An optional prompt to guide the model’s behavior. None will result in the client using the default prompt. (default: None)

function

create_async_ai()

create_async_ai(client, ** kwargs) -> AsyncEdgeDBAI

Creates an instance of AsyncEdgeDBAI w/ the specified client & options.

This function ensures that the client is connected asynchronously before initializing the AI with the specified options.

Parameters

client – An asynchronous EdgeDB client instance.
kwargs – Keyword arguments that are passed to the AIOptions data class to configure AI-specific options. These options are:
- model: The name of the model to be used. (required)
- prompt: An optional prompt to guide the model’s behavior. (default: None)

AI client classes

BaseEdgeDBAI

class

BaseEdgeDBAI

The base class for EdgeDB AI clients.

This class handles the initialization and configuration of AI clients and provides methods to modify their configuration and context dynamically.

Both the blocking and async AI client classes inherit from this one, so these methods are available on an AI client of either type.

Variables

options – An instance of AIOptions, storing the AI options.
context – An instance of QueryContext, storing the context for AI queries.
client_cls – A placeholder for the client class, should be implemented by subclasses.

Parameters

client – An instance of EdgeDB client, which could be either a synchronous or asynchronous client.
options – AI options to be used with the client.
kwargs – Keyword arguments to initialize the query context.

method

with_config()

with_config(** kwargs)

Creates a new instance of the same class with modified configuration options. This method uses the current instance’s configuration as a base and applies the changes specified in kwargs.

Parameters

kwargs – Keyword arguments that specify the changes to the AI configuration. These changes are passed to the derive method of the current configuration options object. Possible keywords include:

model: Specifies the AI model to be used. This must be a string.
prompt: An optional prompt to guide the model’s behavior. This is optional and defaults to None.

method

with_context()

with_context(** kwargs)

Creates a new instance of the same class with a modified context. This method preserves the current AI options and client settings, but uses the modified context specified by kwargs.

Parameters

kwargs – Keyword arguments that specify the changes to the context. These changes are passed to the derive method of the current context object. Possible keywords include:

query: The database query string.
variables: A dictionary of variables used in the query.
globals: A dictionary of global settings affecting the query.
max_object_count: An optional integer to limit the number of objects returned by the query.

EdgeDBAI

class

EdgeDBAI

A synchronous class for creating EdgeDB AI clients.

This class provides methods to send queries and receive responses using both blocking and streaming communication modes synchronously.

Variables

client – An instance of httpx.AsyncClient used for making HTTP requests asynchronously.

method

query_rag()

query_rag(message, context = None) -> str

Sends a request to the AI provider and returns the response as a string.

This method uses a blocking HTTP POST request. It raises an HTTP exception if the request fails.

Parameters

message – The query string to be sent to the AI model.
context – An optional QueryContext object to provide additional context for the query. If not provided, uses the default context of this AI client instance.

method

stream_rag()

stream_rag(message, context = None)

Opens a connection to the AI provider to stream query responses.

This method yields data as it is received, utilizing Server-Sent Events (SSE) to handle streaming data. It raises an HTTP exception if the request fails.

Parameters

message – The query string to be sent to the AI model.
context – An optional QueryContext object to provide additional context for the query. If not provided, uses the default context of this AI client instance.

AsyncEdgeDBAI

class

AsyncEdgeDBAI

An asynchronous class for creating EdgeDB AI clients.

This class provides methods to send queries and receive responses using both blocking and streaming communication modes asynchronously.

Variables

client – An instance of httpx.AsyncClient used for making HTTP requests asynchronously.

method

query_rag()

query_rag(message, context = None) -> str

Sends an async request to the AI provider, returns the response as a string.

This method is asynchronous and should be awaited. It raises an HTTP exception if the request fails.

Parameters

message – The query string to be sent to the AI model.
context – An optional QueryContext object to provide additional context for the query. If not provided, uses the default context of this AI client instance.

method

stream_rag()

stream_rag(message, context = None)

Opens an async connection to the AI provider to stream query responses.

This method yields data as it is received, using asynchronous Server-Sent Events (SSE) to handle streaming data. This is an asynchronous generator method and should be used in an async for loop. It raises an HTTP exception if the connection fails.

Parameters

message – The query string to be sent to the AI model.
context – An optional QueryContext object to provide additional context for the query. If not provided, uses the default context of this AI client instance.

Other classes

class

ChatParticipantRole

An enumeration of roles used when defining a custom text generation prompt.

Variables

SYSTEM – Represents a system-level entity or process.
USER – Represents a human user participating in the chat.
ASSISTANT – Represents an AI assistant.
TOOL – Represents a tool or utility used within the chat context.

class

Custom

A single message in a custom text generation prompt.

Variables

role – The role of the chat participant. Must be an instance of ChatParticipantRole.
content – The content associated with the role, expressed as a string.

class

Prompt

The metadata and content of a text generation prompt.

Variables

name – An optional name identifying the prompt.
id – An optional unique identifier for the prompt.
custom – An optional list of Custom objects, each providing role-specific content within the prompt.

class

AIOptions

A data class for AI options, specifying model and prompt settings.

Variables

model – The name of the AI model.
prompt – An optional Prompt providing additional guiding information for the model.

Method derive(kwargs)

Creates a new instance of AIOptions by merging existing options with provided keyword arguments. Returns a new AIOptions instance with updated attributes.

class

QueryContext

A data class defining the context for a query to an AI model.

Variables

query – The base query string.
variables – An optional dictionary of variables used in the query.
globals – An optional dictionary of global settings affecting the query.
max_object_count – An optional integer specifying the maximum number of objects the query should return.

Method derive(kwargs)

Creates a new instance of QueryContext by merging existing context with provided keyword arguments. Returns a new QueryContext instance with updated attributes.

class

RAGRequest

A data class defining a request to a text generation model.

Variables

model – The name of the AI model to query.
prompt – An optional Prompt associated with the request.
context – The QueryContext defining the query context.
query – The specific query string to be sent to the model.
stream – A boolean indicating whether the response should be streamed (True) or returned in a single response (False).

Method to_httpx_request()

Converts the RAGRequest into a dictionary suitable for making an HTTP request using the httpx library.

EdgeDB AI's Python package​

Usage​

Blocking​

Async​

API reference​

AI client classes​

BaseEdgeDBAI​

EdgeDBAI​

AsyncEdgeDBAI​

Other classes​

EdgeDB AI's Python package

Usage

Blocking

Async

API reference

AI client classes

BaseEdgeDBAI

EdgeDBAI

AsyncEdgeDBAI

Other classes