diff --git a/litellm/proxy/db/sqlmodel/README.md b/litellm/proxy/db/sqlmodel/README.md new file mode 100644 index 000000000000..dff718236b3d --- /dev/null +++ b/litellm/proxy/db/sqlmodel/README.md @@ -0,0 +1,116 @@ +# Prisma -> SQLModel migration -- Phase 1 + +This package is the foundation for migrating `litellm`'s proxy persistence +layer from Prisma (`prisma-client-py==0.11.0`) to SQLModel/SQLAlchemy. The +migration is multi-phase by necessity -- the proxy has ~1,680 Prisma client +call sites across ~147 production files plus ~177 test files, so any +"big-bang" cutover would be unreviewable and unsafe. + +## What this Phase ships + +| Artefact | Purpose | +|---|---| +| `schema_parser.py` | Tiny pure-Python parser for the subset of Prisma DSL actually used by `schema.prisma`. Used by the parity test and the generator. | +| `_generate.py` | Code generator that emits SQLModel class definitions from a parsed schema. Run it manually after schema changes. | +| `models.py` | SQLModel classes for **all 64 models** in `schema.prisma`. Hand-editable; the generator only seeds the file. | +| `tests/test_litellm/proxy/db/sqlmodel_orm/test_schema_parser.py` | Unit tests for the parser. | +| `tests/test_litellm/proxy/db/sqlmodel_orm/test_parity.py` | Parity test: every Prisma model has a matching SQLModel class with matching columns, nullability, primary keys, uniques, and indexes. Also asserts that the committed `models.py` is byte-identical to a fresh generator run. | + +**Nothing in this package is wired into the runtime proxy yet.** Importing +the module has no effect on the existing Prisma-backed code paths -- the +generated classes simply sit alongside the Prisma client and are guarded +against drift by CI. + +## Why a generator at all? + +The schema is the source of truth and changes frequently. Hand-writing 64 +SQLModel classes against a 1,378-line Prisma schema invites typos and +silent drift. The generator gives us one well-tested translation rule per +Prisma construct (`@id`, `@@index`, `String[]`, `@updatedAt`, etc.) and the +parity test catches any regression in either the schema or the generator. + +When subsequent phases need to add SQLAlchemy-only behaviour (custom +relationships, hybrid properties, `Mapped[...]` annotations, etc.), edit +`models.py` by hand. The generator's output should still load and the +parity test should still pass; if they don't, the schema and the SQLModel +layer have diverged. + +## Re-running the generator + +```bash +uv run python -m litellm.proxy.db.sqlmodel._generate \ + --schema schema.prisma \ + --out litellm/proxy/db/sqlmodel/models.py +``` + +The parity test fails CI if a schema change isn't accompanied by a +regenerated `models.py`. + +## Subsequent phases + +The work below is the responsibility of follow-up PRs, in roughly this +order. Each phase is independently testable; do not bundle them. + +1. **Session abstraction.** Introduce a thin `DBSession` interface that + wraps the existing `prisma_client` today and a SQLAlchemy + `AsyncSession` tomorrow. Land with zero behaviour change. This is the + prerequisite for incrementally swapping call sites. +2. **CI: keep `schema.prisma` and `models.py` in sync.** Add a workflow + that runs the parity test on every PR (the test already exists -- this + step is just enabling it as a required check). +3. **Port the raw-SQL hotspots.** ~288 `query_raw` / `execute_raw` calls + across ~37 files (concentrated in `spend_management_endpoints.py`, + `db/create_views.py`, focus/cloudzero exporters). These are the + easiest call sites to migrate -- the SQL is already there; we just + swap the executor to a SQLAlchemy `session.execute(text(...))`. +4. **Migrate per-table call sites.** ~55 tables touched across ~1,680 + Prisma-client call sites. Parallelise by feature area + (keys/teams/users -> spend/logs -> MCP/managed objects -> + adaptive router/workflows). The session abstraction from phase 1 lets + each call site flip independently. +5. **Replace the custom Prisma reliability layer.** The current + `PrismaWrapper` (RDS IAM token rotation), `RoutingPrismaWrapper` + (read/write split), and `PrismaDBExceptionHandler` (~10 distinct + error type classifications) all need SQLAlchemy-native equivalents. +6. **Swap migrations to Alembic.** The current `litellm-proxy-extras` + package bundles 123 Prisma migration files. Establish an Alembic + baseline matching the live schema, with a documented "first-run + after upgrade" path for existing deployments. The 10 / 123 + migrations that contain DML need careful translation; the rest are + pure DDL and can be folded into the baseline for fresh installs. +7. **Tear out Prisma.** Remove `prisma==0.11.0` from `pyproject.toml`, + `prisma generate` from all 7 Dockerfiles, the CI workflows that run + it, the 3 `schema.prisma` copies (with their `check-schema-sync` and + `sync-schema` workflows), and the `litellm-proxy-extras` migration + bundle. + +## Risks and gotchas surfaced during Phase 1 + +* **Reserved attribute names.** SQLModel/SQLAlchemy reserve `metadata` + and `registry` on the mapped class. Several Prisma models have a + `metadata Json` column. The generator emits these as Python attribute + `metadata_` while keeping the on-disk column name `metadata` via + `sa_column_kwargs={'name': 'metadata'}`. Migration of call sites must + use `MyTable.metadata_` in Python. +* **`String[]` (Postgres array columns).** Prisma maps `String[]` to a + Postgres `text[]`. The generator uses + `sqlalchemy.dialects.postgresql.ARRAY(Text())`, which is + Postgres-specific. SQLite-backed test environments will need a + separate fixture path -- this is identical to the current Prisma + situation (`prisma-client-py` on SQLite already requires manual JSON + emulation). +* **`Json` columns default-text quoting.** Prisma's `@default("[]")` and + `@default("{}")` emit `'[]'::jsonb` / `'{}'::jsonb` as the Postgres + `DEFAULT`. The generator preserves both the Python `default_factory` + *and* the `server_default` so migrated rows behave identically when + the column is omitted from an INSERT. +* **`@updatedAt`.** Prisma updates the column from the client. The + generator translates this to a SQLAlchemy `onupdate=lambda: ... + utcnow()` so the behaviour persists when ported off Prisma. +* **`cuid()`.** Only `LiteLLM_CronJob.cronjob_id` uses it. The generator + treats it as opaque-string-equivalent to `uuid()` (which is what every + consumer already assumes). +* **Enums.** The single Prisma enum (`JobStatus`) is emitted as a Python + `str`-Enum and the column is stored as `Text` to match what + `prisma-client-py` already does on Postgres. A real `sa.Enum` can be + introduced later if any call site benefits. diff --git a/litellm/proxy/db/sqlmodel/__init__.py b/litellm/proxy/db/sqlmodel/__init__.py new file mode 100644 index 000000000000..843fd314639d --- /dev/null +++ b/litellm/proxy/db/sqlmodel/__init__.py @@ -0,0 +1,30 @@ +"""SQLModel-based ORM definitions for the LiteLLM proxy database. + +This package is the foundation for migrating proxy persistence from Prisma to +SQLModel/SQLAlchemy. **Phase 1** (this module's current state) introduces: + +* :mod:`schema_parser` -- a small ``schema.prisma`` parser used by the + parity test (and by future code generators). +* :mod:`models` -- hand-maintained, generator-seeded SQLModel classes that + mirror every model in the canonical ``schema.prisma``. +* A parity test (in ``tests/test_litellm/proxy/db/sqlmodel/``) that fails + CI if the SQLModel definitions drift from the Prisma schema. + +Nothing in this package is wired into the runtime proxy yet -- importing it +has no effect on existing Prisma-backed code paths. Subsequent phases will: + +1. introduce a ``DBSession`` abstraction wrapping Prisma today and SQLAlchemy + tomorrow, +2. port raw-SQL call sites (``query_raw`` / ``execute_raw``) onto SQLAlchemy, +3. migrate per-table call sites (~55 tables) behind the abstraction, +4. rebuild the ``PrismaWrapper`` / ``RoutingPrismaWrapper`` / exception + classifier as SQLAlchemy-native components, +5. swap the migration tool from Prisma to Alembic with a baseline derived + from the current schema state. + +See ``litellm/proxy/db/sqlmodel/README.md`` for the full plan. +""" + +from litellm.proxy.db.sqlmodel.models import ALL_MODELS + +__all__ = ["ALL_MODELS"] diff --git a/litellm/proxy/db/sqlmodel/_generate.py b/litellm/proxy/db/sqlmodel/_generate.py new file mode 100644 index 000000000000..9f999058b1eb --- /dev/null +++ b/litellm/proxy/db/sqlmodel/_generate.py @@ -0,0 +1,515 @@ +"""Generator: emit SQLModel class definitions from ``schema.prisma``. + +This is a developer tool, not runtime code. Re-run after schema changes +(or rely on the parity test to flag drift) and copy the output into +:mod:`litellm.proxy.db.sqlmodel.models`. The output is plain Python that +should be reviewed and committed by hand -- this generator is here for +correctness, not for automatic codegen at import time. + +Usage:: + + uv run python -m litellm.proxy.db.sqlmodel._generate \\ + --schema schema.prisma \\ + --out litellm/proxy/db/sqlmodel/models.py + +The generated file is structurally equivalent to ``schema.prisma`` (every +model becomes a SQLModel class with one column per scalar field, plus +table-level constraints and indexes). It does **not** model relations - +those will be added by hand in subsequent migration phases as needed. +""" + +from __future__ import annotations + +import argparse +import sys +from pathlib import Path +from typing import Dict, Iterable, List, Optional, Set, Tuple + +from litellm.proxy.db.sqlmodel.schema_parser import ( + PrismaEnum, + PrismaField, + PrismaModel, + PrismaSchema, + parse_schema_file, +) + +# Map Prisma scalar -> (python annotation, SQLAlchemy column type expression). +# We deliberately use SQLAlchemy types (not SQLModel sugar) to match what the +# Prisma migrations have shipped historically: BigInt -> BigInteger, +# String -> Text (Prisma's default for `String` is unbounded text on Postgres), +# Json -> JSONB, etc. +_SCALAR_PY_TYPE = { + "String": "str", + "Int": "int", + "BigInt": "int", + "Float": "float", + "Decimal": "Decimal", + "Boolean": "bool", + "DateTime": "datetime", + "Json": "Any", + "Bytes": "bytes", +} + +_SCALAR_SA_TYPE = { + "String": "Text()", + "Int": "Integer()", + "BigInt": "BigInteger()", + "Float": "Double()", + "Decimal": "Numeric()", + "Boolean": "Boolean()", + "DateTime": "DateTime(timezone=True)", + "Json": "JSONB()", + "Bytes": "LargeBinary()", +} + +# Default expressions for typical Prisma defaults. Returned strings are +# Python source that produces a SQLAlchemy ``Column(... default=..., server_default=...)`` +# argument value. We bias toward server defaults for ``now()``, scalar defaults +# for booleans/numbers, and Python factories for ``uuid()``/``cuid()`` (so the +# value is set at INSERT time, matching prisma-client-py behavior). + + +def _python_default_for(default_raw: str, base_type: str) -> Optional[str]: + """Return Python ``default=`` argument source, or ``None``.""" + raw = default_raw.strip() + if raw == "uuid()": + return "default_factory=lambda: str(__import__('uuid').uuid4())" + if raw == "cuid()": + # cuid is roughly equivalent to uuid for our purposes; the only + # current user is ``LiteLLM_CronJob.cronjob_id`` and downstream + # consumers treat it as an opaque string. + return "default_factory=lambda: str(__import__('uuid').uuid4())" + if raw == "now()": + return "default_factory=lambda: __import__('datetime').datetime.utcnow()" + if raw == "true": + return "default=True" + if raw == "false": + return "default=False" + if raw == "[]": + return "default_factory=list" + if raw == '"{}"': + return "default_factory=dict" + if raw == '"[]"': + return "default_factory=list" + # Quoted string literal + if raw.startswith('"') and raw.endswith('"'): + return f"default={raw}" + # Numeric literal + try: + float(raw) + return f"default={raw}" + except ValueError: + pass + # Enum reference (e.g. JobStatus value: INACTIVE) + if ( + raw.replace("_", "").isalnum() + and raw[:1].isalpha() + and base_type not in _SCALAR_PY_TYPE + ): + # we can't reference the Python enum here without an import wrangle, + # so fall back to a string default. + return f'default="{raw}"' + return None + + +def _server_default_for(default_raw: str, base_type: str) -> Optional[str]: + """Optional ``server_default=`` to match the existing Postgres DDL.""" + raw = default_raw.strip() + if raw == "now()": + return "server_default=text('CURRENT_TIMESTAMP')" + if base_type == "Json": + if raw == '"{}"': + return "server_default=text(\"'{}'\")" + if raw == '"[]"': + return "server_default=text(\"'[]'\")" + return None + + +def _sa_type_for(field: PrismaField, schema: PrismaSchema) -> str: + """SQLAlchemy column type expression for a Prisma field.""" + if field.is_list: + inner = _SCALAR_SA_TYPE.get(field.base_type, "Text()") + return f"ARRAY({inner})" + if field.base_type in _SCALAR_SA_TYPE: + return _SCALAR_SA_TYPE[field.base_type] + if field.base_type in schema.enums: + # Use a plain Text column; we already index/filter these as strings + # everywhere in production and Prisma's enum type is mostly a + # client-side affair. (Subsequent phases can introduce a real + # ``sa.Enum`` if the call sites benefit from it.) + return "Text()" + return "Text()" + + +def _py_type_for(field: PrismaField, schema: PrismaSchema) -> str: + if field.base_type in _SCALAR_PY_TYPE: + py = _SCALAR_PY_TYPE[field.base_type] + elif field.base_type in schema.enums: + py = "str" + else: + py = "str" + if field.is_list: + py = f"List[{py}]" + if field.is_optional: + py = f"Optional[{py}]" + return py + + +# Names SQLAlchemy's Declarative API reserves on a mapped class. +# When a Prisma column collides with one of these we emit the Python attribute +# with a trailing underscore but keep the on-disk column name unchanged via +# ``sa_column_kwargs={'name': '...'}``. +_RESERVED_PY_ATTRS: Set[str] = {"metadata", "registry"} + + +def _format_field(field: PrismaField, schema: PrismaSchema) -> str: + """Render one ``Foo: = Field(...)`` line for a SQLModel class.""" + py_type = _py_type_for(field, schema) + sa_type = _sa_type_for(field, schema) + + field_kwargs: List[str] = [f"sa_type={sa_type}"] + sa_column_kwargs: List[str] = [] + + py_attr_name = field.name + if field.name in _RESERVED_PY_ATTRS: + py_attr_name = f"{field.name}_" + + if field.column_name != py_attr_name: + sa_column_kwargs.append(f"'name': {field.column_name!r}") + if field.is_id: + field_kwargs.append("primary_key=True") + if field.is_unique and not field.is_id: + field_kwargs.append("unique=True") + + py_default: Optional[str] = None + srv_default: Optional[str] = None + if field.has_default and field.default_raw is not None: + py_default = _python_default_for(field.default_raw, field.base_type) + srv_default = _server_default_for(field.default_raw, field.base_type) + + if py_default is not None: + field_kwargs.append(py_default) + elif field.is_optional: + field_kwargs.append("default=None") + elif field.is_list: + field_kwargs.append("default_factory=list") + + if srv_default is not None: + # ``server_default`` lives on the SA column, not on the SQLModel Field. + # _server_default_for returns ``server_default=text('...')``; rip the + # value off and stuff it into sa_column_kwargs so SQLModel forwards it. + value = srv_default.split("=", 1)[1] + sa_column_kwargs.append(f"'server_default': {value}") + + if field.has_updated_at: + sa_column_kwargs.append( + "'onupdate': lambda: __import__('datetime').datetime.utcnow()" + ) + + if sa_column_kwargs: + joined = ", ".join(sa_column_kwargs) + field_kwargs.append(f"sa_column_kwargs={{{joined}}}") + + field_args = ", ".join(field_kwargs) + return f" {py_attr_name}: {py_type} = Field({field_args})" + + +def _format_index_args(model: PrismaModel) -> List[str]: + args: List[str] = [] + composite_pk: Tuple[str, ...] = ( + model.primary_key if len(model.primary_key) > 1 else () + ) + if composite_pk: + cols = ", ".join(repr(c) for c in composite_pk) + args.append(f"PrimaryKeyConstraint({cols})") + for u in model.uniques: + cols = ", ".join(repr(c) for c in u.fields) + args.append(f"UniqueConstraint({cols})") + for idx in model.indexes: + cols = ", ".join(repr(c) for c in idx.fields) + if idx.map_name: + args.append(f"Index({idx.map_name!r}, {cols})") + else: + # Default index name: ___idx (matches the + # convention Prisma generates so existing DBs stay happy). + default_name = f"{model.table_name}_" + "_".join(idx.fields) + "_idx" + args.append(f"Index({default_name!r}, {cols})") + return args + + +def _model_class_name(model: PrismaModel) -> str: + """Map ``LiteLLM_FooTable`` -> ``LiteLLMFooTable`` (CamelCase, no underscores).""" + parts = model.name.split("_") + return "".join(p[:1].upper() + p[1:] for p in parts if p) + + +def _render_model_class(model: PrismaModel, schema: PrismaSchema) -> str: + cls_name = _model_class_name(model) + lines: List[str] = [] + lines.append(f"class {cls_name}(SQLModel, table=True):") + lines.append(f" __tablename__ = {model.table_name!r}") + index_args = _format_index_args(model) + if index_args: + if len(index_args) == 1: + lines.append(f" __table_args__ = ({index_args[0]},)") + else: + lines.append(" __table_args__ = (") + for arg in index_args: + lines.append(f" {arg},") + lines.append(" )") + lines.append("") + for field in model.fields: + lines.append(_format_field(field, schema)) + lines.append("") + return "\n".join(lines) + + +def _render_enum(enum: PrismaEnum) -> str: + lines = [f"class {enum.name}(str, Enum):"] + for v in enum.values: + lines.append(f" {v} = {v!r}") + lines.append("") + return "\n".join(lines) + + +_DOCSTRING = '''"""SQLModel ORM definitions mirroring ``schema.prisma``. + +THIS FILE IS GENERATED by ``litellm.proxy.db.sqlmodel._generate`` but is +CHECKED IN as ordinary Python source. Hand-edits are allowed -- the parity +test in ``tests/test_litellm/proxy/db/sqlmodel_orm/test_parity.py`` will +fail CI if structural drift from ``schema.prisma`` is introduced (in +either direction). + +Re-generate with:: + + uv run python -m litellm.proxy.db.sqlmodel._generate \\ + --schema schema.prisma \\ + --out litellm/proxy/db/sqlmodel/models.py + +Phase 1 of the Prisma -> SQLModel migration only ships these definitions; +nothing in the runtime proxy currently imports them. See the package +README for the multi-phase plan. +"""''' + + +_FOOTER_TEMPLATE = """ + +ALL_MODELS: List[Type[SQLModel]] = [ +{model_lines} +] +""" + + +# Map Prisma scalar -> (SA import name, sets has_jsonb, sets has_datetime, sets has_decimal, sets has_any) +_SA_IMPORT_FOR_BASE = { + "String": "Text", + "Int": "Integer", + "BigInt": "BigInteger", + "Float": "Double", + "Decimal": "Numeric", + "Boolean": "Boolean", + "DateTime": "DateTime", + "Json": "JSONB", + "Bytes": "LargeBinary", +} + + +def _classify_field( + field: PrismaField, schema: PrismaSchema, flags: Dict[str, bool] +) -> Optional[str]: + """Return the SQLAlchemy import name needed for ``field`` and update ``flags``.""" + base = field.base_type + if base in schema.enums: + return "Text" + sa_name = _SA_IMPORT_FOR_BASE.get(base, "Text") + if base == "Json": + flags["jsonb"] = True + flags["any"] = True + elif base == "DateTime": + flags["datetime"] = True + elif base == "Decimal": + flags["decimal"] = True + return sa_name + + +def _gather_features(schema: PrismaSchema) -> Tuple[Set[str], Dict[str, bool]]: + """Walk the schema once and return (sqlalchemy import names, feature flags).""" + sa_imports: Set[str] = set() + flags: Dict[str, bool] = { + "optional": False, + "any": False, + "datetime": False, + "decimal": False, + "enum_class": bool(schema.enums), + "indexes": False, + "uniques": False, + "composite_pk": False, + "text_default": False, + "array": False, + "jsonb": False, + } + for model in schema.models.values(): + if len(model.primary_key) > 1: + flags["composite_pk"] = True + if model.uniques: + flags["uniques"] = True + if model.indexes: + flags["indexes"] = True + for f in model.fields: + if f.is_optional: + flags["optional"] = True + if f.is_list: + flags["array"] = True + sa = _classify_field(f, schema, flags) + if sa: + sa_imports.add(sa) + if ( + f.has_default + and f.default_raw is not None + and _server_default_for(f.default_raw, f.base_type) is not None + ): + flags["text_default"] = True + return sa_imports, flags + + +def _collect_used_symbols(schema: PrismaSchema) -> Set[str]: + """Return a sentinel-encoded set describing imports needed by the output.""" + sa_imports, flags = _gather_features(schema) + + if flags["indexes"]: + sa_imports.add("Index") + if flags["composite_pk"]: + sa_imports.add("PrimaryKeyConstraint") + if flags["uniques"]: + sa_imports.add("UniqueConstraint") + if flags["text_default"]: + sa_imports.add("text") + + pg_imports: List[str] = [] + if flags["array"]: + pg_imports.append("ARRAY") + if flags["jsonb"]: + pg_imports.append("JSONB") + + typing_imports: List[str] = ["List", "Type"] + if flags["any"]: + typing_imports.append("Any") + if flags["optional"]: + typing_imports.append("Optional") + + stdlib_lines: List[str] = [] + if flags["datetime"]: + stdlib_lines.append("from datetime import datetime") + if flags["decimal"]: + stdlib_lines.append("from decimal import Decimal") + if flags["enum_class"]: + stdlib_lines.append("from enum import Enum") + + used: Set[str] = set(sa_imports) + used.update(f"_pg::{name}" for name in pg_imports) + used.update(f"_typing::{name}" for name in sorted(set(typing_imports))) + used.update(f"_stdlib::{line}" for line in stdlib_lines) + return used + + +def _render_imports(schema: PrismaSchema) -> str: + used = _collect_used_symbols(schema) + sa = sorted( + s + for s in used + if not s.startswith("_") + and s + in { + "BigInteger", + "Boolean", + "DateTime", + "Double", + "Index", + "Integer", + "LargeBinary", + "Numeric", + "PrimaryKeyConstraint", + "Text", + "UniqueConstraint", + "text", + } + ) + pg = sorted(s.split("::", 1)[1] for s in used if s.startswith("_pg::")) + typing = sorted(s.split("::", 1)[1] for s in used if s.startswith("_typing::")) + stdlib = sorted(s.split("::", 1)[1] for s in used if s.startswith("_stdlib::")) + + lines: List[str] = ["from __future__ import annotations", ""] + lines.extend(stdlib) + if stdlib: + lines.append("") + lines.append(f"from typing import {', '.join(typing)}") + lines.append("") + if sa: + if len(sa) == 1: + lines.append(f"from sqlalchemy import {sa[0]}") + else: + lines.append("from sqlalchemy import (") + for s in sa: + lines.append(f" {s},") + lines.append(")") + if pg: + lines.append(f"from sqlalchemy.dialects.postgresql import {', '.join(pg)}") + lines.append("from sqlmodel import Field, SQLModel") + return "\n".join(lines) + + +def _format_with_black(src: str) -> str: + """Run Black over ``src`` so generator output matches the committed style. + + Black is already a hard CI requirement for this repo (see ``CLAUDE.md``), + so we lean on it as the canonical formatter rather than carrying our own + line-wrapping logic. Falls back to the unformatted source if Black is + unavailable -- the parity test will catch the resulting drift. + """ + try: + import black # type: ignore[import-not-found] + except ImportError: + return src + mode = black.Mode(line_length=88) + try: + return black.format_str(src, mode=mode) + except black.InvalidInput: + return src + + +def render_module(schema: PrismaSchema) -> str: + """Render the entire ``models.py`` source for the given schema.""" + out: List[str] = [_DOCSTRING, "", _render_imports(schema), ""] + + if schema.enums: + for name in sorted(schema.enums): + out.append(_render_enum(schema.enums[name])) + + for name in sorted(schema.models): + out.append(_render_model_class(schema.models[name], schema)) + + model_lines = ",\n".join( + f" {_model_class_name(schema.models[n])}" for n in sorted(schema.models) + ) + out.append(_FOOTER_TEMPLATE.format(model_lines=model_lines)) + return _format_with_black("\n".join(out)) + + +def main(argv: Optional[Iterable[str]] = None) -> int: + parser = argparse.ArgumentParser(description=__doc__.splitlines()[0]) + parser.add_argument("--schema", type=Path, default=Path("schema.prisma")) + parser.add_argument("--out", type=Path, required=True) + args = parser.parse_args(list(argv) if argv is not None else None) + + schema = parse_schema_file(args.schema) + src = render_module(schema) + args.out.write_text(src) + sys.stdout.write( + f"wrote {args.out} ({len(schema.models)} models, " + f"{len(schema.enums)} enums)\n" + ) + return 0 + + +if __name__ == "__main__": # pragma: no cover + raise SystemExit(main()) diff --git a/litellm/proxy/db/sqlmodel/models.py b/litellm/proxy/db/sqlmodel/models.py new file mode 100644 index 000000000000..e1396ed069fd --- /dev/null +++ b/litellm/proxy/db/sqlmodel/models.py @@ -0,0 +1,2491 @@ +"""SQLModel ORM definitions mirroring ``schema.prisma``. + +THIS FILE IS GENERATED by ``litellm.proxy.db.sqlmodel._generate`` but is +CHECKED IN as ordinary Python source. Hand-edits are allowed -- the parity +test in ``tests/test_litellm/proxy/db/sqlmodel_orm/test_parity.py`` will +fail CI if structural drift from ``schema.prisma`` is introduced (in +either direction). + +Re-generate with:: + + uv run python -m litellm.proxy.db.sqlmodel._generate \ + --schema schema.prisma \ + --out litellm/proxy/db/sqlmodel/models.py + +Phase 1 of the Prisma -> SQLModel migration only ships these definitions; +nothing in the runtime proxy currently imports them. See the package +README for the multi-phase plan. +""" + +from __future__ import annotations + +from datetime import datetime +from enum import Enum + +from typing import Any, List, Optional, Type + +from sqlalchemy import ( + BigInteger, + Boolean, + DateTime, + Double, + Index, + Integer, + LargeBinary, + PrimaryKeyConstraint, + Text, + UniqueConstraint, + text, +) +from sqlalchemy.dialects.postgresql import ARRAY, JSONB +from sqlmodel import Field, SQLModel + + +class JobStatus(str, Enum): + ACTIVE = "ACTIVE" + INACTIVE = "INACTIVE" + + +class LiteLLMAccessGroupTable(SQLModel, table=True): + __tablename__ = "LiteLLM_AccessGroupTable" + + access_group_id: str = Field( + sa_type=Text(), + primary_key=True, + default_factory=lambda: str(__import__("uuid").uuid4()), + ) + access_group_name: str = Field(sa_type=Text(), unique=True) + description: Optional[str] = Field(sa_type=Text(), default=None) + access_model_names: List[str] = Field(sa_type=ARRAY(Text()), default_factory=list) + access_mcp_server_ids: List[str] = Field( + sa_type=ARRAY(Text()), default_factory=list + ) + access_agent_ids: List[str] = Field(sa_type=ARRAY(Text()), default_factory=list) + assigned_team_ids: List[str] = Field(sa_type=ARRAY(Text()), default_factory=list) + assigned_key_ids: List[str] = Field(sa_type=ARRAY(Text()), default_factory=list) + created_at: datetime = Field( + sa_type=DateTime(timezone=True), + default_factory=lambda: __import__("datetime").datetime.utcnow(), + sa_column_kwargs={"server_default": text("CURRENT_TIMESTAMP")}, + ) + created_by: Optional[str] = Field(sa_type=Text(), default=None) + updated_at: datetime = Field( + sa_type=DateTime(timezone=True), + default_factory=lambda: __import__("datetime").datetime.utcnow(), + sa_column_kwargs={ + "server_default": text("CURRENT_TIMESTAMP"), + "onupdate": lambda: __import__("datetime").datetime.utcnow(), + }, + ) + updated_by: Optional[str] = Field(sa_type=Text(), default=None) + + +class LiteLLMAdaptiveRouterSession(SQLModel, table=True): + __tablename__ = "LiteLLM_AdaptiveRouterSession" + __table_args__ = ( + PrimaryKeyConstraint("session_id", "router_name", "model_name"), + Index("idx_adaptive_router_session_activity", "last_activity_at"), + ) + + session_id: str = Field(sa_type=Text()) + router_name: str = Field(sa_type=Text()) + model_name: str = Field(sa_type=Text()) + classified_type: str = Field(sa_type=Text()) + misalignment_count: int = Field(sa_type=Integer(), default=0) + stagnation_count: int = Field(sa_type=Integer(), default=0) + disengagement_count: int = Field(sa_type=Integer(), default=0) + satisfaction_count: int = Field(sa_type=Integer(), default=0) + failure_count: int = Field(sa_type=Integer(), default=0) + loop_count: int = Field(sa_type=Integer(), default=0) + exhaustion_count: int = Field(sa_type=Integer(), default=0) + last_user_content: Optional[str] = Field(sa_type=Text(), default=None) + last_assistant_content: Optional[str] = Field(sa_type=Text(), default=None) + tool_call_history: Any = Field( + sa_type=JSONB(), + default_factory=list, + sa_column_kwargs={"server_default": text("'[]'")}, + ) + pending_tool_calls: Any = Field( + sa_type=JSONB(), + default_factory=dict, + sa_column_kwargs={"server_default": text("'{}'")}, + ) + turn_count: int = Field(sa_type=Integer(), default=0) + last_processed_turn: int = Field(sa_type=Integer(), default=-1) + clean_credit_awarded: bool = Field(sa_type=Boolean(), default=False) + terminal_status: Optional[int] = Field(sa_type=Integer(), default=None) + last_activity_at: datetime = Field( + sa_type=DateTime(timezone=True), + default_factory=lambda: __import__("datetime").datetime.utcnow(), + sa_column_kwargs={ + "server_default": text("CURRENT_TIMESTAMP"), + "onupdate": lambda: __import__("datetime").datetime.utcnow(), + }, + ) + + +class LiteLLMAdaptiveRouterState(SQLModel, table=True): + __tablename__ = "LiteLLM_AdaptiveRouterState" + __table_args__ = ( + PrimaryKeyConstraint("router_name", "request_type", "model_name"), + ) + + router_name: str = Field(sa_type=Text()) + request_type: str = Field(sa_type=Text()) + model_name: str = Field(sa_type=Text()) + alpha: float = Field(sa_type=Double()) + beta: float = Field(sa_type=Double()) + total_samples: int = Field(sa_type=Integer(), default=0) + last_updated_at: datetime = Field( + sa_type=DateTime(timezone=True), + default_factory=lambda: __import__("datetime").datetime.utcnow(), + sa_column_kwargs={ + "server_default": text("CURRENT_TIMESTAMP"), + "onupdate": lambda: __import__("datetime").datetime.utcnow(), + }, + ) + + +class LiteLLMAgentsTable(SQLModel, table=True): + __tablename__ = "LiteLLM_AgentsTable" + + agent_id: str = Field( + sa_type=Text(), + primary_key=True, + default_factory=lambda: str(__import__("uuid").uuid4()), + ) + agent_name: str = Field(sa_type=Text(), unique=True) + litellm_params: Optional[Any] = Field(sa_type=JSONB(), default=None) + agent_card_params: Any = Field(sa_type=JSONB()) + static_headers: Optional[Any] = Field( + sa_type=JSONB(), + default_factory=dict, + sa_column_kwargs={"server_default": text("'{}'")}, + ) + extra_headers: List[str] = Field(sa_type=ARRAY(Text()), default_factory=list) + agent_access_groups: List[str] = Field(sa_type=ARRAY(Text()), default_factory=list) + object_permission_id: Optional[str] = Field(sa_type=Text(), default=None) + spend: float = Field(sa_type=Double(), default=0.0) + tpm_limit: Optional[int] = Field(sa_type=Integer(), default=None) + rpm_limit: Optional[int] = Field(sa_type=Integer(), default=None) + session_tpm_limit: Optional[int] = Field(sa_type=Integer(), default=None) + session_rpm_limit: Optional[int] = Field(sa_type=Integer(), default=None) + created_at: datetime = Field( + sa_type=DateTime(timezone=True), + default_factory=lambda: __import__("datetime").datetime.utcnow(), + sa_column_kwargs={"server_default": text("CURRENT_TIMESTAMP")}, + ) + created_by: str = Field(sa_type=Text()) + updated_at: datetime = Field( + sa_type=DateTime(timezone=True), + default_factory=lambda: __import__("datetime").datetime.utcnow(), + sa_column_kwargs={ + "server_default": text("CURRENT_TIMESTAMP"), + "onupdate": lambda: __import__("datetime").datetime.utcnow(), + }, + ) + updated_by: str = Field(sa_type=Text()) + + +class LiteLLMAuditLog(SQLModel, table=True): + __tablename__ = "LiteLLM_AuditLog" + + id: str = Field( + sa_type=Text(), + primary_key=True, + default_factory=lambda: str(__import__("uuid").uuid4()), + ) + updated_at: datetime = Field( + sa_type=DateTime(timezone=True), + default_factory=lambda: __import__("datetime").datetime.utcnow(), + sa_column_kwargs={"server_default": text("CURRENT_TIMESTAMP")}, + ) + changed_by: str = Field(sa_type=Text(), default="") + changed_by_api_key: str = Field(sa_type=Text(), default="") + action: str = Field(sa_type=Text()) + table_name: str = Field(sa_type=Text()) + object_id: str = Field(sa_type=Text()) + before_value: Optional[Any] = Field(sa_type=JSONB(), default=None) + updated_values: Optional[Any] = Field(sa_type=JSONB(), default=None) + + +class LiteLLMBudgetTable(SQLModel, table=True): + __tablename__ = "LiteLLM_BudgetTable" + + budget_id: str = Field( + sa_type=Text(), + primary_key=True, + default_factory=lambda: str(__import__("uuid").uuid4()), + ) + max_budget: Optional[float] = Field(sa_type=Double(), default=None) + soft_budget: Optional[float] = Field(sa_type=Double(), default=None) + max_parallel_requests: Optional[int] = Field(sa_type=Integer(), default=None) + tpm_limit: Optional[int] = Field(sa_type=BigInteger(), default=None) + rpm_limit: Optional[int] = Field(sa_type=BigInteger(), default=None) + model_max_budget: Optional[Any] = Field(sa_type=JSONB(), default=None) + budget_duration: Optional[str] = Field(sa_type=Text(), default=None) + budget_reset_at: Optional[datetime] = Field( + sa_type=DateTime(timezone=True), default=None + ) + allowed_models: List[str] = Field(sa_type=ARRAY(Text()), default_factory=list) + created_at: datetime = Field( + sa_type=DateTime(timezone=True), + default_factory=lambda: __import__("datetime").datetime.utcnow(), + sa_column_kwargs={"server_default": text("CURRENT_TIMESTAMP")}, + ) + created_by: str = Field(sa_type=Text()) + updated_at: datetime = Field( + sa_type=DateTime(timezone=True), + default_factory=lambda: __import__("datetime").datetime.utcnow(), + sa_column_kwargs={ + "server_default": text("CURRENT_TIMESTAMP"), + "onupdate": lambda: __import__("datetime").datetime.utcnow(), + }, + ) + updated_by: str = Field(sa_type=Text()) + + +class LiteLLMCacheConfig(SQLModel, table=True): + __tablename__ = "LiteLLM_CacheConfig" + + id: str = Field(sa_type=Text(), primary_key=True, default="cache_config") + cache_settings: Any = Field(sa_type=JSONB()) + created_at: datetime = Field( + sa_type=DateTime(timezone=True), + default_factory=lambda: __import__("datetime").datetime.utcnow(), + sa_column_kwargs={"server_default": text("CURRENT_TIMESTAMP")}, + ) + updated_at: datetime = Field( + sa_type=DateTime(timezone=True), + sa_column_kwargs={"onupdate": lambda: __import__("datetime").datetime.utcnow()}, + ) + + +class LiteLLMClaudeCodePluginTable(SQLModel, table=True): + __tablename__ = "LiteLLM_ClaudeCodePluginTable" + + id: str = Field( + sa_type=Text(), + primary_key=True, + default_factory=lambda: str(__import__("uuid").uuid4()), + ) + name: str = Field(sa_type=Text(), unique=True) + version: Optional[str] = Field(sa_type=Text(), default=None) + description: Optional[str] = Field(sa_type=Text(), default=None) + manifest_json: Optional[str] = Field(sa_type=Text(), default=None) + files_json: Optional[str] = Field(sa_type=Text(), default_factory=dict) + enabled: bool = Field(sa_type=Boolean(), default=True) + created_at: Optional[datetime] = Field( + sa_type=DateTime(timezone=True), + default_factory=lambda: __import__("datetime").datetime.utcnow(), + sa_column_kwargs={"server_default": text("CURRENT_TIMESTAMP")}, + ) + updated_at: Optional[datetime] = Field( + sa_type=DateTime(timezone=True), + default_factory=lambda: __import__("datetime").datetime.utcnow(), + sa_column_kwargs={ + "server_default": text("CURRENT_TIMESTAMP"), + "onupdate": lambda: __import__("datetime").datetime.utcnow(), + }, + ) + created_by: Optional[str] = Field(sa_type=Text(), default=None) + + +class LiteLLMConfig(SQLModel, table=True): + __tablename__ = "LiteLLM_Config" + + param_name: str = Field(sa_type=Text(), primary_key=True) + param_value: Optional[Any] = Field(sa_type=JSONB(), default=None) + + +class LiteLLMConfigOverrides(SQLModel, table=True): + __tablename__ = "LiteLLM_ConfigOverrides" + + config_type: str = Field(sa_type=Text(), primary_key=True) + config_value: Any = Field(sa_type=JSONB()) + created_at: datetime = Field( + sa_type=DateTime(timezone=True), + default_factory=lambda: __import__("datetime").datetime.utcnow(), + sa_column_kwargs={"server_default": text("CURRENT_TIMESTAMP")}, + ) + updated_at: datetime = Field( + sa_type=DateTime(timezone=True), + sa_column_kwargs={"onupdate": lambda: __import__("datetime").datetime.utcnow()}, + ) + + +class LiteLLMCredentialsTable(SQLModel, table=True): + __tablename__ = "LiteLLM_CredentialsTable" + + credential_id: str = Field( + sa_type=Text(), + primary_key=True, + default_factory=lambda: str(__import__("uuid").uuid4()), + ) + credential_name: str = Field(sa_type=Text(), unique=True) + credential_values: Any = Field(sa_type=JSONB()) + credential_info: Optional[Any] = Field(sa_type=JSONB(), default=None) + created_at: datetime = Field( + sa_type=DateTime(timezone=True), + default_factory=lambda: __import__("datetime").datetime.utcnow(), + sa_column_kwargs={"server_default": text("CURRENT_TIMESTAMP")}, + ) + created_by: str = Field(sa_type=Text()) + updated_at: datetime = Field( + sa_type=DateTime(timezone=True), + default_factory=lambda: __import__("datetime").datetime.utcnow(), + sa_column_kwargs={ + "server_default": text("CURRENT_TIMESTAMP"), + "onupdate": lambda: __import__("datetime").datetime.utcnow(), + }, + ) + updated_by: str = Field(sa_type=Text()) + + +class LiteLLMCronJob(SQLModel, table=True): + __tablename__ = "LiteLLM_CronJob" + + cronjob_id: str = Field( + sa_type=Text(), + primary_key=True, + default_factory=lambda: str(__import__("uuid").uuid4()), + ) + pod_id: str = Field(sa_type=Text()) + status: str = Field(sa_type=Text(), default="INACTIVE") + last_updated: datetime = Field( + sa_type=DateTime(timezone=True), + default_factory=lambda: __import__("datetime").datetime.utcnow(), + sa_column_kwargs={"server_default": text("CURRENT_TIMESTAMP")}, + ) + ttl: datetime = Field(sa_type=DateTime(timezone=True)) + + +class LiteLLMDailyAgentSpend(SQLModel, table=True): + __tablename__ = "LiteLLM_DailyAgentSpend" + __table_args__ = ( + UniqueConstraint( + "agent_id", + "date", + "api_key", + "model", + "custom_llm_provider", + "mcp_namespaced_tool_name", + "endpoint", + ), + Index("LiteLLM_DailyAgentSpend_date_idx", "date"), + Index("LiteLLM_DailyAgentSpend_agent_id_date_idx", "agent_id", "date"), + Index("LiteLLM_DailyAgentSpend_api_key_idx", "api_key"), + Index("LiteLLM_DailyAgentSpend_model_idx", "model"), + Index( + "LiteLLM_DailyAgentSpend_mcp_namespaced_tool_name_idx", + "mcp_namespaced_tool_name", + ), + Index("LiteLLM_DailyAgentSpend_endpoint_idx", "endpoint"), + ) + + id: str = Field( + sa_type=Text(), + primary_key=True, + default_factory=lambda: str(__import__("uuid").uuid4()), + ) + agent_id: Optional[str] = Field(sa_type=Text(), default=None) + date: str = Field(sa_type=Text()) + api_key: str = Field(sa_type=Text()) + model: Optional[str] = Field(sa_type=Text(), default=None) + model_group: Optional[str] = Field(sa_type=Text(), default=None) + custom_llm_provider: Optional[str] = Field(sa_type=Text(), default=None) + mcp_namespaced_tool_name: Optional[str] = Field(sa_type=Text(), default=None) + endpoint: Optional[str] = Field(sa_type=Text(), default=None) + prompt_tokens: int = Field(sa_type=BigInteger(), default=0) + completion_tokens: int = Field(sa_type=BigInteger(), default=0) + cache_read_input_tokens: int = Field(sa_type=BigInteger(), default=0) + cache_creation_input_tokens: int = Field(sa_type=BigInteger(), default=0) + spend: float = Field(sa_type=Double(), default=0.0) + api_requests: int = Field(sa_type=BigInteger(), default=0) + successful_requests: int = Field(sa_type=BigInteger(), default=0) + failed_requests: int = Field(sa_type=BigInteger(), default=0) + created_at: datetime = Field( + sa_type=DateTime(timezone=True), + default_factory=lambda: __import__("datetime").datetime.utcnow(), + sa_column_kwargs={"server_default": text("CURRENT_TIMESTAMP")}, + ) + updated_at: datetime = Field( + sa_type=DateTime(timezone=True), + sa_column_kwargs={"onupdate": lambda: __import__("datetime").datetime.utcnow()}, + ) + + +class LiteLLMDailyEndUserSpend(SQLModel, table=True): + __tablename__ = "LiteLLM_DailyEndUserSpend" + __table_args__ = ( + UniqueConstraint( + "end_user_id", + "date", + "api_key", + "model", + "custom_llm_provider", + "mcp_namespaced_tool_name", + "endpoint", + ), + Index("LiteLLM_DailyEndUserSpend_date_idx", "date"), + Index("LiteLLM_DailyEndUserSpend_end_user_id_date_idx", "end_user_id", "date"), + Index("LiteLLM_DailyEndUserSpend_api_key_idx", "api_key"), + Index("LiteLLM_DailyEndUserSpend_model_idx", "model"), + Index( + "LiteLLM_DailyEndUserSpend_mcp_namespaced_tool_name_idx", + "mcp_namespaced_tool_name", + ), + Index("LiteLLM_DailyEndUserSpend_endpoint_idx", "endpoint"), + ) + + id: str = Field( + sa_type=Text(), + primary_key=True, + default_factory=lambda: str(__import__("uuid").uuid4()), + ) + end_user_id: Optional[str] = Field(sa_type=Text(), default=None) + date: str = Field(sa_type=Text()) + api_key: str = Field(sa_type=Text()) + model: Optional[str] = Field(sa_type=Text(), default=None) + model_group: Optional[str] = Field(sa_type=Text(), default=None) + custom_llm_provider: Optional[str] = Field(sa_type=Text(), default=None) + mcp_namespaced_tool_name: Optional[str] = Field(sa_type=Text(), default=None) + endpoint: Optional[str] = Field(sa_type=Text(), default=None) + prompt_tokens: int = Field(sa_type=BigInteger(), default=0) + completion_tokens: int = Field(sa_type=BigInteger(), default=0) + cache_read_input_tokens: int = Field(sa_type=BigInteger(), default=0) + cache_creation_input_tokens: int = Field(sa_type=BigInteger(), default=0) + spend: float = Field(sa_type=Double(), default=0.0) + api_requests: int = Field(sa_type=BigInteger(), default=0) + successful_requests: int = Field(sa_type=BigInteger(), default=0) + failed_requests: int = Field(sa_type=BigInteger(), default=0) + created_at: datetime = Field( + sa_type=DateTime(timezone=True), + default_factory=lambda: __import__("datetime").datetime.utcnow(), + sa_column_kwargs={"server_default": text("CURRENT_TIMESTAMP")}, + ) + updated_at: datetime = Field( + sa_type=DateTime(timezone=True), + sa_column_kwargs={"onupdate": lambda: __import__("datetime").datetime.utcnow()}, + ) + + +class LiteLLMDailyGuardrailMetrics(SQLModel, table=True): + __tablename__ = "LiteLLM_DailyGuardrailMetrics" + __table_args__ = ( + PrimaryKeyConstraint("guardrail_id", "date"), + Index("LiteLLM_DailyGuardrailMetrics_date_idx", "date"), + Index("LiteLLM_DailyGuardrailMetrics_guardrail_id_idx", "guardrail_id"), + ) + + guardrail_id: str = Field(sa_type=Text()) + date: str = Field(sa_type=Text()) + requests_evaluated: int = Field(sa_type=BigInteger(), default=0) + passed_count: int = Field(sa_type=BigInteger(), default=0) + blocked_count: int = Field(sa_type=BigInteger(), default=0) + flagged_count: int = Field(sa_type=BigInteger(), default=0) + avg_score: Optional[float] = Field(sa_type=Double(), default=None) + avg_latency_ms: Optional[float] = Field(sa_type=Double(), default=None) + created_at: datetime = Field( + sa_type=DateTime(timezone=True), + default_factory=lambda: __import__("datetime").datetime.utcnow(), + sa_column_kwargs={"server_default": text("CURRENT_TIMESTAMP")}, + ) + updated_at: datetime = Field( + sa_type=DateTime(timezone=True), + sa_column_kwargs={"onupdate": lambda: __import__("datetime").datetime.utcnow()}, + ) + + +class LiteLLMDailyOrganizationSpend(SQLModel, table=True): + __tablename__ = "LiteLLM_DailyOrganizationSpend" + __table_args__ = ( + UniqueConstraint( + "organization_id", + "date", + "api_key", + "model", + "custom_llm_provider", + "mcp_namespaced_tool_name", + "endpoint", + ), + Index("LiteLLM_DailyOrganizationSpend_date_idx", "date"), + Index( + "LiteLLM_DailyOrganizationSpend_organization_id_date_idx", + "organization_id", + "date", + ), + Index("LiteLLM_DailyOrganizationSpend_api_key_idx", "api_key"), + Index("LiteLLM_DailyOrganizationSpend_model_idx", "model"), + Index( + "LiteLLM_DailyOrganizationSpend_mcp_namespaced_tool_name_idx", + "mcp_namespaced_tool_name", + ), + Index("LiteLLM_DailyOrganizationSpend_endpoint_idx", "endpoint"), + ) + + id: str = Field( + sa_type=Text(), + primary_key=True, + default_factory=lambda: str(__import__("uuid").uuid4()), + ) + organization_id: Optional[str] = Field(sa_type=Text(), default=None) + date: str = Field(sa_type=Text()) + api_key: str = Field(sa_type=Text()) + model: Optional[str] = Field(sa_type=Text(), default=None) + model_group: Optional[str] = Field(sa_type=Text(), default=None) + custom_llm_provider: Optional[str] = Field(sa_type=Text(), default=None) + mcp_namespaced_tool_name: Optional[str] = Field(sa_type=Text(), default=None) + endpoint: Optional[str] = Field(sa_type=Text(), default=None) + prompt_tokens: int = Field(sa_type=BigInteger(), default=0) + completion_tokens: int = Field(sa_type=BigInteger(), default=0) + cache_read_input_tokens: int = Field(sa_type=BigInteger(), default=0) + cache_creation_input_tokens: int = Field(sa_type=BigInteger(), default=0) + spend: float = Field(sa_type=Double(), default=0.0) + api_requests: int = Field(sa_type=BigInteger(), default=0) + successful_requests: int = Field(sa_type=BigInteger(), default=0) + failed_requests: int = Field(sa_type=BigInteger(), default=0) + created_at: datetime = Field( + sa_type=DateTime(timezone=True), + default_factory=lambda: __import__("datetime").datetime.utcnow(), + sa_column_kwargs={"server_default": text("CURRENT_TIMESTAMP")}, + ) + updated_at: datetime = Field( + sa_type=DateTime(timezone=True), + sa_column_kwargs={"onupdate": lambda: __import__("datetime").datetime.utcnow()}, + ) + + +class LiteLLMDailyPolicyMetrics(SQLModel, table=True): + __tablename__ = "LiteLLM_DailyPolicyMetrics" + __table_args__ = ( + PrimaryKeyConstraint("policy_id", "date"), + Index("LiteLLM_DailyPolicyMetrics_date_idx", "date"), + Index("LiteLLM_DailyPolicyMetrics_policy_id_idx", "policy_id"), + ) + + policy_id: str = Field(sa_type=Text()) + date: str = Field(sa_type=Text()) + requests_evaluated: int = Field(sa_type=BigInteger(), default=0) + passed_count: int = Field(sa_type=BigInteger(), default=0) + blocked_count: int = Field(sa_type=BigInteger(), default=0) + flagged_count: int = Field(sa_type=BigInteger(), default=0) + avg_score: Optional[float] = Field(sa_type=Double(), default=None) + avg_latency_ms: Optional[float] = Field(sa_type=Double(), default=None) + created_at: datetime = Field( + sa_type=DateTime(timezone=True), + default_factory=lambda: __import__("datetime").datetime.utcnow(), + sa_column_kwargs={"server_default": text("CURRENT_TIMESTAMP")}, + ) + updated_at: datetime = Field( + sa_type=DateTime(timezone=True), + sa_column_kwargs={"onupdate": lambda: __import__("datetime").datetime.utcnow()}, + ) + + +class LiteLLMDailyTagSpend(SQLModel, table=True): + __tablename__ = "LiteLLM_DailyTagSpend" + __table_args__ = ( + UniqueConstraint( + "tag", + "date", + "api_key", + "model", + "custom_llm_provider", + "mcp_namespaced_tool_name", + "endpoint", + ), + Index("LiteLLM_DailyTagSpend_date_idx", "date"), + Index("LiteLLM_DailyTagSpend_tag_date_idx", "tag", "date"), + Index("LiteLLM_DailyTagSpend_api_key_idx", "api_key"), + Index("LiteLLM_DailyTagSpend_model_idx", "model"), + Index( + "LiteLLM_DailyTagSpend_mcp_namespaced_tool_name_idx", + "mcp_namespaced_tool_name", + ), + Index("LiteLLM_DailyTagSpend_endpoint_idx", "endpoint"), + ) + + id: str = Field( + sa_type=Text(), + primary_key=True, + default_factory=lambda: str(__import__("uuid").uuid4()), + ) + request_id: Optional[str] = Field(sa_type=Text(), default=None) + tag: Optional[str] = Field(sa_type=Text(), default=None) + date: str = Field(sa_type=Text()) + api_key: str = Field(sa_type=Text()) + model: Optional[str] = Field(sa_type=Text(), default=None) + model_group: Optional[str] = Field(sa_type=Text(), default=None) + custom_llm_provider: Optional[str] = Field(sa_type=Text(), default=None) + mcp_namespaced_tool_name: Optional[str] = Field(sa_type=Text(), default=None) + endpoint: Optional[str] = Field(sa_type=Text(), default=None) + prompt_tokens: int = Field(sa_type=BigInteger(), default=0) + completion_tokens: int = Field(sa_type=BigInteger(), default=0) + cache_read_input_tokens: int = Field(sa_type=BigInteger(), default=0) + cache_creation_input_tokens: int = Field(sa_type=BigInteger(), default=0) + spend: float = Field(sa_type=Double(), default=0.0) + api_requests: int = Field(sa_type=BigInteger(), default=0) + successful_requests: int = Field(sa_type=BigInteger(), default=0) + failed_requests: int = Field(sa_type=BigInteger(), default=0) + created_at: datetime = Field( + sa_type=DateTime(timezone=True), + default_factory=lambda: __import__("datetime").datetime.utcnow(), + sa_column_kwargs={"server_default": text("CURRENT_TIMESTAMP")}, + ) + updated_at: datetime = Field( + sa_type=DateTime(timezone=True), + sa_column_kwargs={"onupdate": lambda: __import__("datetime").datetime.utcnow()}, + ) + + +class LiteLLMDailyTeamSpend(SQLModel, table=True): + __tablename__ = "LiteLLM_DailyTeamSpend" + __table_args__ = ( + UniqueConstraint( + "team_id", + "date", + "api_key", + "model", + "custom_llm_provider", + "mcp_namespaced_tool_name", + "endpoint", + ), + Index("LiteLLM_DailyTeamSpend_date_idx", "date"), + Index("LiteLLM_DailyTeamSpend_team_id_date_idx", "team_id", "date"), + Index("LiteLLM_DailyTeamSpend_api_key_idx", "api_key"), + Index("LiteLLM_DailyTeamSpend_model_idx", "model"), + Index( + "LiteLLM_DailyTeamSpend_mcp_namespaced_tool_name_idx", + "mcp_namespaced_tool_name", + ), + Index("LiteLLM_DailyTeamSpend_endpoint_idx", "endpoint"), + ) + + id: str = Field( + sa_type=Text(), + primary_key=True, + default_factory=lambda: str(__import__("uuid").uuid4()), + ) + team_id: Optional[str] = Field(sa_type=Text(), default=None) + date: str = Field(sa_type=Text()) + api_key: str = Field(sa_type=Text()) + model: Optional[str] = Field(sa_type=Text(), default=None) + model_group: Optional[str] = Field(sa_type=Text(), default=None) + custom_llm_provider: Optional[str] = Field(sa_type=Text(), default=None) + mcp_namespaced_tool_name: Optional[str] = Field(sa_type=Text(), default=None) + endpoint: Optional[str] = Field(sa_type=Text(), default=None) + prompt_tokens: int = Field(sa_type=BigInteger(), default=0) + completion_tokens: int = Field(sa_type=BigInteger(), default=0) + cache_read_input_tokens: int = Field(sa_type=BigInteger(), default=0) + cache_creation_input_tokens: int = Field(sa_type=BigInteger(), default=0) + spend: float = Field(sa_type=Double(), default=0.0) + api_requests: int = Field(sa_type=BigInteger(), default=0) + successful_requests: int = Field(sa_type=BigInteger(), default=0) + failed_requests: int = Field(sa_type=BigInteger(), default=0) + created_at: datetime = Field( + sa_type=DateTime(timezone=True), + default_factory=lambda: __import__("datetime").datetime.utcnow(), + sa_column_kwargs={"server_default": text("CURRENT_TIMESTAMP")}, + ) + updated_at: datetime = Field( + sa_type=DateTime(timezone=True), + sa_column_kwargs={"onupdate": lambda: __import__("datetime").datetime.utcnow()}, + ) + + +class LiteLLMDailyUserSpend(SQLModel, table=True): + __tablename__ = "LiteLLM_DailyUserSpend" + __table_args__ = ( + UniqueConstraint( + "user_id", + "date", + "api_key", + "model", + "custom_llm_provider", + "mcp_namespaced_tool_name", + "endpoint", + ), + Index("LiteLLM_DailyUserSpend_date_idx", "date"), + Index("LiteLLM_DailyUserSpend_user_id_date_idx", "user_id", "date"), + Index("LiteLLM_DailyUserSpend_api_key_idx", "api_key"), + Index("LiteLLM_DailyUserSpend_model_idx", "model"), + Index( + "LiteLLM_DailyUserSpend_mcp_namespaced_tool_name_idx", + "mcp_namespaced_tool_name", + ), + Index("LiteLLM_DailyUserSpend_endpoint_idx", "endpoint"), + ) + + id: str = Field( + sa_type=Text(), + primary_key=True, + default_factory=lambda: str(__import__("uuid").uuid4()), + ) + user_id: Optional[str] = Field(sa_type=Text(), default=None) + date: str = Field(sa_type=Text()) + api_key: str = Field(sa_type=Text()) + model: Optional[str] = Field(sa_type=Text(), default=None) + model_group: Optional[str] = Field(sa_type=Text(), default=None) + custom_llm_provider: Optional[str] = Field(sa_type=Text(), default=None) + mcp_namespaced_tool_name: Optional[str] = Field(sa_type=Text(), default=None) + endpoint: Optional[str] = Field(sa_type=Text(), default=None) + prompt_tokens: int = Field(sa_type=BigInteger(), default=0) + completion_tokens: int = Field(sa_type=BigInteger(), default=0) + cache_read_input_tokens: int = Field(sa_type=BigInteger(), default=0) + cache_creation_input_tokens: int = Field(sa_type=BigInteger(), default=0) + spend: float = Field(sa_type=Double(), default=0.0) + api_requests: int = Field(sa_type=BigInteger(), default=0) + successful_requests: int = Field(sa_type=BigInteger(), default=0) + failed_requests: int = Field(sa_type=BigInteger(), default=0) + created_at: datetime = Field( + sa_type=DateTime(timezone=True), + default_factory=lambda: __import__("datetime").datetime.utcnow(), + sa_column_kwargs={"server_default": text("CURRENT_TIMESTAMP")}, + ) + updated_at: datetime = Field( + sa_type=DateTime(timezone=True), + sa_column_kwargs={"onupdate": lambda: __import__("datetime").datetime.utcnow()}, + ) + + +class LiteLLMDeletedTeamTable(SQLModel, table=True): + __tablename__ = "LiteLLM_DeletedTeamTable" + __table_args__ = ( + Index("LiteLLM_DeletedTeamTable_team_id_idx", "team_id"), + Index("LiteLLM_DeletedTeamTable_deleted_at_idx", "deleted_at"), + Index("LiteLLM_DeletedTeamTable_organization_id_idx", "organization_id"), + Index("LiteLLM_DeletedTeamTable_team_alias_idx", "team_alias"), + Index("LiteLLM_DeletedTeamTable_created_at_idx", "created_at"), + ) + + id: str = Field( + sa_type=Text(), + primary_key=True, + default_factory=lambda: str(__import__("uuid").uuid4()), + ) + team_id: str = Field(sa_type=Text()) + team_alias: Optional[str] = Field(sa_type=Text(), default=None) + organization_id: Optional[str] = Field(sa_type=Text(), default=None) + object_permission_id: Optional[str] = Field(sa_type=Text(), default=None) + admins: List[str] = Field(sa_type=ARRAY(Text()), default_factory=list) + members: List[str] = Field(sa_type=ARRAY(Text()), default_factory=list) + members_with_roles: Any = Field( + sa_type=JSONB(), + default_factory=dict, + sa_column_kwargs={"server_default": text("'{}'")}, + ) + metadata_: Any = Field( + sa_type=JSONB(), + default_factory=dict, + sa_column_kwargs={"name": "metadata", "server_default": text("'{}'")}, + ) + max_budget: Optional[float] = Field(sa_type=Double(), default=None) + soft_budget: Optional[float] = Field(sa_type=Double(), default=None) + spend: float = Field(sa_type=Double(), default=0.0) + models: List[str] = Field(sa_type=ARRAY(Text()), default_factory=list) + max_parallel_requests: Optional[int] = Field(sa_type=Integer(), default=None) + tpm_limit: Optional[int] = Field(sa_type=BigInteger(), default=None) + rpm_limit: Optional[int] = Field(sa_type=BigInteger(), default=None) + budget_duration: Optional[str] = Field(sa_type=Text(), default=None) + budget_reset_at: Optional[datetime] = Field( + sa_type=DateTime(timezone=True), default=None + ) + blocked: bool = Field(sa_type=Boolean(), default=False) + model_spend: Any = Field( + sa_type=JSONB(), + default_factory=dict, + sa_column_kwargs={"server_default": text("'{}'")}, + ) + model_max_budget: Any = Field( + sa_type=JSONB(), + default_factory=dict, + sa_column_kwargs={"server_default": text("'{}'")}, + ) + router_settings: Optional[Any] = Field( + sa_type=JSONB(), + default_factory=dict, + sa_column_kwargs={"server_default": text("'{}'")}, + ) + team_member_permissions: List[str] = Field( + sa_type=ARRAY(Text()), default_factory=list + ) + access_group_ids: List[str] = Field(sa_type=ARRAY(Text()), default_factory=list) + policies: List[str] = Field(sa_type=ARRAY(Text()), default_factory=list) + model_id: Optional[int] = Field(sa_type=Integer(), default=None) + allow_team_guardrail_config: bool = Field(sa_type=Boolean(), default=False) + created_at: Optional[datetime] = Field( + sa_type=DateTime(timezone=True), default=None + ) + updated_at: Optional[datetime] = Field( + sa_type=DateTime(timezone=True), default=None + ) + deleted_at: datetime = Field( + sa_type=DateTime(timezone=True), + default_factory=lambda: __import__("datetime").datetime.utcnow(), + sa_column_kwargs={"server_default": text("CURRENT_TIMESTAMP")}, + ) + deleted_by: Optional[str] = Field(sa_type=Text(), default=None) + deleted_by_api_key: Optional[str] = Field(sa_type=Text(), default=None) + litellm_changed_by: Optional[str] = Field(sa_type=Text(), default=None) + + +class LiteLLMDeletedVerificationToken(SQLModel, table=True): + __tablename__ = "LiteLLM_DeletedVerificationToken" + __table_args__ = ( + Index("LiteLLM_DeletedVerificationToken_token_idx", "token"), + Index("LiteLLM_DeletedVerificationToken_deleted_at_idx", "deleted_at"), + Index("LiteLLM_DeletedVerificationToken_user_id_idx", "user_id"), + Index("LiteLLM_DeletedVerificationToken_team_id_idx", "team_id"), + Index( + "LiteLLM_DeletedVerificationToken_organization_id_idx", "organization_id" + ), + Index("LiteLLM_DeletedVerificationToken_key_alias_idx", "key_alias"), + Index("LiteLLM_DeletedVerificationToken_created_at_idx", "created_at"), + ) + + id: str = Field( + sa_type=Text(), + primary_key=True, + default_factory=lambda: str(__import__("uuid").uuid4()), + ) + token: str = Field(sa_type=Text()) + key_name: Optional[str] = Field(sa_type=Text(), default=None) + key_alias: Optional[str] = Field(sa_type=Text(), default=None) + soft_budget_cooldown: bool = Field(sa_type=Boolean(), default=False) + spend: float = Field(sa_type=Double(), default=0.0) + expires: Optional[datetime] = Field(sa_type=DateTime(timezone=True), default=None) + models: List[str] = Field(sa_type=ARRAY(Text()), default_factory=list) + aliases: Any = Field( + sa_type=JSONB(), + default_factory=dict, + sa_column_kwargs={"server_default": text("'{}'")}, + ) + config: Any = Field( + sa_type=JSONB(), + default_factory=dict, + sa_column_kwargs={"server_default": text("'{}'")}, + ) + user_id: Optional[str] = Field(sa_type=Text(), default=None) + team_id: Optional[str] = Field(sa_type=Text(), default=None) + agent_id: Optional[str] = Field(sa_type=Text(), default=None) + project_id: Optional[str] = Field(sa_type=Text(), default=None) + permissions: Any = Field( + sa_type=JSONB(), + default_factory=dict, + sa_column_kwargs={"server_default": text("'{}'")}, + ) + max_parallel_requests: Optional[int] = Field(sa_type=Integer(), default=None) + metadata_: Any = Field( + sa_type=JSONB(), + default_factory=dict, + sa_column_kwargs={"name": "metadata", "server_default": text("'{}'")}, + ) + blocked: Optional[bool] = Field(sa_type=Boolean(), default=None) + tpm_limit: Optional[int] = Field(sa_type=BigInteger(), default=None) + rpm_limit: Optional[int] = Field(sa_type=BigInteger(), default=None) + max_budget: Optional[float] = Field(sa_type=Double(), default=None) + budget_duration: Optional[str] = Field(sa_type=Text(), default=None) + budget_reset_at: Optional[datetime] = Field( + sa_type=DateTime(timezone=True), default=None + ) + allowed_cache_controls: List[str] = Field( + sa_type=ARRAY(Text()), default_factory=list + ) + allowed_routes: List[str] = Field(sa_type=ARRAY(Text()), default_factory=list) + policies: List[str] = Field(sa_type=ARRAY(Text()), default_factory=list) + access_group_ids: List[str] = Field(sa_type=ARRAY(Text()), default_factory=list) + model_spend: Any = Field( + sa_type=JSONB(), + default_factory=dict, + sa_column_kwargs={"server_default": text("'{}'")}, + ) + model_max_budget: Any = Field( + sa_type=JSONB(), + default_factory=dict, + sa_column_kwargs={"server_default": text("'{}'")}, + ) + router_settings: Optional[Any] = Field( + sa_type=JSONB(), + default_factory=dict, + sa_column_kwargs={"server_default": text("'{}'")}, + ) + budget_id: Optional[str] = Field(sa_type=Text(), default=None) + organization_id: Optional[str] = Field(sa_type=Text(), default=None) + object_permission_id: Optional[str] = Field(sa_type=Text(), default=None) + created_at: Optional[datetime] = Field( + sa_type=DateTime(timezone=True), default=None + ) + created_by: Optional[str] = Field(sa_type=Text(), default=None) + updated_at: Optional[datetime] = Field( + sa_type=DateTime(timezone=True), default=None + ) + updated_by: Optional[str] = Field(sa_type=Text(), default=None) + last_active: Optional[datetime] = Field( + sa_type=DateTime(timezone=True), default=None + ) + rotation_count: Optional[int] = Field(sa_type=Integer(), default=0) + auto_rotate: Optional[bool] = Field(sa_type=Boolean(), default=False) + rotation_interval: Optional[str] = Field(sa_type=Text(), default=None) + last_rotation_at: Optional[datetime] = Field( + sa_type=DateTime(timezone=True), default=None + ) + key_rotation_at: Optional[datetime] = Field( + sa_type=DateTime(timezone=True), default=None + ) + deleted_at: datetime = Field( + sa_type=DateTime(timezone=True), + default_factory=lambda: __import__("datetime").datetime.utcnow(), + sa_column_kwargs={"server_default": text("CURRENT_TIMESTAMP")}, + ) + deleted_by: Optional[str] = Field(sa_type=Text(), default=None) + deleted_by_api_key: Optional[str] = Field(sa_type=Text(), default=None) + litellm_changed_by: Optional[str] = Field(sa_type=Text(), default=None) + + +class LiteLLMDeprecatedVerificationToken(SQLModel, table=True): + __tablename__ = "LiteLLM_DeprecatedVerificationToken" + __table_args__ = ( + UniqueConstraint("token"), + Index( + "LiteLLM_DeprecatedVerificationToken_token_revoke_at_idx", + "token", + "revoke_at", + ), + Index("LiteLLM_DeprecatedVerificationToken_revoke_at_idx", "revoke_at"), + ) + + id: str = Field( + sa_type=Text(), + primary_key=True, + default_factory=lambda: str(__import__("uuid").uuid4()), + ) + token: str = Field(sa_type=Text()) + active_token_id: str = Field(sa_type=Text()) + revoke_at: datetime = Field(sa_type=DateTime(timezone=True)) + created_at: datetime = Field( + sa_type=DateTime(timezone=True), + default_factory=lambda: __import__("datetime").datetime.utcnow(), + sa_column_kwargs={"server_default": text("CURRENT_TIMESTAMP")}, + ) + + +class LiteLLMEndUserTable(SQLModel, table=True): + __tablename__ = "LiteLLM_EndUserTable" + + user_id: str = Field(sa_type=Text(), primary_key=True) + alias: Optional[str] = Field(sa_type=Text(), default=None) + spend: float = Field(sa_type=Double(), default=0.0) + allowed_model_region: Optional[str] = Field(sa_type=Text(), default=None) + default_model: Optional[str] = Field(sa_type=Text(), default=None) + budget_id: Optional[str] = Field(sa_type=Text(), default=None) + object_permission_id: Optional[str] = Field(sa_type=Text(), default=None) + blocked: bool = Field(sa_type=Boolean(), default=False) + + +class LiteLLMErrorLogs(SQLModel, table=True): + __tablename__ = "LiteLLM_ErrorLogs" + + request_id: str = Field( + sa_type=Text(), + primary_key=True, + default_factory=lambda: str(__import__("uuid").uuid4()), + ) + startTime: datetime = Field(sa_type=DateTime(timezone=True)) + endTime: datetime = Field(sa_type=DateTime(timezone=True)) + api_base: str = Field(sa_type=Text(), default="") + model_group: str = Field(sa_type=Text(), default="") + litellm_model_name: str = Field(sa_type=Text(), default="") + model_id: str = Field(sa_type=Text(), default="") + request_kwargs: Any = Field( + sa_type=JSONB(), + default_factory=dict, + sa_column_kwargs={"server_default": text("'{}'")}, + ) + exception_type: str = Field(sa_type=Text(), default="") + exception_string: str = Field(sa_type=Text(), default="") + status_code: str = Field(sa_type=Text(), default="") + + +class LiteLLMGuardrailsTable(SQLModel, table=True): + __tablename__ = "LiteLLM_GuardrailsTable" + __table_args__ = (Index("LiteLLM_GuardrailsTable_status_idx", "status"),) + + guardrail_id: str = Field( + sa_type=Text(), + primary_key=True, + default_factory=lambda: str(__import__("uuid").uuid4()), + ) + guardrail_name: str = Field(sa_type=Text(), unique=True) + litellm_params: Any = Field(sa_type=JSONB()) + guardrail_info: Optional[Any] = Field(sa_type=JSONB(), default=None) + team_id: Optional[str] = Field(sa_type=Text(), default=None) + created_at: datetime = Field( + sa_type=DateTime(timezone=True), + default_factory=lambda: __import__("datetime").datetime.utcnow(), + sa_column_kwargs={"server_default": text("CURRENT_TIMESTAMP")}, + ) + updated_at: datetime = Field( + sa_type=DateTime(timezone=True), + sa_column_kwargs={"onupdate": lambda: __import__("datetime").datetime.utcnow()}, + ) + status: str = Field(sa_type=Text(), default="active") + submitted_at: Optional[datetime] = Field( + sa_type=DateTime(timezone=True), default=None + ) + reviewed_at: Optional[datetime] = Field( + sa_type=DateTime(timezone=True), default=None + ) + + +class LiteLLMHealthCheckTable(SQLModel, table=True): + __tablename__ = "LiteLLM_HealthCheckTable" + __table_args__ = ( + Index("LiteLLM_HealthCheckTable_model_name_idx", "model_name"), + Index("LiteLLM_HealthCheckTable_checked_at_idx", "checked_at"), + Index("LiteLLM_HealthCheckTable_status_idx", "status"), + Index( + "LiteLLM_HealthCheckTable_model_id_model_name_checked_at_idx", + "model_id", + "model_name", + "checked_at", + ), + ) + + health_check_id: str = Field( + sa_type=Text(), + primary_key=True, + default_factory=lambda: str(__import__("uuid").uuid4()), + ) + model_name: str = Field(sa_type=Text()) + model_id: Optional[str] = Field(sa_type=Text(), default=None) + status: str = Field(sa_type=Text()) + healthy_count: int = Field(sa_type=Integer(), default=0) + unhealthy_count: int = Field(sa_type=Integer(), default=0) + error_message: Optional[str] = Field(sa_type=Text(), default=None) + response_time_ms: Optional[float] = Field(sa_type=Double(), default=None) + details: Optional[Any] = Field(sa_type=JSONB(), default=None) + checked_by: Optional[str] = Field(sa_type=Text(), default=None) + checked_at: datetime = Field( + sa_type=DateTime(timezone=True), + default_factory=lambda: __import__("datetime").datetime.utcnow(), + sa_column_kwargs={"server_default": text("CURRENT_TIMESTAMP")}, + ) + created_at: datetime = Field( + sa_type=DateTime(timezone=True), + default_factory=lambda: __import__("datetime").datetime.utcnow(), + sa_column_kwargs={"server_default": text("CURRENT_TIMESTAMP")}, + ) + updated_at: datetime = Field( + sa_type=DateTime(timezone=True), + sa_column_kwargs={"onupdate": lambda: __import__("datetime").datetime.utcnow()}, + ) + + +class LiteLLMInvitationLink(SQLModel, table=True): + __tablename__ = "LiteLLM_InvitationLink" + + id: str = Field( + sa_type=Text(), + primary_key=True, + default_factory=lambda: str(__import__("uuid").uuid4()), + ) + user_id: str = Field(sa_type=Text()) + is_accepted: bool = Field(sa_type=Boolean(), default=False) + accepted_at: Optional[datetime] = Field( + sa_type=DateTime(timezone=True), default=None + ) + expires_at: datetime = Field(sa_type=DateTime(timezone=True)) + created_at: datetime = Field(sa_type=DateTime(timezone=True)) + created_by: str = Field(sa_type=Text()) + updated_at: datetime = Field(sa_type=DateTime(timezone=True)) + updated_by: str = Field(sa_type=Text()) + + +class LiteLLMJWTKeyMapping(SQLModel, table=True): + __tablename__ = "LiteLLM_JWTKeyMapping" + __table_args__ = ( + UniqueConstraint("jwt_claim_name", "jwt_claim_value"), + Index( + "LiteLLM_JWTKeyMapping_jwt_claim_name_jwt_claim_value_is_active_idx", + "jwt_claim_name", + "jwt_claim_value", + "is_active", + ), + ) + + id: str = Field( + sa_type=Text(), + primary_key=True, + default_factory=lambda: str(__import__("uuid").uuid4()), + ) + jwt_claim_name: str = Field(sa_type=Text()) + jwt_claim_value: str = Field(sa_type=Text()) + token: str = Field(sa_type=Text()) + description: Optional[str] = Field(sa_type=Text(), default=None) + is_active: bool = Field(sa_type=Boolean(), default=True) + created_at: datetime = Field( + sa_type=DateTime(timezone=True), + default_factory=lambda: __import__("datetime").datetime.utcnow(), + sa_column_kwargs={"server_default": text("CURRENT_TIMESTAMP")}, + ) + created_by: Optional[str] = Field(sa_type=Text(), default=None) + updated_at: datetime = Field( + sa_type=DateTime(timezone=True), + default_factory=lambda: __import__("datetime").datetime.utcnow(), + sa_column_kwargs={ + "server_default": text("CURRENT_TIMESTAMP"), + "onupdate": lambda: __import__("datetime").datetime.utcnow(), + }, + ) + updated_by: Optional[str] = Field(sa_type=Text(), default=None) + + +class LiteLLMMCPServerTable(SQLModel, table=True): + __tablename__ = "LiteLLM_MCPServerTable" + __table_args__ = ( + Index("LiteLLM_MCPServerTable_approval_status_idx", "approval_status"), + ) + + server_id: str = Field( + sa_type=Text(), + primary_key=True, + default_factory=lambda: str(__import__("uuid").uuid4()), + ) + server_name: Optional[str] = Field(sa_type=Text(), default=None) + alias: Optional[str] = Field(sa_type=Text(), default=None) + description: Optional[str] = Field(sa_type=Text(), default=None) + instructions: Optional[str] = Field(sa_type=Text(), default=None) + url: Optional[str] = Field(sa_type=Text(), default=None) + spec_path: Optional[str] = Field(sa_type=Text(), default=None) + transport: str = Field(sa_type=Text(), default="sse") + auth_type: Optional[str] = Field(sa_type=Text(), default=None) + credentials: Optional[Any] = Field( + sa_type=JSONB(), + default_factory=dict, + sa_column_kwargs={"server_default": text("'{}'")}, + ) + created_at: Optional[datetime] = Field( + sa_type=DateTime(timezone=True), + default_factory=lambda: __import__("datetime").datetime.utcnow(), + sa_column_kwargs={"server_default": text("CURRENT_TIMESTAMP")}, + ) + created_by: Optional[str] = Field(sa_type=Text(), default=None) + updated_at: Optional[datetime] = Field( + sa_type=DateTime(timezone=True), + default_factory=lambda: __import__("datetime").datetime.utcnow(), + sa_column_kwargs={ + "server_default": text("CURRENT_TIMESTAMP"), + "onupdate": lambda: __import__("datetime").datetime.utcnow(), + }, + ) + updated_by: Optional[str] = Field(sa_type=Text(), default=None) + mcp_info: Optional[Any] = Field( + sa_type=JSONB(), + default_factory=dict, + sa_column_kwargs={"server_default": text("'{}'")}, + ) + mcp_access_groups: List[str] = Field(sa_type=ARRAY(Text()), default_factory=list) + allowed_tools: List[str] = Field(sa_type=ARRAY(Text()), default_factory=list) + tool_name_to_display_name: Optional[Any] = Field( + sa_type=JSONB(), + default_factory=dict, + sa_column_kwargs={"server_default": text("'{}'")}, + ) + tool_name_to_description: Optional[Any] = Field( + sa_type=JSONB(), + default_factory=dict, + sa_column_kwargs={"server_default": text("'{}'")}, + ) + extra_headers: List[str] = Field(sa_type=ARRAY(Text()), default_factory=list) + static_headers: Optional[Any] = Field( + sa_type=JSONB(), + default_factory=dict, + sa_column_kwargs={"server_default": text("'{}'")}, + ) + status: Optional[str] = Field(sa_type=Text(), default="unknown") + last_health_check: Optional[datetime] = Field( + sa_type=DateTime(timezone=True), default=None + ) + health_check_error: Optional[str] = Field(sa_type=Text(), default=None) + command: Optional[str] = Field(sa_type=Text(), default=None) + args: List[str] = Field(sa_type=ARRAY(Text()), default_factory=list) + env: Optional[Any] = Field( + sa_type=JSONB(), + default_factory=dict, + sa_column_kwargs={"server_default": text("'{}'")}, + ) + authorization_url: Optional[str] = Field(sa_type=Text(), default=None) + token_url: Optional[str] = Field(sa_type=Text(), default=None) + registration_url: Optional[str] = Field(sa_type=Text(), default=None) + allow_all_keys: bool = Field(sa_type=Boolean(), default=False) + available_on_public_internet: bool = Field(sa_type=Boolean(), default=True) + delegate_auth_to_upstream: bool = Field(sa_type=Boolean(), default=False) + is_byok: bool = Field(sa_type=Boolean(), default=False) + byok_description: List[str] = Field(sa_type=ARRAY(Text()), default_factory=list) + byok_api_key_help_url: Optional[str] = Field(sa_type=Text(), default=None) + source_url: Optional[str] = Field(sa_type=Text(), default=None) + approval_status: Optional[str] = Field(sa_type=Text(), default="active") + submitted_by: Optional[str] = Field(sa_type=Text(), default=None) + submitted_at: Optional[datetime] = Field( + sa_type=DateTime(timezone=True), default=None + ) + reviewed_at: Optional[datetime] = Field( + sa_type=DateTime(timezone=True), default=None + ) + review_notes: Optional[str] = Field(sa_type=Text(), default=None) + + +class LiteLLMMCPToolsetTable(SQLModel, table=True): + __tablename__ = "LiteLLM_MCPToolsetTable" + + toolset_id: str = Field( + sa_type=Text(), + primary_key=True, + default_factory=lambda: str(__import__("uuid").uuid4()), + ) + toolset_name: str = Field(sa_type=Text(), unique=True) + description: Optional[str] = Field(sa_type=Text(), default=None) + tools: Any = Field( + sa_type=JSONB(), + default_factory=list, + sa_column_kwargs={"server_default": text("'[]'")}, + ) + created_at: datetime = Field( + sa_type=DateTime(timezone=True), + default_factory=lambda: __import__("datetime").datetime.utcnow(), + sa_column_kwargs={"server_default": text("CURRENT_TIMESTAMP")}, + ) + created_by: Optional[str] = Field(sa_type=Text(), default=None) + updated_at: datetime = Field( + sa_type=DateTime(timezone=True), + default_factory=lambda: __import__("datetime").datetime.utcnow(), + sa_column_kwargs={ + "server_default": text("CURRENT_TIMESTAMP"), + "onupdate": lambda: __import__("datetime").datetime.utcnow(), + }, + ) + updated_by: Optional[str] = Field(sa_type=Text(), default=None) + + +class LiteLLMMCPUserCredentials(SQLModel, table=True): + __tablename__ = "LiteLLM_MCPUserCredentials" + __table_args__ = (UniqueConstraint("user_id", "server_id"),) + + id: str = Field( + sa_type=Text(), + primary_key=True, + default_factory=lambda: str(__import__("uuid").uuid4()), + ) + user_id: str = Field(sa_type=Text()) + server_id: str = Field(sa_type=Text()) + credential_b64: str = Field(sa_type=Text()) + created_at: datetime = Field( + sa_type=DateTime(timezone=True), + default_factory=lambda: __import__("datetime").datetime.utcnow(), + sa_column_kwargs={"server_default": text("CURRENT_TIMESTAMP")}, + ) + updated_at: datetime = Field( + sa_type=DateTime(timezone=True), + default_factory=lambda: __import__("datetime").datetime.utcnow(), + sa_column_kwargs={ + "server_default": text("CURRENT_TIMESTAMP"), + "onupdate": lambda: __import__("datetime").datetime.utcnow(), + }, + ) + + +class LiteLLMManagedFileTable(SQLModel, table=True): + __tablename__ = "LiteLLM_ManagedFileTable" + __table_args__ = ( + Index("LiteLLM_ManagedFileTable_unified_file_id_idx", "unified_file_id"), + Index( + "LiteLLM_ManagedFileTable_team_id_created_at_idx", "team_id", "created_at" + ), + ) + + id: str = Field( + sa_type=Text(), + primary_key=True, + default_factory=lambda: str(__import__("uuid").uuid4()), + ) + unified_file_id: str = Field(sa_type=Text(), unique=True) + file_object: Optional[Any] = Field(sa_type=JSONB(), default=None) + model_mappings: Any = Field(sa_type=JSONB()) + flat_model_file_ids: List[str] = Field(sa_type=ARRAY(Text()), default_factory=list) + storage_backend: Optional[str] = Field(sa_type=Text(), default=None) + storage_url: Optional[str] = Field(sa_type=Text(), default=None) + created_at: datetime = Field( + sa_type=DateTime(timezone=True), + default_factory=lambda: __import__("datetime").datetime.utcnow(), + sa_column_kwargs={"server_default": text("CURRENT_TIMESTAMP")}, + ) + created_by: Optional[str] = Field(sa_type=Text(), default=None) + team_id: Optional[str] = Field(sa_type=Text(), default=None) + updated_at: datetime = Field( + sa_type=DateTime(timezone=True), + sa_column_kwargs={"onupdate": lambda: __import__("datetime").datetime.utcnow()}, + ) + updated_by: Optional[str] = Field(sa_type=Text(), default=None) + + +class LiteLLMManagedObjectTable(SQLModel, table=True): + __tablename__ = "LiteLLM_ManagedObjectTable" + __table_args__ = ( + Index("LiteLLM_ManagedObjectTable_unified_object_id_idx", "unified_object_id"), + Index("LiteLLM_ManagedObjectTable_model_object_id_idx", "model_object_id"), + Index( + "LiteLLM_ManagedObjectTable_team_id_created_at_idx", "team_id", "created_at" + ), + ) + + id: str = Field( + sa_type=Text(), + primary_key=True, + default_factory=lambda: str(__import__("uuid").uuid4()), + ) + unified_object_id: str = Field(sa_type=Text(), unique=True) + model_object_id: str = Field(sa_type=Text(), unique=True) + file_object: Any = Field(sa_type=JSONB()) + file_purpose: str = Field(sa_type=Text()) + status: Optional[str] = Field(sa_type=Text(), default=None) + batch_processed: bool = Field(sa_type=Boolean(), default=False) + created_at: datetime = Field( + sa_type=DateTime(timezone=True), + default_factory=lambda: __import__("datetime").datetime.utcnow(), + sa_column_kwargs={"server_default": text("CURRENT_TIMESTAMP")}, + ) + created_by: Optional[str] = Field(sa_type=Text(), default=None) + team_id: Optional[str] = Field(sa_type=Text(), default=None) + updated_at: datetime = Field( + sa_type=DateTime(timezone=True), + sa_column_kwargs={"onupdate": lambda: __import__("datetime").datetime.utcnow()}, + ) + updated_by: Optional[str] = Field(sa_type=Text(), default=None) + + +class LiteLLMManagedVectorStoreIndexTable(SQLModel, table=True): + __tablename__ = "LiteLLM_ManagedVectorStoreIndexTable" + + id: str = Field( + sa_type=Text(), + primary_key=True, + default_factory=lambda: str(__import__("uuid").uuid4()), + ) + index_name: str = Field(sa_type=Text(), unique=True) + litellm_params: Any = Field(sa_type=JSONB()) + index_info: Optional[Any] = Field(sa_type=JSONB(), default=None) + created_at: datetime = Field( + sa_type=DateTime(timezone=True), + default_factory=lambda: __import__("datetime").datetime.utcnow(), + sa_column_kwargs={"server_default": text("CURRENT_TIMESTAMP")}, + ) + created_by: Optional[str] = Field(sa_type=Text(), default=None) + updated_at: datetime = Field( + sa_type=DateTime(timezone=True), + sa_column_kwargs={"onupdate": lambda: __import__("datetime").datetime.utcnow()}, + ) + updated_by: Optional[str] = Field(sa_type=Text(), default=None) + + +class LiteLLMManagedVectorStoreTable(SQLModel, table=True): + __tablename__ = "LiteLLM_ManagedVectorStoreTable" + __table_args__ = ( + Index( + "LiteLLM_ManagedVectorStoreTable_unified_resource_id_idx", + "unified_resource_id", + ), + Index( + "LiteLLM_ManagedVectorStoreTable_team_id_created_at_idx", + "team_id", + "created_at", + ), + ) + + id: str = Field( + sa_type=Text(), + primary_key=True, + default_factory=lambda: str(__import__("uuid").uuid4()), + ) + unified_resource_id: str = Field(sa_type=Text(), unique=True) + resource_object: Optional[Any] = Field(sa_type=JSONB(), default=None) + model_mappings: Any = Field(sa_type=JSONB()) + flat_model_resource_ids: List[str] = Field( + sa_type=ARRAY(Text()), default_factory=list + ) + storage_backend: Optional[str] = Field(sa_type=Text(), default=None) + storage_url: Optional[str] = Field(sa_type=Text(), default=None) + created_at: datetime = Field( + sa_type=DateTime(timezone=True), + default_factory=lambda: __import__("datetime").datetime.utcnow(), + sa_column_kwargs={"server_default": text("CURRENT_TIMESTAMP")}, + ) + created_by: Optional[str] = Field(sa_type=Text(), default=None) + team_id: Optional[str] = Field(sa_type=Text(), default=None) + updated_at: datetime = Field( + sa_type=DateTime(timezone=True), + sa_column_kwargs={"onupdate": lambda: __import__("datetime").datetime.utcnow()}, + ) + updated_by: Optional[str] = Field(sa_type=Text(), default=None) + + +class LiteLLMManagedVectorStoresTable(SQLModel, table=True): + __tablename__ = "LiteLLM_ManagedVectorStoresTable" + __table_args__ = ( + Index("LiteLLM_ManagedVectorStoresTable_team_id_idx", "team_id"), + Index("LiteLLM_ManagedVectorStoresTable_user_id_idx", "user_id"), + ) + + vector_store_id: str = Field(sa_type=Text(), primary_key=True) + custom_llm_provider: str = Field(sa_type=Text()) + vector_store_name: Optional[str] = Field(sa_type=Text(), default=None) + vector_store_description: Optional[str] = Field(sa_type=Text(), default=None) + vector_store_metadata: Optional[Any] = Field(sa_type=JSONB(), default=None) + created_at: datetime = Field( + sa_type=DateTime(timezone=True), + default_factory=lambda: __import__("datetime").datetime.utcnow(), + sa_column_kwargs={"server_default": text("CURRENT_TIMESTAMP")}, + ) + updated_at: datetime = Field( + sa_type=DateTime(timezone=True), + sa_column_kwargs={"onupdate": lambda: __import__("datetime").datetime.utcnow()}, + ) + litellm_credential_name: Optional[str] = Field(sa_type=Text(), default=None) + litellm_params: Optional[Any] = Field(sa_type=JSONB(), default=None) + team_id: Optional[str] = Field(sa_type=Text(), default=None) + user_id: Optional[str] = Field(sa_type=Text(), default=None) + + +class LiteLLMMemoryTable(SQLModel, table=True): + __tablename__ = "LiteLLM_MemoryTable" + __table_args__ = ( + Index("LiteLLM_MemoryTable_user_id_idx", "user_id"), + Index("LiteLLM_MemoryTable_team_id_idx", "team_id"), + ) + + memory_id: str = Field( + sa_type=Text(), + primary_key=True, + default_factory=lambda: str(__import__("uuid").uuid4()), + ) + key: str = Field(sa_type=Text(), unique=True) + value: str = Field(sa_type=Text()) + metadata_: Optional[Any] = Field( + sa_type=JSONB(), default=None, sa_column_kwargs={"name": "metadata"} + ) + user_id: Optional[str] = Field(sa_type=Text(), default=None) + team_id: Optional[str] = Field(sa_type=Text(), default=None) + created_at: datetime = Field( + sa_type=DateTime(timezone=True), + default_factory=lambda: __import__("datetime").datetime.utcnow(), + sa_column_kwargs={"server_default": text("CURRENT_TIMESTAMP")}, + ) + created_by: Optional[str] = Field(sa_type=Text(), default=None) + updated_at: datetime = Field( + sa_type=DateTime(timezone=True), + default_factory=lambda: __import__("datetime").datetime.utcnow(), + sa_column_kwargs={ + "server_default": text("CURRENT_TIMESTAMP"), + "onupdate": lambda: __import__("datetime").datetime.utcnow(), + }, + ) + updated_by: Optional[str] = Field(sa_type=Text(), default=None) + + +class LiteLLMModelTable(SQLModel, table=True): + __tablename__ = "LiteLLM_ModelTable" + + id: int = Field(sa_type=Integer(), primary_key=True) + model_aliases: Optional[Any] = Field( + sa_type=JSONB(), default=None, sa_column_kwargs={"name": "aliases"} + ) + created_at: datetime = Field( + sa_type=DateTime(timezone=True), + default_factory=lambda: __import__("datetime").datetime.utcnow(), + sa_column_kwargs={"server_default": text("CURRENT_TIMESTAMP")}, + ) + created_by: str = Field(sa_type=Text()) + updated_at: datetime = Field( + sa_type=DateTime(timezone=True), + default_factory=lambda: __import__("datetime").datetime.utcnow(), + sa_column_kwargs={ + "server_default": text("CURRENT_TIMESTAMP"), + "onupdate": lambda: __import__("datetime").datetime.utcnow(), + }, + ) + updated_by: str = Field(sa_type=Text()) + team: Optional[str] = Field(sa_type=Text(), default=None) + + +class LiteLLMObjectPermissionTable(SQLModel, table=True): + __tablename__ = "LiteLLM_ObjectPermissionTable" + + object_permission_id: str = Field( + sa_type=Text(), + primary_key=True, + default_factory=lambda: str(__import__("uuid").uuid4()), + ) + mcp_servers: List[str] = Field(sa_type=ARRAY(Text()), default_factory=list) + mcp_access_groups: List[str] = Field(sa_type=ARRAY(Text()), default_factory=list) + mcp_tool_permissions: Optional[Any] = Field(sa_type=JSONB(), default=None) + vector_stores: List[str] = Field(sa_type=ARRAY(Text()), default_factory=list) + agents: List[str] = Field(sa_type=ARRAY(Text()), default_factory=list) + agent_access_groups: List[str] = Field(sa_type=ARRAY(Text()), default_factory=list) + models: List[str] = Field(sa_type=ARRAY(Text()), default_factory=list) + blocked_tools: List[str] = Field(sa_type=ARRAY(Text()), default_factory=list) + mcp_toolsets: List[str] = Field(sa_type=ARRAY(Text()), default_factory=list) + search_tools: List[str] = Field(sa_type=ARRAY(Text()), default_factory=list) + + +class LiteLLMOrganizationMembership(SQLModel, table=True): + __tablename__ = "LiteLLM_OrganizationMembership" + __table_args__ = ( + PrimaryKeyConstraint("user_id", "organization_id"), + UniqueConstraint("user_id", "organization_id"), + ) + + user_id: str = Field(sa_type=Text()) + organization_id: str = Field(sa_type=Text()) + user_role: Optional[str] = Field(sa_type=Text(), default=None) + spend: Optional[float] = Field(sa_type=Double(), default=0.0) + budget_id: Optional[str] = Field(sa_type=Text(), default=None) + created_at: Optional[datetime] = Field( + sa_type=DateTime(timezone=True), + default_factory=lambda: __import__("datetime").datetime.utcnow(), + sa_column_kwargs={"server_default": text("CURRENT_TIMESTAMP")}, + ) + updated_at: Optional[datetime] = Field( + sa_type=DateTime(timezone=True), + default_factory=lambda: __import__("datetime").datetime.utcnow(), + sa_column_kwargs={ + "server_default": text("CURRENT_TIMESTAMP"), + "onupdate": lambda: __import__("datetime").datetime.utcnow(), + }, + ) + + +class LiteLLMOrganizationTable(SQLModel, table=True): + __tablename__ = "LiteLLM_OrganizationTable" + + organization_id: str = Field( + sa_type=Text(), + primary_key=True, + default_factory=lambda: str(__import__("uuid").uuid4()), + ) + organization_alias: str = Field(sa_type=Text()) + budget_id: str = Field(sa_type=Text()) + metadata_: Any = Field( + sa_type=JSONB(), + default_factory=dict, + sa_column_kwargs={"name": "metadata", "server_default": text("'{}'")}, + ) + models: List[str] = Field(sa_type=ARRAY(Text()), default_factory=list) + spend: float = Field(sa_type=Double(), default=0.0) + model_spend: Any = Field( + sa_type=JSONB(), + default_factory=dict, + sa_column_kwargs={"server_default": text("'{}'")}, + ) + object_permission_id: Optional[str] = Field(sa_type=Text(), default=None) + created_at: datetime = Field( + sa_type=DateTime(timezone=True), + default_factory=lambda: __import__("datetime").datetime.utcnow(), + sa_column_kwargs={"server_default": text("CURRENT_TIMESTAMP")}, + ) + created_by: str = Field(sa_type=Text()) + updated_at: datetime = Field( + sa_type=DateTime(timezone=True), + default_factory=lambda: __import__("datetime").datetime.utcnow(), + sa_column_kwargs={ + "server_default": text("CURRENT_TIMESTAMP"), + "onupdate": lambda: __import__("datetime").datetime.utcnow(), + }, + ) + updated_by: str = Field(sa_type=Text()) + + +class LiteLLMPolicyAttachmentTable(SQLModel, table=True): + __tablename__ = "LiteLLM_PolicyAttachmentTable" + + attachment_id: str = Field( + sa_type=Text(), + primary_key=True, + default_factory=lambda: str(__import__("uuid").uuid4()), + ) + policy_name: str = Field(sa_type=Text()) + scope: Optional[str] = Field(sa_type=Text(), default=None) + teams: List[str] = Field(sa_type=ARRAY(Text()), default_factory=list) + keys: List[str] = Field(sa_type=ARRAY(Text()), default_factory=list) + models: List[str] = Field(sa_type=ARRAY(Text()), default_factory=list) + tags: List[str] = Field(sa_type=ARRAY(Text()), default_factory=list) + created_at: datetime = Field( + sa_type=DateTime(timezone=True), + default_factory=lambda: __import__("datetime").datetime.utcnow(), + sa_column_kwargs={"server_default": text("CURRENT_TIMESTAMP")}, + ) + created_by: Optional[str] = Field(sa_type=Text(), default=None) + updated_at: datetime = Field( + sa_type=DateTime(timezone=True), + default_factory=lambda: __import__("datetime").datetime.utcnow(), + sa_column_kwargs={ + "server_default": text("CURRENT_TIMESTAMP"), + "onupdate": lambda: __import__("datetime").datetime.utcnow(), + }, + ) + updated_by: Optional[str] = Field(sa_type=Text(), default=None) + + +class LiteLLMPolicyTable(SQLModel, table=True): + __tablename__ = "LiteLLM_PolicyTable" + __table_args__ = ( + UniqueConstraint("policy_name", "version_number"), + Index( + "LiteLLM_PolicyTable_policy_name_version_status_idx", + "policy_name", + "version_status", + ), + ) + + policy_id: str = Field( + sa_type=Text(), + primary_key=True, + default_factory=lambda: str(__import__("uuid").uuid4()), + ) + policy_name: str = Field(sa_type=Text()) + version_number: int = Field(sa_type=Integer(), default=1) + version_status: str = Field(sa_type=Text(), default="production") + parent_version_id: Optional[str] = Field(sa_type=Text(), default=None) + is_latest: bool = Field(sa_type=Boolean(), default=True) + published_at: Optional[datetime] = Field( + sa_type=DateTime(timezone=True), default=None + ) + production_at: Optional[datetime] = Field( + sa_type=DateTime(timezone=True), default=None + ) + inherit: Optional[str] = Field(sa_type=Text(), default=None) + description: Optional[str] = Field(sa_type=Text(), default=None) + guardrails_add: List[str] = Field(sa_type=ARRAY(Text()), default_factory=list) + guardrails_remove: List[str] = Field(sa_type=ARRAY(Text()), default_factory=list) + condition: Optional[Any] = Field( + sa_type=JSONB(), + default_factory=dict, + sa_column_kwargs={"server_default": text("'{}'")}, + ) + pipeline: Optional[Any] = Field(sa_type=JSONB(), default=None) + created_at: datetime = Field( + sa_type=DateTime(timezone=True), + default_factory=lambda: __import__("datetime").datetime.utcnow(), + sa_column_kwargs={"server_default": text("CURRENT_TIMESTAMP")}, + ) + created_by: Optional[str] = Field(sa_type=Text(), default=None) + updated_at: datetime = Field( + sa_type=DateTime(timezone=True), + default_factory=lambda: __import__("datetime").datetime.utcnow(), + sa_column_kwargs={ + "server_default": text("CURRENT_TIMESTAMP"), + "onupdate": lambda: __import__("datetime").datetime.utcnow(), + }, + ) + updated_by: Optional[str] = Field(sa_type=Text(), default=None) + + +class LiteLLMProjectTable(SQLModel, table=True): + __tablename__ = "LiteLLM_ProjectTable" + + project_id: str = Field( + sa_type=Text(), + primary_key=True, + default_factory=lambda: str(__import__("uuid").uuid4()), + ) + project_alias: Optional[str] = Field(sa_type=Text(), default=None) + description: Optional[str] = Field(sa_type=Text(), default=None) + team_id: Optional[str] = Field(sa_type=Text(), default=None) + budget_id: Optional[str] = Field(sa_type=Text(), default=None) + metadata_: Any = Field( + sa_type=JSONB(), + default_factory=dict, + sa_column_kwargs={"name": "metadata", "server_default": text("'{}'")}, + ) + models: List[str] = Field(sa_type=ARRAY(Text()), default_factory=list) + spend: float = Field(sa_type=Double(), default=0.0) + model_spend: Any = Field( + sa_type=JSONB(), + default_factory=dict, + sa_column_kwargs={"server_default": text("'{}'")}, + ) + model_rpm_limit: Any = Field( + sa_type=JSONB(), + default_factory=dict, + sa_column_kwargs={"server_default": text("'{}'")}, + ) + model_tpm_limit: Any = Field( + sa_type=JSONB(), + default_factory=dict, + sa_column_kwargs={"server_default": text("'{}'")}, + ) + blocked: bool = Field(sa_type=Boolean(), default=False) + object_permission_id: Optional[str] = Field(sa_type=Text(), default=None) + created_at: datetime = Field( + sa_type=DateTime(timezone=True), + default_factory=lambda: __import__("datetime").datetime.utcnow(), + sa_column_kwargs={"server_default": text("CURRENT_TIMESTAMP")}, + ) + created_by: str = Field(sa_type=Text()) + updated_at: datetime = Field( + sa_type=DateTime(timezone=True), + default_factory=lambda: __import__("datetime").datetime.utcnow(), + sa_column_kwargs={ + "server_default": text("CURRENT_TIMESTAMP"), + "onupdate": lambda: __import__("datetime").datetime.utcnow(), + }, + ) + updated_by: str = Field(sa_type=Text()) + + +class LiteLLMPromptTable(SQLModel, table=True): + __tablename__ = "LiteLLM_PromptTable" + __table_args__ = ( + UniqueConstraint("prompt_id", "version", "environment"), + Index( + "LiteLLM_PromptTable_prompt_id_environment_idx", "prompt_id", "environment" + ), + Index("LiteLLM_PromptTable_prompt_id_idx", "prompt_id"), + ) + + id: str = Field( + sa_type=Text(), + primary_key=True, + default_factory=lambda: str(__import__("uuid").uuid4()), + ) + prompt_id: str = Field(sa_type=Text()) + version: int = Field(sa_type=Integer(), default=1) + environment: str = Field(sa_type=Text(), default="development") + created_by: Optional[str] = Field(sa_type=Text(), default=None) + litellm_params: Any = Field(sa_type=JSONB()) + prompt_info: Optional[Any] = Field(sa_type=JSONB(), default=None) + created_at: datetime = Field( + sa_type=DateTime(timezone=True), + default_factory=lambda: __import__("datetime").datetime.utcnow(), + sa_column_kwargs={"server_default": text("CURRENT_TIMESTAMP")}, + ) + updated_at: datetime = Field( + sa_type=DateTime(timezone=True), + sa_column_kwargs={"onupdate": lambda: __import__("datetime").datetime.utcnow()}, + ) + + +class LiteLLMProxyModelTable(SQLModel, table=True): + __tablename__ = "LiteLLM_ProxyModelTable" + + model_id: str = Field( + sa_type=Text(), + primary_key=True, + default_factory=lambda: str(__import__("uuid").uuid4()), + ) + model_name: str = Field(sa_type=Text()) + litellm_params: Any = Field(sa_type=JSONB()) + model_info: Optional[Any] = Field(sa_type=JSONB(), default=None) + blocked: bool = Field(sa_type=Boolean(), default=False) + created_at: datetime = Field( + sa_type=DateTime(timezone=True), + default_factory=lambda: __import__("datetime").datetime.utcnow(), + sa_column_kwargs={"server_default": text("CURRENT_TIMESTAMP")}, + ) + created_by: str = Field(sa_type=Text()) + updated_at: datetime = Field( + sa_type=DateTime(timezone=True), + default_factory=lambda: __import__("datetime").datetime.utcnow(), + sa_column_kwargs={ + "server_default": text("CURRENT_TIMESTAMP"), + "onupdate": lambda: __import__("datetime").datetime.utcnow(), + }, + ) + updated_by: str = Field(sa_type=Text()) + + +class LiteLLMSSOConfig(SQLModel, table=True): + __tablename__ = "LiteLLM_SSOConfig" + + id: str = Field(sa_type=Text(), primary_key=True, default="sso_config") + sso_settings: Any = Field(sa_type=JSONB()) + created_at: datetime = Field( + sa_type=DateTime(timezone=True), + default_factory=lambda: __import__("datetime").datetime.utcnow(), + sa_column_kwargs={"server_default": text("CURRENT_TIMESTAMP")}, + ) + updated_at: datetime = Field( + sa_type=DateTime(timezone=True), + sa_column_kwargs={"onupdate": lambda: __import__("datetime").datetime.utcnow()}, + ) + + +class LiteLLMSearchToolsTable(SQLModel, table=True): + __tablename__ = "LiteLLM_SearchToolsTable" + + search_tool_id: str = Field( + sa_type=Text(), + primary_key=True, + default_factory=lambda: str(__import__("uuid").uuid4()), + ) + search_tool_name: str = Field(sa_type=Text(), unique=True) + litellm_params: Any = Field(sa_type=JSONB()) + search_tool_info: Optional[Any] = Field(sa_type=JSONB(), default=None) + created_at: datetime = Field( + sa_type=DateTime(timezone=True), + default_factory=lambda: __import__("datetime").datetime.utcnow(), + sa_column_kwargs={"server_default": text("CURRENT_TIMESTAMP")}, + ) + updated_at: datetime = Field( + sa_type=DateTime(timezone=True), + sa_column_kwargs={"onupdate": lambda: __import__("datetime").datetime.utcnow()}, + ) + + +class LiteLLMSkillsTable(SQLModel, table=True): + __tablename__ = "LiteLLM_SkillsTable" + + skill_id: str = Field( + sa_type=Text(), + primary_key=True, + default_factory=lambda: str(__import__("uuid").uuid4()), + ) + display_title: Optional[str] = Field(sa_type=Text(), default=None) + description: Optional[str] = Field(sa_type=Text(), default=None) + instructions: Optional[str] = Field(sa_type=Text(), default=None) + source: str = Field(sa_type=Text(), default="custom") + latest_version: Optional[str] = Field(sa_type=Text(), default=None) + file_content: Optional[bytes] = Field(sa_type=LargeBinary(), default=None) + file_name: Optional[str] = Field(sa_type=Text(), default=None) + file_type: Optional[str] = Field(sa_type=Text(), default=None) + metadata_: Optional[Any] = Field( + sa_type=JSONB(), + default_factory=dict, + sa_column_kwargs={"name": "metadata", "server_default": text("'{}'")}, + ) + created_at: datetime = Field( + sa_type=DateTime(timezone=True), + default_factory=lambda: __import__("datetime").datetime.utcnow(), + sa_column_kwargs={"server_default": text("CURRENT_TIMESTAMP")}, + ) + created_by: Optional[str] = Field(sa_type=Text(), default=None) + updated_at: datetime = Field( + sa_type=DateTime(timezone=True), + default_factory=lambda: __import__("datetime").datetime.utcnow(), + sa_column_kwargs={ + "server_default": text("CURRENT_TIMESTAMP"), + "onupdate": lambda: __import__("datetime").datetime.utcnow(), + }, + ) + updated_by: Optional[str] = Field(sa_type=Text(), default=None) + + +class LiteLLMSpendLogGuardrailIndex(SQLModel, table=True): + __tablename__ = "LiteLLM_SpendLogGuardrailIndex" + __table_args__ = ( + PrimaryKeyConstraint("request_id", "guardrail_id"), + Index( + "LiteLLM_SpendLogGuardrailIndex_guardrail_id_start_time_idx", + "guardrail_id", + "start_time", + ), + Index( + "LiteLLM_SpendLogGuardrailIndex_policy_id_start_time_idx", + "policy_id", + "start_time", + ), + ) + + request_id: str = Field(sa_type=Text()) + guardrail_id: str = Field(sa_type=Text()) + policy_id: Optional[str] = Field(sa_type=Text(), default=None) + start_time: datetime = Field(sa_type=DateTime(timezone=True)) + + +class LiteLLMSpendLogToolIndex(SQLModel, table=True): + __tablename__ = "LiteLLM_SpendLogToolIndex" + __table_args__ = ( + PrimaryKeyConstraint("request_id", "tool_name"), + Index( + "LiteLLM_SpendLogToolIndex_tool_name_start_time_idx", + "tool_name", + "start_time", + ), + ) + + request_id: str = Field(sa_type=Text()) + tool_name: str = Field(sa_type=Text()) + start_time: datetime = Field(sa_type=DateTime(timezone=True)) + + +class LiteLLMSpendLogs(SQLModel, table=True): + __tablename__ = "LiteLLM_SpendLogs" + __table_args__ = ( + Index("LiteLLM_SpendLogs_startTime_idx", "startTime"), + Index("LiteLLM_SpendLogs_startTime_request_id_idx", "startTime", "request_id"), + Index("LiteLLM_SpendLogs_end_user_idx", "end_user"), + Index("LiteLLM_SpendLogs_session_id_idx", "session_id"), + ) + + request_id: str = Field(sa_type=Text(), primary_key=True) + call_type: str = Field(sa_type=Text()) + api_key: str = Field(sa_type=Text()) + spend: float = Field(sa_type=Double(), default=0.0) + total_tokens: int = Field(sa_type=Integer(), default=0) + prompt_tokens: int = Field(sa_type=Integer(), default=0) + completion_tokens: int = Field(sa_type=Integer(), default=0) + startTime: datetime = Field(sa_type=DateTime(timezone=True)) + endTime: datetime = Field(sa_type=DateTime(timezone=True)) + request_duration_ms: Optional[int] = Field(sa_type=Integer(), default=None) + completionStartTime: Optional[datetime] = Field( + sa_type=DateTime(timezone=True), default=None + ) + model: str = Field(sa_type=Text(), default="") + model_id: Optional[str] = Field(sa_type=Text(), default="") + model_group: Optional[str] = Field(sa_type=Text(), default="") + custom_llm_provider: Optional[str] = Field(sa_type=Text(), default="") + api_base: Optional[str] = Field(sa_type=Text(), default="") + user: Optional[str] = Field(sa_type=Text(), default="") + metadata_: Optional[Any] = Field( + sa_type=JSONB(), + default_factory=dict, + sa_column_kwargs={"name": "metadata", "server_default": text("'{}'")}, + ) + cache_hit: Optional[str] = Field(sa_type=Text(), default="") + cache_key: Optional[str] = Field(sa_type=Text(), default="") + request_tags: Optional[Any] = Field( + sa_type=JSONB(), + default_factory=list, + sa_column_kwargs={"server_default": text("'[]'")}, + ) + team_id: Optional[str] = Field(sa_type=Text(), default=None) + organization_id: Optional[str] = Field(sa_type=Text(), default=None) + end_user: Optional[str] = Field(sa_type=Text(), default=None) + requester_ip_address: Optional[str] = Field(sa_type=Text(), default=None) + messages: Optional[Any] = Field( + sa_type=JSONB(), + default_factory=dict, + sa_column_kwargs={"server_default": text("'{}'")}, + ) + response: Optional[Any] = Field( + sa_type=JSONB(), + default_factory=dict, + sa_column_kwargs={"server_default": text("'{}'")}, + ) + session_id: Optional[str] = Field(sa_type=Text(), default=None) + status: Optional[str] = Field(sa_type=Text(), default=None) + mcp_namespaced_tool_name: Optional[str] = Field(sa_type=Text(), default=None) + agent_id: Optional[str] = Field(sa_type=Text(), default=None) + proxy_server_request: Optional[Any] = Field( + sa_type=JSONB(), + default_factory=dict, + sa_column_kwargs={"server_default": text("'{}'")}, + ) + + +class LiteLLMTagTable(SQLModel, table=True): + __tablename__ = "LiteLLM_TagTable" + + tag_name: str = Field(sa_type=Text(), primary_key=True) + description: Optional[str] = Field(sa_type=Text(), default=None) + models: List[str] = Field(sa_type=ARRAY(Text()), default_factory=list) + model_info: Optional[Any] = Field(sa_type=JSONB(), default=None) + spend: float = Field(sa_type=Double(), default=0.0) + budget_id: Optional[str] = Field(sa_type=Text(), default=None) + created_at: datetime = Field( + sa_type=DateTime(timezone=True), + default_factory=lambda: __import__("datetime").datetime.utcnow(), + sa_column_kwargs={"server_default": text("CURRENT_TIMESTAMP")}, + ) + created_by: Optional[str] = Field(sa_type=Text(), default=None) + updated_at: datetime = Field( + sa_type=DateTime(timezone=True), + default_factory=lambda: __import__("datetime").datetime.utcnow(), + sa_column_kwargs={ + "server_default": text("CURRENT_TIMESTAMP"), + "onupdate": lambda: __import__("datetime").datetime.utcnow(), + }, + ) + + +class LiteLLMTeamMembership(SQLModel, table=True): + __tablename__ = "LiteLLM_TeamMembership" + __table_args__ = (PrimaryKeyConstraint("user_id", "team_id"),) + + user_id: str = Field(sa_type=Text()) + team_id: str = Field(sa_type=Text()) + spend: float = Field(sa_type=Double(), default=0.0) + total_spend: float = Field(sa_type=Double(), default=0.0) + budget_id: Optional[str] = Field(sa_type=Text(), default=None) + + +class LiteLLMTeamTable(SQLModel, table=True): + __tablename__ = "LiteLLM_TeamTable" + __table_args__ = ( + Index("LiteLLM_TeamTable_organization_id_idx", "organization_id"), + Index("LiteLLM_TeamTable_team_alias_idx", "team_alias"), + Index("LiteLLM_TeamTable_created_at_idx", "created_at"), + ) + + team_id: str = Field( + sa_type=Text(), + primary_key=True, + default_factory=lambda: str(__import__("uuid").uuid4()), + ) + team_alias: Optional[str] = Field(sa_type=Text(), default=None) + organization_id: Optional[str] = Field(sa_type=Text(), default=None) + object_permission_id: Optional[str] = Field(sa_type=Text(), default=None) + admins: List[str] = Field(sa_type=ARRAY(Text()), default_factory=list) + members: List[str] = Field(sa_type=ARRAY(Text()), default_factory=list) + members_with_roles: Any = Field( + sa_type=JSONB(), + default_factory=dict, + sa_column_kwargs={"server_default": text("'{}'")}, + ) + metadata_: Any = Field( + sa_type=JSONB(), + default_factory=dict, + sa_column_kwargs={"name": "metadata", "server_default": text("'{}'")}, + ) + max_budget: Optional[float] = Field(sa_type=Double(), default=None) + soft_budget: Optional[float] = Field(sa_type=Double(), default=None) + spend: float = Field(sa_type=Double(), default=0.0) + models: List[str] = Field(sa_type=ARRAY(Text()), default_factory=list) + max_parallel_requests: Optional[int] = Field(sa_type=Integer(), default=None) + tpm_limit: Optional[int] = Field(sa_type=BigInteger(), default=None) + rpm_limit: Optional[int] = Field(sa_type=BigInteger(), default=None) + budget_duration: Optional[str] = Field(sa_type=Text(), default=None) + budget_reset_at: Optional[datetime] = Field( + sa_type=DateTime(timezone=True), default=None + ) + blocked: bool = Field(sa_type=Boolean(), default=False) + created_at: datetime = Field( + sa_type=DateTime(timezone=True), + default_factory=lambda: __import__("datetime").datetime.utcnow(), + sa_column_kwargs={"server_default": text("CURRENT_TIMESTAMP")}, + ) + updated_at: datetime = Field( + sa_type=DateTime(timezone=True), + default_factory=lambda: __import__("datetime").datetime.utcnow(), + sa_column_kwargs={ + "server_default": text("CURRENT_TIMESTAMP"), + "onupdate": lambda: __import__("datetime").datetime.utcnow(), + }, + ) + model_spend: Any = Field( + sa_type=JSONB(), + default_factory=dict, + sa_column_kwargs={"server_default": text("'{}'")}, + ) + model_max_budget: Any = Field( + sa_type=JSONB(), + default_factory=dict, + sa_column_kwargs={"server_default": text("'{}'")}, + ) + router_settings: Optional[Any] = Field( + sa_type=JSONB(), + default_factory=dict, + sa_column_kwargs={"server_default": text("'{}'")}, + ) + team_member_permissions: List[str] = Field( + sa_type=ARRAY(Text()), default_factory=list + ) + access_group_ids: List[str] = Field(sa_type=ARRAY(Text()), default_factory=list) + policies: List[str] = Field(sa_type=ARRAY(Text()), default_factory=list) + default_team_member_models: List[str] = Field( + sa_type=ARRAY(Text()), default_factory=list + ) + budget_limits: Optional[Any] = Field(sa_type=JSONB(), default=None) + model_id: Optional[int] = Field(sa_type=Integer(), unique=True, default=None) + allow_team_guardrail_config: bool = Field(sa_type=Boolean(), default=False) + + +class LiteLLMToolTable(SQLModel, table=True): + __tablename__ = "LiteLLM_ToolTable" + __table_args__ = ( + Index("LiteLLM_ToolTable_input_policy_idx", "input_policy"), + Index("LiteLLM_ToolTable_output_policy_idx", "output_policy"), + Index("LiteLLM_ToolTable_team_id_idx", "team_id"), + ) + + tool_id: str = Field( + sa_type=Text(), + primary_key=True, + default_factory=lambda: str(__import__("uuid").uuid4()), + ) + tool_name: str = Field(sa_type=Text(), unique=True) + origin: Optional[str] = Field(sa_type=Text(), default=None) + input_policy: str = Field(sa_type=Text(), default="untrusted") + output_policy: str = Field(sa_type=Text(), default="untrusted") + call_count: int = Field(sa_type=Integer(), default=0) + assignments: Optional[Any] = Field( + sa_type=JSONB(), + default_factory=dict, + sa_column_kwargs={"server_default": text("'{}'")}, + ) + key_hash: Optional[str] = Field(sa_type=Text(), default=None) + team_id: Optional[str] = Field(sa_type=Text(), default=None) + key_alias: Optional[str] = Field(sa_type=Text(), default=None) + user_agent: Optional[str] = Field(sa_type=Text(), default=None) + last_used_at: Optional[datetime] = Field( + sa_type=DateTime(timezone=True), default=None + ) + created_at: datetime = Field( + sa_type=DateTime(timezone=True), + default_factory=lambda: __import__("datetime").datetime.utcnow(), + sa_column_kwargs={"server_default": text("CURRENT_TIMESTAMP")}, + ) + created_by: Optional[str] = Field(sa_type=Text(), default=None) + updated_at: datetime = Field( + sa_type=DateTime(timezone=True), + default_factory=lambda: __import__("datetime").datetime.utcnow(), + sa_column_kwargs={ + "server_default": text("CURRENT_TIMESTAMP"), + "onupdate": lambda: __import__("datetime").datetime.utcnow(), + }, + ) + updated_by: Optional[str] = Field(sa_type=Text(), default=None) + + +class LiteLLMUISettings(SQLModel, table=True): + __tablename__ = "LiteLLM_UISettings" + + id: str = Field(sa_type=Text(), primary_key=True, default="ui_settings") + ui_settings: Any = Field(sa_type=JSONB()) + created_at: datetime = Field( + sa_type=DateTime(timezone=True), + default_factory=lambda: __import__("datetime").datetime.utcnow(), + sa_column_kwargs={"server_default": text("CURRENT_TIMESTAMP")}, + ) + updated_at: datetime = Field( + sa_type=DateTime(timezone=True), + sa_column_kwargs={"onupdate": lambda: __import__("datetime").datetime.utcnow()}, + ) + + +class LiteLLMUserNotifications(SQLModel, table=True): + __tablename__ = "LiteLLM_UserNotifications" + + request_id: str = Field(sa_type=Text(), primary_key=True) + user_id: str = Field(sa_type=Text()) + models: List[str] = Field(sa_type=ARRAY(Text()), default_factory=list) + justification: str = Field(sa_type=Text()) + status: str = Field(sa_type=Text()) + + +class LiteLLMUserTable(SQLModel, table=True): + __tablename__ = "LiteLLM_UserTable" + + user_id: str = Field(sa_type=Text(), primary_key=True) + user_alias: Optional[str] = Field(sa_type=Text(), default=None) + team_id: Optional[str] = Field(sa_type=Text(), default=None) + sso_user_id: Optional[str] = Field(sa_type=Text(), unique=True, default=None) + organization_id: Optional[str] = Field(sa_type=Text(), default=None) + object_permission_id: Optional[str] = Field(sa_type=Text(), default=None) + password: Optional[str] = Field(sa_type=Text(), default=None) + teams: List[str] = Field(sa_type=ARRAY(Text()), default_factory=list) + user_role: Optional[str] = Field(sa_type=Text(), default=None) + max_budget: Optional[float] = Field(sa_type=Double(), default=None) + spend: float = Field(sa_type=Double(), default=0.0) + user_email: Optional[str] = Field(sa_type=Text(), default=None) + models: List[str] = Field(sa_type=ARRAY(Text()), default_factory=list) + metadata_: Any = Field( + sa_type=JSONB(), + default_factory=dict, + sa_column_kwargs={"name": "metadata", "server_default": text("'{}'")}, + ) + max_parallel_requests: Optional[int] = Field(sa_type=Integer(), default=None) + tpm_limit: Optional[int] = Field(sa_type=BigInteger(), default=None) + rpm_limit: Optional[int] = Field(sa_type=BigInteger(), default=None) + budget_duration: Optional[str] = Field(sa_type=Text(), default=None) + budget_reset_at: Optional[datetime] = Field( + sa_type=DateTime(timezone=True), default=None + ) + allowed_cache_controls: List[str] = Field( + sa_type=ARRAY(Text()), default_factory=list + ) + policies: List[str] = Field(sa_type=ARRAY(Text()), default_factory=list) + model_spend: Any = Field( + sa_type=JSONB(), + default_factory=dict, + sa_column_kwargs={"server_default": text("'{}'")}, + ) + model_max_budget: Any = Field( + sa_type=JSONB(), + default_factory=dict, + sa_column_kwargs={"server_default": text("'{}'")}, + ) + created_at: Optional[datetime] = Field( + sa_type=DateTime(timezone=True), + default_factory=lambda: __import__("datetime").datetime.utcnow(), + sa_column_kwargs={"server_default": text("CURRENT_TIMESTAMP")}, + ) + updated_at: Optional[datetime] = Field( + sa_type=DateTime(timezone=True), + default_factory=lambda: __import__("datetime").datetime.utcnow(), + sa_column_kwargs={ + "server_default": text("CURRENT_TIMESTAMP"), + "onupdate": lambda: __import__("datetime").datetime.utcnow(), + }, + ) + + +class LiteLLMVerificationToken(SQLModel, table=True): + __tablename__ = "LiteLLM_VerificationToken" + __table_args__ = ( + Index("LiteLLM_VerificationToken_user_id_team_id_idx", "user_id", "team_id"), + Index("LiteLLM_VerificationToken_team_id_idx", "team_id"), + Index( + "LiteLLM_VerificationToken_budget_reset_at_expires_idx", + "budget_reset_at", + "expires", + ), + ) + + token: str = Field(sa_type=Text(), primary_key=True) + key_name: Optional[str] = Field(sa_type=Text(), default=None) + key_alias: Optional[str] = Field(sa_type=Text(), default=None) + soft_budget_cooldown: bool = Field(sa_type=Boolean(), default=False) + spend: float = Field(sa_type=Double(), default=0.0) + expires: Optional[datetime] = Field(sa_type=DateTime(timezone=True), default=None) + models: List[str] = Field(sa_type=ARRAY(Text()), default_factory=list) + aliases: Any = Field( + sa_type=JSONB(), + default_factory=dict, + sa_column_kwargs={"server_default": text("'{}'")}, + ) + config: Any = Field( + sa_type=JSONB(), + default_factory=dict, + sa_column_kwargs={"server_default": text("'{}'")}, + ) + router_settings: Optional[Any] = Field( + sa_type=JSONB(), + default_factory=dict, + sa_column_kwargs={"server_default": text("'{}'")}, + ) + user_id: Optional[str] = Field(sa_type=Text(), default=None) + team_id: Optional[str] = Field(sa_type=Text(), default=None) + agent_id: Optional[str] = Field(sa_type=Text(), default=None) + project_id: Optional[str] = Field(sa_type=Text(), default=None) + permissions: Any = Field( + sa_type=JSONB(), + default_factory=dict, + sa_column_kwargs={"server_default": text("'{}'")}, + ) + max_parallel_requests: Optional[int] = Field(sa_type=Integer(), default=None) + metadata_: Any = Field( + sa_type=JSONB(), + default_factory=dict, + sa_column_kwargs={"name": "metadata", "server_default": text("'{}'")}, + ) + blocked: Optional[bool] = Field(sa_type=Boolean(), default=None) + tpm_limit: Optional[int] = Field(sa_type=BigInteger(), default=None) + rpm_limit: Optional[int] = Field(sa_type=BigInteger(), default=None) + max_budget: Optional[float] = Field(sa_type=Double(), default=None) + budget_duration: Optional[str] = Field(sa_type=Text(), default=None) + budget_reset_at: Optional[datetime] = Field( + sa_type=DateTime(timezone=True), default=None + ) + allowed_cache_controls: List[str] = Field( + sa_type=ARRAY(Text()), default_factory=list + ) + allowed_routes: List[str] = Field(sa_type=ARRAY(Text()), default_factory=list) + policies: List[str] = Field(sa_type=ARRAY(Text()), default_factory=list) + access_group_ids: List[str] = Field(sa_type=ARRAY(Text()), default_factory=list) + model_spend: Any = Field( + sa_type=JSONB(), + default_factory=dict, + sa_column_kwargs={"server_default": text("'{}'")}, + ) + model_max_budget: Any = Field( + sa_type=JSONB(), + default_factory=dict, + sa_column_kwargs={"server_default": text("'{}'")}, + ) + budget_id: Optional[str] = Field(sa_type=Text(), default=None) + organization_id: Optional[str] = Field(sa_type=Text(), default=None) + object_permission_id: Optional[str] = Field(sa_type=Text(), default=None) + created_at: Optional[datetime] = Field( + sa_type=DateTime(timezone=True), + default_factory=lambda: __import__("datetime").datetime.utcnow(), + sa_column_kwargs={"server_default": text("CURRENT_TIMESTAMP")}, + ) + created_by: Optional[str] = Field(sa_type=Text(), default=None) + updated_at: Optional[datetime] = Field( + sa_type=DateTime(timezone=True), + default_factory=lambda: __import__("datetime").datetime.utcnow(), + sa_column_kwargs={ + "server_default": text("CURRENT_TIMESTAMP"), + "onupdate": lambda: __import__("datetime").datetime.utcnow(), + }, + ) + updated_by: Optional[str] = Field(sa_type=Text(), default=None) + last_active: Optional[datetime] = Field( + sa_type=DateTime(timezone=True), default=None + ) + rotation_count: Optional[int] = Field(sa_type=Integer(), default=0) + auto_rotate: Optional[bool] = Field(sa_type=Boolean(), default=False) + rotation_interval: Optional[str] = Field(sa_type=Text(), default=None) + last_rotation_at: Optional[datetime] = Field( + sa_type=DateTime(timezone=True), default=None + ) + key_rotation_at: Optional[datetime] = Field( + sa_type=DateTime(timezone=True), default=None + ) + budget_limits: Optional[Any] = Field(sa_type=JSONB(), default=None) + + +class LiteLLMWorkflowEvent(SQLModel, table=True): + __tablename__ = "LiteLLM_WorkflowEvent" + __table_args__ = ( + UniqueConstraint("run_id", "sequence_number"), + Index("LiteLLM_WorkflowEvent_run_id_idx", "run_id"), + ) + + event_id: str = Field( + sa_type=Text(), + primary_key=True, + default_factory=lambda: str(__import__("uuid").uuid4()), + ) + run_id: str = Field(sa_type=Text()) + event_type: str = Field(sa_type=Text()) + step_name: str = Field(sa_type=Text()) + sequence_number: int = Field(sa_type=Integer()) + data: Optional[Any] = Field(sa_type=JSONB(), default=None) + created_at: datetime = Field( + sa_type=DateTime(timezone=True), + default_factory=lambda: __import__("datetime").datetime.utcnow(), + sa_column_kwargs={"server_default": text("CURRENT_TIMESTAMP")}, + ) + + +class LiteLLMWorkflowMessage(SQLModel, table=True): + __tablename__ = "LiteLLM_WorkflowMessage" + __table_args__ = ( + UniqueConstraint("run_id", "sequence_number"), + Index("LiteLLM_WorkflowMessage_run_id_idx", "run_id"), + ) + + message_id: str = Field( + sa_type=Text(), + primary_key=True, + default_factory=lambda: str(__import__("uuid").uuid4()), + ) + run_id: str = Field(sa_type=Text()) + role: str = Field(sa_type=Text()) + content: str = Field(sa_type=Text()) + sequence_number: int = Field(sa_type=Integer()) + session_id: Optional[str] = Field(sa_type=Text(), default=None) + created_at: datetime = Field( + sa_type=DateTime(timezone=True), + default_factory=lambda: __import__("datetime").datetime.utcnow(), + sa_column_kwargs={"server_default": text("CURRENT_TIMESTAMP")}, + ) + + +class LiteLLMWorkflowRun(SQLModel, table=True): + __tablename__ = "LiteLLM_WorkflowRun" + __table_args__ = ( + Index( + "LiteLLM_WorkflowRun_workflow_type_status_idx", "workflow_type", "status" + ), + Index("LiteLLM_WorkflowRun_session_id_idx", "session_id"), + Index("LiteLLM_WorkflowRun_created_at_idx", "created_at"), + Index("LiteLLM_WorkflowRun_created_by_idx", "created_by"), + ) + + run_id: str = Field( + sa_type=Text(), + primary_key=True, + default_factory=lambda: str(__import__("uuid").uuid4()), + ) + session_id: str = Field( + sa_type=Text(), + unique=True, + default_factory=lambda: str(__import__("uuid").uuid4()), + ) + workflow_type: str = Field(sa_type=Text()) + status: str = Field(sa_type=Text(), default="pending") + created_by: Optional[str] = Field(sa_type=Text(), default=None) + created_at: datetime = Field( + sa_type=DateTime(timezone=True), + default_factory=lambda: __import__("datetime").datetime.utcnow(), + sa_column_kwargs={"server_default": text("CURRENT_TIMESTAMP")}, + ) + updated_at: datetime = Field( + sa_type=DateTime(timezone=True), + sa_column_kwargs={"onupdate": lambda: __import__("datetime").datetime.utcnow()}, + ) + input: Optional[Any] = Field(sa_type=JSONB(), default=None) + output: Optional[Any] = Field(sa_type=JSONB(), default=None) + metadata_: Optional[Any] = Field( + sa_type=JSONB(), default=None, sa_column_kwargs={"name": "metadata"} + ) + + +ALL_MODELS: List[Type[SQLModel]] = [ + LiteLLMAccessGroupTable, + LiteLLMAdaptiveRouterSession, + LiteLLMAdaptiveRouterState, + LiteLLMAgentsTable, + LiteLLMAuditLog, + LiteLLMBudgetTable, + LiteLLMCacheConfig, + LiteLLMClaudeCodePluginTable, + LiteLLMConfig, + LiteLLMConfigOverrides, + LiteLLMCredentialsTable, + LiteLLMCronJob, + LiteLLMDailyAgentSpend, + LiteLLMDailyEndUserSpend, + LiteLLMDailyGuardrailMetrics, + LiteLLMDailyOrganizationSpend, + LiteLLMDailyPolicyMetrics, + LiteLLMDailyTagSpend, + LiteLLMDailyTeamSpend, + LiteLLMDailyUserSpend, + LiteLLMDeletedTeamTable, + LiteLLMDeletedVerificationToken, + LiteLLMDeprecatedVerificationToken, + LiteLLMEndUserTable, + LiteLLMErrorLogs, + LiteLLMGuardrailsTable, + LiteLLMHealthCheckTable, + LiteLLMInvitationLink, + LiteLLMJWTKeyMapping, + LiteLLMMCPServerTable, + LiteLLMMCPToolsetTable, + LiteLLMMCPUserCredentials, + LiteLLMManagedFileTable, + LiteLLMManagedObjectTable, + LiteLLMManagedVectorStoreIndexTable, + LiteLLMManagedVectorStoreTable, + LiteLLMManagedVectorStoresTable, + LiteLLMMemoryTable, + LiteLLMModelTable, + LiteLLMObjectPermissionTable, + LiteLLMOrganizationMembership, + LiteLLMOrganizationTable, + LiteLLMPolicyAttachmentTable, + LiteLLMPolicyTable, + LiteLLMProjectTable, + LiteLLMPromptTable, + LiteLLMProxyModelTable, + LiteLLMSSOConfig, + LiteLLMSearchToolsTable, + LiteLLMSkillsTable, + LiteLLMSpendLogGuardrailIndex, + LiteLLMSpendLogToolIndex, + LiteLLMSpendLogs, + LiteLLMTagTable, + LiteLLMTeamMembership, + LiteLLMTeamTable, + LiteLLMToolTable, + LiteLLMUISettings, + LiteLLMUserNotifications, + LiteLLMUserTable, + LiteLLMVerificationToken, + LiteLLMWorkflowEvent, + LiteLLMWorkflowMessage, + LiteLLMWorkflowRun, +] diff --git a/litellm/proxy/db/sqlmodel/schema_parser.py b/litellm/proxy/db/sqlmodel/schema_parser.py new file mode 100644 index 000000000000..cf13aa682250 --- /dev/null +++ b/litellm/proxy/db/sqlmodel/schema_parser.py @@ -0,0 +1,473 @@ +"""Minimal ``schema.prisma`` parser used by the SQLModel parity test. + +This is intentionally **not** a full Prisma parser. It targets only the +constructs that actually appear in ``litellm``'s ``schema.prisma`` (as of the +start of the Prisma -> SQLModel migration) and is exercised by the parity +test in ``tests/test_litellm/proxy/db/sqlmodel/``. + +The parser produces a structured representation that is easy to compare +against the SQLAlchemy ``MetaData`` of the generated SQLModel classes: + +* Top-level ``PrismaSchema`` with ``models`` (dict by model name) and + ``enums`` (dict by enum name). +* Each ``PrismaModel`` carries its **scalar** fields, primary key, + uniqueness constraints, and indexes. +* Relation fields (``Foo[]`` / ``Foo? @relation(...)``) are recorded + separately in ``relations`` and are explicitly ignored by the column + parity check -- relations are not columns. + +The parser is pure-Python (no third-party deps) so it can run in any test +environment and serve as a building block for future code generators. +""" + +from __future__ import annotations + +import re +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any, Dict, List, Optional, Tuple + +# --------------------------------------------------------------------------- +# Public dataclasses +# --------------------------------------------------------------------------- + + +@dataclass +class PrismaField: + """A single scalar (or scalar-array) column on a Prisma model.""" + + name: str # field name as written in schema.prisma + column_name: str # column name on disk (respects @map(...)) + base_type: str # e.g. "String", "Int", "BigInt", "DateTime", "Json", "Bytes", "Float", "Boolean", or an enum name + is_optional: bool # True if `?` + is_list: bool # True if `[]` + is_id: bool # True if marked `@id` + is_unique: bool # True if marked `@unique` + has_default: bool + default_raw: Optional[str] # raw text inside `@default(...)` + has_updated_at: bool # True if marked `@updatedAt` + attributes: List[str] = field(default_factory=list) # raw `@...` attributes + + +@dataclass +class PrismaRelation: + """A relation field (``Foo[]`` or ``Foo? @relation(...)``) -- not a column.""" + + name: str + target_model: str + is_optional: bool + is_list: bool + relation_attributes: List[str] = field(default_factory=list) + + +@dataclass +class PrismaIndex: + """A ``@@index([...])`` declaration.""" + + fields: Tuple[str, ...] + map_name: Optional[str] = None + + +@dataclass +class PrismaUnique: + """A ``@@unique([...])`` declaration.""" + + fields: Tuple[str, ...] + + +@dataclass +class PrismaModel: + """A Prisma ``model`` block, scalar columns + constraints only.""" + + name: str + table_name: str # respects ``@@map("...")``; defaults to model name + fields: List[PrismaField] = field(default_factory=list) + relations: List[PrismaRelation] = field(default_factory=list) + primary_key: Tuple[str, ...] = () # field names (not column names) + uniques: List[PrismaUnique] = field(default_factory=list) + indexes: List[PrismaIndex] = field(default_factory=list) + raw_attributes: List[str] = field(default_factory=list) + + def field_by_name(self, name: str) -> Optional[PrismaField]: + for f in self.fields: + if f.name == name: + return f + return None + + +@dataclass +class PrismaEnum: + name: str + values: Tuple[str, ...] + + +@dataclass +class PrismaSchema: + models: Dict[str, PrismaModel] = field(default_factory=dict) + enums: Dict[str, PrismaEnum] = field(default_factory=dict) + + +# --------------------------------------------------------------------------- +# Parser +# --------------------------------------------------------------------------- + +# Built-in Prisma scalar types we know how to map. +_SCALAR_TYPES = { + "String", + "Int", + "BigInt", + "Float", + "Decimal", + "Boolean", + "DateTime", + "Json", + "Bytes", +} + + +_MODEL_RE = re.compile(r"^\s*model\s+(\w+)\s*\{\s*$") +_ENUM_RE = re.compile(r"^\s*enum\s+(\w+)\s*\{\s*$") +_DATASOURCE_RE = re.compile(r"^\s*(datasource|generator)\s+\w+\s*\{\s*$") +_TABLE_ATTR_RE = re.compile(r"^\s*@@(\w+)\s*\((.*)\)\s*$") +_TABLE_MAP_RE = re.compile(r"^\s*@@map\s*\(\s*\"([^\"]+)\"\s*\)\s*$") + + +def _strip_comment(line: str) -> str: + """Remove a trailing ``// ...`` comment, ignoring `//` inside quotes.""" + out: List[str] = [] + in_str = False + i = 0 + while i < len(line): + ch = line[i] + if ch == '"' and (i == 0 or line[i - 1] != "\\"): + in_str = not in_str + out.append(ch) + i += 1 + continue + if not in_str and ch == "/" and i + 1 < len(line) and line[i + 1] == "/": + break + out.append(ch) + i += 1 + return "".join(out).rstrip() + + +def _split_top_level_commas(s: str) -> List[str]: + """Split a parenthesized argument list on top-level commas only.""" + parts: List[str] = [] + depth = 0 + in_str = False + buf: List[str] = [] + for ch in s: + if ch == '"': + in_str = not in_str + buf.append(ch) + elif in_str: + buf.append(ch) + elif ch in "([{": + depth += 1 + buf.append(ch) + elif ch in ")]}": + depth -= 1 + buf.append(ch) + elif ch == "," and depth == 0: + parts.append("".join(buf).strip()) + buf = [] + else: + buf.append(ch) + tail = "".join(buf).strip() + if tail: + parts.append(tail) + return parts + + +def _extract_attributes(rest: str) -> List[str]: + """Extract ``@foo(...)`` / ``@foo`` attribute substrings from a field tail.""" + attrs: List[str] = [] + i = 0 + while i < len(rest): + if rest[i] == "@": + j = i + 1 + while j < len(rest) and (rest[j].isalnum() or rest[j] in "._"): + j += 1 + if j < len(rest) and rest[j] == "(": + depth = 1 + k = j + 1 + in_str = False + while k < len(rest) and depth > 0: + ch = rest[k] + if ch == '"' and rest[k - 1] != "\\": + in_str = not in_str + elif not in_str: + if ch == "(": + depth += 1 + elif ch == ")": + depth -= 1 + k += 1 + attrs.append(rest[i:k]) + i = k + continue + attrs.append(rest[i:j]) + i = j + continue + i += 1 + return attrs + + +def _parse_default_value(attr: str) -> Optional[str]: + m = re.match(r"^@default\((.*)\)$", attr) + if not m: + return None + return m.group(1).strip() + + +def _parse_map_value(attr: str) -> Optional[str]: + m = re.match(r"^@map\(\s*\"([^\"]+)\"\s*\)$", attr) + if not m: + return None + return m.group(1) + + +def _parse_field_line(line: str) -> Optional[Any]: + """Parse a single field line inside a model block. + + Returns either a ``PrismaField``, a ``PrismaRelation``, or ``None`` if the + line is blank/comment-only. + """ + stripped = _strip_comment(line).strip() + if not stripped: + return None + if stripped.startswith("@@"): + return None # handled separately + + parts = stripped.split(None, 2) + if len(parts) < 2: + return None + name = parts[0] + type_token = parts[1] + rest = parts[2] if len(parts) == 3 else "" + + is_list = type_token.endswith("[]") + if is_list: + base = type_token[:-2] + is_optional = False + elif type_token.endswith("?"): + base = type_token[:-1] + is_optional = True + else: + base = type_token + is_optional = False + + attributes = _extract_attributes(rest) + + is_relation = base not in _SCALAR_TYPES and any( + a.startswith("@relation") for a in attributes + ) + is_relation = is_relation or ( + base not in _SCALAR_TYPES and is_list # `Foo[]` back-reference + ) + + if is_relation: + return PrismaRelation( + name=name, + target_model=base, + is_optional=is_optional, + is_list=is_list, + relation_attributes=attributes, + ) + + column_name = name + has_default = False + default_raw: Optional[str] = None + has_updated_at = False + is_id = False + is_unique = False + + for attr in attributes: + if attr == "@id": + is_id = True + elif attr == "@unique": + is_unique = True + elif attr == "@updatedAt": + has_updated_at = True + elif attr.startswith("@default("): + has_default = True + default_raw = _parse_default_value(attr) + elif attr.startswith("@map("): + mapped = _parse_map_value(attr) + if mapped is not None: + column_name = mapped + + return PrismaField( + name=name, + column_name=column_name, + base_type=base, + is_optional=is_optional, + is_list=is_list, + is_id=is_id, + is_unique=is_unique, + has_default=has_default, + default_raw=default_raw, + has_updated_at=has_updated_at, + attributes=attributes, + ) + + +def _parse_field_list(arg: str) -> Tuple[str, ...]: + """Parse the field list inside ``@@id([...])`` / ``@@index([...])``. + + Field expressions like ``checked_at(sort: Desc)`` are reduced to the bare + field name, which is what we need for parity (SQLAlchemy index objects + don't capture sort direction in the simple comparison we do). + """ + m = re.match(r"^\s*\[(.*)\]\s*(?:,\s*map\s*:\s*\"([^\"]+)\")?\s*$", arg) + if not m: + return () + inner = m.group(1) + pieces = _split_top_level_commas(inner) + out: List[str] = [] + for p in pieces: + # strip ``(sort: Desc)`` etc. + bare = re.sub(r"\(.*\)", "", p).strip() + if bare: + out.append(bare) + return tuple(out) + + +def _parse_index_attr(arg: str) -> PrismaIndex: + map_name = None + m = re.search(r"map\s*:\s*\"([^\"]+)\"", arg) + if m: + map_name = m.group(1) + fields = _parse_field_list(arg) + return PrismaIndex(fields=fields, map_name=map_name) + + +def parse_schema(text: str) -> PrismaSchema: + """Parse a ``schema.prisma`` source string.""" + schema = PrismaSchema() + lines = text.splitlines() + i = 0 + n = len(lines) + while i < n: + line = _strip_comment(lines[i]) + m_model = _MODEL_RE.match(line) + m_enum = _ENUM_RE.match(line) + m_ds = _DATASOURCE_RE.match(line) + if m_ds: + i = _skip_block(lines, i) + continue + if m_enum: + name = m_enum.group(1) + values, i = _consume_enum(lines, i + 1) + schema.enums[name] = PrismaEnum(name=name, values=values) + continue + if m_model: + name = m_model.group(1) + model, i = _consume_model(lines, i + 1, name) + schema.models[name] = model + continue + i += 1 + return schema + + +def parse_schema_file(path: Path) -> PrismaSchema: + return parse_schema(Path(path).read_text()) + + +def _skip_block(lines: List[str], i: int) -> int: + """Skip a balanced ``{ ... }`` block starting at ``lines[i]``.""" + depth = 0 + while i < len(lines): + for ch in lines[i]: + if ch == "{": + depth += 1 + elif ch == "}": + depth -= 1 + if depth == 0: + return i + 1 + i += 1 + return i + + +def _consume_enum(lines: List[str], i: int) -> Tuple[Tuple[str, ...], int]: + values: List[str] = [] + while i < len(lines): + stripped = _strip_comment(lines[i]).strip() + if stripped == "}": + return tuple(values), i + 1 + if stripped: + # one identifier per line + tok = stripped.split()[0] + values.append(tok) + i += 1 + return tuple(values), i + + +def _consume_model( + lines: List[str], i: int, model_name: str +) -> Tuple[PrismaModel, int]: + model = PrismaModel(name=model_name, table_name=model_name) + while i < len(lines): + raw = lines[i] + stripped_no_comment = _strip_comment(raw).strip() + if stripped_no_comment == "}": + i += 1 + break + if not stripped_no_comment: + i += 1 + continue + + # @@map / @@id / @@unique / @@index / other table-level attrs + m_map = _TABLE_MAP_RE.match(raw) + if m_map: + model.table_name = m_map.group(1) + i += 1 + continue + m_attr = _TABLE_ATTR_RE.match(raw) + if m_attr: + kind = m_attr.group(1) + arg = m_attr.group(2).strip() + model.raw_attributes.append(stripped_no_comment) + if kind == "id": + model.primary_key = _parse_field_list(arg) + elif kind == "unique": + model.uniques.append(PrismaUnique(fields=_parse_field_list(arg))) + elif kind == "index": + model.indexes.append(_parse_index_attr(arg)) + i += 1 + continue + + parsed = _parse_field_line(raw) + if parsed is None: + i += 1 + continue + if isinstance(parsed, PrismaField): + model.fields.append(parsed) + if parsed.is_id and not model.primary_key: + model.primary_key = (parsed.name,) + elif isinstance(parsed, PrismaRelation): + model.relations.append(parsed) + i += 1 + return model, i + + +# --------------------------------------------------------------------------- +# Convenience helpers used by the parity test +# --------------------------------------------------------------------------- + + +def column_specs_for(model: PrismaModel) -> Dict[str, Dict[str, Any]]: + """Return a normalized ``{column_name: spec}`` for parity comparison.""" + specs: Dict[str, Dict[str, Any]] = {} + for f in model.fields: + specs[f.column_name] = { + "field_name": f.name, + "base_type": f.base_type, + "is_optional": f.is_optional, + "is_list": f.is_list, + "is_id": f.is_id, + "is_unique": f.is_unique, + "has_default": f.has_default, + "has_updated_at": f.has_updated_at, + } + return specs diff --git a/pyproject.toml b/pyproject.toml index 70681c4ed6c4..e3a68e2e421c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -76,6 +76,7 @@ extra_proxy = [ "resend==2.23.0", "redisvl==0.4.1; python_version < '3.14'", "a2a-sdk==0.3.24", + "sqlmodel>=0.0.22,<1.0", ] utils = [ # Not in Docker or PyPI proxy extra. diff --git a/tests/test_litellm/proxy/db/sqlmodel_orm/__init__.py b/tests/test_litellm/proxy/db/sqlmodel_orm/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/tests/test_litellm/proxy/db/sqlmodel_orm/test_parity.py b/tests/test_litellm/proxy/db/sqlmodel_orm/test_parity.py new file mode 100644 index 000000000000..82eceb4eb77e --- /dev/null +++ b/tests/test_litellm/proxy/db/sqlmodel_orm/test_parity.py @@ -0,0 +1,350 @@ +"""Parity test: SQLModel definitions must match ``schema.prisma``. + +If this test fails, either: + +* ``schema.prisma`` was changed and ``litellm/proxy/db/sqlmodel/models.py`` + was not regenerated, OR +* ``models.py`` was hand-edited in a way that no longer reflects the Prisma + schema (which is still the source of truth during the migration). + +Re-run the generator and commit the diff:: + + uv run python -m litellm.proxy.db.sqlmodel._generate \\ + --schema schema.prisma \\ + --out litellm/proxy/db/sqlmodel/models.py + +The test only enforces structural parity that matters for behavioural +equivalence at the database layer: + +* every Prisma model has exactly one SQLModel class, +* every scalar Prisma field has a column with the same on-disk name and + nullability, +* primary keys, ``@@unique`` and ``@@index`` clauses match, +* table names (``@@map``) match. + +It deliberately does *not* check Python attribute names, type granularity +beyond the broad SQL category, default values, or relation back-refs -- +those are implementation details of the SQLModel layer that may diverge +once we hand-tune for SQLAlchemy idioms in later phases. +""" + +from __future__ import annotations + +from pathlib import Path +from typing import Dict, Set, Tuple + +import pytest +from sqlalchemy import Index, PrimaryKeyConstraint, Table, UniqueConstraint + +from litellm.proxy.db.sqlmodel.models import ALL_MODELS +from litellm.proxy.db.sqlmodel.schema_parser import ( + PrismaModel, + PrismaSchema, + parse_schema_file, +) + + +# --------------------------------------------------------------------------- +# Fixtures +# --------------------------------------------------------------------------- + + +def _find_repo_root() -> Path: + p = Path(__file__).resolve() + while not (p / "schema.prisma").exists(): + if p.parent == p: + raise RuntimeError("schema.prisma not found in any ancestor directory") + p = p.parent + return p + + +@pytest.fixture(scope="module") +def prisma_schema() -> PrismaSchema: + return parse_schema_file(_find_repo_root() / "schema.prisma") + + +@pytest.fixture(scope="module") +def sqlmodel_tables() -> Dict[str, Table]: + return {cls.__tablename__: cls.__table__ for cls in ALL_MODELS} + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _index_signatures(table: Table) -> Set[Tuple[str, ...]]: + """Set of ``(col1, col2, ...)`` tuples from non-unique SQLAlchemy indexes.""" + sigs: Set[Tuple[str, ...]] = set() + for ix in table.indexes: + if ix.unique: + continue + sigs.add(tuple(c.name for c in ix.columns)) + return sigs + + +def _unique_signatures(table: Table) -> Set[Tuple[str, ...]]: + sigs: Set[Tuple[str, ...]] = set() + for cons in table.constraints: + if isinstance(cons, UniqueConstraint): + sigs.add(tuple(c.name for c in cons.columns)) + for col in table.columns: + if col.unique and not col.primary_key: + sigs.add((col.name,)) + return sigs + + +def _pk_signature(table: Table) -> Tuple[str, ...]: + return tuple(c.name for c in table.primary_key.columns) + + +def _prisma_pk_columns(model: PrismaModel) -> Tuple[str, ...]: + """Map field-name PK to column-name PK (respects ``@map``).""" + cols: list[str] = [] + for fname in model.primary_key: + f = model.field_by_name(fname) + cols.append(f.column_name if f is not None else fname) + return tuple(cols) + + +def _prisma_unique_signatures(model: PrismaModel) -> Set[Tuple[str, ...]]: + sigs: Set[Tuple[str, ...]] = set() + for u in model.uniques: + sigs.add(tuple(_field_to_column(model, fn) for fn in u.fields)) + for f in model.fields: + if f.is_unique and not f.is_id: + sigs.add((f.column_name,)) + return sigs + + +def _prisma_index_signatures(model: PrismaModel) -> Set[Tuple[str, ...]]: + sigs: Set[Tuple[str, ...]] = set() + for idx in model.indexes: + sigs.add(tuple(_field_to_column(model, fn) for fn in idx.fields)) + return sigs + + +def _field_to_column(model: PrismaModel, fname: str) -> str: + f = model.field_by_name(fname) + return f.column_name if f is not None else fname + + +# Prisma scalar -> coarse SQL category we expect on the generated column. +_EXPECTED_TYPE_CATEGORIES = { + "String": {"text", "varchar"}, + "Int": {"integer"}, + "BigInt": {"biginteger", "bigint"}, + "Float": {"double", "double_precision", "float"}, + "Decimal": {"numeric", "decimal"}, + "Boolean": {"boolean"}, + "DateTime": {"datetime", "timestamp"}, + "Json": {"json", "jsonb"}, + "Bytes": {"largebinary", "bytea"}, +} + + +def _column_type_category(col) -> str: + return type(col.type).__name__.lower() + + +# --------------------------------------------------------------------------- +# Tests +# --------------------------------------------------------------------------- + + +def test_one_sqlmodel_class_per_prisma_model(prisma_schema, sqlmodel_tables): + prisma_table_names = {m.table_name for m in prisma_schema.models.values()} + sqlmodel_table_names = set(sqlmodel_tables) + missing_in_sqlmodel = prisma_table_names - sqlmodel_table_names + extra_in_sqlmodel = sqlmodel_table_names - prisma_table_names + assert not missing_in_sqlmodel, ( + f"Prisma tables with no SQLModel class: {sorted(missing_in_sqlmodel)}. " + "Did you forget to regenerate models.py?" + ) + assert not extra_in_sqlmodel, ( + f"SQLModel classes with no Prisma model: {sorted(extra_in_sqlmodel)}. " + "Did you forget to update schema.prisma?" + ) + + +def test_columns_match_for_every_table(prisma_schema, sqlmodel_tables): + failures: list[str] = [] + for prisma_model in prisma_schema.models.values(): + table = sqlmodel_tables[prisma_model.table_name] + prisma_cols = {f.column_name: f for f in prisma_model.fields} + sqlmodel_cols = {c.name: c for c in table.columns} + + missing = set(prisma_cols) - set(sqlmodel_cols) + extra = set(sqlmodel_cols) - set(prisma_cols) + if missing: + failures.append( + f"{prisma_model.table_name}: missing columns in SQLModel: {sorted(missing)}" + ) + if extra: + failures.append( + f"{prisma_model.table_name}: unexpected columns in SQLModel: {sorted(extra)}" + ) + assert not failures, "\n".join(failures) + + +def test_column_nullability_matches(prisma_schema, sqlmodel_tables): + failures: list[str] = [] + for prisma_model in prisma_schema.models.values(): + table = sqlmodel_tables[prisma_model.table_name] + sqlmodel_cols = {c.name: c for c in table.columns} + for f in prisma_model.fields: + col = sqlmodel_cols.get(f.column_name) + if col is None: + continue + expected_nullable = f.is_optional + if col.nullable != expected_nullable: + failures.append( + f"{prisma_model.table_name}.{f.column_name}: " + f"prisma optional={f.is_optional} but SQLModel nullable={col.nullable}" + ) + assert not failures, "\n".join(failures) + + +def test_column_types_in_expected_category(prisma_schema, sqlmodel_tables): + """Coarse type check: e.g. ``BigInt`` -> a BigInteger-class type, not Integer. + + We deliberately do not enforce exact ``server_default`` or precision -- those + are implementation details that can drift without behavioural impact, and + they are guarded separately by the migration tests. + """ + failures: list[str] = [] + for prisma_model in prisma_schema.models.values(): + table = sqlmodel_tables[prisma_model.table_name] + sqlmodel_cols = {c.name: c for c in table.columns} + for f in prisma_model.fields: + col = sqlmodel_cols.get(f.column_name) + if col is None: + continue + expected = _EXPECTED_TYPE_CATEGORIES.get(f.base_type) + if expected is None: + # enum reference or unknown scalar -> skip + continue + actual_kind = _column_type_category(col) + ok = any(token in actual_kind for token in expected) + # ARRAY columns wrap an inner type; check the item type instead. + if not ok and "array" in actual_kind and f.is_list: + inner = type(col.type.item_type).__name__.lower() + ok = any(token in inner for token in expected) + if not ok: + failures.append( + f"{prisma_model.table_name}.{f.column_name}: " + f"prisma type={f.base_type}{'[]' if f.is_list else ''} " + f"but SQLModel column type is {actual_kind}" + ) + assert not failures, "\n".join(failures) + + +def test_array_columns_match(prisma_schema, sqlmodel_tables): + failures: list[str] = [] + for prisma_model in prisma_schema.models.values(): + table = sqlmodel_tables[prisma_model.table_name] + sqlmodel_cols = {c.name: c for c in table.columns} + for f in prisma_model.fields: + col = sqlmodel_cols.get(f.column_name) + if col is None: + continue + actual_is_array = "array" in type(col.type).__name__.lower() + if f.is_list != actual_is_array: + failures.append( + f"{prisma_model.table_name}.{f.column_name}: " + f"prisma is_list={f.is_list} but SQLModel ARRAY={actual_is_array}" + ) + assert not failures, "\n".join(failures) + + +def test_primary_keys_match(prisma_schema, sqlmodel_tables): + failures: list[str] = [] + for prisma_model in prisma_schema.models.values(): + table = sqlmodel_tables[prisma_model.table_name] + prisma_pk = _prisma_pk_columns(prisma_model) + sqlmodel_pk = _pk_signature(table) + if set(prisma_pk) != set(sqlmodel_pk): + failures.append( + f"{prisma_model.table_name}: prisma PK={prisma_pk} but SQLModel PK={sqlmodel_pk}" + ) + assert not failures, "\n".join(failures) + + +def test_unique_constraints_match(prisma_schema, sqlmodel_tables): + failures: list[str] = [] + for prisma_model in prisma_schema.models.values(): + table = sqlmodel_tables[prisma_model.table_name] + prisma_uniques = _prisma_unique_signatures(prisma_model) + sqlmodel_uniques = _unique_signatures(table) + # Set comparison ignores ordering of the unique-constraint columns, + # which matches what Postgres treats as logically equivalent. + prisma_norm = {tuple(sorted(s)) for s in prisma_uniques} + sqlmodel_norm = {tuple(sorted(s)) for s in sqlmodel_uniques} + missing = prisma_norm - sqlmodel_norm + extra = sqlmodel_norm - prisma_norm + if missing: + failures.append( + f"{prisma_model.table_name}: missing unique constraints in SQLModel: {sorted(missing)}" + ) + if extra: + failures.append( + f"{prisma_model.table_name}: unexpected unique constraints in SQLModel: {sorted(extra)}" + ) + assert not failures, "\n".join(failures) + + +def test_indexes_match(prisma_schema, sqlmodel_tables): + failures: list[str] = [] + for prisma_model in prisma_schema.models.values(): + table = sqlmodel_tables[prisma_model.table_name] + prisma_idx = _prisma_index_signatures(prisma_model) + sqlmodel_idx = _index_signatures(table) + # We compare ordered tuples here because index column order + # affects which queries the index can serve. + missing = prisma_idx - sqlmodel_idx + extra = sqlmodel_idx - prisma_idx + if missing: + failures.append( + f"{prisma_model.table_name}: missing indexes in SQLModel: {sorted(missing)}" + ) + if extra: + failures.append( + f"{prisma_model.table_name}: unexpected indexes in SQLModel: {sorted(extra)}" + ) + assert not failures, "\n".join(failures) + + +def test_generator_output_is_committed(tmp_path): + """Re-run the generator and assert the result matches the checked-in file. + + This is the strongest guard: it catches any drift in either the schema + or the generator (or hand-edits to ``models.py`` that don't roundtrip). + """ + from litellm.proxy.db.sqlmodel import _generate + + schema = parse_schema_file(_find_repo_root() / "schema.prisma") + expected = _generate.render_module(schema) + actual = ( + _find_repo_root() / "litellm" / "proxy" / "db" / "sqlmodel" / "models.py" + ).read_text() + if expected != actual: + # Surface a small diff so the failure message is actionable. + import difflib + + diff = "\n".join( + difflib.unified_diff( + actual.splitlines(), + expected.splitlines(), + fromfile="models.py (committed)", + tofile="models.py (regenerated)", + lineterm="", + n=3, + ) + ) + pytest.fail( + "litellm/proxy/db/sqlmodel/models.py is out of sync with " + "schema.prisma. Run:\n" + " uv run python -m litellm.proxy.db.sqlmodel._generate " + "--schema schema.prisma --out litellm/proxy/db/sqlmodel/models.py\n\n" + f"Diff (truncated to first 60 lines):\n{chr(10).join(diff.splitlines()[:60])}" + ) diff --git a/tests/test_litellm/proxy/db/sqlmodel_orm/test_schema_parser.py b/tests/test_litellm/proxy/db/sqlmodel_orm/test_schema_parser.py new file mode 100644 index 000000000000..b764660aea4c --- /dev/null +++ b/tests/test_litellm/proxy/db/sqlmodel_orm/test_schema_parser.py @@ -0,0 +1,228 @@ +"""Unit tests for the ``schema.prisma`` parser. + +Run with:: + + uv run pytest tests/test_litellm/proxy/db/sqlmodel/test_schema_parser.py -vv +""" + +from __future__ import annotations + +import textwrap + +import pytest + +from litellm.proxy.db.sqlmodel.schema_parser import ( + PrismaField, + PrismaRelation, + parse_schema, +) + + +def test_parse_simple_model(): + src = textwrap.dedent( + """ + model Foo { + id String @id @default(uuid()) + name String @unique + } + """ + ) + schema = parse_schema(src) + assert "Foo" in schema.models + foo = schema.models["Foo"] + assert foo.table_name == "Foo" + assert foo.primary_key == ("id",) + assert [f.name for f in foo.fields] == ["id", "name"] + assert foo.fields[0].is_id + assert foo.fields[0].has_default + assert foo.fields[0].default_raw == "uuid()" + assert foo.fields[1].is_unique + + +def test_optional_and_array_fields(): + src = textwrap.dedent( + """ + model Foo { + id String @id + tags String[] @default([]) + note String? + } + """ + ) + foo = parse_schema(src).models["Foo"] + f_tags = foo.field_by_name("tags") + assert f_tags is not None + assert f_tags.is_list and not f_tags.is_optional + assert f_tags.has_default and f_tags.default_raw == "[]" + f_note = foo.field_by_name("note") + assert f_note is not None + assert f_note.is_optional and not f_note.is_list + + +def test_composite_primary_key_and_index(): + src = textwrap.dedent( + """ + model Foo { + a String + b String + c Int @default(0) + + @@id([a, b]) + @@index([c]) + @@unique([a, c]) + } + """ + ) + foo = parse_schema(src).models["Foo"] + assert foo.primary_key == ("a", "b") + assert len(foo.indexes) == 1 + assert foo.indexes[0].fields == ("c",) + assert len(foo.uniques) == 1 + assert foo.uniques[0].fields == ("a", "c") + + +def test_index_with_map_and_sort(): + src = textwrap.dedent( + """ + model Foo { + a String @id + b DateTime + c String + + @@index([a, b, c(sort: Desc)], map: "Foo_custom_idx") + } + """ + ) + foo = parse_schema(src).models["Foo"] + assert len(foo.indexes) == 1 + idx = foo.indexes[0] + assert idx.map_name == "Foo_custom_idx" + assert idx.fields == ("a", "b", "c") + + +def test_at_map_renames_column(): + src = textwrap.dedent( + """ + model Foo { + id String @id + created String @map("created_at") + } + """ + ) + foo = parse_schema(src).models["Foo"] + f = foo.field_by_name("created") + assert f is not None + assert f.column_name == "created_at" + + +def test_at_at_map_renames_table(): + src = textwrap.dedent( + """ + model Foo { + id String @id + @@map("foo_table") + } + """ + ) + foo = parse_schema(src).models["Foo"] + assert foo.table_name == "foo_table" + + +def test_relations_are_separated_from_fields(): + src = textwrap.dedent( + """ + model Bar { + id String @id + } + + model Foo { + id String @id + bar_id String? + bar Bar? @relation(fields: [bar_id], references: [id]) + many Bar[] + } + """ + ) + foo = parse_schema(src).models["Foo"] + field_names = {f.name for f in foo.fields} + rel_names = {r.name for r in foo.relations} + assert field_names == {"id", "bar_id"} + assert rel_names == {"bar", "many"} + rel_bar = next(r for r in foo.relations if r.name == "bar") + assert rel_bar.target_model == "Bar" + assert rel_bar.is_optional and not rel_bar.is_list + + +def test_enum_parsed(): + src = textwrap.dedent( + """ + enum Status { + ACTIVE + INACTIVE + } + + model Foo { + id String @id + status Status @default(INACTIVE) + } + """ + ) + schema = parse_schema(src) + assert schema.enums["Status"].values == ("ACTIVE", "INACTIVE") + f = schema.models["Foo"].field_by_name("status") + assert f is not None + assert f.base_type == "Status" + assert f.default_raw == "INACTIVE" + + +def test_handles_trailing_block_comment_on_model_line(): + """``model Foo { // comment`` should still be recognized as a model.""" + src = textwrap.dedent( + """ + model Foo { // a trailing comment after the brace + id String @id + } + """ + ) + schema = parse_schema(src) + assert "Foo" in schema.models + + +def test_strip_comment_handles_quoted_double_slash(): + """A `//` inside a quoted string default must not be treated as a comment.""" + src = textwrap.dedent( + """ + model Foo { + id String @id + url String @default("https://example.com") + } + """ + ) + foo = parse_schema(src).models["Foo"] + f = foo.field_by_name("url") + assert f is not None + assert f.default_raw == '"https://example.com"' + + +def test_real_schema_round_trip(tmp_path): + """Parse the actual repository ``schema.prisma`` and assert basic shape. + + This is a smoke test -- the deeper structural parity check lives in + ``test_parity.py``. + """ + from pathlib import Path + + repo_root = Path(__file__).resolve() + while not (repo_root / "schema.prisma").exists(): + if repo_root.parent == repo_root: + pytest.skip("schema.prisma not found in any ancestor directory") + repo_root = repo_root.parent + schema = parse_schema((repo_root / "schema.prisma").read_text()) + assert len(schema.models) >= 60 + assert "LiteLLM_VerificationToken" in schema.models + assert "LiteLLM_TeamMembership" in schema.models + # composite PK on TeamMembership + assert schema.models["LiteLLM_TeamMembership"].primary_key == ( + "user_id", + "team_id", + ) diff --git a/uv.lock b/uv.lock index cafb6664958d..5ddbd0ee5692 100644 --- a/uv.lock +++ b/uv.lock @@ -9,7 +9,7 @@ resolution-markers = [ ] [options] -exclude-newer = "0001-01-01T00:00:00Z" # This has no effect and is included for backwards compatibility when using relative exclude-newer values. +exclude-newer = "2026-05-17T16:09:35.533374581Z" exclude-newer-span = "P3D" [manifest] @@ -3219,6 +3219,7 @@ extra-proxy = [ { name = "prisma" }, { name = "redisvl" }, { name = "resend" }, + { name = "sqlmodel" }, ] google = [ { name = "google-cloud-aiplatform" }, @@ -3444,6 +3445,7 @@ requires-dist = [ { name = "sentry-sdk", marker = "extra == 'proxy-runtime'", specifier = "==2.21.0" }, { name = "soundfile", marker = "extra == 'proxy'", specifier = "==0.12.1" }, { name = "soundfile", marker = "extra == 'stt-nvidia-riva'", specifier = ">=0.12.1" }, + { name = "sqlmodel", marker = "extra == 'extra-proxy'", specifier = ">=0.0.22,<1.0" }, { name = "tiktoken", specifier = ">=0.8.0,<1.0" }, { name = "tokenizers", specifier = ">=0.21.0,<1.0" }, { name = "uvicorn", marker = "extra == 'proxy'", specifier = "==0.33.0" }, @@ -7183,6 +7185,20 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/e5/30/8519fdde58a7bdf155b714359791ad1dc018b47d60269d5d160d311fdc36/sqlalchemy-2.0.49-py3-none-any.whl", hash = "sha256:ec44cfa7ef1a728e88ad41674de50f6db8cfdb3e2af84af86e0041aaf02d43d0", size = 1942158, upload-time = "2026-04-03T16:53:44.135Z" }, ] +[[package]] +name = "sqlmodel" +version = "0.0.38" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "pydantic" }, + { name = "sqlalchemy" }, + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/64/0d/26ec1329960ea9430131fe63f63a95ea4cb8971d49c891ff7e1f3255421c/sqlmodel-0.0.38.tar.gz", hash = "sha256:d583ec237b14103809f74e8630032bc40ab68cd6b754a610f0813c56911a547b", size = 86710, upload-time = "2026-04-02T21:03:55.571Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/72/c7/10c60af0607ab6fa136264f7f39d205932218516226d38585324ffda705d/sqlmodel-0.0.38-py3-none-any.whl", hash = "sha256:84e3fa990a77395461ded72a6c73173438ce8449d5c1c4d97fbff1b1df692649", size = 27294, upload-time = "2026-04-02T21:03:56.406Z" }, +] + [[package]] name = "sqlparse" version = "0.5.5"