Agents API¶

The agent runtime lives directly in the top-level kavalai package. The headline class is Agent — a multi-step reasoning loop that calls tools through a FunctionKernel until it produces a final, optionally structured, answer.

Agent¶

class kavalai.agent.ToolCall(*, name: str, literal_args: str = '{}', planner_context_args: str = '{}', input_args: str = '{}', call_id: str | None = None)[source]¶

Bases: BaseModel

This data structure represents tool call requests.

Arguments are expected to be JSON encoded to help LLM models encode the data.

model_config : ClassVar[ConfigDict] = {'extra': 'forbid'}¶

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

name : str¶

literal_args : str¶

planner_context_args : str¶

input_args : str¶

call_id : str | None¶

kavalai.agent.get_step_output_type(ResponseModel=typing.Type[pydantic.main.BaseModel])[source]¶

class kavalai.agent.StepStreamDemuxer(step_idx: int, *, stream_output: bool = False, stream_instructions: bool = False, stream_partials: bool = False, stream_delta: bool = False)[source]¶

Bases: object

Routes one agent step’s raw LLM stream into named sub-streams.

Consumes the step’s StreamContent chunks (the step is always streamed in raw-delta mode) and returns the chunks to emit, gated by the stream flags:

• the step’s raw model output under step<N> (stream_partials),

• the instructions field under instructions, with a per-step complete (stream_instructions),

• the output field under response as full safe-parsed JSON (stream_output),

• pass-through of auxiliary provider streams (e.g. Gemini thoughts) and of restart chunks, which also reset the accumulated state.

One instance handles exactly one step; the raw text accumulates in buffer and the completed step’s JSON is exposed as step_json.

on_chunk(chunk: StreamContent) → list[StreamContent][source]¶

Process one incoming chunk and return the chunks to emit.

step_completed(step_output) → list[StreamContent][source]¶

Per-step closing chunks, once the parsed StepOutput is known.

class kavalai.agent.Agent(llm_client: BaseLlmClient, *, kernel: FunctionKernel | None = None, run_context: RunContext | None = None, prompt_template: Template | None = None, allowed_tools: list[str] | None = None, debug: bool = False)[source]¶

Bases: object

async prompt_stream(prompt: str, response_model: type[BaseModel] | None = None, max_steps: int = 10, *, stream_output: bool = False, stream_instructions: bool = False, stream_partials: bool = False, stream_delta: bool = False) → AsyncGenerator[StreamContent, None][source]¶

Run the agent loop, streaming progress as StreamContent.

The agent iterates up to max_steps times. On each step the LLM returns a StepOutput with optional tool_calls and an optional final output. Tool calls are executed through the FunctionKernel and their results are fed back into the prompt so the model can reason over them on the next step. The loop stops once the model returns an output without requesting further tool calls, or when max_steps is reached.

The final chunk is always complete/response carrying the final output (JSON for structured outputs, the plain string otherwise, or None when no output was produced). The flags gate the optional progress streams; their semantics (including the naming and the full-JSON structured output stream) are documented on kavalai.workflow.models.AgentNode. Stream names here are unscoped (response, instructions, step<N>); the workflow engine prefixes them with the node name.

Parameters:¶

prompt: str¶

The task description for the agent.

response_model: type[BaseModel] | None = None¶

Optional Pydantic model describing the structured final output. When omitted, a plain string is produced.

max_steps: int = 10¶

Maximum number of reasoning/tool-calling iterations.

stream_output: bool = False¶

Stream the step’s output field as it is written.

stream_instructions: bool = False¶

Stream each step’s instructions field.

stream_partials: bool = False¶

Stream each step’s raw model output.

stream_delta: bool = False¶

Emit deltas instead of full accumulated values on the instructions and step<N> streams.

async prompt(prompt: str, response_model: type[BaseModel] | None = None, max_steps: int = 10) → str | BaseModel[source]¶

Run the agent loop and return the final output (blocking wrapper).

Drains prompt_stream() with all progress streams disabled.

Parameters:¶

prompt: str¶

The task description for the agent.

response_model: type[BaseModel] | None = None¶

Optional Pydantic model describing the structured final output. When omitted, a plain string is returned.

max_steps: int = 10¶

Maximum number of reasoning/tool-calling iterations.

Returns:¶

The structured response_model instance, or a string when no response_model is provided. None if no output was produced.

Run context¶

Copyright 2026 OÜ KAVAL AI (registry code 17393877)

Licensed under the Apache License, Version 2.0 (the “License”); you may not use this file except in compliance with the License. You may obtain a copy of the License at

http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an “AS IS” BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and limitations under the License.

class kavalai.run_context.RunContext(*, agent_id: UUID | None = None, session_id: UUID | None = None, run_id: UUID | None = None, data: dict = {}, templates: dict[str, str] = {}, agent_service: Any | None = None)[source]¶

Bases: BaseModel

Runtime data for a single interaction.

model_config : ClassVar[ConfigDict] = {'arbitrary_types_allowed': True}¶

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

agent_id : UUID | None¶

session_id : UUID | None¶

run_id : UUID | None¶

data : dict¶

templates : Dict[str, str]¶

agent_service : Any | None¶

resolve_context_value(path: str)[source]¶

Resolve a dotted path like ‘input.user_message’ from context data.

async resolve_history_value(path: str)[source]¶

Resolve a value from session history.

async resolve_template_value(name: str)[source]¶

Resolve a template value by name.

async render_prompt(prompt: str) → str[source]¶

Render a prompt string by replacing {{ templates.NAME }}, {{ context.PATH }}, and {{ history.PATH }} with their resolved values.

async resolve_input_info(info: ArgumentInfo)[source]¶

Resolve a TypeInputInfo to its actual value.

async prepare_tool_inputs(task: Any) → dict[source]¶

Resolve a task/node’s inputs mapping into plain values.

Agent service & persistence¶

Copyright 2026 OÜ KAVAL AI (registry code 17393877)

Licensed under the Apache License, Version 2.0 (the “License”); you may not use this file except in compliance with the License. You may obtain a copy of the License at

http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an “AS IS” BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and limitations under the License.

class kavalai.agent_service.AgentService(session_maker: async_sessionmaker[AsyncSession])[source]¶

Bases: object

Database operations for the agent runtime.

Manages the core entities (agents, sessions, runs) as well as the history data recorded while they execute (chat messages, tasks, model-call stats). Works against Postgres and SQLite alike (the models are dialect-agnostic and schema-less; the schema comes from the engine’s schema_translate_map).

async get_or_create_agent(name: str, description: str | None = None, input_schema: dict | None = None, output_schema: dict | None = None, workflow: dict | None = None) → Agent[source]¶

Finds an agent by name or creates a new one if not found.

async get_or_create_session(agent_id: UUID, session_id: UUID | None = None, external_id: UUID | None = None) → Session | None[source]¶

async create_run(session_id: UUID, input_data: dict | None = None, context: dict | None = None) → Run[source]¶

Creates a new run entry for a specific session.

async initialize_workflow_run(agent_name: str, agent_description: str | None = None, input_schema: dict | None = None, output_schema: dict | None = None, workflow: dict | None = None, session_id: UUID | None = None, external_id: str | None = None, input_data: dict | None = None) → tuple[Agent, Session, Run][source]¶

Initialize agent, session, and run in a single database transaction.

This is an optimized batch operation that reduces 3 DB roundtrips to 1, improving performance especially for remote databases.

session_id selects an existing session by primary id (raises ValueError if absent). Without it, external_id reuses the agent’s most recent session carrying that caller-supplied id — letting clients pin a conversation to their own identifier — and a new session is created when neither matches.

Returns:¶

tuple of (agent, session, run)

async update_run(run_id: UUID, *, output_data: dict | None = None, context: dict | None = None) → Run[source]¶

Updates an existing run with final output_data and/or context.

async get_history_value(session_id: UUID, key: str) → Any | None[source]¶

Retrieves a value from the context of previous runs in the same session.

• If key is a dotted path (e.g., “output.search_results”), resolves it as such.

• If key is a plain name (e.g., “search_results”), searches recursively for the first matching key in the context dicts of previous runs (newest first).

Returns the most recent value found for the given key.

async add_chat_message(agent_id: UUID, session_id: UUID, role: str, content: str | None, run_id: UUID | None = None) → ChatMessage[source]¶

Helper to append messages to the chat history.

async get_chat_history(session_id: UUID, limit: int = 50) → list[ChatMessage][source]¶

Retrieves the conversation history for a session, ordered from oldest to newest.

async add_task(session_id: UUID, run_id: UUID, name: str | None = None, agent_id: UUID | None = None, inputs: dict | None = None, output: dict | None = None, prompt: str | None = None, errors: list[str] | None = None, duration_seconds: float | None = None, node_type: str | None = None) → Task[source]¶

Records a specific unit of work (Task) performed within a run.

async add_model_call_stats(stats: ModelCallStat, agent_id: UUID | None = None) → ModelCallStat[source]¶

Records LLM/Embedding call statistics.

async get_model_call_stats(call_type: str | None = None, limit: int = 50, offset: int = 0) → list[ModelCallStat][source]¶

Retrieves paginated model call stats, optionally filtered by call type.

async delete_history_for_session(session_id: UUID) → None[source]¶

Delete all history (chat, tasks) belonging to a session.

async delete_history_for_agent(agent_id: UUID) → None[source]¶

Delete all history (chat, tasks, stats) belonging to an agent.

Copyright 2026 OÜ KAVAL AI (registry code 17393877)

Licensed under the Apache License, Version 2.0 (the “License”); you may not use this file except in compliance with the License. You may obtain a copy of the License at

http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an “AS IS” BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and limitations under the License.

class kavalai.backoffice.sessions.SessionSummary(*, session_id: UUID, agent_id: UUID, agent_name: str, runs_count: int, tasks_count: int, messages_count: int, first_message: str | None, last_message: str | None, errors_count: int, created_at: datetime, updated_at: datetime)[source]¶

Bases: BaseModel

Row-level summary of a session for the Conversations list.

Aggregates a session’s owning agent, its run/task/message and error counts, and a preview of its first and last messages.

session_id : UUID¶

agent_id : UUID¶

agent_name : str¶

runs_count : int¶

tasks_count : int¶

messages_count : int¶

first_message : str | None¶

last_message : str | None¶

errors_count : int¶

created_at : datetime¶

updated_at : datetime¶

model_config : ClassVar[ConfigDict] = {}¶

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

class kavalai.backoffice.sessions.TaskSummary(*, id: UUID, agent_id: UUID | None, session_id: UUID, run_id: UUID, inputs: Any | None, output: Any | None, name: str | None = None, prompt: str | None = None, errors: list[str] | None = None, duration_seconds: float | None = None, created_at: datetime, updated_at: datetime)[source]¶

Bases: BaseModel

Summary of a single task (workflow-node execution) for the Tasks view.

Exposes the task’s inputs, output, name, prompt, any errors and duration.

model_config : ClassVar[ConfigDict] = {'from_attributes': True}¶

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

id : UUID¶

agent_id : UUID | None¶

session_id : UUID¶

run_id : UUID¶

inputs : Any | None¶

output : Any | None¶

name : str | None¶

prompt : str | None¶

errors : list[str] | None¶

duration_seconds : float | None¶

created_at : datetime¶

updated_at : datetime¶

class kavalai.backoffice.sessions.RunSummary(*, id: UUID, session_id: UUID, input_data: Any | None, output_data: Any | None, context: Any | None, tasks_count: int, created_at: datetime, updated_at: datetime)[source]¶

Bases: BaseModel

Summary of a single workflow run for the Runs view.

Exposes the run’s input/output data, resolved context and the number of tasks it executed.

model_config : ClassVar[ConfigDict] = {'from_attributes': True}¶

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

id : UUID¶

session_id : UUID¶

input_data : Any | None¶

output_data : Any | None¶

context : Any | None¶

tasks_count : int¶

created_at : datetime¶

updated_at : datetime¶

class kavalai.backoffice.sessions.ChatMessageSummary(*, id: UUID, agent_id: UUID, session_id: UUID, run_id: UUID | None, role: str, content: str, created_at: datetime, updated_at: datetime)[source]¶

Bases: BaseModel

Summary of a single chat message for the conversation transcript.

Exposes the message’s role, content and the run it is associated with.

model_config : ClassVar[ConfigDict] = {'from_attributes': True}¶

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

id : UUID¶

agent_id : UUID¶

session_id : UUID¶

run_id : UUID | None¶

role : str¶

content : str¶

created_at : datetime¶

updated_at : datetime¶

class kavalai.backoffice.sessions.SessionDetails(*, session_id: UUID, messages: list[ChatMessageSummary], runs: list[RunSummary], tasks: list[TaskSummary])[source]¶

Bases: BaseModel

Full detail of one session: its messages, runs and tasks.

Powers the per-conversation detail view in the backoffice, bundling the session’s chat transcript together with all of its runs and tasks.

session_id : UUID¶

messages : list[ChatMessageSummary]¶

runs : list[RunSummary]¶

tasks : list[TaskSummary]¶

model_config : ClassVar[ConfigDict] = {}¶

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

class kavalai.backoffice.sessions.SessionsResponse[source]¶

Bases: TypedDict

sessions : list[SessionSummary]¶

total_count : int¶

async kavalai.backoffice.sessions.get_sessions_summary(session: AsyncSession, agent_id: UUID | None = None, search: str | None = None, start_date: datetime | None = None, end_date: datetime | None = None, limit: int = 50, offset: int = 0) → SessionsResponse[source]¶

async kavalai.backoffice.sessions.get_session_details(session: AsyncSession, session_id: UUID) → SessionDetails[source]¶

Remote agent client¶

Copyright 2026 OÜ KAVAL AI (registry code 17393877)

Licensed under the Apache License, Version 2.0 (the “License”); you may not use this file except in compliance with the License. You may obtain a copy of the License at

http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an “AS IS” BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and limitations under the License.

class kavalai.client.AgentClient(base_url: str, username: str | None = None, password: str | None = None, timeout: float = 60.0, transport: AsyncBaseTransport | None = None)[source]¶

Bases: object

Async HTTP client for invoking a remote Kaval.AI agent server.

Wraps the agent server’s /run_agent and /stream_agent endpoints, discovering the agent’s input/output schemas from its OpenAPI spec and transparently maintaining the conversation session_id across calls so successive invocations share the same session. Optional HTTP Basic Auth is used when both username and password are provided.

Parameters:¶

base_url: str¶

Base URL of the agent server.

username: str | None = None¶

Optional HTTP Basic Auth username.

password: str | None = None¶

Optional HTTP Basic Auth password.

timeout: float = 60.0¶

Per-request timeout in seconds.

transport: AsyncBaseTransport | None = None¶

Optional httpx transport, e.g. to route requests through a proxy or, in tests, to serve them without a network.

async discover_schemas()[source]¶

Fetch the server’s OpenAPI spec and derive the agent’s schemas.

Populates self.input_schema and self.output_schema with the Pydantic models for the agent’s request and response payloads. Called automatically by run_agent() and stream_agent() on first use, but may be invoked directly to inspect the schemas up front.

async run_agent(data: BaseModel, external_id: str | None = None) → BaseModel[source]¶

Run the agent once and return its complete response.

Sends data to the server’s /run_agent endpoint, blocking until the run finishes. Updates self.session_id from the response so the next call continues the same conversation.

Parameters:¶

data: BaseModel¶

The request payload (an instance matching the agent’s input schema).

external_id: str | None = None¶

Optional caller-side identifier to correlate the session with an external system.

Returns:¶

An instance of the agent’s output schema with the run’s result.

async stream_agent(data: BaseModel, external_id: str | None = None)[source]¶

Run the agent and stream its output incrementally.

Sends data to the server’s /stream_agent (Server-Sent Events) endpoint and yields each data: chunk as a string as it arrives, letting callers consume partial output before the run completes.

Parameters:¶

data: BaseModel¶

The request payload (an instance matching the agent’s input schema).

external_id: str | None = None¶

Optional caller-side identifier to correlate the session with an external system.

Yields:¶

str – Successive content chunks from the streamed response.

RAG service¶

Copyright 2026 OÜ KAVAL AI (registry code 17393877)

Licensed under the Apache License, Version 2.0 (the “License”); you may not use this file except in compliance with the License. You may obtain a copy of the License at

http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an “AS IS” BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and limitations under the License.

class kavalai.rag.base.RagServiceResult(*, id: UUID, model: str, collection_name: str, source_id: str, content: str | None = None, embedding_size: int, rag_metadata: dict, similarity: float, created_at: datetime | None = None, updated_at: datetime | None = None, query_index: int | None = None)[source]¶

Bases: BaseModel

Represents a single result from a RAG query.

Variables:¶

id : UUID¶

Unique identifier of the indexed item.

model : str¶

The embedding model used for this item.

collection_name : str¶

The name of the collection this item belongs to.

source_id : str¶

An external identifier for the source of this item.

content : Optional[str]¶

The original text content that was indexed.

embedding_size : int¶

The dimension of the embedding vector.

rag_metadata : dict¶

Additional metadata associated with the item.

similarity : float¶

The similarity score (1.0 - distance) relative to the query.

created_at : Optional[datetime]¶

Timestamp when the item was created.

updated_at : Optional[datetime]¶

Timestamp when the item was last updated.

query_index : Optional[int]¶

Index of the query in batch queries (for query_batch results).

id : UUID¶

model : str¶

collection_name : str¶

source_id : str¶

content : str | None¶

embedding_size : int¶

rag_metadata : dict¶

similarity : float¶

created_at : datetime | None¶

updated_at : datetime | None¶

query_index : int | None¶

model_config : ClassVar[ConfigDict] = {}¶

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

class kavalai.rag.base.BaseRagService[source]¶

Bases: ABC

Interface for RAG (Retrieval-Augmented Generation) storage backends.

A RAG service indexes text documents as embeddings and answers similarity queries against them. Concrete backends (e.g. PostgresRagService) implement the abstract methods; compute_similarity_matrix and learn_normalizer have generic default implementations that backends may override with more efficient or exact versions.

similarity_matrix_candidates_per_source : int = 100¶

abstractmethod async index(text: str, source_metadata: dict | None = None, collection_name: str = 'default', source_id: str = 'default')[source]¶

Index a single text blob with metadata.

Parameters:¶

text : str¶

The text content to index.

source_metadata : Optional[dict]¶

Metadata to associate with the text.

collection_name : str¶

Name of the collection. Defaults to “default”.

source_id : str¶

Source identifier. Defaults to “default”.

Returns:¶

The created index entry (backend-specific type).

abstractmethod async index_batch(texts: list[str], metadata_list: list[dict], source_ids: list[str] | None = None, collection_name: str = 'default')[source]¶

Index multiple text items in a single batch.

Batch indexing can be significantly more efficient than repeated index() calls with certain backends.

Parameters:¶

texts : list[str]¶

List of text strings to index.

metadata_list : list[dict]¶

List of metadata dictionaries for each text.

source_ids : Optional[list[str]]¶

Optional list of source identifiers. If not provided, “default” is used.

collection_name : str¶

Name of the collection to add items to. Defaults to “default”.

Returns:¶

List of created index entries (backend-specific type).

Raises:¶

ValueError – If the lengths of texts, metadata_list, or source_ids do not match.

abstractmethod async query(text: str, top_k: int = 5, collection_name: str | None = None, source_ids: list[str] | None = None, keep_best: bool = False) → list[RagServiceResult][source]¶

Query the indexed items for similarities to the input text.

Parameters:¶

text : str¶

The query text.

top_k : int¶

Number of top results to return. Defaults to 5.

collection_name : Optional[str]¶

If provided, filter by collection name.

source_ids : Optional[list[str]]¶

If provided, filter by source identifiers.

keep_best : bool¶

If True, only the best result per source_id is returned. Useful when a single source is split into multiple indexed items.

Returns:¶

List of results with similarity scores.

Return type:¶

list[RagServiceResult]

abstractmethod async query_batch(texts: list[str], top_k: int = 5, collection_name: str | None = None, source_ids: list[str] | None = None) → list[list[RagServiceResult]][source]¶

Query the indexed items for similarities to multiple input texts.

Batch querying can be significantly more efficient than repeated query() calls with certain backends.

Parameters:¶

texts : list[str]¶

List of query texts to search for.

top_k : int¶

Number of top results to return per query. Defaults to 5.

collection_name : Optional[str]¶

If provided, filter by collection name.

source_ids : Optional[list[str]]¶

If provided, filter by source identifiers.

Returns:¶

A list of result lists, where each inner list contains

the top_k results for the corresponding query text.

Return type:¶

list[list[RagServiceResult]]

abstractmethod async delete(item_id: UUID, collection_name: str | None = None) → None[source]¶

Delete a single indexed item by its identifier.

Parameters:¶

item_id : UUID¶

Identifier of the indexed item to delete.

collection_name : Optional[str]¶

Collection the item belongs to. Backends that store collections separately search all collections when omitted.

abstractmethod async delete_by_source_id(collection_name: str, source_id: str | list[str]) → None[source]¶

Delete all items in a collection that match the given source identifier(s).

Parameters:¶

collection_name : str¶

The name of the collection.

source_id : Union[str, list[str]]¶

A source identifier, or a list of them.

async count_entries(collection_name: str) → int[source]¶

Number of entries in a collection (0 if it doesn’t exist).

iter_entries(collection_name: str, batch_size: int = 500) → AsyncIterator[dict][source]¶

Iterate all entries of a collection (including embeddings).

Yields dicts with keys: id, source_id, content, embedding, rag_metadata, created_at, updated_at. Used for bulk export (e.g. the backoffice embedding projector) so no caller needs to touch backend storage directly.

async compute_similarity_matrix(texts: list[str], source_ids: list[str], method: str = 'min', collection_name: str = 'default') → list[list[float]][source]¶

Compute a similarity matrix between multiple texts and multiple source identifiers.

Default implementation built on query_batch(): it retrieves up to similarity_matrix_candidates_per_source candidates per source and aggregates similarities per source_id. Sources with more indexed items than that may yield approximate “avg” aggregates; backends can override this with an exact implementation.

Parameters:¶

texts : list[str]¶

List of query texts (rows in the matrix).

source_ids : list[str]¶

List of source identifiers to compare against (columns in the matrix).

method : str¶

Aggregate method to use when multiple items exist for a source_id. “min” (default) uses the shortest distance (highest similarity). “avg” uses the average distance.

Returns:¶

A 2D matrix where matrix[i][j] is the similarity between

texts[i] and source_ids[j]. Missing sources score 0.0.

Return type:¶

list[list[float]]

async learn_normalizer(collection_name: str | None = None) → Normalizer[source]¶

Learn a normalizer from the indexed data.

Default implementation returns the process-wide default normalizer; backends with access to the stored embeddings should override this to learn (e.g.) a centering vector from the index.

Parameters:¶

collection_name : Optional[str]¶

If provided, learn only from this collection.

Returns:¶

The learned (or default) normalizer.

Return type:¶

Normalizer

Copyright 2026 OÜ KAVAL AI (registry code 17393877)

Licensed under the Apache License, Version 2.0 (the “License”); you may not use this file except in compliance with the License. You may obtain a copy of the License at

http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an “AS IS” BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and limitations under the License.

PostgreSQL (pgvector) backed RAG service — self-provisioning.

This backend owns its schema entirely; no Alembic migration set covers it. It maintains a small registry table (rag_collections) plus one table per collection with a typed vector(N) column, a real HNSW index for the collection’s exact dimension, and a GIN index on the metadata column. Dropping a collection is DROP TABLE. The registry row carries a schema_version used for in-code upgrades of collection tables.

All SQL here is raw and therefore bypasses schema_translate_map — every statement qualifies the configured schema explicitly.

kavalai.rag.postgres.RAG_COLLECTION_SCHEMA_VERSION = 1¶

Version of the per-collection table layout. Bump when the layout changes and register an upgrade step in _COLLECTION_UPGRADES.

class kavalai.rag.postgres.CollectionInfo(name: str, table_name: str, model: str, embedding_size: int, schema_version: int)[source]¶

Bases: object

Registry entry for one RAG collection.

class kavalai.rag.postgres.PostgresRagService(session_maker: async_sessionmaker[AsyncSession] | Callable[[], AbstractAsyncContextManager[AsyncSession]], model: str | None = None, agent: Agent | None = None, normalizer: Normalizer | None = None, schema: str | None = None)[source]¶

Bases: BaseRagService

PostgreSQL (pgvector) backed RAG service with backend-owned DDL.

Each collection lives in its own table (typed vector(N) column, per-collection HNSW + GIN indexes) registered in rag_collections. Collections are provisioned lazily on first index (the embedding dimension is taken from the first batch) or explicitly via create_collection().

All operations are scoped to a single collection ("default" unless specified) — collection_name is a logical handle in the BaseRagService interface; here it maps to a table.

REGISTRY_TABLE = 'rag_collections'¶

property embedding_client¶

Embedding client, created lazily so model-less usage works.

classmethod from_uri(uri: str, model: str, agent: Agent | None = None, normalizer: Normalizer | None = None, schema: str | None = None) → PostgresRagService[source]¶

Create a PostgresRagService from a database URI.

classmethod from_session_maker(session_maker: async_sessionmaker[AsyncSession], model: str, agent: Agent | None = None, normalizer: Normalizer | None = None, schema: str | None = None) → PostgresRagService[source]¶

Create a PostgresRagService from a session maker.

static table_name_for_collection(collection_name: str) → str[source]¶

Deterministic, SQL-safe table name for a collection.

A sanitized slug keeps the name readable; a short hash of the exact collection name guarantees uniqueness across names that sanitize to the same slug.

async create_collection(collection_name: str, embedding_size: int) → None[source]¶

Explicitly provision a collection with a known embedding dimension.

async drop_collection(collection_name: str) → None[source]¶

Drop a collection: its table and registry entry.

async list_collections() → list[dict][source]¶

List registered collections with entry counts.

async get_stats() → dict[source]¶

Aggregate stats across collections (for e.g. the backoffice).

async index(text: str, source_metadata: dict | None = None, collection_name: str = 'default', source_id: str = 'default') → dict[source]¶

Index a single text blob with metadata. Returns the created row dict.

async index_batch(texts: list[str], metadata_list: list[dict], source_ids: list[str] | None = None, collection_name: str = 'default') → list[dict][source]¶

Index multiple text items in a single batch.

The collection is provisioned on first use, taking its embedding dimension from the computed embeddings.

Returns:¶

Created rows (id, model, collection_name, source_id,

content, embedding_size, rag_metadata, created_at, updated_at).

Return type:¶

list[dict]

async delete(item_id: UUID, collection_name: str | None = None) → None[source]¶

Delete a single indexed item by its identifier.

Parameters:¶

item_id : UUID¶

Identifier of the indexed item to delete.

collection_name : Optional[str]¶

Collection the item belongs to. If omitted, all registered collections are searched.

async delete_by_source_id(collection_name: str, source_id: str | list[str]) → None[source]¶

Delete all items in a collection matching the source identifier(s).

async query(text: str, top_k: int = 5, collection_name: str | None = None, source_ids: list[str] | None = None, keep_best: bool = False) → list[RagServiceResult][source]¶

Query one collection for similarities to the input text.

collection_name defaults to "default" — with table-per-collection storage there is no cross-collection search; query each collection explicitly if needed.

async query_batch(texts: list[str], top_k: int = 5, collection_name: str | None = None, source_ids: list[str] | None = None, keep_best: bool = False) → list[list[RagServiceResult]][source]¶

Query one collection for similarities to multiple input texts in a single database call (CROSS JOIN LATERAL over an unnested vector array).

async compute_similarity_matrix(texts: list[str], source_ids: list[str], method: str = 'min', collection_name: str = 'default') → list[list[float]][source]¶

Compute a similarity matrix between texts and source identifiers within one collection, in a single database query.

async learn_normalizer(collection_name: str | None = None) → Normalizer[source]¶

Learn a centering normalizer from one collection’s embeddings.

async iter_entries(collection_name: str, batch_size: int = 500) → AsyncIterator[dict][source]¶

Iterate all entries of a collection (including embeddings) in stable id order using keyset pagination. Yields dicts with keys: id, source_id, content, embedding, rag_metadata, created_at, updated_at.

async count_entries(collection_name: str) → int[source]¶

Number of entries in a collection (0 if it doesn’t exist).

async get_embeddings_by_ids(collection_name: str, ids: list[UUID]) → dict[UUID, list[float]][source]¶

Fetch embeddings for specific entry ids within a collection.