Skip to content

Commit 39cf5b7

Browse files
committed
Add database model for model muxing
This introduces the `models` and `muxes` database tables to store models with their configuration as well as muxing configurations that allow us to appropriately route prompts to the relevant models. Closes: #752 Signed-off-by: Juan Antonio Osorio <[email protected]>
1 parent 0a97a60 commit 39cf5b7

File tree

1 file changed

+94
-0
lines changed

1 file changed

+94
-0
lines changed
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,94 @@
1+
"""add llm and mux
2+
3+
Revision ID: 0f9b8edc8e46
4+
Revises: 90d5471db49a
5+
Create Date: 2025-01-24 07:58:34.907908+00:00
6+
7+
"""
8+
9+
from typing import Sequence, Union
10+
11+
from alembic import context, op
12+
13+
# revision identifiers, used by Alembic.
14+
revision: str = "0f9b8edc8e46"
15+
down_revision: Union[str, None] = "90d5471db49a"
16+
branch_labels: Union[str, Sequence[str], None] = None
17+
depends_on: Union[str, Sequence[str], None] = None
18+
19+
20+
def upgrade() -> None:
21+
with context.begin_transaction():
22+
# This table is used to store the models that are available
23+
# for references, e.g. in Muxing. The `auth_blob` field is
24+
# used to store the credentials for the model, which can be
25+
# a JSON object or a string, depending on the `auth_type`.
26+
# The `auth_type` field is used to determine how to interpret
27+
# the `auth_blob` field. If `auth_type` is `none`, then the
28+
# `auth_blob` field is ignored.
29+
# The `endpoint` field is used to store the endpoint of the
30+
# model.
31+
# NOTE: This resource is not namespaced by a workspace; that is
32+
# because the models are shared across workspaces.
33+
# NOTE: The lack of `deleted_at` is intentional. This resource
34+
# is not soft-deleted.
35+
# TODO: Do we need a display name here? An option is to
36+
# use the `name` field as the display name and normalize
37+
# the `name` field to be a slug when used as a reference.
38+
op.execute(
39+
"""
40+
CREATE TABLE IF NOT EXISTS models (
41+
id TEXT PRIMARY KEY, -- UUID stored as TEXT
42+
provider TEXT NOT NULL,
43+
name TEXT NOT NULL UNIQUE,
44+
description TEXT NOT NULL DEFAULT '',
45+
endpoint TEXT NOT NULL DEFAULT '',
46+
auth_type TEXT NOT NULL DEFAULT 'none',
47+
auth_blob TEXT NOT NULL DEFAULT '',
48+
created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP,
49+
updated_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP
50+
);
51+
"""
52+
)
53+
54+
# This table is used to store the Muxing configuration. The
55+
# `destination_model_id` field is used to reference the model that the
56+
# Muxing configuration is for.
57+
# The `matcher_type` field is used to determine the type of the
58+
# matcher that is used in the Muxing configuration. e.g. `file_glob` would
59+
# be a matcher that uses file globbing to match files if a file is
60+
# detected in the prompt. The `matcher_blob` field is used to store the
61+
# configuration for the matcher, which can be a JSON object or a string,
62+
# depending on the `matcher_type`. On an initial implementation, the
63+
# `matcher_blob` field will simply be a string that is used to match
64+
# the prompt file name (if a file is detected in the prompt).
65+
# The `priority` field is used to determine the priority of the Muxing
66+
# configuration. The lower the number, the higher the priority. Note that
67+
# prompts will be matched against the Muxing configurations in ascending
68+
# order of priority.
69+
op.execute(
70+
"""
71+
CREATE TABLE IF NOT EXISTS muxes (
72+
id TEXT PRIMARY KEY, -- UUID stored as TEXT
73+
destination_model_id TEXT NOT NULL REFERENCES models(id) ON DELETE CASCADE,
74+
workspace_id TEXT NOT NULL REFERENCES workspaces(id) ON DELETE CASCADE,
75+
matcher_type TEXT NOT NULL,
76+
matcher_blob TEXT NOT NULL,
77+
priority INTEGER NOT NULL,
78+
created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP,
79+
updated_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP
80+
);
81+
"""
82+
)
83+
84+
# In terms of access patterns, the `muxes` table will be queried
85+
# to find the Muxing configuration for a given prompt. On initial search,
86+
# the `muxes` table will be queried by the `workspace_id`.
87+
op.execute("CREATE INDEX IF NOT EXISTS idx_muxes_workspace_id ON muxes (workspace_id);")
88+
89+
90+
def downgrade() -> None:
91+
with context.begin_transaction():
92+
op.execute("DROP INDEX IF EXISTS idx_muxes_workspace_id;")
93+
op.execute("DROP TABLE IF EXISTS muxes;")
94+
op.execute("DROP TABLE IF EXISTS models;")

0 commit comments

Comments
 (0)