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: type[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, step_template: Template | None = None, allowed_tools: list[str] | None = None, on_step: Callable[[dict], None] | 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, history: Sequence[ChatMessage] | None = None, inputs: Mapping[str, Any] | None = None, request: str | None = None) → 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.
history: Sequence[ChatMessage] | None = None¶: Earlier turns of the conversation, sent with every step between the system prompt and the step message. The current user message belongs in prompt, not here — the model would read it twice.
inputs: Mapping[str, Any] | None = None¶: The context variables of this invocation: listed in the system prompt and the values input_args of a tool call resolve against. Defaults to the run context’s data. The workflow engine passes the node’s resolved inputs, so an agent node sees what its author declared and not every value the run has produced so far.
request: str | None = None¶: The request this invocation answers, stated in every step message. history ends on the turn before it, so without it a model reading the conversation answers that turn. The workflow engine passes the run’s user message.

async prompt(prompt: str, response_model: type[BaseModel] | None = None, max_steps: int = 10, *, history: Sequence[ChatMessage] | None = None, inputs: Mapping[str, Any] | None = None, request: str | None = None) → 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.
history: Sequence[ChatMessage] | None = None¶: Earlier turns of the conversation, as on prompt_stream().
inputs: Mapping[str, Any] | None = None¶: The context variables, as on prompt_stream().
request: str | None = None¶: The request this invocation answers, as on prompt_stream().

Returns:¶

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

Run context¶

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.

kavalai.run_context.PLACEHOLDER_PATTERN = re.compile('\\{\\{\\s*(templates|context|history)\\.(.+?)\\s*\\}\\}')¶: One {{ templates.NAME }}, {{ context.PATH }} or {{ history.PATH }} placeholder; the only three prefixes a template may reference.

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] = {}, template_overrides: dict[str, str] = {}, agent_service: Any | None = None, token_stats: Any | None = None, seq_counter: Any | None = None, started_at: float | None = None, current_seq: int | None = None, task_logger: Any | None = None, user_message: str | None = None)[source]¶

Bases: BaseModel

Runtime data for a single interaction.

token_stats is the run’s TokenAccumulator. Every LLM client built during the run reports into it, and parallel branches share the parent’s, so the totals cover the whole run.

seq_counter is the run’s task-sequence counter, shared with parallel branches in the same way so that every task row of the run draws from one sequence and the interleaving of concurrent branches is recorded rather than lost. current_seq is the number handed to the node this context is currently executing, and is per-branch.

task_logger overrides the engine’s TaskLogger for this run only. One engine serves many concurrent runs, so a caller that wants this run’s trajectory on its own — the evaluation runner, a notebook debugging one call — cannot swap the engine’s logger.

template_overrides holds the template values passed to this run only (run_stream(templates=...)); they are already merged into templates and are kept separately so the run record can say which values came from the caller rather than from the document.

started_at is the monotonic-clock reading taken when the run started, shared with parallel branches, from which the run’s duration_seconds is measured when it completes or fails.

user_message is the run’s input as the chat history records it (the input’s user_message field, else the input itself as text). An agent node states it in its step message: the history sent with each step ends on the previous turn, so without it the model answers that turn.

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]¶

template_overrides : Dict[str, str]¶

agent_service : Any | None¶

token_stats : Any | None¶

seq_counter : Any | None¶

started_at : float | None¶

current_seq : int | None¶

task_logger : Any | None¶

user_message : str | None¶

next_seq() → int | None[source]¶

Take the next number from the run’s task sequence.

Returns None when the run has no counter, which is the case for a bare RunContext() built outside the engine.

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 resolve_placeholder(prefix: str, path: str)[source]¶

The value behind one {{ prefix.path }} placeholder.

prefix is templates, context or history. A reference that resolves to nothing raises ValueError: a prompt or filter silently missing a value is worse than a failed run.

async render_value(template: str)[source]¶

Render a template that may stand for a single typed value.

A template that is exactly one placeholder resolves to the referenced value itself — a number stays a number, a boolean a boolean — so a value an earlier node extracted can be compared without passing through text. Anything else is rendered with render_prompt().

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 an ArgumentInfo 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¶

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

class kavalai.agent_service.SessionSummary(*, session_id: UUID, agent_id: UUID, agent_name: str, external_id: str | None = None, 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

One row of a session list.

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

external_id is the caller-supplied key for the conversation. Evaluation runs record eval:{tag}:{case}, so filtering on that prefix lands on exactly the conversation a failing case produced.

session_id : UUID¶

agent_id : UUID¶

agent_name : str¶

external_id : str | None¶

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.agent_service.SessionsResponse[source]¶

Bases: TypedDict

sessions : list[SessionSummary]¶

total_count : int¶

async kavalai.agent_service.summarise_sessions(db_session: AsyncSession, agent_ids: Sequence[UUID] | None = None, *, search: str | None = None, external_id_prefix: str | None = None, exclude_external_id_prefix: str | None = None, start: datetime | None = None, end: datetime | None = None, limit: int = 50, offset: int = 0) → SessionsResponse[source]¶

List sessions, most recently active first, with their counts.

The query behind AgentService.list_sessions() and the backoffice conversation list, on a session the caller already holds.

Parameters:¶

agent_ids: Sequence[UUID] | None = None¶: Sessions of these agents only. None lists every agent’s; an empty list lists nothing.
search: str | None = None¶: Case-insensitive substring of any message in the session.
external_id_prefix: str | None = None¶: The session’s caller-supplied key starts with this, case-insensitively. Pasting eval:pr-412: shows every conversation that experiment produced.
exclude_external_id_prefix: str | None = None¶: Leave out sessions whose key starts with this — a host’s preview conversations, for example.
start: datetime | None = None¶: Sessions created at or after this time.
end: datetime | None = None¶: Sessions created at or before this time.

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: str | None = None) → Session | None[source]¶: Look up a session by id, or create a new one when no id is given.

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.

One session and one commit instead of three, which matters against a remote database.

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. A continued session’s updated_at is moved to now, so it records the session’s last activity.

Returns:¶: tuple of (agent, session, run)

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

Updates an existing run with its output, context and/or duration.

Each argument left None keeps the stored value, so the engine can record the output at the end node and the duration in the same call, while a failure recording carries the context and duration only.

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, max_chars: int | None = None, exclude_run_id: UUID | None = None) → list[ChatMessage][source]¶

The most recent messages of a session, ordered oldest to newest.

Parameters:¶

limit: int = 50¶: At most this many messages.
max_chars: int | None = None¶: At most this many characters of content. Whole messages are dropped from the oldest end until the rest fits; a message is never cut, and one that does not fit ends the window.
exclude_run_id: UUID | None = None¶: Leave out the messages of this run. The engine records the user message before the nodes run, so a node that already carries the current message in its prompt passes the current run and gets the earlier turns only; limit and max_chars then measure those turns alone.

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

seq is the run’s execution order, parent_task_name names the node a tool-call row belongs to, and tool_uri identifies the tool that ran. See kavalai.db.Task for what they are for.

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

Record one LLM or embedding call.

Parameters:¶

stats: Any¶: The Pydantic ModelCallStat a client reported, or a model_call_stats row built by hand.
agent_id: UUID | None = None¶: The agent the call is attributed to.
session_id: UUID | None = None¶: The session of the run that made the call.
run_id: UUID | None = None¶: The run that made the call.

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

Model call records, newest first, optionally filtered.

Parameters:¶

call_type: str | None = None¶: "llm" or "embedding".
agent_id: UUID | None = None¶: Calls attributed to this agent.
session_id: UUID | None = None¶: Calls made by runs of this session.
run_id: UUID | None = None¶: Calls made by this run.

async list_sessions(agent_ids: Sequence[UUID] | None = None, *, search: str | None = None, external_id_prefix: str | None = None, exclude_external_id_prefix: str | None = None, start: datetime | None = None, end: datetime | None = None, limit: int = 50, offset: int = 0) → SessionsResponse[source]¶

Sessions, most recently active first, with their counts.

agent_ids scopes the list to several agents at once — the agents of one tenant, for example — and an empty list lists nothing. See summarise_sessions() for the other filters.

async purge_sessions(before: datetime, agent_ids: Sequence[UUID] | None = None, batch_size: int = 1000) → AsyncIterator[list[UUID]][source]¶

Delete sessions whose last activity is older than before.

Deletes in batches of batch_size, one transaction each, and yields each batch’s session ids once it is committed, so a caller can delete rows of its own keyed by session in step. A session’s runs, tasks and chat messages go with it through the foreign keys. Its model calls stay, with their request and response payloads cleared: the token counts are a cost record, the payloads are the conversation again.

Parameters:¶

before: datetime¶: Sessions whose updated_at (last run) is earlier.
agent_ids: Sequence[UUID] | None = None¶: Only these agents’ sessions. None means every agent; an empty list purges nothing.
batch_size: int = 1000¶: Sessions per transaction.

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

Delete a session’s chat messages and tasks.

Its model calls keep their token counts and lose their payloads, as in purge_sessions().

async delete_history_for_agent(agent_id: UUID) → None[source]¶: Delete all history (chat, tasks, stats) belonging to an agent.

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

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.RunSummary(*, id: UUID, session_id: UUID, input_data: Any | None, output_data: Any | None, context: Any | None, tasks_count: int, duration_seconds: float | None = None, 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, the number of tasks it executed and how long it took (None for a run still in progress or recorded before the runs table carried a duration).

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¶

duration_seconds : float | None¶

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.SessionSummary(*, session_id: UUID, agent_id: UUID, agent_name: str, external_id: str | None = None, 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

One row of a session list.

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

external_id is the caller-supplied key for the conversation. Evaluation runs record eval:{tag}:{case}, so filtering on that prefix lands on exactly the conversation a failing case produced.

model_config : ClassVar[ConfigDict] = {}¶: Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

session_id : UUID¶

agent_id : UUID¶

agent_name : str¶

external_id : str | None¶

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¶

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

Bases: TypedDict

sessions : list[SessionSummary]¶

total_count : int¶

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.

seq is the run’s execution order — order by it and you have the path the run actually took. parent_task_name is set on the tool calls an agent node made, so they render indented under it, and tool_uri names the tool that ran.

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¶

node_type : str | None¶

prompt : str | None¶

errors : list[str] | None¶

duration_seconds : float | None¶

seq : int | None¶

parent_task_name : str | None¶

tool_uri : str | None¶

created_at : datetime¶

updated_at : datetime¶

async kavalai.backoffice.sessions.get_session_details(session: AsyncSession, session_id: UUID) → SessionDetails[source]¶: The transcript, runs and tasks of one session.

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

List sessions, most recently active first, with their counts.

The conversation list’s query, which is kavalai.agent_service.summarise_sessions() scoped to at most one agent. search matches message content; external_id matches the session’s caller-supplied key as a prefix. The two are separate because they answer different questions — “what did someone ask about” versus “show me this exact conversation”.

Database models¶

The ORM models below define the agent runtime schema. Their purpose and relationships are described in The database and data model.

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

kavalai.db.uuid_column()[source]¶

UUID column type.

Uses PostgreSQL’s native uuid on Postgres and SQLAlchemy’s generic Uuid on SQLite, so the same ORM models work against both the production Postgres backend and the SQLite/IndexedDB backend used under Pyodide.

kavalai.db.json_column()[source]¶: JSON column type: JSONB on Postgres, generic JSON on SQLite.

class kavalai.db.json_typeof(*args: _ColumnExpressionOrLiteralArgument[Any], **kwargs: Any)[source]¶

Bases: GenericFunction

Type name of a JSON value: array, object, null, …

Renders as jsonb_typeof on Postgres and json_type on SQLite, so a query over a json_column() reads the same against both. Written for the backoffice session list, which counts the tasks whose errors hold something.

type : TypeEngine[_T] = TEXT()¶: The SQL type of the result.

inherit_cache : bool | None = True¶: The compiled form depends on nothing but the arguments, so SQLAlchemy may cache it.

identifier : str = 'json_typeof'¶

name : str = 'json_typeof'¶

class kavalai.db.json_array_length(*args: _ColumnExpressionOrLiteralArgument[Any], **kwargs: Any)[source]¶

Bases: GenericFunction

Number of elements in a JSON array: jsonb_array_length on Postgres, json_array_length on SQLite.

type : TypeEngine[_T] = Integer()¶: The SQL type of the result.

inherit_cache : bool | None = True¶: The compiled form depends on nothing but the arguments, so SQLAlchemy may cache it.

identifier : str = 'json_array_length'¶

name : str = 'json_array_length'¶

kavalai.db.ensure_async_scheme(uri: str) → str[source]¶

Return uri with the async driver for its backend.

postgresql (any driver) becomes postgresql+asyncpg and sqlite becomes sqlite+aiosqlite; anything else is returned unchanged.

kavalai.db.is_sqlite_uri(uri: str) → bool[source]¶: Whether uri names a SQLite database (any driver spelling).

kavalai.db.sqlite_path_from_uri(uri: str) → str[source]¶: The file path a sqlite:///... URI points at (:memory: included).

kavalai.db.build_db_uri(user, password, host, port, db_name, scheme='postgresql+asyncpg') → str[source]¶: Builds a database URI from components.

exception kavalai.db.EngineOptionsConflictError[source]¶

Bases: ValueError

Raised when a cached engine is requested with conflicting options.

Engines are cached per (url, schema) and their options (pool_size, max_overflow, echo) are fixed by whichever call creates them. A later call that explicitly asks for different options would otherwise silently receive an engine configured differently than requested — raising surfaces the disagreeing call sites so they can be aligned.

class kavalai.db.DatabaseManager[source]¶

Bases: object

Manages dynamic engine creation and session factories.

SQLITE_PYODIDE_DB_PATH = '/kavalai/kavalai.db'¶

get_sessionmaker(*, user=None, password=None, host=None, port=None, db_name=None, uri=None, schema=None, echo=None, pool_size=None, max_overflow=None, read_only=False)[source]¶

Return an async_sessionmaker for the given database and schema.

A sqlite:// URI is served by get_sqlite_sessionmaker(): the file named by the URI, foreign keys on, one shared connection, and schema ignored because SQLite has none.

read_only=True returns an engine of its own whose connections refuse every write at the database: on PostgreSQL each connection is opened with default_transaction_read_only=on, so an INSERT, UPDATE, DELETE or DDL statement fails with “cannot execute … in a read-only transaction” whatever the role may do; on SQLite each connection sets PRAGMA query_only. The guarantee is the connection’s, not the caller’s discipline, which is what the backoffice needs when it is pointed at a production database.

schema selects the schema the ORM tables live in. The models are defined schema-less; the schema is applied per-engine via SQLAlchemy’s schema_translate_map, so the same models can target any schema at runtime. schema=None leaves table names unqualified (Postgres then resolves them via the connection’s search_path, i.e. public).

echo, pool_size and max_overflow are engine-level options, fixed by the call that first creates the engine for a given (url, schema) (unspecified options fall back to echo=False, pool_size=1, max_overflow=0). A later call may repeat the effective options or leave them unspecified; explicitly requesting different ones raises EngineOptionsConflictError instead of silently returning an engine configured otherwise.

get_sqlite_engine(*, db_path=None, echo=False, read_only=False)[source]¶

Return a cached async SQLite engine backed by db_path.

Requires greenlet (SQLAlchemy’s async engine dependency), so this is for CPython. Under pyodide use get_sqlite_sync_engine() instead.

db_path defaults to the IndexedDB-backed file under Pyodide and to an in-memory database otherwise. SQLite foreign keys are enabled. A StaticPool keeps a single shared connection so an in-memory database survives across sessions. read_only=True is a second engine on the same file whose connection has PRAGMA query_only set, so a write fails with “attempt to write a readonly database”.

get_sqlite_sessionmaker(*, db_path=None, echo=False, read_only=False)[source]¶: Return an async_sessionmaker bound to the async SQLite engine.

async init_sqlite(*, db_path=None, echo=False)[source]¶

Prepare the async SQLite/IndexedDB backend (mount, create tables, persist).

Stamps PRAGMA user_version; a database written by a different schema version is dropped and recreated (browser stores are ephemeral, there is no in-place upgrade path).

get_sqlite_sync_engine(*, db_path=None, echo=False)[source]¶

Return a cached sync SQLite engine backed by db_path.

This is the pyodide-friendly variant: SQLAlchemy’s synchronous engine does not need greenlet (which has no pyodide build). It behaves like get_sqlite_engine() otherwise – foreign keys enabled and a shared connection.

get_sqlite_sync_sessionmaker(*, db_path=None, echo=False)[source]¶: Return a sync sessionmaker bound to the sync SQLite engine.

async init_sqlite_sync(*, db_path=None, echo=False)[source]¶

Prepare the sync SQLite/IndexedDB backend (mount, create tables, persist).

Mounts IndexedDB (under Pyodide), creates all ORM tables and persists the resulting database. The table creation itself is synchronous; only the IndexedDB mount/flush are awaited. Safe to call repeatedly. Stamps PRAGMA user_version and drops a stale-version database (see init_sqlite()).

get_sqlite_compat_sessionmaker(*, db_path=':memory:', echo=False)[source]¶

Return an async-interface sessionmaker over the sync SQLite engine.

For platforms without greenlet (Pyodide/WASM), where neither SQLAlchemy’s async engine nor aiosqlite can run: sessions produced here expose the awaitable surface AgentService uses (execute/commit/refresh/flush/rollback) but run each call synchronously, briefly blocking the event loop — fine for local SQLite queries.

Bootstraps the schema on first use. db_path defaults to a private in-memory database; for a persistent Pyodide/IndexedDB store call init_sqlite_sync() first and pass its (default) file path.

async flush()[source]¶: Persist pending SQLite changes to IndexedDB (no-op outside Pyodide).

class kavalai.db.AsyncSessionShim(session)[source]¶

Bases: object

Async facade over a synchronous ORM Session.

Produced by DatabaseManager.get_sqlite_compat_sessionmaker() for greenlet-less platforms (Pyodide/WASM). The awaitable methods execute synchronously on the calling thread.

The surface has to match what AgentService and kavalai.crud actually call — a method missing here does not fail at import, it fails in the browser at the moment a user hits that code path (crud.get_one did exactly that). tests/test_agent_service.py runs the service against this shim as well as the real async session so the two cannot drift apart again.

add(obj)[source]¶

async execute(statement)[source]¶

async flush()[source]¶

async commit()[source]¶

async rollback()[source]¶

async refresh(obj)[source]¶

async get(entity, ident, **kwargs)[source]¶

async delete(obj)[source]¶

async merge(obj, **kwargs)[source]¶

async scalar(statement, *args, **kwargs)[source]¶

async scalars(statement, *args, **kwargs)[source]¶

async close()[source]¶

class kavalai.db.Agent(**kwargs)[source]¶

Bases: Base

A configured agent.

One row holds an agent’s definition: its name, optional description, the JSON input/output schemas it exposes, and the workflow document that drives its execution. Sessions, tasks and chat messages all reference the agent they belong to.

id : Mapped[UUID]¶

name : Mapped[str]¶

description : Mapped[str | None]¶

input_schema : Mapped[dict | None]¶

output_schema : Mapped[dict | None]¶

workflow : Mapped[dict | None]¶

created_at : Mapped[datetime]¶

updated_at : Mapped[datetime]¶

sessions : Mapped[list[Session]]¶

tasks : Mapped[list[Task]]¶

chat_messages : Mapped[list[ChatMessage]]¶

class kavalai.db.ModelCallStat(**kwargs)[source]¶

Bases: Base

Token usage and timing for a single LLM or embedding call.

Each row records one model call (call_type distinguishes e.g. chat completion from embedding): the model used, the request/response payloads, token counts, batch size and duration.

Cost is deliberately not recorded. Providers do not return it — only tokens — so any figure here would come from a price table the library would have to keep current, and cached input tokens are priced differently enough that a total derived from prompt_tokens alone would be wrong by a multiple. cached_prompt_tokens and reasoning_tokens are stored instead, which is what makes cost computable downstream (e.g. with genai-prices) from a row like this one.

agent_id, session_id and run_id attribute the call to the run that made it, so cost per run and per conversation is a query. None of them is a foreign key: a cost record outlives the conversation it belongs to. Purging a session blanks its calls’ payloads — the prompt is the conversation again — and keeps the token counts.

id : Mapped[UUID]¶

call_type : Mapped[str]¶

model : Mapped[str]¶

agent_id : Mapped[UUID | None]¶

session_id : Mapped[UUID | None]¶

run_id : Mapped[UUID | None]¶

request_data : Mapped[dict | None]¶

response_data : Mapped[dict | None]¶

response_code : Mapped[int | None]¶

prompt_tokens : Mapped[int | None]¶

completion_tokens : Mapped[int | None]¶

total_tokens : Mapped[int | None]¶

batch_size : Mapped[int | None]¶

duration_seconds : Mapped[float | None]¶

cached_prompt_tokens : Mapped[int | None]¶

reasoning_tokens : Mapped[int | None]¶

created_at : Mapped[datetime]¶

updated_at : Mapped[datetime]¶

class kavalai.db.Session(**kwargs)[source]¶

Bases: Base

A user conversation/session with an agent.

A session groups together everything exchanged with one agent over a conversation: its runs, tasks and chat messages. external_id lets a caller correlate the session with an identifier in their own system.

updated_at is the time of the session’s last run: every run moves it. Retention and the recent-first lists sort on it, which is why it is indexed.

id : Mapped[UUID]¶

agent_id : Mapped[UUID]¶

external_id : Mapped[str | None]¶

created_at : Mapped[datetime]¶

updated_at : Mapped[datetime]¶

agent : Mapped[Agent]¶

runs : Mapped[list[Run]]¶

tasks : Mapped[list[Task]]¶

chat_messages : Mapped[list[ChatMessage]]¶

class kavalai.db.Run(**kwargs)[source]¶

Bases: Base

One workflow run within a session.

A run captures a single invocation of the agent’s workflow: the input it was called with, the output it produced, the resolved run context and how long it took. Each run belongs to a session and owns the tasks executed during it.

duration_seconds is the wall-clock time of the invocation, from the moment the engine starts the run until it records the result — a success or a failure alike — measured with a monotonic clock. It is NULL while the run is in progress, for a run whose process died before it could record anything, and for every run written before this column existed. updated_at - created_at approximates it for those.

id : Mapped[UUID]¶

session_id : Mapped[UUID]¶

input_data : Mapped[dict | None]¶

output_data : Mapped[dict | None]¶

context : Mapped[dict | None]¶

duration_seconds : Mapped[float | None]¶

created_at : Mapped[datetime]¶

updated_at : Mapped[datetime]¶

session : Mapped[Session]¶

tasks : Mapped[list[Task]]¶

chat_messages : Mapped[list[ChatMessage]]¶

class kavalai.db.Task(**kwargs)[source]¶

Bases: Base

One workflow-node execution within a run.

A task records the execution of a single node in the workflow: its name and node type, the inputs it received, the output it produced, the prompt used (if any), any errors raised and how long it took. Tasks belong to a run (and its session) and provide the step-by-step trace of a run.

id : Mapped[UUID]¶

agent_id : Mapped[UUID | None]¶

session_id : Mapped[UUID]¶

run_id : Mapped[UUID]¶

inputs : Mapped[dict | None]¶

output : Mapped[dict | None]¶

name : Mapped[str | None]¶

node_type : Mapped[str | None]¶

prompt : Mapped[str | None]¶

errors : Mapped[list[str] | None]¶

duration_seconds : Mapped[float | None]¶

seq : Mapped[int | None]¶

parent_task_name : Mapped[str | None]¶

tool_uri : Mapped[str | None]¶

created_at : Mapped[datetime]¶

updated_at : Mapped[datetime]¶

agent : Mapped[Agent]¶

session : Mapped[Session]¶

run : Mapped[Run]¶

class kavalai.db.ChatMessage(**kwargs)[source]¶

Bases: Base

One message in a session’s chat history.

Each row is a single message (role such as user or assistant with its text content) belonging to a session. It optionally references the run that produced it, forming the ordered conversation transcript.

id : Mapped[UUID]¶

agent_id : Mapped[UUID]¶

session_id : Mapped[UUID]¶

run_id : Mapped[UUID | None]¶

role : Mapped[str]¶

content : Mapped[str]¶

created_at : Mapped[datetime]¶

updated_at : Mapped[datetime]¶

agent : Mapped[Agent]¶

session : Mapped[Session]¶

run : Mapped[Run]¶

Remote agent client¶

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

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¶

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

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. Kaval.AI mints these, and backends must be able to store one: a store whose identifiers are integers or bounded strings has to accept a UUID (as text if necessary) and give the same value back.
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¶: How close the item is to the query, higher is better, comparable across backends for the same embedding model. Backends index with cosine distance and report 1.0 - distance, so a perfect match scores 1.0. A backend whose store reports a metric-dependent score must normalise to this convention rather than passing its own number through.
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. The interface is generic; the implementations are deliberately not. PostgresRagService pushes similarity into pgvector and carries a dozen methods this interface never mentions, and that is the intended shape — a backend is free to expose whatever its store does well, as long as the declared surface below behaves identically everywhere.

The surface comes in three tiers:

Required. index(), index_batch(), query(), query_batch(), delete() and delete_by_source_id() are abstract. Every backend implements all six.

Optional. count_entries(), iter_entries(), delete_by_metadata() and replace() raise NotImplementedError by default. A backend that overrides them lists their names in capabilities; ask supports() before calling them rather than catching the exception.

Defaulted. compute_similarity_matrix(), learn_normalizer() and delete_many() work on every backend as written, and may be overridden with an exact or cheaper implementation.

source_ids means the same on every method that takes it: None does not filter, and an empty list matches nothing. The distinction is a safety property — a filter computed to be empty (a tenant with no sources, a permission check that passed nothing) must not widen into a search of everything.

The SQL backends additionally take min_similarity and stats_receiver on their query and index methods; a rag_query node passes the run’s receiver to any service whose query accepts one.

tests/rag/test_conformance.py runs the declared contract against every backend, so “implements the interface correctly” is checked rather than assumed.

capabilities : frozenset = frozenset({})¶

supports(capability: str) → bool[source]¶

Whether an optional method is implemented by this backend.

count_entries and iter_entries are optional, so a caller that needs one (the backoffice embedding projector needs iter_entries) should ask rather than catch NotImplementedError.

Parameters:¶

capability: str¶: An optional method name, e.g. "iter_entries".

Returns:¶

True if this backend implements it.

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') → dict[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 entry. It always carries id, model, collection_name, source_id, content, embedding_size, rag_metadata, created_at and updated_at; a backend may add its own keys (Postgres also returns embedding). The guaranteed keys are part of the contract, not a backend detail.

Return type:¶

dict

abstractmethod 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.

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:¶

The created entries, each shaped as in index().

Return type:¶

list[dict]

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, include_content: bool = True, match: dict | None = None) → 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]]¶: None searches every source; a list restricts the search to those sources, and an empty list returns no results.
match : Optional[dict]¶: Return only items whose metadata has every key of match equal — top-level keys and scalar values, as delete_by_metadata() takes them. None does not filter. The filter is part of the nearest-neighbour search, not applied to its result, so top_k matching items come back when the collection holds that many.
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. Backends whose store cannot group server-side over-fetch and reduce in the client, which makes the result approximate when a source has more indexed items than the over-fetch window — the same caveat compute_similarity_matrix() carries.
include_content : bool¶: When False, content is omitted from the results. Everything else is unchanged. Free on a local database; on a remote store it avoids transferring text that the caller is only going to discard.

Returns:¶

Results ordered by descending similarity.

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, include_content: bool = True, match: dict | 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]]¶: As in query(); an empty list returns an empty result list for every text.
include_content : bool¶: When False, content is omitted from the results. See query().
match : Optional[dict]¶: Metadata equality filter, as in query(); it applies to every text.

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 delete_many(ids: list[UUID], collection_name: str | None = None) → None[source]¶

Delete several indexed items by identifier.

The default deletes them one at a time through delete(); a backend overrides it with one statement.

Parameters:¶

ids : list[UUID]¶: Identifiers of the items to delete.
collection_name : Optional[str]¶: Collection the items belong to, as in delete().

async delete_by_metadata(collection_name: str, match: dict) → None[source]¶

Delete the items whose metadata has every key of match equal.

match holds top-level keys and scalar values (strings, numbers, booleans). Optional: check supports("delete_by_metadata").

async replace(collection_name: str, texts: list[str], metadata_list: list[dict], source_ids: list[str] | None = None, *, match: dict | None = None, source_id: str | None = None) → list[dict][source]¶

Replace the items selected by match or source_id atomically.

The new texts are embedded before anything is deleted, and the delete and the insert commit together, so a failure leaves the old items in place. Optional: check supports("replace").

Returns:¶: The created entries, shaped as in index().
Return type:¶: list[dict]

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

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

PostgreSQL (pgvector) backed RAG service.

This backend owns its schema; no Alembic migration set covers it. The storage model is the one described in kavalai.rag.collections: a rag_collections registry plus one table per collection with a typed vector(N) or halfvec(N) column, an HNSW index for the collection’s exact dimension, and a GIN index on the metadata column. Dropping a collection is DROP TABLE. With provision=False the service issues no DDL and the tables are created by ensure_registry() and create_collection() from a role that may.

All SQL here is raw and therefore bypasses schema_translate_map — every statement qualifies the configured schema explicitly, and the vector types are qualified with public, where the extension is created.

HALFVEC_MIN_VERSION and ITERATIVE_SCAN_MIN_VERSION are the pgvector releases that introduced halfvec and iterative index scans. _COLLECTION_UPGRADES maps a from_version to an async callable (session, service, collection_info) -> None that brings a collection from from_version to from_version + 1; see RAG_COLLECTION_SCHEMA_VERSION in kavalai.rag.collections.

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

Bases: object

Registry entry for one RAG collection.

vector_type is the column type of the collection’s embeddings — "vector" (32-bit floats) or, on PostgreSQL, "halfvec" (16-bit). It is read from the table itself rather than stored in the registry.

class kavalai.rag.postgres.PostgresRagService(session_maker: async_sessionmaker[AsyncSession] | Callable[[], AbstractAsyncContextManager[AsyncSession]], model: str | None = None, *, normalizer: Normalizer | None = None, schema: str | None = None, provision: bool = True, stats_receiver: Any = None, vector_type: str = 'vector')[source]¶

Bases: CollectionRagService

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

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

vector_type="halfvec" stores new collections’ embeddings as 16-bit floats: half the memory, and HNSW indexes for up to 4,000 dimensions instead of 2,000, so a 3,072-dimension model can be indexed at all. It needs pgvector 0.7 or later. An existing collection keeps the type it was created with.

A query restricted by source_ids or match enables pgvector’s iterative index scans (0.8 or later) for its transaction. Without them an HNSW search filters its candidates after the search, so a filter whose rows lie away from the query can return fewer than top_k rows, or none. match is metadata @> :match, the containment the collection’s GIN index serves.

collection_upgrades : dict[int, Callable] = {}¶

vector_types : frozenset = frozenset({'halfvec', 'vector'})¶

classmethod from_uri(uri: str, model: str | None = None, *, schema: str | None = None, **options: Any) → PostgresRagService[source]¶

Create a PostgresRagService from a database URI.

options are the constructor’s keyword arguments.

classmethod from_session_maker(session_maker: async_sessionmaker[AsyncSession], model: str | None = None, **options: Any) → PostgresRagService[source]¶: Create a PostgresRagService from a session maker.

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.

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

SQLite (sqlite-vector) backed RAG service.

The storage model is the one described in kavalai.rag.collections — a rag_collections registry and one table per collection — in a single ordinary SQLite file. Ids are TEXT UUIDs, metadata is JSON text, timestamps are ISO-8601 text and embeddings are FLOAT32 blobs (vector_as_f32) searched through the sqlite-vector extension under cosine distance, so similarity scores match the Postgres backend (similarity = 1 - distance).

_COLLECTION_UPGRADES maps a from_version to an async callable (connection, service, collection_info) -> None; see RAG_COLLECTION_SCHEMA_VERSION in kavalai.rag.collections.

class kavalai.rag.sqllite.SqliteRagService(filename: str, model: str | None = None, auto_create: bool = True, normalizer: Normalizer | None = None, *, provision: bool = True, stats_receiver: Any = None, vector_type: str = 'vector', read_only: bool = False)[source]¶

Bases: CollectionRagService

SQLite backed RAG service using the sqlite-vector extension (https://github.com/sqliteai/sqlite-vector).

No server: the whole index — registry and collection tables — lives in one file that can be built offline and copied wherever it is needed, ":memory:" included. Each collection carries its own embedding dimension, as on Postgres.

All methods run the SQLite work synchronously on the calling loop — no worker threads — so the service also works under Pyodide / WebAssembly. Embeddings are stored as 32-bit floats; halfvec is a PostgreSQL option.

collection_upgrades : dict[int, Callable] = {}¶

close() → None[source]¶: Close the underlying SQLite connection.