xtrace_sdk.integrations.xtrace

Classes

XTraceIntegration

XTrace integration for retrievers.

Module Contents

class xtrace_sdk.integrations.xtrace.XTraceIntegration(org_id, api_key=None, api_url='https://api.production.xtrace.ai')

Bases: xtrace_sdk.integrations.base.IntegrationBase

XTrace integration for retrievers.

Parameters:

  • XtraceIntegrationBase (class) – Interface for XTrace integrations

  • org_id (Any)

  • api_key (Optional[Any])

  • api_url (str)

org_id
api_url = 'https://api.production.xtrace.ai'
batch_size = 25
session: aiohttp.ClientSession | None = None
_loop: asyncio.AbstractEventLoop | None = None
_lock
api_key = None
async init_session()

Ensure there is an open ClientSession usable on the current event loop.

  • Create once on first use.

  • If the session was closed or the event loop changed, (re)create it.

Return type:

aiohttp.ClientSession

async close()
Return type:

None

async __aenter__()
Return type:

Any

async __aexit__(exc_type, exc, tb)
Parameters:

  • exc_type (Any)

  • exc (Any)

  • tb (Any)

Return type:

None

async get_chunk_by_id(chunk_ids, kb_id)

Get a chunk from the XTrace storage.

Parameters:

  • chunk_id (str) – chunk id of the chunk to be retrieved

  • kb_id (str) – knowledge base id of the chunk to be retrieved

  • chunk_ids (list[int])

Returns:

the chunk

Return type:

Chunk

store_index(index)

Store the index in the storage.

Parameters:

index (Any)

Return type:

Any

_normalize_meta(md)

Ensure meta_data has all required tags, filling missing ones with None.

Parameters:

md (dict | None)

Return type:

dict

preprocess_chunk(db, index, update=False)
Parameters:

  • db (Any)

  • index (Any)

  • update (bool)

Return type:

list[list[dict]]

async store_db(db, index, kb_id, context_id, concurrent=False)

Store the database in the XTrace storage. Assumes index and db are built together.

Parameters:

  • db (EncryptedDB) – the database to be stored

  • index (Index) – the index to be stored

  • kb_id (str) – knowledge base id

  • org_id (str) – organization id

  • concurrent (bool, optional) – whether to insert chunks concurrently, defaults to False

  • meta_data (list[MetaData], optional) – metadata for each chunk, defaults to None

  • context_id (str)

Raises:

AssertionError – if the length of db and index are not the same, or if the length of db and meta_data are not the same

Returns:

the response from the API

Return type:

requests.Response

async insert_chunks(kb_id, chunks, context_id)

Insert a chunk into the XTrace storage.

Parameters:

  • kb_id (str) – the id of the knowledge base

  • chunk_id (ChunkID) – the id of the chunk

  • chunk_content (str) – the content of the chunk

  • meta_data (MetaData) – the metadata of the chunk

  • index (int) – the index of the chunk

  • chunks (list[dict])

  • context_id (str)

Returns:

the response from the API

Return type:

requests.Response

load_db()

Load the database from the storage.

Return type:

None

load_index()

Load the index from the storage.

Return type:

None

async compute_hamming_distances_bench(query, pk, **kwargs)
Parameters:

  • query (list[gmpy2.mpz])

  • pk (xtrace_sdk.utils.xtrace_types.HomomorphicPublicKey)

  • kwargs (Any)

Return type:

None

async compute_hamming_distances(query, context_id, **kwargs)

Compute the hamming distances between the query and the data in the XTrace database.

Parameters:

  • query (int) – encrypted query to be used for comparison

  • pk (HomomorphicPublicKey) – the public key used for encryption

  • kb_id (str) – the id of the knowledge base to query

  • context_id (str)

  • kwargs (Any)

Returns:

list of tuples with the id and the encrypted hamming distance

Return type:

List[Tuple[ChunkID,int]]

async get_serialized_exec_context(context_id)

Get the execution context from the XTrace storage.

Parameters:

context_id (str) – the id of the execution context to be retrieved

Returns:

the serialized execution context

Return type:

str

async store_exec_context(exec_context_dict, context_id)

Store the execution context in the XTrace storage.

Parameters:
Return type:

str

async list_exec_contexts()

List all execution contexts for the organization.

Returns:

list of context ids

Return type:

list[str]

async update_chunks(index_updates, chunk_updates, context_id, kb_id)

Update chunks in the XTrace storage.

Parameters:

  • kb_id (str) – the id of the knowledge base

  • chunks (list[dict]) – list of chunks to be updated

  • context_id (str) – the id of the execution context

  • index_updates (Any)

  • chunk_updates (list[dict])

Raises:

AssertionError – if any chunk does not have a chunk_id

Return type:

Any

async delete_chunks(chunk_ids, kb_id)

Delete chunks from the XTrace storage.

Parameters:

  • kb_id (str) – the id of the knowledge base

  • chunk_ids (list[int]) – list of chunk ids to be deleted

Return type:

Any

Search chunks in the XTrace storage based on metadata. Only Equality operator is supported for now.

Parameters:

  • kb_id (str) – the id of the knowledge base

  • meta_filter (str) – the metadata filter as a JSON string

  • context_id (str)

Returns:

list of chunk ids matching the metadata filter

Return type:

list[int]