Skip to content

Commit ec0cb34

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 e6372d7 commit ec0cb34

File tree

1 file changed

+90
-0
lines changed

1 file changed

+90
-0
lines changed
Lines changed: 90 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,90 @@
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+
"""CREATE TABLE IF NOT EXISTS models (
40+
id TEXT PRIMARY KEY, -- UUID stored as TEXT
41+
provider TEXT NOT NULL,
42+
name TEXT NOT NULL UNIQUE,
43+
description TEXT NOT NULL DEFAULT '',
44+
endpoint TEXT NOT NULL DEFAULT '',
45+
auth_type TEXT NOT NULL DEFAULT 'none',
46+
auth_blob TEXT NOT NULL DEFAULT '',
47+
created_at TIMESTAMP NOT NULL DEFAULT NOW(),
48+
updated_at TIMESTAMP NOT NULL DEFAULT NOW()
49+
);"""
50+
)
51+
52+
# This table is used to store the Muxing configuration. The
53+
# `destination_model_id` field is used to reference the model that the
54+
# Muxing configuration is for.
55+
# The `matcher_type` field is used to determine the type of the
56+
# matcher that is used in the Muxing configuration. e.g. `file_glob` would
57+
# be a matcher that uses file globbing to match files if a file is
58+
# detected in the prompt. The `matcher_blob` field is used to store the
59+
# configuration for the matcher, which can be a JSON object or a string,
60+
# depending on the `matcher_type`. On an initial implementation, the
61+
# `matcher_blob` field will simply be a string that is used to match
62+
# the prompt file name (if a file is detected in the prompt).
63+
# The `priority` field is used to determine the priority of the Muxing
64+
# configuration. The lower the number, the higher the priority. Note that
65+
# prompts will be matched against the Muxing configurations in ascending
66+
# order of priority.
67+
op.execute(
68+
"""CREATE TABLE IF NOT EXISTS muxes (
69+
id TEXT PRIMARY KEY, -- UUID stored as TEXT
70+
destination_model_id TEXT NOT NULL REFERENCES models(id) ON DELETE CASCADE,
71+
workspace_id TEXT NOT NULL REFERENCES workspaces(id) ON DELETE CASCADE,
72+
matcher_type TEXT NOT NULL,
73+
matcher_blob TEXT NOT NULL,
74+
priority INTEGER NOT NULL,
75+
created_at TIMESTAMP NOT NULL DEFAULT NOW(),
76+
updated_at TIMESTAMP NOT NULL DEFAULT NOW()
77+
);"""
78+
)
79+
80+
# In terms of access patterns, the `muxes` table will be queried
81+
# to find the Muxing configuration for a given prompt. On initial search,
82+
# the `muxes` table will be queried by the `workspace_id`.
83+
op.execute("CREATE INDEX IF NOT EXISTS idx_muxes_workspace_id ON muxes (workspace_id);")
84+
85+
86+
def downgrade() -> None:
87+
with context.begin_transaction():
88+
op.execute("DROP INDEX IF EXISTS idx_muxes_workspace_id;")
89+
op.execute("DROP TABLE IF EXISTS muxes;")
90+
op.execute("DROP TABLE IF EXISTS models;")

0 commit comments

Comments
 (0)