diff --git a/conftest.py b/conftest.py new file mode 100644 index 00000000000..03d2e64a193 --- /dev/null +++ b/conftest.py @@ -0,0 +1,29 @@ +"""Repository-root pytest configuration. + +After the SQLAlchemy big-bang migration, the ``prisma`` PyPI package is +no longer installed but a handful of test modules still do ``from prisma +... import ...``. We install a stand-in for the ``prisma`` namespace at +plugin-load time -- *before* pytest collects test modules -- so those +imports resolve without rewriting every test file in this PR. + +The stand-in lives in ``tests/_prisma_compat.py`` and re-exports the +LiteLLM-native error classes from ``litellm.proxy.db.sqlmodel.errors``. +""" + +from __future__ import annotations + +import os +import sys + +# Make ``tests/_prisma_compat`` importable regardless of where pytest is +# invoked from. +_TESTS_DIR = os.path.join(os.path.dirname(__file__), "tests") +if _TESTS_DIR not in sys.path: + sys.path.insert(0, _TESTS_DIR) + +try: + from _prisma_compat import install as _install_prisma_compat + + _install_prisma_compat() +except ImportError: # pragma: no cover -- shim is best-effort + pass diff --git a/litellm/proxy/_experimental/mcp_server/toolset_db.py b/litellm/proxy/_experimental/mcp_server/toolset_db.py index 08ac7dbd33b..5301416a6fe 100644 --- a/litellm/proxy/_experimental/mcp_server/toolset_db.py +++ b/litellm/proxy/_experimental/mcp_server/toolset_db.py @@ -92,7 +92,7 @@ async def update_mcp_toolset( data=data_dict, ) except Exception as e: - from prisma.errors import RecordNotFoundError + from litellm.proxy.db.sqlmodel.errors import RecordNotFoundError if isinstance(e, RecordNotFoundError): return None @@ -109,7 +109,7 @@ async def delete_mcp_toolset( where={"toolset_id": toolset_id} ) except Exception as e: - from prisma.errors import RecordNotFoundError + from litellm.proxy.db.sqlmodel.errors import RecordNotFoundError if isinstance(e, RecordNotFoundError): return None diff --git a/litellm/proxy/db/exception_handler.py b/litellm/proxy/db/exception_handler.py index ab9d341aa51..07467eb8a81 100644 --- a/litellm/proxy/db/exception_handler.py +++ b/litellm/proxy/db/exception_handler.py @@ -43,24 +43,24 @@ def is_database_connection_error(e: Exception) -> bool: to True so genuine outages that don't match a specific subclass still trigger the fallback. """ - import prisma + from litellm.proxy.db.sqlmodel import errors as prisma_errors # Explicit data-layer exclusion: DB IS reachable, fallback must # NOT fire. data_layer_errors = ( - prisma.errors.DataError, - prisma.errors.UniqueViolationError, - prisma.errors.ForeignKeyViolationError, - prisma.errors.MissingRequiredValueError, - prisma.errors.RawQueryError, - prisma.errors.TableNotFoundError, - prisma.errors.RecordNotFoundError, + prisma_errors.DataError, + prisma_errors.UniqueViolationError, + prisma_errors.ForeignKeyViolationError, + prisma_errors.MissingRequiredValueError, + prisma_errors.RawQueryError, + prisma_errors.TableNotFoundError, + prisma_errors.RecordNotFoundError, ) if isinstance(e, data_layer_errors): return False if isinstance(e, DB_CONNECTION_ERROR_TYPES): return True - if isinstance(e, prisma.errors.PrismaError): + if isinstance(e, prisma_errors.PrismaError): return True if isinstance(e, ProxyException) and e.type == ProxyErrorTypes.no_db_connection: return True @@ -75,19 +75,19 @@ def is_database_transport_error(e: Exception) -> bool: Use this for reconnect logic — data-layer errors like UniqueViolationError mean the DB IS reachable, so reconnecting would be pointless. """ - import prisma + from litellm.proxy.db.sqlmodel import errors as prisma_errors if isinstance(e, DB_CONNECTION_ERROR_TYPES): return True if isinstance( e, ( - prisma.errors.ClientNotConnectedError, - prisma.errors.HTTPClientClosedError, + prisma_errors.ClientNotConnectedError, + prisma_errors.HTTPClientClosedError, ), ): return True - if isinstance(e, prisma.errors.PrismaError): + if isinstance(e, prisma_errors.PrismaError): error_message = str(e).lower() connection_keywords = ( "can't reach database server", diff --git a/litellm/proxy/db/log_db_metrics.py b/litellm/proxy/db/log_db_metrics.py index 5c795155324..faf601cb1e6 100644 --- a/litellm/proxy/db/log_db_metrics.py +++ b/litellm/proxy/db/log_db_metrics.py @@ -107,7 +107,7 @@ def _is_exception_related_to_db(e: Exception) -> bool: """ import httpx - from prisma.errors import PrismaError + from litellm.proxy.db.sqlmodel.errors import PrismaError return isinstance(e, (PrismaError, httpx.ConnectError, httpx.TimeoutException)) diff --git a/litellm/proxy/db/prisma_client.py b/litellm/proxy/db/prisma_client.py index af5a58802bb..6263f664f66 100644 --- a/litellm/proxy/db/prisma_client.py +++ b/litellm/proxy/db/prisma_client.py @@ -328,7 +328,19 @@ async def recreate_prisma_client( new URL is passed explicitly via `datasource={"url": ...}` (Prisma does not auto-read alternate env vars like DATABASE_URL_READ_REPLICA). """ - from prisma import Prisma # type: ignore + # NOTE: ``PrismaWrapper`` is dead code after the SQLAlchemy big-bang + # migration -- ``PrismaClient.__init__`` no longer instantiates one. + # We keep the class for backwards-compat with tests that import it + # directly, and guard the prisma-client-py import so module load + # succeeds when the ``prisma`` PyPI package is no longer installed. + try: + from prisma import Prisma # type: ignore[import-not-found] + except ImportError as exc: # pragma: no cover - migration transitional + raise RuntimeError( + "PrismaWrapper.recreate_prisma_client is no longer supported " + "after the SQLAlchemy migration. Use LiteLLMDB / " + "PrismaCompatClient from litellm.proxy.db.sqlmodel." + ) from exc old_engine_pid = self._get_engine_pid() if old_engine_pid > 0: diff --git a/litellm/proxy/db/sqlmodel/README.md b/litellm/proxy/db/sqlmodel/README.md new file mode 100644 index 00000000000..dff718236b3 --- /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 00000000000..843fd314639 --- /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 00000000000..7a68e4a1e03 --- /dev/null +++ b/litellm/proxy/db/sqlmodel/_generate.py @@ -0,0 +1,523 @@ +"""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)", + # ``JSONB`` is Postgres-specific; fall back to portable ``JSON`` on + # SQLite so the test environment can run without a real database. + "Json": "JSONB().with_variant(JSON(), 'sqlite')", + "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") + + if flags["jsonb"]: + # JSONB columns use ``with_variant(JSON(), 'sqlite')`` -- pull JSON + # in from sqlalchemy.types so SQLite test envs compile cleanly. + sa_imports.add("JSON") + + 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", + "JSON", + "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/compat.py b/litellm/proxy/db/sqlmodel/compat.py new file mode 100644 index 00000000000..0541ce85c69 --- /dev/null +++ b/litellm/proxy/db/sqlmodel/compat.py @@ -0,0 +1,1004 @@ +"""Prisma-client-py compatibility shim backed by SQLAlchemy. + +The existing LiteLLM proxy has ~1,680 call sites of the form:: + + await prisma_client.db.litellm_verificationtoken.find_unique( + where={"token": hashed_token}, + include={"litellm_budget_table": True}, + ) + +Rewriting every call site at once is the diff size that blocks any +big-bang migration. Instead, this module exposes the same +``db.litellm_
.(...)`` surface and translates each call +into a SQLAlchemy statement. Call sites continue to work unchanged. + +What is implemented (and tested in the unit suite): + +* per-table accessors auto-derived from :data:`models.ALL_MODELS` +* ``find_unique``, ``find_first``, ``find_many``, ``count``, + ``create``, ``update``, ``upsert``, ``delete``, + ``update_many``, ``delete_many``, ``create_many`` +* ``where=`` translation for ``in``, ``not``, ``equals``, ``contains`` + (with ``mode: 'insensitive'``), ``startswith``, ``endswith``, + ``gt``, ``lt``, ``gte``, ``lte``, ``has``, ``AND``, ``OR``, ``NOT`` +* ``data=`` translation including ``{"increment": N}`` +* ``order=``, ``take=``, ``skip=`` +* ``query_raw(sql, *params)`` and ``query_raw(query=sql, *params)`` +* ``execute_raw(sql, *params)`` +* ``batch_()`` returning a batcher with ``await batcher.commit()`` +* ``async with db.tx() as tx:`` + +What is **not** implemented (raises ``NotImplementedError`` or logs a +warning and returns ``None``): + +* ``include=`` for relation eager-loading -- SQLModel classes do not + carry SQLAlchemy ``relationship()`` definitions yet. Existing call + sites that rely on ``include`` will see ``None`` for the included + attribute. Adding ``relationship()`` declarations and selectinload + wiring is a follow-up of comparable scope to a CRUD-feature port. +* ``select=`` for column projection -- only 3 production call sites + use it; they should be ported to direct attribute access. +* ``group_by`` -- 4 production call sites; port to raw SQL. +* Nested ``some``/``every`` relation filters -- 2 production call + sites; port to explicit JOIN via raw SQL. + +The tests under ``tests/test_litellm/proxy/db/sqlmodel_orm/`` exercise +the implemented features against an in-memory SQLite DB. +""" + +from __future__ import annotations + +import asyncio +import logging +from contextlib import asynccontextmanager +from typing import ( + Any, + AsyncIterator, + Dict, + Iterable, + List, + Mapping, + Optional, + Sequence, + Type, +) + +from sqlalchemy import ( + and_, + delete, + func, + not_, + or_, + select, + text, + update, +) +from sqlalchemy.exc import SQLAlchemyError +from sqlalchemy.ext.asyncio import AsyncSession +from sqlmodel import SQLModel + +from litellm.proxy.db.sqlmodel import errors +from litellm.proxy.db.sqlmodel.engine import LiteLLMDB +from litellm.proxy.db.sqlmodel.models import ALL_MODELS + +logger = logging.getLogger(__name__) + + +# --------------------------------------------------------------------------- +# Table registry +# --------------------------------------------------------------------------- + + +def _build_table_index() -> Dict[str, Type[SQLModel]]: + """Map ``litellm_`` -> SQLModel class. + + Mirrors prisma-client-py's accessor naming: the Prisma model + ``LiteLLM_TeamTable`` becomes ``litellm_teamtable``. + """ + index: Dict[str, Type[SQLModel]] = {} + for cls in ALL_MODELS: + # prisma-client-py exposes a model accessor by lowercasing the + # original Prisma model name verbatim -- underscores included. + # ``LiteLLM_VerificationToken`` -> ``litellm_verificationtoken`` + # ``LiteLLM_Config`` -> ``litellm_config`` + prisma_table = getattr(cls, "__tablename__", cls.__name__) + accessor = prisma_table.lower() + index[accessor] = cls + return index + + +_TABLE_INDEX: Dict[str, Type[SQLModel]] = _build_table_index() + + +def _model_for_accessor(name: str) -> Type[SQLModel]: + cls = _TABLE_INDEX.get(name) + if cls is None: + raise errors.TableNotFoundError( + f"No SQLModel class registered for prisma accessor '{name}'. " + "Did you regenerate models.py after a schema change?" + ) + return cls + + +# --------------------------------------------------------------------------- +# Filter / data translators +# --------------------------------------------------------------------------- + + +_FIELD_OPERATORS = { + "equals", + "not", + "in", + "notIn", + "not_in", + "contains", + "startswith", + "endswith", + "lt", + "lte", + "gt", + "gte", + "has", + "mode", +} + + +def _resolve_attr(model: Type[SQLModel], name: str) -> Any: + """Return the SQLAlchemy column for ``name`` on ``model``. + + Honours the Python-attribute-rename trick the generator uses for + SQLAlchemy reserved names (``metadata`` -> ``metadata_``). Callers + pass the on-disk column name (``metadata``); we look it up in + ``__table__.columns`` first, then fall back to attribute lookup. + """ + table = model.__table__ # type: ignore[attr-defined] + if name in table.columns: + return table.columns[name] + attr = getattr(model, name, None) + if attr is None: + raise AttributeError( + f"Column '{name}' not found on {model.__name__}; " + "shim cannot translate this filter." + ) + return attr + + +def _translate_field_filter(column: Any, operand: Any) -> Any: + if not isinstance(operand, Mapping): + return column == operand + + mode = operand.get("mode") + case_insensitive = mode == "insensitive" + clauses: List[Any] = [] + for op, val in operand.items(): + if op == "mode": + continue + if op == "equals": + if case_insensitive and isinstance(val, str): + clauses.append(func.lower(column) == val.lower()) + else: + clauses.append(column == val) + elif op == "not": + clauses.append(column != val) + elif op == "in": + clauses.append(column.in_(list(val) if val is not None else [])) + elif op in ("notIn", "not_in"): + clauses.append(~column.in_(list(val) if val is not None else [])) + elif op == "contains": + pattern = f"%{val}%" + clauses.append( + column.ilike(pattern) if case_insensitive else column.like(pattern) + ) + elif op == "startswith": + pattern = f"{val}%" + clauses.append( + column.ilike(pattern) if case_insensitive else column.like(pattern) + ) + elif op == "endswith": + pattern = f"%{val}" + clauses.append( + column.ilike(pattern) if case_insensitive else column.like(pattern) + ) + elif op == "lt": + clauses.append(column < val) + elif op == "lte": + clauses.append(column <= val) + elif op == "gt": + clauses.append(column > val) + elif op == "gte": + clauses.append(column >= val) + elif op == "has": + # Postgres array contains -- ``column @> ARRAY[val]``. + clauses.append(column.contains([val])) + else: + raise NotImplementedError( + f"Field operator '{op}' not implemented in compat shim. " + "Either port the call site to raw SQL or extend " + "_translate_field_filter." + ) + if not clauses: + return None + if len(clauses) == 1: + return clauses[0] + return and_(*clauses) + + +def _translate_where(model: Type[SQLModel], where: Optional[Mapping[str, Any]]) -> Any: + """Translate a Prisma ``where=`` mapping to a SQLAlchemy clause. + + Returns ``None`` when ``where`` is empty (caller should emit no + WHERE clause). + """ + if not where: + return None + clauses: List[Any] = [] + for key, val in where.items(): + if key == "AND": + assert isinstance(val, list), "AND requires a list of sub-filters" + sub = [_translate_where(model, item) for item in val] + clauses.append(and_(*[c for c in sub if c is not None])) + elif key == "OR": + assert isinstance(val, list), "OR requires a list of sub-filters" + sub = [_translate_where(model, item) for item in val] + clauses.append(or_(*[c for c in sub if c is not None])) + elif key == "NOT": + if isinstance(val, list): + sub = [_translate_where(model, item) for item in val] + clauses.append(not_(and_(*[c for c in sub if c is not None]))) + else: + inner = _translate_where(model, val) + if inner is not None: + clauses.append(not_(inner)) + elif key in ("some", "every", "is", "isNot", "is_not"): + raise NotImplementedError( + f"Relation filter '{key}' is not yet supported by the shim. " + "Port this call site to an explicit JOIN via raw SQL." + ) + else: + column = _resolve_attr(model, key) + translated = _translate_field_filter(column, val) + if translated is not None: + clauses.append(translated) + if not clauses: + return None + if len(clauses) == 1: + return clauses[0] + return and_(*clauses) + + +def _translate_data(model: Type[SQLModel], data: Mapping[str, Any]) -> Dict[str, Any]: + """Prisma ``data=`` -> kwargs/values for ``insert``/``update`` statements. + + Handles ``{"increment": N}`` by emitting a SQL expression; everything + else passes through verbatim. + """ + out: Dict[str, Any] = {} + table = model.__table__ # type: ignore[attr-defined] + for key, val in data.items(): + if isinstance(val, Mapping) and "increment" in val: + column = table.columns[key] + out[key] = column + val["increment"] + elif isinstance(val, Mapping) and "decrement" in val: + column = table.columns[key] + out[key] = column - val["decrement"] + elif isinstance(val, Mapping) and "set" in val: + out[key] = val["set"] + else: + out[key] = val + return out + + +def _translate_order(model: Type[SQLModel], order: Any) -> List[Any]: + """Prisma ``order=`` -> list of SQLAlchemy ORDER BY expressions.""" + if order is None: + return [] + if isinstance(order, Mapping): + items = list(order.items()) + elif isinstance(order, list): + items = [] + for entry in order: + if isinstance(entry, Mapping): + items.extend(entry.items()) + # Keep deterministic iteration order + else: + return [] + out: List[Any] = [] + for col_name, direction in items: + column = _resolve_attr(model, col_name) + if str(direction).lower().startswith("desc"): + out.append(column.desc()) + else: + out.append(column.asc()) + return out + + +# --------------------------------------------------------------------------- +# Result helpers +# --------------------------------------------------------------------------- + + +class CountResult: + """Mimic prisma-client-py's ``BatchPayload`` for ``update_many`` etc. + + Some call sites read ``result.count`` and others compare to an int; + we support both by also implementing ``__int__`` and ``__index__``. + """ + + __slots__ = ("count",) + + def __init__(self, count: int) -> None: + self.count = count + + def __int__(self) -> int: + return self.count + + def __index__(self) -> int: + return self.count + + def __eq__(self, other: object) -> bool: + if isinstance(other, CountResult): + return self.count == other.count + if isinstance(other, int): + return self.count == other + return NotImplemented + + def __repr__(self) -> str: + return f"CountResult(count={self.count})" + + +def _row_to_dict(row: Any) -> Dict[str, Any]: + if hasattr(row, "_mapping"): + return dict(row._mapping) + if isinstance(row, dict): + return row + return dict(row) # last resort + + +# --------------------------------------------------------------------------- +# Per-table accessor +# --------------------------------------------------------------------------- + + +def _warn_unsupported(kwarg: str, table: str) -> None: + if kwarg in ("include",): + logger.warning( + "compat shim: '%s=' on %s is not yet implemented; the relation " + "attribute will be unset on returned rows.", + kwarg, + table, + ) + elif kwarg in ("select",): + logger.warning( + "compat shim: 'select=' on %s is not yet implemented; returning " + "the full row.", + table, + ) + + +class TableAccessor: + """The thing returned by ``client.db.litellm_
``. + + Two construction modes: + + * **standalone** (``owned_session=False``): each method opens its own + session, performs the work, commits, and closes. This matches the + historical Prisma-client-py call pattern where each ``await + db.foo.bar()`` is its own implicit transaction. + * **bound** (``owned_session=True``): the accessor reuses an existing + :class:`AsyncSession` opened by a surrounding ``tx()`` context. The + session is *not* closed and *not* committed -- the surrounding + transaction owns lifetime. + """ + + def __init__( + self, + model: Type[SQLModel], + db: "LiteLLMDB | None" = None, + bound_session: "AsyncSession | None" = None, + ) -> None: + if (db is None) == (bound_session is None): + raise ValueError( + "TableAccessor requires exactly one of db= or bound_session=" + ) + self._model = model + self._db = db + self._bound = bound_session + + @asynccontextmanager + async def _session(self) -> AsyncIterator[AsyncSession]: + if self._bound is not None: + yield self._bound + return + assert self._db is not None + async with self._db.session_ctx() as sess: + yield sess + + @property + def _is_bound(self) -> bool: + return self._bound is not None + + async def _maybe_commit(self, sess: AsyncSession) -> None: + if not self._is_bound: + await sess.commit() + + # ------------------------------------------------------------------ + # Queries + # ------------------------------------------------------------------ + + async def find_unique( + self, + *, + where: Optional[Mapping[str, Any]] = None, + include: Any = None, + select: Any = None, # noqa: A002 -- prisma kwarg name + ) -> Optional[SQLModel]: + if include is not None: + _warn_unsupported("include", self._model.__tablename__) + if select is not None: + _warn_unsupported("select", self._model.__tablename__) + clause = _translate_where(self._model, where) + stmt = self._select_stmt().limit(1) + if clause is not None: + stmt = stmt.where(clause) + return await self._scalar_one_or_none(stmt) + + async def find_first( + self, + *, + where: Optional[Mapping[str, Any]] = None, + order: Any = None, + include: Any = None, + skip: Optional[int] = None, + ) -> Optional[SQLModel]: + if include is not None: + _warn_unsupported("include", self._model.__tablename__) + clause = _translate_where(self._model, where) + stmt = self._select_stmt() + if clause is not None: + stmt = stmt.where(clause) + for clause_obj in _translate_order(self._model, order): + stmt = stmt.order_by(clause_obj) + if skip: + stmt = stmt.offset(skip) + stmt = stmt.limit(1) + return await self._scalar_one_or_none(stmt) + + async def find_many( + self, + *, + where: Optional[Mapping[str, Any]] = None, + order: Any = None, + take: Optional[int] = None, + skip: Optional[int] = None, + include: Any = None, + select: Any = None, # noqa: A002 + ) -> List[SQLModel]: + if include is not None: + _warn_unsupported("include", self._model.__tablename__) + if select is not None: + _warn_unsupported("select", self._model.__tablename__) + clause = _translate_where(self._model, where) + stmt = self._select_stmt() + if clause is not None: + stmt = stmt.where(clause) + for clause_obj in _translate_order(self._model, order): + stmt = stmt.order_by(clause_obj) + if take: + stmt = stmt.limit(take) + if skip: + stmt = stmt.offset(skip) + return await self._scalars_all(stmt) + + async def count( + self, *, where: Optional[Mapping[str, Any]] = None, **_: Any + ) -> int: + clause = _translate_where(self._model, where) + stmt = select(func.count()).select_from(self._model) + if clause is not None: + stmt = stmt.where(clause) + async with self._session() as sess: + try: + result = await sess.execute(stmt) + return int(result.scalar_one()) + except SQLAlchemyError as exc: + raise errors.map_sqlalchemy_error(exc) from exc + + # ------------------------------------------------------------------ + # Mutations + # ------------------------------------------------------------------ + + async def create(self, *, data: Mapping[str, Any], include: Any = None) -> SQLModel: + if include is not None: + _warn_unsupported("include", self._model.__tablename__) + translated = _translate_data(self._model, data) + instance = self._model(**translated) + async with self._session() as sess: + try: + sess.add(instance) + await sess.flush() + await self._maybe_commit(sess) + except SQLAlchemyError as exc: + if not self._is_bound: + await sess.rollback() + raise errors.map_sqlalchemy_error(exc) from exc + return instance + + async def update( + self, + *, + where: Mapping[str, Any], + data: Mapping[str, Any], + include: Any = None, + ) -> Optional[SQLModel]: + if include is not None: + _warn_unsupported("include", self._model.__tablename__) + clause = _translate_where(self._model, where) + translated = _translate_data(self._model, data) + async with self._session() as sess: + try: + stmt = self._select_stmt() + if clause is not None: + stmt = stmt.where(clause) + obj = (await sess.execute(stmt.limit(1))).scalar_one_or_none() + if obj is None: + raise errors.RecordNotFoundError( + f"No {self._model.__name__} matched where={dict(where)}" + ) + for col, val in translated.items(): + setattr(obj, col, val) + await sess.flush() + await self._maybe_commit(sess) + return obj + except SQLAlchemyError as exc: + if not self._is_bound: + await sess.rollback() + raise errors.map_sqlalchemy_error(exc) from exc + + async def upsert( + self, + *, + where: Mapping[str, Any], + data: Mapping[str, Any], + include: Any = None, + ) -> SQLModel: + """Prisma's ``upsert`` takes ``data={"create": {...}, "update": {...}}``. + + We mirror that contract -- the existing call sites all use the + nested ``create``/``update`` shape. + """ + if include is not None: + _warn_unsupported("include", self._model.__tablename__) + if ( + not isinstance(data, Mapping) + or "create" not in data + or "update" not in data + ): + raise ValueError( + "upsert(data=...) must be {'create': {...}, 'update': {...}}" + ) + existing = await self.find_unique(where=where) + if existing is None: + create_payload = dict(data["create"]) + return await self.create(data=create_payload) + return await self.update(where=where, data=data["update"]) # type: ignore[return-value] + + async def delete(self, *, where: Mapping[str, Any]) -> Optional[SQLModel]: + clause = _translate_where(self._model, where) + async with self._session() as sess: + try: + stmt = self._select_stmt() + if clause is not None: + stmt = stmt.where(clause) + obj = (await sess.execute(stmt.limit(1))).scalar_one_or_none() + if obj is None: + raise errors.RecordNotFoundError( + f"No {self._model.__name__} matched where={dict(where)}" + ) + await sess.delete(obj) + await self._maybe_commit(sess) + return obj + except SQLAlchemyError as exc: + if not self._is_bound: + await sess.rollback() + raise errors.map_sqlalchemy_error(exc) from exc + + async def update_many( + self, *, where: Mapping[str, Any], data: Mapping[str, Any] + ) -> CountResult: + clause = _translate_where(self._model, where) + translated = _translate_data(self._model, data) + async with self._session() as sess: + try: + stmt = update(self._model).values(**translated) + if clause is not None: + stmt = stmt.where(clause) + result = await sess.execute(stmt) + await self._maybe_commit(sess) + return CountResult(count=result.rowcount or 0) + except SQLAlchemyError as exc: + if not self._is_bound: + await sess.rollback() + raise errors.map_sqlalchemy_error(exc) from exc + + async def delete_many( + self, *, where: Optional[Mapping[str, Any]] = None + ) -> CountResult: + clause = _translate_where(self._model, where) + async with self._session() as sess: + try: + stmt = delete(self._model) + if clause is not None: + stmt = stmt.where(clause) + result = await sess.execute(stmt) + await self._maybe_commit(sess) + return CountResult(count=result.rowcount or 0) + except SQLAlchemyError as exc: + if not self._is_bound: + await sess.rollback() + raise errors.map_sqlalchemy_error(exc) from exc + + async def create_many( + self, *, data: Sequence[Mapping[str, Any]], skip_duplicates: bool = False + ) -> CountResult: + if not data: + return CountResult(count=0) + instances = [self._model(**_translate_data(self._model, item)) for item in data] + async with self._session() as sess: + try: + sess.add_all(instances) + await sess.flush() + await self._maybe_commit(sess) + return CountResult(count=len(instances)) + except SQLAlchemyError as exc: + if not self._is_bound: + await sess.rollback() + if skip_duplicates and isinstance( + errors.map_sqlalchemy_error(exc), errors.UniqueViolationError + ): + return CountResult(count=0) + raise errors.map_sqlalchemy_error(exc) from exc + + # ------------------------------------------------------------------ + # Internals + # ------------------------------------------------------------------ + + def _select_stmt(self): + return select(self._model) + + async def _scalar_one_or_none(self, stmt) -> Optional[SQLModel]: + async with self._session() as sess: + try: + result = await sess.execute(stmt) + return result.scalar_one_or_none() + except SQLAlchemyError as exc: + raise errors.map_sqlalchemy_error(exc) from exc + + async def _scalars_all(self, stmt) -> List[SQLModel]: + async with self._session() as sess: + try: + result = await sess.execute(stmt) + return list(result.scalars().all()) + except SQLAlchemyError as exc: + raise errors.map_sqlalchemy_error(exc) from exc + + +# --------------------------------------------------------------------------- +# Batch / transaction +# --------------------------------------------------------------------------- + + +class _BatchOp: + __slots__ = ("kind", "args", "kwargs", "model") + + def __init__(self, kind: str, model: Type[SQLModel], args, kwargs) -> None: + self.kind = kind + self.model = model + self.args = args + self.kwargs = kwargs + + +class _BatchTableAccessor: + """Records ops queued against a particular table for later flush.""" + + def __init__(self, model: Type[SQLModel], ops: List[_BatchOp]) -> None: + self._model = model + self._ops = ops + + def update(self, *args, **kwargs) -> None: + self._ops.append(_BatchOp("update", self._model, args, kwargs)) + + def upsert(self, *args, **kwargs) -> None: + self._ops.append(_BatchOp("upsert", self._model, args, kwargs)) + + def update_many(self, *args, **kwargs) -> None: + self._ops.append(_BatchOp("update_many", self._model, args, kwargs)) + + def create(self, *args, **kwargs) -> None: + self._ops.append(_BatchOp("create", self._model, args, kwargs)) + + def delete(self, *args, **kwargs) -> None: + self._ops.append(_BatchOp("delete", self._model, args, kwargs)) + + +class BatchAccessor: + """Object returned by ``db.batch_()``. + + Mimics prisma-client-py's batcher: queue mutations against any table + via ``batcher.litellm_
.(...)``, then flush all of them + in a single transaction with ``await batcher.commit()``. + """ + + def __init__(self, db: LiteLLMDB) -> None: + self._db = db + self._ops: List[_BatchOp] = [] + + def __getattr__(self, name: str) -> _BatchTableAccessor: + if name.startswith("_") or name in ("commit",): + raise AttributeError(name) + return _BatchTableAccessor(_model_for_accessor(name), self._ops) + + async def commit(self) -> None: + if not self._ops: + return + async with self._db.session_ctx() as sess: + try: + async with sess.begin(): + for op in self._ops: + await _execute_op_in_session(sess, op) + except SQLAlchemyError as exc: + raise errors.map_sqlalchemy_error(exc) from exc + + +class TransactionContext: + """Object yielded by ``async with db.tx():``. + + Exposes table accessors that share a single session so that all the + ops inside the ``with`` block execute in the same transaction. + """ + + def __init__(self, session: AsyncSession) -> None: + self._session = session + + def __getattr__(self, name: str) -> "TableAccessor": + if name.startswith("_") or name in ("batch_",): + raise AttributeError(name) + model = _model_for_accessor(name) + return TableAccessor(model, bound_session=self._session) + + def batch_(self) -> "TransactionBatch": + return TransactionBatch(self._session) + + +class TransactionBatch: + """``batch_`` inside an open transaction shares the same session.""" + + def __init__(self, session: AsyncSession) -> None: + self._session = session + self._ops: List[_BatchOp] = [] + + def __getattr__(self, name: str) -> _BatchTableAccessor: + if name.startswith("_") or name in ("commit",): + raise AttributeError(name) + return _BatchTableAccessor(_model_for_accessor(name), self._ops) + + async def commit(self) -> None: + for op in self._ops: + await _execute_op_in_session(self._session, op) + # No commit() here -- the surrounding tx() context manager owns + # transaction lifetime. + + +async def _execute_op_in_session(sess: AsyncSession, op: _BatchOp) -> None: + """Execute a queued batch op against an open session. + + Mirrors :class:`TableAccessor` but never opens its own transaction. + """ + model = op.model + kwargs = op.kwargs + if op.kind == "update": + clause = _translate_where(model, kwargs.get("where")) + translated = _translate_data(model, kwargs.get("data") or {}) + stmt = update(model).values(**translated) + if clause is not None: + stmt = stmt.where(clause) + await sess.execute(stmt) + elif op.kind == "update_many": + clause = _translate_where(model, kwargs.get("where")) + translated = _translate_data(model, kwargs.get("data") or {}) + stmt = update(model).values(**translated) + if clause is not None: + stmt = stmt.where(clause) + await sess.execute(stmt) + elif op.kind == "upsert": + # Without ON CONFLICT awareness in the generic shim, fall back to + # SELECT-then-update-or-insert. Sufficient for the call sites we + # see, all of which use unique-keyed where clauses. + clause = _translate_where(model, kwargs.get("where")) + existing_stmt = select(model).limit(1) + if clause is not None: + existing_stmt = existing_stmt.where(clause) + existing = (await sess.execute(existing_stmt)).scalar_one_or_none() + data = kwargs.get("data") or {} + if existing is None: + create_payload = _translate_data(model, dict(data.get("create") or {})) + sess.add(model(**create_payload)) + else: + update_payload = _translate_data(model, dict(data.get("update") or {})) + for col, val in update_payload.items(): + setattr(existing, col, val) + elif op.kind == "create": + translated = _translate_data(model, kwargs.get("data") or {}) + sess.add(model(**translated)) + elif op.kind == "delete": + clause = _translate_where(model, kwargs.get("where")) + stmt = delete(model) + if clause is not None: + stmt = stmt.where(clause) + await sess.execute(stmt) + else: + raise NotImplementedError(f"Batch op kind '{op.kind}' is not implemented.") + + +# --------------------------------------------------------------------------- +# Top-level Prisma-compatible client +# --------------------------------------------------------------------------- + + +class PrismaCompatClient: + """Replacement for prisma-client-py's ``Prisma`` class. + + Exposes a ``.db`` attribute (returns ``self`` -- the historical + ``prisma_client.db.litellm_x`` shape collapses cleanly because the + table accessors live on this object) plus connection lifecycle and + raw-SQL methods. + """ + + def __init__(self, db: LiteLLMDB) -> None: + self._db = db + + # The historical surface is ``prisma_client.db.litellm_x`` -- callers + # do ``client.db.litellm_x`` *and* ``client.litellm_x`` interchangeably + # in some places. We expose ``.db`` as ``self`` to match both. + @property + def db(self) -> "PrismaCompatClient": + return self + + # ------------------------------------------------------------------ + # Lifecycle + # ------------------------------------------------------------------ + + async def connect(self) -> None: + await self._db.connect() + + async def disconnect(self) -> None: + await self._db.disconnect() + + def is_connected(self) -> bool: + return self._db.is_connected() + + async def start_token_refresh_task(self) -> None: + await self._db.start_token_refresh_task() + + async def stop_token_refresh_task(self) -> None: + await self._db.stop_token_refresh_task() + + # ------------------------------------------------------------------ + # Table accessor lookup + # ------------------------------------------------------------------ + + def __getattr__(self, name: str) -> TableAccessor: + # Accessor names are lowercase; the property ``db`` is intentionally + # not routed here. Names starting with ``_`` are protected (so + # SQLAlchemy / pydantic introspection doesn't trip the lookup). + if name.startswith("_"): + raise AttributeError(name) + if name in _TABLE_INDEX: + return TableAccessor(_TABLE_INDEX[name], db=self._db) + raise AttributeError( + f"PrismaCompatClient has no attribute '{name}'. " + f"Known accessors: {sorted(_TABLE_INDEX.keys())[:5]}..." + ) + + # ------------------------------------------------------------------ + # Raw SQL + # ------------------------------------------------------------------ + + async def query_raw( + self, *args: Any, query: Optional[str] = None + ) -> List[Dict[str, Any]]: + sql = query + params: Iterable[Any] = () + if sql is None: + if not args: + raise ValueError("query_raw() requires either positional SQL or query=") + sql = args[0] + params = args[1:] + else: + params = args + async with self._db.session_ctx() as sess: + try: + result = await sess.execute(_to_text(sql, params), _params_dict(params)) + rows = result.fetchall() + return [_row_to_dict(r) for r in rows] + except SQLAlchemyError as exc: + raise errors.map_sqlalchemy_error(exc) from exc + + async def execute_raw(self, *args: Any, query: Optional[str] = None) -> int: + sql = query + params: Iterable[Any] = () + if sql is None: + if not args: + raise ValueError( + "execute_raw() requires either positional SQL or query=" + ) + sql = args[0] + params = args[1:] + else: + params = args + async with self._db.session_ctx() as sess: + try: + result = await sess.execute(_to_text(sql, params), _params_dict(params)) + await sess.commit() + return result.rowcount or 0 + except SQLAlchemyError as exc: + await sess.rollback() + raise errors.map_sqlalchemy_error(exc) from exc + + # ------------------------------------------------------------------ + # Batch / tx + # ------------------------------------------------------------------ + + def batch_(self) -> BatchAccessor: + return BatchAccessor(self._db) + + @asynccontextmanager + async def tx( + self, *, timeout: Optional[Any] = None, max_wait: Optional[Any] = None + ) -> AsyncIterator[TransactionContext]: + # ``timeout`` / ``max_wait`` are honoured by the underlying engine + # statement_timeout (set on the connection) when configured; we do + # not enforce them inside Python. + del timeout, max_wait + async with self._db.session_ctx() as sess: + async with sess.begin(): + yield TransactionContext(sess) + + +# --------------------------------------------------------------------------- +# Raw-SQL parameter handling (Postgres ``$1``-style -> SQLAlchemy named) +# --------------------------------------------------------------------------- + + +def _to_text(sql: str, params: Iterable[Any]): + """Rewrite ``$1``/``$2`` placeholders to SQLAlchemy ``:p1``/``:p2``. + + Prisma raw queries use Postgres-style positional placeholders. SQLAlchemy + needs ``text(":p1")`` with bound params. We rewrite numerically-stable + placeholders so the same SQL works on either backend. + """ + if not params: + return text(sql) + rewritten = sql + for i, _ in enumerate(params, start=1): + rewritten = rewritten.replace(f"${i}", f":p{i}") + return text(rewritten) + + +def _params_dict(params: Iterable[Any]) -> Dict[str, Any]: + return {f"p{i}": v for i, v in enumerate(params, start=1)} + + +# --------------------------------------------------------------------------- +# Convenience: build the global client +# --------------------------------------------------------------------------- + + +def create_client(db: LiteLLMDB) -> PrismaCompatClient: + """Factory used by ``litellm.proxy.utils.PrismaClient``.""" + return PrismaCompatClient(db) + + +# Silence unused-import warning for asyncio (used implicitly in cancellation +# semantics inside engine.py; keeping the import here documents that this +# module is async-aware end-to-end). +_ = asyncio diff --git a/litellm/proxy/db/sqlmodel/engine.py b/litellm/proxy/db/sqlmodel/engine.py new file mode 100644 index 00000000000..6f5f1e3fbbb --- /dev/null +++ b/litellm/proxy/db/sqlmodel/engine.py @@ -0,0 +1,326 @@ +"""SQLAlchemy async engine + session factory for the LiteLLM proxy. + +Replaces the Prisma engine subprocess that ``PrismaWrapper`` previously +managed. A handful of features the Prisma layer carried over remain the +responsibility of the layers above this module: + +* RDS IAM token rotation (``IAM_TOKEN_DB_AUTH=True``) -- handled here by + recreating the engine on token refresh. +* Read-replica routing (``DATABASE_URL_READ_REPLICA``) -- handled by + exposing a separate read-only engine alongside the writer. +* Postgres ``DATETIME`` / ``JSONB`` semantics -- the SQLModel classes in + :mod:`models` already pin Postgres-specific types via + ``sqlalchemy.dialects.postgresql``. + +This module is intentionally narrow: engine + sessionmaker + lifecycle. +The Prisma-compatible query API lives in :mod:`compat`. +""" + +from __future__ import annotations + +import asyncio +import logging +import os +from contextlib import asynccontextmanager +from typing import AsyncIterator, Optional + +from sqlalchemy.exc import SQLAlchemyError +from sqlalchemy.ext.asyncio import ( + AsyncEngine, + AsyncSession, + async_sessionmaker, + create_async_engine, +) + +logger = logging.getLogger(__name__) + + +def _normalize_url(url: str) -> str: + """Return an ``asyncpg``-compatible Postgres URL. + + LiteLLM accepts ``postgres://`` and ``postgresql://`` URLs (Prisma + accepted both). SQLAlchemy needs ``postgresql+asyncpg://`` for the + async driver. We rewrite the scheme conservatively and leave the rest + of the URL (host, port, query string) untouched so that connection + parameters such as ``sslmode`` survive. + """ + if url.startswith("postgresql+asyncpg://"): + return url + if url.startswith("postgresql://"): + return "postgresql+asyncpg://" + url[len("postgresql://") :] + if url.startswith("postgres://"): + return "postgresql+asyncpg://" + url[len("postgres://") :] + return url + + +class LiteLLMDB: + """Owns the writer (and optional reader) async engines for the proxy. + + The class is deliberately simple. We do not attempt to replicate the + subprocess-watchdog / engine-PID-reaping behaviour that ``PrismaWrapper`` + inherited from prisma-client-py: SQLAlchemy uses an in-process + connection pool and exposes the same recovery surface via + ``engine.dispose()``. + """ + + def __init__( + self, + database_url: str, + *, + read_replica_url: Optional[str] = None, + echo: bool = False, + pool_size: int = 10, + max_overflow: int = 20, + pool_recycle: int = 3600, + ) -> None: + if not database_url: + raise ValueError("LiteLLMDB requires a non-empty DATABASE_URL") + self._writer_url = database_url + self._reader_url = read_replica_url + self._echo = echo + self._pool_size = pool_size + self._max_overflow = max_overflow + self._pool_recycle = pool_recycle + + self._writer: AsyncEngine = self._build_engine(database_url) + self._reader: Optional[AsyncEngine] = ( + self._build_engine(read_replica_url) if read_replica_url else None + ) + self._writer_sm: async_sessionmaker[AsyncSession] = async_sessionmaker( + self._writer, expire_on_commit=False, class_=AsyncSession + ) + self._reader_sm: Optional[async_sessionmaker[AsyncSession]] = ( + async_sessionmaker( + self._reader, expire_on_commit=False, class_=AsyncSession + ) + if self._reader is not None + else None + ) + + self._iam_refresh_task: Optional[asyncio.Task[None]] = None + self._iam_refresh_interval_seconds: Optional[int] = None + self._iam_token_provider = None # callable[[], str] + self._connected: bool = False + + # ------------------------------------------------------------------ + # Engine plumbing + # ------------------------------------------------------------------ + + def _build_engine(self, url: str) -> AsyncEngine: + return create_async_engine( + _normalize_url(url), + echo=self._echo, + pool_size=self._pool_size, + max_overflow=self._max_overflow, + pool_recycle=self._pool_recycle, + pool_pre_ping=True, + future=True, + ) + + @property + def writer(self) -> AsyncEngine: + return self._writer + + @property + def reader(self) -> AsyncEngine: + return self._reader if self._reader is not None else self._writer + + # ------------------------------------------------------------------ + # Session factories + # ------------------------------------------------------------------ + + def session(self) -> AsyncSession: + """A new writer session (caller is responsible for ``await session.close()``).""" + return self._writer_sm() + + def reader_session(self) -> AsyncSession: + """A new reader session, or a writer session if no replica is configured.""" + if self._reader_sm is not None: + return self._reader_sm() + return self._writer_sm() + + @asynccontextmanager + async def session_ctx(self) -> AsyncIterator[AsyncSession]: + """Context manager that opens, yields, and closes a writer session.""" + session = self._writer_sm() + try: + yield session + finally: + await session.close() + + @asynccontextmanager + async def reader_session_ctx(self) -> AsyncIterator[AsyncSession]: + if self._reader_sm is None: + async with self.session_ctx() as session: + yield session + return + session = self._reader_sm() + try: + yield session + finally: + await session.close() + + # ------------------------------------------------------------------ + # Lifecycle (Prisma-compatible surface) + # ------------------------------------------------------------------ + + async def connect(self) -> None: + """Eagerly establish a single connection on each engine. + + SQLAlchemy lazily connects on first query, so this is mostly a + smoke test that catches misconfigured URLs / unreachable hosts at + startup time -- which is what callers expect from the previous + ``await prisma_client.db.connect()`` semantics. + """ + try: + async with self._writer.connect() as conn: + await conn.execute(_select_one()) + if self._reader is not None: + async with self._reader.connect() as conn: + await conn.execute(_select_one()) + except SQLAlchemyError as exc: + logger.error("LiteLLMDB.connect() failed: %s", exc) + raise + self._connected = True + + async def disconnect(self) -> None: + await self.stop_token_refresh_task() + try: + await self._writer.dispose() + finally: + if self._reader is not None: + await self._reader.dispose() + self._connected = False + + def is_connected(self) -> bool: + return self._connected + + # ------------------------------------------------------------------ + # IAM token rotation + # ------------------------------------------------------------------ + + def configure_iam_token_refresh( + self, *, token_provider, interval_seconds: int + ) -> None: + """Wire up a periodic engine recreate using the given token provider. + + ``token_provider`` is a callable returning a fresh password; we + rebuild the URL and dispose+recreate the engines on each tick. + """ + if not callable(token_provider): + raise TypeError("token_provider must be callable") + if interval_seconds <= 0: + raise ValueError("interval_seconds must be positive") + self._iam_token_provider = token_provider + self._iam_refresh_interval_seconds = interval_seconds + + async def start_token_refresh_task(self) -> None: + if self._iam_refresh_task is not None and not self._iam_refresh_task.done(): + return + if ( + self._iam_token_provider is None + or self._iam_refresh_interval_seconds is None + ): + return # IAM auth not configured + loop = asyncio.get_event_loop() + self._iam_refresh_task = loop.create_task(self._iam_refresh_loop()) + + async def stop_token_refresh_task(self) -> None: + task = self._iam_refresh_task + if task is None: + return + task.cancel() + try: + await task + except (asyncio.CancelledError, Exception): # noqa: BLE001 + pass + self._iam_refresh_task = None + + async def _iam_refresh_loop(self) -> None: + assert self._iam_refresh_interval_seconds is not None + assert self._iam_token_provider is not None + while True: + try: + await asyncio.sleep(self._iam_refresh_interval_seconds) + await self._rotate_iam_token() + except asyncio.CancelledError: + raise + except Exception: # noqa: BLE001 + logger.exception("IAM token rotation failed; will retry") + + async def _rotate_iam_token(self) -> None: + """Rebuild engines with a freshly minted IAM token.""" + assert self._iam_token_provider is not None + new_password = self._iam_token_provider() + new_writer_url = _replace_password(self._writer_url, new_password) + old_writer = self._writer + self._writer = self._build_engine(new_writer_url) + self._writer_sm = async_sessionmaker( + self._writer, expire_on_commit=False, class_=AsyncSession + ) + await old_writer.dispose() + + if self._reader is not None and self._reader_url is not None: + new_reader_url = _replace_password(self._reader_url, new_password) + old_reader = self._reader + self._reader = self._build_engine(new_reader_url) + self._reader_sm = async_sessionmaker( + self._reader, expire_on_commit=False, class_=AsyncSession + ) + await old_reader.dispose() + + +def _select_one(): + """Lazy import so the module loads without sqlalchemy installed.""" + from sqlalchemy import text + + return text("SELECT 1") + + +def _replace_password(url: str, new_password: str) -> str: + """Return ``url`` with its password component replaced. + + Avoids importing ``urllib`` at module import time because some + deployments swap in a different URL parser. + """ + from urllib.parse import urlparse, urlunparse + + parsed = urlparse(url) + if parsed.username is None: + return url + netloc_user = parsed.username + netloc_host = parsed.hostname or "" + port = f":{parsed.port}" if parsed.port else "" + new_netloc = f"{netloc_user}:{new_password}@{netloc_host}{port}" + return urlunparse(parsed._replace(netloc=new_netloc)) + + +# --------------------------------------------------------------------------- +# Module-level singleton +# --------------------------------------------------------------------------- + +_DB: Optional[LiteLLMDB] = None + + +def configure(database_url: str, **kwargs) -> LiteLLMDB: + """Initialise (or replace) the module-level :class:`LiteLLMDB` singleton.""" + global _DB + if _DB is not None: + # Best-effort: caller is reconfiguring (e.g. after credential rotation). + # We do not auto-dispose the previous engine because the caller may + # still hold sessions; PrismaClient.connect handles lifecycle there. + pass + _DB = LiteLLMDB(database_url, **kwargs) + return _DB + + +def get_db() -> LiteLLMDB: + if _DB is None: + url = os.getenv("DATABASE_URL") + if not url: + raise RuntimeError( + "LiteLLMDB has not been configured -- call engine.configure(...) " + "or set DATABASE_URL before any DB query." + ) + return configure(url, read_replica_url=os.getenv("DATABASE_URL_READ_REPLICA")) + return _DB diff --git a/litellm/proxy/db/sqlmodel/errors.py b/litellm/proxy/db/sqlmodel/errors.py new file mode 100644 index 00000000000..1bdfd327a9c --- /dev/null +++ b/litellm/proxy/db/sqlmodel/errors.py @@ -0,0 +1,106 @@ +"""Prisma-compatible exception types backed by SQLAlchemy errors. + +Existing call sites do ``except prisma.errors.UniqueViolationError`` and +similar. The shim raises *these* classes (re-exported from this module) +so call sites need not change. The Prisma package is no longer a runtime +dependency once the rip-and-replace lands; without these stand-ins, +``except prisma.errors.X`` would itself fail to import. +""" + +from __future__ import annotations + + +class PrismaError(Exception): + """Base class mirroring ``prisma.errors.PrismaError``. + + prisma-client-py's exception classes were Pydantic-ish: they accepted + arbitrary keyword arguments to capture structured error metadata + (table name, query, the upstream Prisma JSON envelope, etc.). We + don't model that envelope, but we do accept and ignore arbitrary + kwargs so legacy ``raise UniqueViolationError(data={"...": ...})`` + call sites and exception fixtures keep working. + """ + + def __init__(self, *args: object, **kwargs: object) -> None: + super().__init__(*args) + self.data = kwargs.get("data") + self._extra = kwargs + + +class DataError(PrismaError): + """Generic data-layer error (e.g. constraint, type cast).""" + + +class UniqueViolationError(DataError): + """Raised on UNIQUE constraint violations (Postgres SQLSTATE 23505).""" + + +class ForeignKeyViolationError(DataError): + """Raised on FOREIGN KEY constraint violations (Postgres SQLSTATE 23503).""" + + +class RecordNotFoundError(DataError): + """Raised when ``find_unique`` / ``update`` targets a non-existent row.""" + + +class MissingRequiredValueError(DataError): + """Raised when a non-null column is omitted from ``data=``.""" + + +class TableNotFoundError(DataError): + """Raised when a table referenced by the shim does not exist in the schema.""" + + +class RawQueryError(DataError): + """Raised when a ``query_raw`` / ``execute_raw`` statement fails.""" + + +class ClientNotConnectedError(PrismaError): + """Raised when the engine has been disposed but a query is attempted.""" + + +class HTTPClientClosedError(ClientNotConnectedError): + """Compatibility alias -- prisma-client-py distinguishes these.""" + + +def map_sqlalchemy_error(exc: Exception) -> Exception: + """Translate a SQLAlchemy ``IntegrityError`` (etc.) to a Prisma-compat error. + + The mapping is conservative: anything we cannot classify is wrapped in + :class:`PrismaError` so callers' ``except PrismaError`` blocks still + behave. Connection-layer errors are intentionally **not** mapped here + -- they propagate untranslated so the existing reconnect logic in + :class:`litellm.proxy.db.exception_handler.PrismaDBExceptionHandler` + can recognise them. + """ + from sqlalchemy.exc import ( # local import; SQLAlchemy is optional + IntegrityError, + NoResultFound, + ProgrammingError, + StatementError, + ) + + if isinstance(exc, NoResultFound): + return RecordNotFoundError(str(exc)) + + if isinstance(exc, IntegrityError): + msg = str(exc.orig) if getattr(exc, "orig", None) is not None else str(exc) + lowered = msg.lower() + if "unique" in lowered or "duplicate key" in lowered: + return UniqueViolationError(msg) + if "foreign key" in lowered: + return ForeignKeyViolationError(msg) + if "not-null" in lowered or "violates not-null constraint" in lowered: + return MissingRequiredValueError(msg) + return DataError(msg) + + if isinstance(exc, ProgrammingError): + msg = str(exc.orig) if getattr(exc, "orig", None) is not None else str(exc) + if "does not exist" in msg.lower() and "relation" in msg.lower(): + return TableNotFoundError(msg) + return RawQueryError(msg) + + if isinstance(exc, StatementError): + return DataError(str(exc)) + + return exc diff --git a/litellm/proxy/db/sqlmodel/models.py b/litellm/proxy/db/sqlmodel/models.py new file mode 100644 index 00000000000..8e52b8795ba --- /dev/null +++ b/litellm/proxy/db/sqlmodel/models.py @@ -0,0 +1,2546 @@ +"""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, + JSON, + 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().with_variant(JSON(), "sqlite"), + default_factory=list, + sa_column_kwargs={"server_default": text("'[]'")}, + ) + pending_tool_calls: Any = Field( + sa_type=JSONB().with_variant(JSON(), "sqlite"), + 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().with_variant(JSON(), "sqlite"), default=None + ) + agent_card_params: Any = Field(sa_type=JSONB().with_variant(JSON(), "sqlite")) + static_headers: Optional[Any] = Field( + sa_type=JSONB().with_variant(JSON(), "sqlite"), + 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().with_variant(JSON(), "sqlite"), default=None + ) + updated_values: Optional[Any] = Field( + sa_type=JSONB().with_variant(JSON(), "sqlite"), 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().with_variant(JSON(), "sqlite"), 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().with_variant(JSON(), "sqlite")) + 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().with_variant(JSON(), "sqlite"), 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().with_variant(JSON(), "sqlite")) + 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().with_variant(JSON(), "sqlite")) + credential_info: Optional[Any] = Field( + sa_type=JSONB().with_variant(JSON(), "sqlite"), 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().with_variant(JSON(), "sqlite"), + default_factory=dict, + sa_column_kwargs={"server_default": text("'{}'")}, + ) + metadata_: Any = Field( + sa_type=JSONB().with_variant(JSON(), "sqlite"), + 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().with_variant(JSON(), "sqlite"), + default_factory=dict, + sa_column_kwargs={"server_default": text("'{}'")}, + ) + model_max_budget: Any = Field( + sa_type=JSONB().with_variant(JSON(), "sqlite"), + default_factory=dict, + sa_column_kwargs={"server_default": text("'{}'")}, + ) + router_settings: Optional[Any] = Field( + sa_type=JSONB().with_variant(JSON(), "sqlite"), + 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().with_variant(JSON(), "sqlite"), + default_factory=dict, + sa_column_kwargs={"server_default": text("'{}'")}, + ) + config: Any = Field( + sa_type=JSONB().with_variant(JSON(), "sqlite"), + 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().with_variant(JSON(), "sqlite"), + 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().with_variant(JSON(), "sqlite"), + 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().with_variant(JSON(), "sqlite"), + default_factory=dict, + sa_column_kwargs={"server_default": text("'{}'")}, + ) + model_max_budget: Any = Field( + sa_type=JSONB().with_variant(JSON(), "sqlite"), + default_factory=dict, + sa_column_kwargs={"server_default": text("'{}'")}, + ) + router_settings: Optional[Any] = Field( + sa_type=JSONB().with_variant(JSON(), "sqlite"), + 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().with_variant(JSON(), "sqlite"), + 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().with_variant(JSON(), "sqlite")) + guardrail_info: Optional[Any] = Field( + sa_type=JSONB().with_variant(JSON(), "sqlite"), 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().with_variant(JSON(), "sqlite"), 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().with_variant(JSON(), "sqlite"), + 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().with_variant(JSON(), "sqlite"), + 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().with_variant(JSON(), "sqlite"), + default_factory=dict, + sa_column_kwargs={"server_default": text("'{}'")}, + ) + tool_name_to_description: Optional[Any] = Field( + sa_type=JSONB().with_variant(JSON(), "sqlite"), + 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().with_variant(JSON(), "sqlite"), + 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().with_variant(JSON(), "sqlite"), + 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().with_variant(JSON(), "sqlite"), + 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().with_variant(JSON(), "sqlite"), default=None + ) + model_mappings: Any = Field(sa_type=JSONB().with_variant(JSON(), "sqlite")) + 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().with_variant(JSON(), "sqlite")) + 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().with_variant(JSON(), "sqlite")) + index_info: Optional[Any] = Field( + sa_type=JSONB().with_variant(JSON(), "sqlite"), 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().with_variant(JSON(), "sqlite"), default=None + ) + model_mappings: Any = Field(sa_type=JSONB().with_variant(JSON(), "sqlite")) + 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().with_variant(JSON(), "sqlite"), 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().with_variant(JSON(), "sqlite"), 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().with_variant(JSON(), "sqlite"), + 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().with_variant(JSON(), "sqlite"), + 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().with_variant(JSON(), "sqlite"), 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().with_variant(JSON(), "sqlite"), + 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().with_variant(JSON(), "sqlite"), + 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().with_variant(JSON(), "sqlite"), + default_factory=dict, + sa_column_kwargs={"server_default": text("'{}'")}, + ) + pipeline: Optional[Any] = Field( + sa_type=JSONB().with_variant(JSON(), "sqlite"), 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().with_variant(JSON(), "sqlite"), + 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().with_variant(JSON(), "sqlite"), + default_factory=dict, + sa_column_kwargs={"server_default": text("'{}'")}, + ) + model_rpm_limit: Any = Field( + sa_type=JSONB().with_variant(JSON(), "sqlite"), + default_factory=dict, + sa_column_kwargs={"server_default": text("'{}'")}, + ) + model_tpm_limit: Any = Field( + sa_type=JSONB().with_variant(JSON(), "sqlite"), + 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().with_variant(JSON(), "sqlite")) + prompt_info: Optional[Any] = Field( + sa_type=JSONB().with_variant(JSON(), "sqlite"), 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().with_variant(JSON(), "sqlite")) + model_info: Optional[Any] = Field( + sa_type=JSONB().with_variant(JSON(), "sqlite"), 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().with_variant(JSON(), "sqlite")) + 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().with_variant(JSON(), "sqlite")) + search_tool_info: Optional[Any] = Field( + sa_type=JSONB().with_variant(JSON(), "sqlite"), 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().with_variant(JSON(), "sqlite"), + 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().with_variant(JSON(), "sqlite"), + 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().with_variant(JSON(), "sqlite"), + 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().with_variant(JSON(), "sqlite"), + default_factory=dict, + sa_column_kwargs={"server_default": text("'{}'")}, + ) + response: Optional[Any] = Field( + sa_type=JSONB().with_variant(JSON(), "sqlite"), + 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().with_variant(JSON(), "sqlite"), + 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().with_variant(JSON(), "sqlite"), 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().with_variant(JSON(), "sqlite"), + default_factory=dict, + sa_column_kwargs={"server_default": text("'{}'")}, + ) + metadata_: Any = Field( + sa_type=JSONB().with_variant(JSON(), "sqlite"), + 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().with_variant(JSON(), "sqlite"), + default_factory=dict, + sa_column_kwargs={"server_default": text("'{}'")}, + ) + model_max_budget: Any = Field( + sa_type=JSONB().with_variant(JSON(), "sqlite"), + default_factory=dict, + sa_column_kwargs={"server_default": text("'{}'")}, + ) + router_settings: Optional[Any] = Field( + sa_type=JSONB().with_variant(JSON(), "sqlite"), + 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().with_variant(JSON(), "sqlite"), 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().with_variant(JSON(), "sqlite"), + 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().with_variant(JSON(), "sqlite")) + 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().with_variant(JSON(), "sqlite"), + 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().with_variant(JSON(), "sqlite"), + default_factory=dict, + sa_column_kwargs={"server_default": text("'{}'")}, + ) + model_max_budget: Any = Field( + sa_type=JSONB().with_variant(JSON(), "sqlite"), + 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().with_variant(JSON(), "sqlite"), + default_factory=dict, + sa_column_kwargs={"server_default": text("'{}'")}, + ) + config: Any = Field( + sa_type=JSONB().with_variant(JSON(), "sqlite"), + default_factory=dict, + sa_column_kwargs={"server_default": text("'{}'")}, + ) + router_settings: Optional[Any] = Field( + sa_type=JSONB().with_variant(JSON(), "sqlite"), + 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().with_variant(JSON(), "sqlite"), + 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().with_variant(JSON(), "sqlite"), + 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().with_variant(JSON(), "sqlite"), + default_factory=dict, + sa_column_kwargs={"server_default": text("'{}'")}, + ) + model_max_budget: Any = Field( + sa_type=JSONB().with_variant(JSON(), "sqlite"), + 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().with_variant(JSON(), "sqlite"), 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().with_variant(JSON(), "sqlite"), 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().with_variant(JSON(), "sqlite"), default=None + ) + output: Optional[Any] = Field( + sa_type=JSONB().with_variant(JSON(), "sqlite"), default=None + ) + metadata_: Optional[Any] = Field( + sa_type=JSONB().with_variant(JSON(), "sqlite"), + 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 00000000000..cf13aa68225 --- /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/litellm/proxy/management_endpoints/budget_management_endpoints.py b/litellm/proxy/management_endpoints/budget_management_endpoints.py index 60dc7827a6f..740b5c221a3 100644 --- a/litellm/proxy/management_endpoints/budget_management_endpoints.py +++ b/litellm/proxy/management_endpoints/budget_management_endpoints.py @@ -48,7 +48,7 @@ async def new_budget( - model_max_budget: Optional[dict] - Specify max budget for a given model. Example: {"openai/gpt-4o-mini": {"max_budget": 100.0, "budget_duration": "1d", "tpm_limit": 100000, "rpm_limit": 100000}} - budget_reset_at: Optional[datetime] - Datetime when the initial budget is reset. Default is now. """ - from prisma.errors import UniqueViolationError + from litellm.proxy.db.sqlmodel.errors import UniqueViolationError from litellm.proxy.proxy_server import litellm_proxy_admin_name, prisma_client diff --git a/litellm/proxy/management_endpoints/config_override_endpoints.py b/litellm/proxy/management_endpoints/config_override_endpoints.py index 7f7aa485fb3..92df64fdf62 100644 --- a/litellm/proxy/management_endpoints/config_override_endpoints.py +++ b/litellm/proxy/management_endpoints/config_override_endpoints.py @@ -12,7 +12,7 @@ from litellm.litellm_core_utils.safe_json_loads import safe_json_loads try: - from prisma.errors import RecordNotFoundError + from litellm.proxy.db.sqlmodel.errors import RecordNotFoundError except ImportError: RecordNotFoundError = Exception # type: ignore diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 2ab147043d9..5eeaa7f195a 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -3994,7 +3994,17 @@ async def _rotate_master_key( # noqa: PLR0915 3. Encrypt the values with the new master key 4. Update the values in the DB """ - import prisma + + # ``prisma.Json(...)`` was a prisma-client-py marker that wrapped Python + # dicts before insertion into Prisma ``Json`` columns. SQLAlchemy's + # ``JSONB`` type accepts ``dict`` / ``list`` directly, so we provide a + # local identity shim and skip importing the obsolete prisma package. + class _PrismaJsonShim: + @staticmethod + def Json(value): # noqa: N802 -- preserves prisma-client-py call shape + return value + + prisma = _PrismaJsonShim() from litellm.proxy.proxy_server import proxy_config diff --git a/litellm/proxy/management_endpoints/mcp_management_endpoints.py b/litellm/proxy/management_endpoints/mcp_management_endpoints.py index e9d9c243e7c..97725a2bc8e 100644 --- a/litellm/proxy/management_endpoints/mcp_management_endpoints.py +++ b/litellm/proxy/management_endpoints/mcp_management_endpoints.py @@ -37,7 +37,10 @@ from fastapi.responses import JSONResponse try: - from prisma.errors import RecordNotFoundError, UniqueViolationError + from litellm.proxy.db.sqlmodel.errors import ( + RecordNotFoundError, + UniqueViolationError, + ) except ImportError: RecordNotFoundError = Exception # type: ignore UniqueViolationError = Exception # type: ignore diff --git a/litellm/proxy/management_endpoints/workflow_management_endpoints.py b/litellm/proxy/management_endpoints/workflow_management_endpoints.py index a19af4dd484..20af1759085 100644 --- a/litellm/proxy/management_endpoints/workflow_management_endpoints.py +++ b/litellm/proxy/management_endpoints/workflow_management_endpoints.py @@ -19,7 +19,7 @@ from fastapi import APIRouter, Depends, HTTPException, Query try: - from prisma.errors import UniqueViolationError + from litellm.proxy.db.sqlmodel.errors import UniqueViolationError except ImportError: UniqueViolationError = None # type: ignore from pydantic import BaseModel diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 32c887f17b2..07e75c571a9 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -2738,115 +2738,85 @@ def __init__( self.iam_token_db_auth: Optional[bool] = str_to_bool( os.getenv("IAM_TOKEN_DB_AUTH") ) - verbose_proxy_logger.debug("Creating Prisma Client..") - try: - from prisma import Prisma # type: ignore - except Exception as e: - verbose_proxy_logger.error(f"Failed to import Prisma client: {e}") - verbose_proxy_logger.error( - "This usually means 'prisma generate' hasn't been run yet." - ) - verbose_proxy_logger.error( - "Please run 'prisma generate' to generate the Prisma client." - ) - raise Exception( - "Unable to find Prisma binaries. Please run 'prisma generate' first." + verbose_proxy_logger.debug( + "Creating SQLAlchemy-backed prisma-compatibility client.." + ) + # NOTE: as part of the Prisma -> SQLAlchemy big-bang migration, this + # constructor no longer instantiates a Prisma client. Instead it + # builds a :class:`LiteLLMDB` (engine + sessionmaker) and exposes + # the historical ``prisma_client.db.
.(...)`` surface + # via :class:`PrismaCompatClient` (see + # ``litellm/proxy/db/sqlmodel/compat.py``). All ~1,680 existing + # call sites continue to work without modification. + from litellm.proxy.db.sqlmodel.compat import ( + PrismaCompatClient, + create_client, + ) + from litellm.proxy.db.sqlmodel.engine import LiteLLMDB + + # ``http_client`` is preserved in the signature for compatibility but + # is unused -- prisma-client-py forwarded it to its HTTP engine, and + # the SQLAlchemy backend has no equivalent. We log a one-line + # warning so deployment scripts that previously passed a custom + # client see why the kwarg is being ignored. + if http_client is not None: + verbose_proxy_logger.warning( + "PrismaClient(http_client=...) is ignored after the SQLAlchemy " + "migration; SQLAlchemy manages its own connection pool." ) + iam_flag = ( self.iam_token_db_auth if self.iam_token_db_auth is not None else False ) - # When read-replica routing is on, tag log lines with [writer]/[reader] - # so the two wrappers' interleaved IAM refresh logs can be told apart. - # Single-DB deployments get an empty prefix (logs unchanged). read_replica_url = os.getenv("DATABASE_URL_READ_REPLICA") - writer_log_prefix = "[writer]" if read_replica_url else "" - if http_client is not None: - writer_wrapper = PrismaWrapper( - original_prisma=Prisma(http=http_client), - iam_token_db_auth=iam_flag, - log_prefix=writer_log_prefix, - ) - else: - writer_wrapper = PrismaWrapper( - original_prisma=Prisma(), - iam_token_db_auth=iam_flag, - log_prefix=writer_log_prefix, - ) - - # Optional read-replica routing. When DATABASE_URL_READ_REPLICA is set, - # reads (find_*, count, group_by, query_raw/_first) are routed to the - # reader endpoint and writes stay on the writer. Falls back to the - # writer-only wrapper when the env var is unset, preserving existing - # single-DB deployments. - self.db: Union[PrismaWrapper, RoutingPrismaWrapper] - if read_replica_url: + litellm_db = LiteLLMDB( + database_url=database_url, + read_replica_url=read_replica_url, + ) + # IAM token rotation is wired in :class:`LiteLLMDB`; callers that + # need it should configure ``iam_token_db_auth`` on the client. + if iam_flag: try: - # If IAM auth is enabled, the reader refreshes its own token on - # the same cadence as the writer. We parse the static endpoint - # pieces (host/port/user/db) once from the reader URL — only - # the IAM token rotates after that. - reader_iam_endpoint = ( - parse_iam_endpoint_from_url(read_replica_url) if iam_flag else None - ) - # Mint a fresh IAM token for the reader BEFORE constructing the - # Prisma client. Mirrors what `proxy_cli.py` already does for - # the writer (proxy_cli.py:812-832) — without this, the reader - # Prisma is built with whatever placeholder URL the user - # supplied (no real token), and the first query falls through - # to the synchronous fallback path in - # `PrismaWrapper.__getattr__`, which deadlocks the event loop - # and times out after 30s. - if iam_flag and reader_iam_endpoint is not None: - from litellm.proxy.auth.rds_iam_token import ( - generate_iam_auth_token, - ) + from litellm.proxy.auth.rds_iam_token import generate_iam_auth_token + + parsed = urllib.parse.urlparse(database_url) - reader_token = generate_iam_auth_token( - db_host=reader_iam_endpoint.host, - db_port=reader_iam_endpoint.port, - db_user=reader_iam_endpoint.user, + def _provider() -> str: + return generate_iam_auth_token( + db_host=parsed.hostname or "", + db_port=str(parsed.port or 5432), + db_user=parsed.username or "", ) - read_replica_url = reader_iam_endpoint.build_url(reader_token) - os.environ["DATABASE_URL_READ_REPLICA"] = read_replica_url - reader_kwargs: Dict[str, Any] = { - "datasource": {"url": read_replica_url} - } - if http_client is not None: - reader_prisma = Prisma(http=http_client, **reader_kwargs) - else: - reader_prisma = Prisma(**reader_kwargs) - reader_wrapper = PrismaWrapper( - original_prisma=reader_prisma, - iam_token_db_auth=iam_flag, - db_url_env_var="DATABASE_URL_READ_REPLICA", - iam_endpoint=reader_iam_endpoint, - recreate_uses_datasource=True, - log_prefix="[reader]", - ) - self.db = RoutingPrismaWrapper( - writer=writer_wrapper, reader=reader_wrapper - ) - verbose_proxy_logger.info( - "PrismaClient: read-replica routing enabled via DATABASE_URL_READ_REPLICA" - + (" (with IAM token auto-refresh)" if iam_flag else "") + + # 12-minute refresh -- matches the historical default + # cadence (RDS tokens expire after 15 minutes). + litellm_db.configure_iam_token_refresh( + token_provider=_provider, interval_seconds=12 * 60 ) - except Exception as e: - # Reader is opt-in; never let its construction fail proxy - # startup. Mirrors the runtime contract from - # `RoutingPrismaWrapper.connect`: reader-side failures are - # logged and we keep serving traffic via the writer alone. - # This recovers from transient AWS STS hiccups during the - # reader IAM token mint, malformed DATABASE_URL_READ_REPLICA, - # and Prisma construction errors. Operator restart is required - # to retry read-routing once the underlying issue is resolved. + except Exception as exc: # noqa: BLE001 verbose_proxy_logger.warning( - "Failed to initialize read replica Prisma client: %s. " - "Falling back to writer-only mode (no read routing) until proxy restart.", - e, + "Failed to wire IAM token refresh: %s. Continuing without " + "automatic rotation -- the proxy will only succeed for the " + "lifetime of the initial IAM token.", + exc, ) - self.db = writer_wrapper - else: - self.db = writer_wrapper # Client to connect to Prisma db + + compat_client: PrismaCompatClient = create_client(litellm_db) + # ``self.db`` historically pointed at PrismaWrapper / RoutingPrismaWrapper + # which wrapped a real ``prisma.Prisma`` instance. PrismaCompatClient + # exposes the same surface (``litellm_
.``, + # ``query_raw``, ``execute_raw``, ``batch_``, ``tx``, + # ``connect/disconnect/is_connected``, IAM hooks) so call sites + # continue to work unchanged. Read-replica routing is honoured at + # the engine level (see :class:`LiteLLMDB`) rather than via a + # wrapper that dispatches per-call -- the two-engine + per-session + # routing approach is more idiomatic for SQLAlchemy. + self.db = compat_client # type: ignore[assignment] + if read_replica_url: + verbose_proxy_logger.info( + "PrismaClient: read-replica routing enabled via DATABASE_URL_READ_REPLICA " + "(handled at the SQLAlchemy engine layer)." + ) self._db_reconnect_lock = asyncio.Lock() self._db_health_watchdog_task: Optional[asyncio.Task] = None self._db_last_reconnect_attempt_ts: float = 0.0 @@ -4182,13 +4152,10 @@ async def disconnect(self): raise e def _get_engine_pid(self) -> int: - try: - engine = self.db._original_prisma._engine # type: ignore[attr-defined] - process = getattr(engine, "process", None) if engine is not None else None - if process is not None: - return process.pid - except (AttributeError, TypeError): - pass + # SQLAlchemy uses an in-process connection pool; there is no engine + # subprocess to track. Returning 0 short-circuits ``_is_engine_alive`` + # to ``True`` and disables the zombie-reaping watchdog that + # prisma-client-py needed. return 0 def _is_engine_alive(self) -> bool: diff --git a/pyproject.toml b/pyproject.toml index 70681c4ed6c..005f314cfb1 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -66,7 +66,9 @@ proxy = [ "pydantic-settings>=2.14.1", ] extra_proxy = [ - "prisma==0.11.0", + # prisma==0.11.0 was removed in the SQLAlchemy migration. The + # ``litellm.proxy.db.sqlmodel`` package provides a Prisma-compatible + # shim backed by SQLAlchemy + asyncpg. "azure-identity==1.25.2", "azure-keyvault-secrets==4.10.0", # Not in PyPI proxy extra. @@ -76,6 +78,9 @@ 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", + "asyncpg>=0.29,<1.0", + "alembic>=1.13,<2.0", ] utils = [ # Not in Docker or PyPI proxy extra. @@ -122,6 +127,9 @@ proxy-runtime = [ "llm-sandbox==0.3.39", "detect-secrets==1.5.0", ] +proxy-dev = [ + "aiosqlite>=0.19,<1.0", +] [project.scripts] litellm = "litellm:run_server" @@ -166,7 +174,8 @@ dev = [ "pytest-recording==0.13.4", ] proxy-dev = [ - "prisma==0.11.0", + # prisma==0.11.0 was removed in the SQLAlchemy migration. See + # ``litellm.proxy.db.sqlmodel`` for the SQLAlchemy-backed replacement. "hypercorn==0.17.3", "prometheus-client==0.20.0", "opentelemetry-api==1.28.0", diff --git a/tests/_prisma_compat.py b/tests/_prisma_compat.py new file mode 100644 index 00000000000..a5d2dc4e88b --- /dev/null +++ b/tests/_prisma_compat.py @@ -0,0 +1,84 @@ +"""Test-only prisma stand-in. + +After the SQLAlchemy big-bang migration, the ``prisma`` PyPI package is +no longer installed. A handful of tests still do +``from prisma.errors import RecordNotFoundError`` etc.; this module +registers thin sys.modules entries so those imports resolve to the +LiteLLM-native error classes defined in +:mod:`litellm.proxy.db.sqlmodel.errors`. + +The shim is intentionally only loaded from ``tests/test_litellm/conftest.py`` +and ``tests/conftest.py`` -- it does **not** ship in the ``litellm`` +wheel, and production code that previously imported from ``prisma`` has +been ported to the native errors module. +""" + +from __future__ import annotations + +import sys +import types + +from litellm.proxy.db.sqlmodel import errors as _native_errors + + +def _build_prisma_errors_module() -> types.ModuleType: + mod = types.ModuleType("prisma.errors") + for name in ( + "PrismaError", + "DataError", + "UniqueViolationError", + "ForeignKeyViolationError", + "RecordNotFoundError", + "MissingRequiredValueError", + "TableNotFoundError", + "RawQueryError", + "ClientNotConnectedError", + "HTTPClientClosedError", + ): + setattr(mod, name, getattr(_native_errors, name)) + return mod + + +def _build_prisma_module(errors_mod: types.ModuleType) -> types.ModuleType: + mod = types.ModuleType("prisma") + mod.errors = errors_mod # type: ignore[attr-defined] + + # ``prisma.Json`` was a marker callable for inserting Python dicts into + # Prisma ``Json`` columns. SQLAlchemy's JSONB accepts dicts directly, + # so we re-export an identity function. + def _json(value): # noqa: N802 -- preserves prisma-client-py call shape + return value + + mod.Json = _json # type: ignore[attr-defined] + + class _PrismaUnusable: + """Stand-in for the obsolete ``prisma.Prisma`` class. + + Tests that imported it for ``isinstance`` checks or to construct + a real engine subprocess no longer apply -- we raise on construction + so any forgotten production usage surfaces loudly rather than + silently passing. + """ + + def __init__(self, *args, **kwargs): + raise RuntimeError( + "prisma.Prisma is no longer available -- the SQLAlchemy " + "migration removed the prisma-client-py dependency. Use " + "litellm.proxy.db.sqlmodel.compat.PrismaCompatClient instead." + ) + + mod.Prisma = _PrismaUnusable # type: ignore[attr-defined] + return mod + + +def install() -> None: + """Register the stub modules in ``sys.modules`` if real prisma is absent.""" + if "prisma" in sys.modules: + return + try: + import prisma # type: ignore[import-not-found] # noqa: F401 + except ImportError: + errors_mod = _build_prisma_errors_module() + prisma_mod = _build_prisma_module(errors_mod) + sys.modules["prisma"] = prisma_mod + sys.modules["prisma.errors"] = errors_mod diff --git a/tests/test_litellm/conftest.py b/tests/test_litellm/conftest.py index f4aa1926d21..8f724124ddf 100644 --- a/tests/test_litellm/conftest.py +++ b/tests/test_litellm/conftest.py @@ -16,8 +16,20 @@ sys.path.insert( 0, os.path.abspath("../..") ) # Adds the parent directory to the system path +sys.path.insert(0, os.path.abspath("../")) # tests/_prisma_compat import asyncio +# Install the prisma stand-in **before** importing litellm. After the +# SQLAlchemy big-bang migration, prisma is no longer a runtime dep but +# a handful of tests still ``from prisma.errors import X``; the shim +# registers stub modules in ``sys.modules`` so those imports resolve. +try: + from _prisma_compat import install as _install_prisma_compat + + _install_prisma_compat() +except ImportError: + pass + import litellm from litellm._logging import ALL_LOGGERS from litellm.litellm_core_utils.prompt_templates import ( 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 00000000000..e69de29bb2d diff --git a/tests/test_litellm/proxy/db/sqlmodel_orm/test_compat.py b/tests/test_litellm/proxy/db/sqlmodel_orm/test_compat.py new file mode 100644 index 00000000000..e001382971d --- /dev/null +++ b/tests/test_litellm/proxy/db/sqlmodel_orm/test_compat.py @@ -0,0 +1,324 @@ +"""End-to-end tests for the Prisma-compatibility shim against SQLite. + +These exercise the shim through its public surface so that breakage in +filter translation, mutations, batch_, or tx surfaces here rather than +in the proxy at runtime. + +Run with:: + + uv run pytest tests/test_litellm/proxy/db/sqlmodel_orm/test_compat.py -vv +""" + +from __future__ import annotations + +import os +from typing import AsyncIterator + +import pytest +import pytest_asyncio +from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine +from sqlmodel import SQLModel + +from litellm.proxy.db.sqlmodel import errors +from litellm.proxy.db.sqlmodel.compat import PrismaCompatClient +from litellm.proxy.db.sqlmodel.engine import LiteLLMDB +from litellm.proxy.db.sqlmodel.models import ( + LiteLLMConfig, + LiteLLMSpendLogs, + LiteLLMTeamMembership, +) + + +class _SQLiteLiteLLMDB(LiteLLMDB): + """LiteLLMDB subclass pinned to in-memory SQLite for tests. + + The shipped LiteLLMDB pins ``pool_size``/``max_overflow`` which are + NullPool-incompatible on SQLite. We override the engine builder for + the test environment. + """ + + def _build_engine(self, url: str): # type: ignore[override] + return create_async_engine(url, future=True, echo=False) + + +@pytest_asyncio.fixture +async def db() -> AsyncIterator[_SQLiteLiteLLMDB]: + db = _SQLiteLiteLLMDB("sqlite+aiosqlite:///:memory:") + # Create the subset of tables we exercise in this test file. + async with db.writer.begin() as conn: + for cls in (LiteLLMConfig, LiteLLMSpendLogs, LiteLLMTeamMembership): + await conn.run_sync(cls.__table__.create) + yield db + await db.disconnect() + + +@pytest_asyncio.fixture +async def client(db) -> PrismaCompatClient: + return PrismaCompatClient(db) + + +# --------------------------------------------------------------------------- +# Basic CRUD +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_create_and_find_unique(client: PrismaCompatClient): + await client.litellm_config.create( + data={"param_name": "foo", "param_value": {"x": 1}} + ) + row = await client.litellm_config.find_unique(where={"param_name": "foo"}) + assert row is not None + assert row.param_name == "foo" + assert row.param_value == {"x": 1} + + +@pytest.mark.asyncio +async def test_find_unique_returns_none_when_missing(client: PrismaCompatClient): + row = await client.litellm_config.find_unique(where={"param_name": "absent"}) + assert row is None + + +@pytest.mark.asyncio +async def test_update_round_trip(client: PrismaCompatClient): + await client.litellm_config.create( + data={"param_name": "k", "param_value": {"v": 1}} + ) + await client.litellm_config.update( + where={"param_name": "k"}, data={"param_value": {"v": 2}} + ) + row = await client.litellm_config.find_unique(where={"param_name": "k"}) + assert row is not None and row.param_value == {"v": 2} + + +@pytest.mark.asyncio +async def test_update_missing_raises_record_not_found(client: PrismaCompatClient): + with pytest.raises(errors.RecordNotFoundError): + await client.litellm_config.update( + where={"param_name": "absent"}, data={"param_value": {}} + ) + + +@pytest.mark.asyncio +async def test_upsert_creates_then_updates(client: PrismaCompatClient): + await client.litellm_config.upsert( + where={"param_name": "u"}, + data={ + "create": {"param_name": "u", "param_value": {"v": 1}}, + "update": {"param_value": {"v": 2}}, + }, + ) + row = await client.litellm_config.find_unique(where={"param_name": "u"}) + assert row is not None and row.param_value == {"v": 1} + + await client.litellm_config.upsert( + where={"param_name": "u"}, + data={ + "create": {"param_name": "u", "param_value": {"v": 99}}, + "update": {"param_value": {"v": 2}}, + }, + ) + row2 = await client.litellm_config.find_unique(where={"param_name": "u"}) + assert row2 is not None and row2.param_value == {"v": 2} + + +@pytest.mark.asyncio +async def test_delete_returns_row(client: PrismaCompatClient): + await client.litellm_config.create( + data={"param_name": "d", "param_value": {"v": 1}} + ) + deleted = await client.litellm_config.delete(where={"param_name": "d"}) + assert deleted is not None and deleted.param_name == "d" + assert (await client.litellm_config.find_unique(where={"param_name": "d"})) is None + + +# --------------------------------------------------------------------------- +# Bulk + count +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_find_many_with_take_skip_order(client: PrismaCompatClient): + for i in range(5): + await client.litellm_config.create( + data={"param_name": f"x{i}", "param_value": {"i": i}} + ) + rows = await client.litellm_config.find_many( + order={"param_name": "asc"}, take=2, skip=1 + ) + assert [r.param_name for r in rows] == ["x1", "x2"] + + +@pytest.mark.asyncio +async def test_count_with_where(client: PrismaCompatClient): + for name in ("a", "b", "c"): + await client.litellm_config.create(data={"param_name": name, "param_value": {}}) + assert ( + await client.litellm_config.count(where={"param_name": {"in": ["a", "b"]}}) + ) == 2 + + +@pytest.mark.asyncio +async def test_update_many_returns_count(client: PrismaCompatClient): + for name in ("a1", "a2", "a3"): + await client.litellm_config.create( + data={"param_name": name, "param_value": {"v": 0}} + ) + res = await client.litellm_config.update_many( + where={"param_name": {"in": ["a1", "a2"]}}, + data={"param_value": {"v": 9}}, + ) + assert int(res) == 2 + assert res == 2 # CountResult __eq__ with int + + +@pytest.mark.asyncio +async def test_delete_many_returns_count(client: PrismaCompatClient): + for name in ("d1", "d2", "d3"): + await client.litellm_config.create(data={"param_name": name, "param_value": {}}) + res = await client.litellm_config.delete_many( + where={"param_name": {"in": ["d1", "d2"]}} + ) + assert int(res) == 2 + assert (await client.litellm_config.count()) == 1 + + +# --------------------------------------------------------------------------- +# Filter translation +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_or_and_not_filters(client: PrismaCompatClient): + for name in ("alpha", "beta", "gamma"): + await client.litellm_config.create(data={"param_name": name, "param_value": {}}) + rows = await client.litellm_config.find_many( + where={"OR": [{"param_name": "alpha"}, {"param_name": "gamma"}]}, + order={"param_name": "asc"}, + ) + assert [r.param_name for r in rows] == ["alpha", "gamma"] + + rows = await client.litellm_config.find_many( + where={"NOT": {"param_name": "beta"}}, + order={"param_name": "asc"}, + ) + assert [r.param_name for r in rows] == ["alpha", "gamma"] + + +@pytest.mark.asyncio +async def test_contains_insensitive(client: PrismaCompatClient): + for name in ("FooBar", "BarBaz", "Quux"): + await client.litellm_config.create(data={"param_name": name, "param_value": {}}) + rows = await client.litellm_config.find_many( + where={"param_name": {"contains": "bar", "mode": "insensitive"}}, + order={"param_name": "asc"}, + ) + names = [r.param_name for r in rows] + assert "FooBar" in names and "BarBaz" in names + assert "Quux" not in names + + +# --------------------------------------------------------------------------- +# Increment / data translation +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_update_with_increment(client: PrismaCompatClient): + await client.litellm_teammembership.create( + data={"user_id": "u1", "team_id": "t1", "spend": 1.0, "total_spend": 0.0} + ) + await client.litellm_teammembership.update_many( + where={"user_id": "u1", "team_id": "t1"}, + data={"spend": {"increment": 4.0}, "total_spend": {"increment": 1.0}}, + ) + row = await client.litellm_teammembership.find_first( + where={"user_id": "u1", "team_id": "t1"} + ) + assert row is not None + assert row.spend == 5.0 + assert row.total_spend == 1.0 + + +# --------------------------------------------------------------------------- +# Raw SQL +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_query_raw_positional_and_kw(client: PrismaCompatClient): + await client.litellm_config.create( + data={"param_name": "raw1", "param_value": {"v": 1}} + ) + rows = await client.query_raw( + 'SELECT param_name FROM "LiteLLM_Config" WHERE param_name = $1', "raw1" + ) + assert rows and rows[0]["param_name"] == "raw1" + + rows2 = await client.query_raw( + query='SELECT param_name FROM "LiteLLM_Config" WHERE param_name = $1', + *("raw1",), + ) + assert rows2 and rows2[0]["param_name"] == "raw1" + + +@pytest.mark.asyncio +async def test_execute_raw_returns_rowcount(client: PrismaCompatClient): + await client.litellm_config.create(data={"param_name": "ex1", "param_value": {}}) + n = await client.execute_raw( + 'DELETE FROM "LiteLLM_Config" WHERE param_name = $1', "ex1" + ) + assert n == 1 + + +# --------------------------------------------------------------------------- +# batch_ + tx +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_batch_commits_all_ops(client: PrismaCompatClient): + await client.litellm_config.create( + data={"param_name": "b1", "param_value": {"v": 0}} + ) + batcher = client.batch_() + batcher.litellm_config.update( + where={"param_name": "b1"}, data={"param_value": {"v": 1}} + ) + batcher.litellm_config.create(data={"param_name": "b2", "param_value": {"v": 2}}) + await batcher.commit() + + a = await client.litellm_config.find_unique(where={"param_name": "b1"}) + b = await client.litellm_config.find_unique(where={"param_name": "b2"}) + assert a is not None and a.param_value == {"v": 1} + assert b is not None and b.param_value == {"v": 2} + + +@pytest.mark.asyncio +async def test_tx_rolls_back_on_exception(client: PrismaCompatClient): + with pytest.raises(RuntimeError): + async with client.tx() as tx: + await tx.litellm_config.create( + data={"param_name": "tx1", "param_value": {}} + ) + raise RuntimeError("boom") + assert ( + await client.litellm_config.find_unique(where={"param_name": "tx1"}) + ) is None + + +@pytest.mark.asyncio +async def test_tx_with_inner_batch(client: PrismaCompatClient): + async with client.tx() as tx: + batch = tx.batch_() + batch.litellm_config.create( + data={"param_name": "tx_b1", "param_value": {"v": 1}} + ) + batch.litellm_config.create( + data={"param_name": "tx_b2", "param_value": {"v": 2}} + ) + await batch.commit() + rows = await client.litellm_config.find_many( + where={"param_name": {"in": ["tx_b1", "tx_b2"]}} + ) + assert {r.param_name for r in rows} == {"tx_b1", "tx_b2"} 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 00000000000..82eceb4eb77 --- /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 00000000000..b764660aea4 --- /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 cafb6664958..49ac69a2023 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-17T21:39:40.457296356Z" exclude-newer-span = "P3D" [manifest] @@ -164,6 +164,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/fb/76/641ae371508676492379f16e2fa48f4e2c11741bd63c48be4b12a6b09cba/aiosignal-1.4.0-py3-none-any.whl", hash = "sha256:053243f8b92b990551949e63930a839ff0cf0b0ebbe0597b0f3fb19e1a0fe82e", size = 7490, upload-time = "2025-07-03T22:54:42.156Z" }, ] +[[package]] +name = "aiosqlite" +version = "0.22.1" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/4e/8a/64761f4005f17809769d23e518d915db74e6310474e733e3593cfc854ef1/aiosqlite-0.22.1.tar.gz", hash = "sha256:043e0bd78d32888c0a9ca90fc788b38796843360c855a7262a532813133a0650", size = 14821, upload-time = "2025-12-23T19:25:43.997Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/00/b7/e3bf5133d697a08128598c8d0abc5e16377b51465a33756de24fa7dee953/aiosqlite-0.22.1-py3-none-any.whl", hash = "sha256:21c002eb13823fad740196c5a2e9d8e62f6243bd9e7e4a1f87fb5e44ecb4fceb", size = 17405, upload-time = "2025-12-23T19:25:42.139Z" }, +] + [[package]] name = "alabaster" version = "1.0.0" @@ -330,6 +339,49 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/fe/ba/e2081de779ca30d473f21f5b30e0e737c438205440784c7dfc81efc2b029/async_timeout-5.0.1-py3-none-any.whl", hash = "sha256:39e3809566ff85354557ec2398b55e096c8364bacac9405a7a1fa429e77fe76c", size = 6233, upload-time = "2024-11-06T16:41:37.9Z" }, ] +[[package]] +name = "asyncpg" +version = "0.31.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "async-timeout", marker = "python_full_version < '3.11'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/fe/cc/d18065ce2380d80b1bcce927c24a2642efd38918e33fd724bc4bca904877/asyncpg-0.31.0.tar.gz", hash = "sha256:c989386c83940bfbd787180f2b1519415e2d3d6277a70d9d0f0145ac73500735", size = 993667, upload-time = "2025-11-24T23:27:00.812Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/c3/d9/507c80bdac2e95e5a525644af94b03fa7f9a44596a84bd48a6e80f854f92/asyncpg-0.31.0-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:831712dd3cf117eec68575a9b50da711893fd63ebe277fc155ecae1c6c9f0f61", size = 644865, upload-time = "2025-11-24T23:25:23.527Z" }, + { url = "https://files.pythonhosted.org/packages/ea/03/f93b5e543f65c5f504e91405e8d21bb9e600548be95032951a754781a41d/asyncpg-0.31.0-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:0b17c89312c2f4ccea222a3a6571f7df65d4ba2c0e803339bfc7bed46a96d3be", size = 639297, upload-time = "2025-11-24T23:25:25.192Z" }, + { url = "https://files.pythonhosted.org/packages/e5/1e/de2177e57e03a06e697f6c1ddf2a9a7fcfdc236ce69966f54ffc830fd481/asyncpg-0.31.0-cp310-cp310-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:3faa62f997db0c9add34504a68ac2c342cfee4d57a0c3062fcf0d86c7f9cb1e8", size = 2816679, upload-time = "2025-11-24T23:25:26.718Z" }, + { url = "https://files.pythonhosted.org/packages/d0/98/1a853f6870ac7ad48383a948c8ff3c85dc278066a4d69fc9af7d3d4b1106/asyncpg-0.31.0-cp310-cp310-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:8ea599d45c361dfbf398cb67da7fd052affa556a401482d3ff1ee99bd68808a1", size = 2867087, upload-time = "2025-11-24T23:25:28.399Z" }, + { url = "https://files.pythonhosted.org/packages/11/29/7e76f2a51f2360a7c90d2cf6d0d9b210c8bb0ae342edebd16173611a55c2/asyncpg-0.31.0-cp310-cp310-musllinux_1_2_aarch64.whl", hash = "sha256:795416369c3d284e1837461909f58418ad22b305f955e625a4b3a2521d80a5f3", size = 2747631, upload-time = "2025-11-24T23:25:30.154Z" }, + { url = "https://files.pythonhosted.org/packages/5d/3f/716e10cb57c4f388248db46555e9226901688fbfabd0afb85b5e1d65d5a7/asyncpg-0.31.0-cp310-cp310-musllinux_1_2_x86_64.whl", hash = "sha256:a8d758dac9d2e723e173d286ef5e574f0b350ec00e9186fce84d0fc5f6a8e6b8", size = 2855107, upload-time = "2025-11-24T23:25:31.888Z" }, + { url = "https://files.pythonhosted.org/packages/7e/ec/3ebae9dfb23a1bd3f68acfd4f795983b65b413291c0e2b0d982d6ae6c920/asyncpg-0.31.0-cp310-cp310-win32.whl", hash = "sha256:2d076d42eb583601179efa246c5d7ae44614b4144bc1c7a683ad1222814ed095", size = 521990, upload-time = "2025-11-24T23:25:33.402Z" }, + { url = "https://files.pythonhosted.org/packages/20/b4/9fbb4b0af4e36d96a61d026dd37acab3cf521a70290a09640b215da5ab7c/asyncpg-0.31.0-cp310-cp310-win_amd64.whl", hash = "sha256:9ea33213ac044171f4cac23740bed9a3805abae10e7025314cfbd725ec670540", size = 581629, upload-time = "2025-11-24T23:25:34.846Z" }, + { url = "https://files.pythonhosted.org/packages/08/17/cc02bc49bc350623d050fa139e34ea512cd6e020562f2a7312a7bcae4bc9/asyncpg-0.31.0-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:eee690960e8ab85063ba93af2ce128c0f52fd655fdff9fdb1a28df01329f031d", size = 643159, upload-time = "2025-11-24T23:25:36.443Z" }, + { url = "https://files.pythonhosted.org/packages/a4/62/4ded7d400a7b651adf06f49ea8f73100cca07c6df012119594d1e3447aa6/asyncpg-0.31.0-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:2657204552b75f8288de08ca60faf4a99a65deef3a71d1467454123205a88fab", size = 638157, upload-time = "2025-11-24T23:25:37.89Z" }, + { url = "https://files.pythonhosted.org/packages/d6/5b/4179538a9a72166a0bf60ad783b1ef16efb7960e4d7b9afe9f77a5551680/asyncpg-0.31.0-cp311-cp311-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:a429e842a3a4b4ea240ea52d7fe3f82d5149853249306f7ff166cb9948faa46c", size = 2918051, upload-time = "2025-11-24T23:25:39.461Z" }, + { url = "https://files.pythonhosted.org/packages/e6/35/c27719ae0536c5b6e61e4701391ffe435ef59539e9360959240d6e47c8c8/asyncpg-0.31.0-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:c0807be46c32c963ae40d329b3a686356e417f674c976c07fa49f1b30303f109", size = 2972640, upload-time = "2025-11-24T23:25:41.512Z" }, + { url = "https://files.pythonhosted.org/packages/43/f4/01ebb9207f29e645a64699b9ce0eefeff8e7a33494e1d29bb53736f7766b/asyncpg-0.31.0-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:e5d5098f63beeae93512ee513d4c0c53dc12e9aa2b7a1af5a81cddf93fe4e4da", size = 2851050, upload-time = "2025-11-24T23:25:43.153Z" }, + { url = "https://files.pythonhosted.org/packages/3e/f4/03ff1426acc87be0f4e8d40fa2bff5c3952bef0080062af9efc2212e3be8/asyncpg-0.31.0-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:37fc6c00a814e18eef51833545d1891cac9aa69140598bb076b4cd29b3e010b9", size = 2962574, upload-time = "2025-11-24T23:25:44.942Z" }, + { url = "https://files.pythonhosted.org/packages/c7/39/cc788dfca3d4060f9d93e67be396ceec458dfc429e26139059e58c2c244d/asyncpg-0.31.0-cp311-cp311-win32.whl", hash = "sha256:5a4af56edf82a701aece93190cc4e094d2df7d33f6e915c222fb09efbb5afc24", size = 521076, upload-time = "2025-11-24T23:25:46.486Z" }, + { url = "https://files.pythonhosted.org/packages/28/fc/735af5384c029eb7f1ca60ccb8fa95521dbdaeef788edf4cecfc604c3cab/asyncpg-0.31.0-cp311-cp311-win_amd64.whl", hash = "sha256:480c4befbdf079c14c9ca43c8c5e1fe8b6296c96f1f927158d4f1e750aacc047", size = 584980, upload-time = "2025-11-24T23:25:47.938Z" }, + { url = "https://files.pythonhosted.org/packages/2a/a6/59d0a146e61d20e18db7396583242e32e0f120693b67a8de43f1557033e2/asyncpg-0.31.0-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:b44c31e1efc1c15188ef183f287c728e2046abb1d26af4d20858215d50d91fad", size = 662042, upload-time = "2025-11-24T23:25:49.578Z" }, + { url = "https://files.pythonhosted.org/packages/36/01/ffaa189dcb63a2471720615e60185c3f6327716fdc0fc04334436fbb7c65/asyncpg-0.31.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:0c89ccf741c067614c9b5fc7f1fc6f3b61ab05ae4aaa966e6fd6b93097c7d20d", size = 638504, upload-time = "2025-11-24T23:25:51.501Z" }, + { url = "https://files.pythonhosted.org/packages/9f/62/3f699ba45d8bd24c5d65392190d19656d74ff0185f42e19d0bbd973bb371/asyncpg-0.31.0-cp312-cp312-manylinux_2_28_aarch64.whl", hash = "sha256:12b3b2e39dc5470abd5e98c8d3373e4b1d1234d9fbdedf538798b2c13c64460a", size = 3426241, upload-time = "2025-11-24T23:25:53.278Z" }, + { url = "https://files.pythonhosted.org/packages/8c/d1/a867c2150f9c6e7af6462637f613ba67f78a314b00db220cd26ff559d532/asyncpg-0.31.0-cp312-cp312-manylinux_2_28_x86_64.whl", hash = "sha256:aad7a33913fb8bcb5454313377cc330fbb19a0cd5faa7272407d8a0c4257b671", size = 3520321, upload-time = "2025-11-24T23:25:54.982Z" }, + { url = "https://files.pythonhosted.org/packages/7a/1a/cce4c3f246805ecd285a3591222a2611141f1669d002163abef999b60f98/asyncpg-0.31.0-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:3df118d94f46d85b2e434fd62c84cb66d5834d5a890725fe625f498e72e4d5ec", size = 3316685, upload-time = "2025-11-24T23:25:57.43Z" }, + { url = "https://files.pythonhosted.org/packages/40/ae/0fc961179e78cc579e138fad6eb580448ecae64908f95b8cb8ee2f241f67/asyncpg-0.31.0-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:bd5b6efff3c17c3202d4b37189969acf8927438a238c6257f66be3c426beba20", size = 3471858, upload-time = "2025-11-24T23:25:59.636Z" }, + { url = "https://files.pythonhosted.org/packages/52/b2/b20e09670be031afa4cbfabd645caece7f85ec62d69c312239de568e058e/asyncpg-0.31.0-cp312-cp312-win32.whl", hash = "sha256:027eaa61361ec735926566f995d959ade4796f6a49d3bde17e5134b9964f9ba8", size = 527852, upload-time = "2025-11-24T23:26:01.084Z" }, + { url = "https://files.pythonhosted.org/packages/b5/f0/f2ed1de154e15b107dc692262395b3c17fc34eafe2a78fc2115931561730/asyncpg-0.31.0-cp312-cp312-win_amd64.whl", hash = "sha256:72d6bdcbc93d608a1158f17932de2321f68b1a967a13e014998db87a72ed3186", size = 597175, upload-time = "2025-11-24T23:26:02.564Z" }, + { url = "https://files.pythonhosted.org/packages/95/11/97b5c2af72a5d0b9bc3fa30cd4b9ce22284a9a943a150fdc768763caf035/asyncpg-0.31.0-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:c204fab1b91e08b0f47e90a75d1b3c62174dab21f670ad6c5d0f243a228f015b", size = 661111, upload-time = "2025-11-24T23:26:04.467Z" }, + { url = "https://files.pythonhosted.org/packages/1b/71/157d611c791a5e2d0423f09f027bd499935f0906e0c2a416ce712ba51ef3/asyncpg-0.31.0-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:54a64f91839ba59008eccf7aad2e93d6e3de688d796f35803235ea1c4898ae1e", size = 636928, upload-time = "2025-11-24T23:26:05.944Z" }, + { url = "https://files.pythonhosted.org/packages/2e/fc/9e3486fb2bbe69d4a867c0b76d68542650a7ff1574ca40e84c3111bb0c6e/asyncpg-0.31.0-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:c0e0822b1038dc7253b337b0f3f676cadc4ac31b126c5d42691c39691962e403", size = 3424067, upload-time = "2025-11-24T23:26:07.957Z" }, + { url = "https://files.pythonhosted.org/packages/12/c6/8c9d076f73f07f995013c791e018a1cd5f31823c2a3187fc8581706aa00f/asyncpg-0.31.0-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:bef056aa502ee34204c161c72ca1f3c274917596877f825968368b2c33f585f4", size = 3518156, upload-time = "2025-11-24T23:26:09.591Z" }, + { url = "https://files.pythonhosted.org/packages/ae/3b/60683a0baf50fbc546499cfb53132cb6835b92b529a05f6a81471ab60d0c/asyncpg-0.31.0-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:0bfbcc5b7ffcd9b75ab1558f00db2ae07db9c80637ad1b2469c43df79d7a5ae2", size = 3319636, upload-time = "2025-11-24T23:26:11.168Z" }, + { url = "https://files.pythonhosted.org/packages/50/dc/8487df0f69bd398a61e1792b3cba0e47477f214eff085ba0efa7eac9ce87/asyncpg-0.31.0-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:22bc525ebbdc24d1261ecbf6f504998244d4e3be1721784b5f64664d61fbe602", size = 3472079, upload-time = "2025-11-24T23:26:13.164Z" }, + { url = "https://files.pythonhosted.org/packages/13/a1/c5bbeeb8531c05c89135cb8b28575ac2fac618bcb60119ee9696c3faf71c/asyncpg-0.31.0-cp313-cp313-win32.whl", hash = "sha256:f890de5e1e4f7e14023619399a471ce4b71f5418cd67a51853b9910fdfa73696", size = 527606, upload-time = "2025-11-24T23:26:14.78Z" }, + { url = "https://files.pythonhosted.org/packages/91/66/b25ccb84a246b470eb943b0107c07edcae51804912b824054b3413995a10/asyncpg-0.31.0-cp313-cp313-win_amd64.whl", hash = "sha256:dc5f2fa9916f292e5c5c8b2ac2813763bcd7f58e130055b4ad8a0531314201ab", size = 596569, upload-time = "2025-11-24T23:26:16.189Z" }, +] + [[package]] name = "attrs" version = "26.1.0" @@ -3212,13 +3264,15 @@ caching = [ ] extra-proxy = [ { name = "a2a-sdk" }, + { name = "alembic" }, + { name = "asyncpg" }, { name = "azure-identity" }, { name = "azure-keyvault-secrets" }, { name = "google-cloud-iam" }, { name = "google-cloud-kms" }, - { name = "prisma" }, { name = "redisvl" }, { name = "resend" }, + { name = "sqlmodel" }, ] google = [ { name = "google-cloud-aiplatform" }, @@ -3258,6 +3312,9 @@ proxy = [ { name = "uvloop", marker = "sys_platform != 'win32'" }, { name = "websockets" }, ] +proxy-dev = [ + { name = "aiosqlite" }, +] proxy-runtime = [ { name = "anthropic", extra = ["vertex"] }, { name = "azure-ai-contentsafety" }, @@ -3368,7 +3425,6 @@ proxy-dev = [ { name = "opentelemetry-api" }, { name = "opentelemetry-exporter-otlp" }, { name = "opentelemetry-sdk" }, - { name = "prisma" }, { name = "prometheus-client" }, ] @@ -3376,8 +3432,11 @@ proxy-dev = [ requires-dist = [ { name = "a2a-sdk", marker = "extra == 'extra-proxy'", specifier = "==0.3.24" }, { name = "aiohttp", specifier = ">=3.10,<4.0" }, + { name = "aiosqlite", marker = "extra == 'proxy-dev'", specifier = ">=0.19,<1.0" }, + { name = "alembic", marker = "extra == 'extra-proxy'", specifier = ">=1.13,<2.0" }, { name = "anthropic", extras = ["vertex"], marker = "extra == 'proxy-runtime'", specifier = "==0.84.0" }, { name = "apscheduler", marker = "extra == 'proxy'", specifier = "==3.11.2" }, + { name = "asyncpg", marker = "extra == 'extra-proxy'", specifier = ">=0.29,<1.0" }, { name = "audioread", marker = "extra == 'stt-nvidia-riva'", specifier = ">=3.0.1" }, { name = "aurelio-sdk", marker = "python_full_version < '3.14' and extra == 'semantic-router'", specifier = "==0.0.19" }, { name = "azure-ai-contentsafety", marker = "extra == 'proxy-runtime'", specifier = "==1.0.0" }, @@ -3424,7 +3483,6 @@ requires-dist = [ { name = "opentelemetry-sdk", marker = "extra == 'proxy-runtime'", specifier = "==1.28.0" }, { name = "orjson", marker = "extra == 'proxy'", specifier = "==3.11.6" }, { name = "polars", marker = "extra == 'proxy'", specifier = "==1.38.1" }, - { name = "prisma", marker = "extra == 'extra-proxy'", specifier = "==0.11.0" }, { name = "prometheus-client", marker = "extra == 'proxy-runtime'", specifier = "==0.20.0" }, { name = "pydantic", specifier = ">=2.10.0,<3.0.0" }, { name = "pydantic-settings", marker = "extra == 'proxy'", specifier = ">=2.14.1" }, @@ -3444,13 +3502,14 @@ 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" }, { name = "uvloop", marker = "sys_platform != 'win32' and extra == 'proxy'", specifier = "==0.21.0" }, { name = "websockets", marker = "extra == 'proxy'", specifier = "==15.0.1" }, ] -provides-extras = ["proxy", "extra-proxy", "utils", "caching", "semantic-router", "mlflow", "grpc", "stt-nvidia-riva", "google", "proxy-runtime"] +provides-extras = ["proxy", "extra-proxy", "utils", "caching", "semantic-router", "mlflow", "grpc", "stt-nvidia-riva", "google", "proxy-runtime", "proxy-dev"] [package.metadata.requires-dev] ci = [ @@ -3528,7 +3587,6 @@ proxy-dev = [ { name = "opentelemetry-api", specifier = "==1.28.0" }, { name = "opentelemetry-exporter-otlp", specifier = "==1.28.0" }, { name = "opentelemetry-sdk", specifier = "==1.28.0" }, - { name = "prisma", specifier = "==0.11.0" }, { name = "prometheus-client", specifier = "==0.20.0" }, ] @@ -5321,25 +5379,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/5e/5f/82c8074f7e84978129347c2c6ec8b6c59f3584ff1a20bc3c940a3e061790/priority-2.0.0-py3-none-any.whl", hash = "sha256:6f8eefce5f3ad59baf2c080a664037bb4725cd0a790d53d59ab4059288faf6aa", size = 8946, upload-time = "2021-06-27T10:15:03.856Z" }, ] -[[package]] -name = "prisma" -version = "0.11.0" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "click" }, - { name = "httpx" }, - { name = "jinja2" }, - { name = "nodeenv" }, - { name = "pydantic" }, - { name = "python-dotenv" }, - { name = "tomlkit" }, - { name = "typing-extensions" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/2b/62/23a2638aa14d8eefd16bc4b230a6e65d8de5f69bfaeeaf954654c923f02e/prisma-0.11.0.tar.gz", hash = "sha256:3f2f2fd2361e1ec5ff655f2a04c7860c2f2a5bc4c91f78ca9c5c6349735bf693", size = 94181, upload-time = "2023-10-22T22:59:55.817Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/85/a1/a8734bb5668bb47eb777ac90176ad8135580d4f676c67416d468e03d3eaa/prisma-0.11.0-py3-none-any.whl", hash = "sha256:22bb869e59a2968b99f3483bb417717273ffbc569fd1e9ceed95e5614cbaf53a", size = 109302, upload-time = "2023-10-22T22:59:53.897Z" }, -] - [[package]] name = "prometheus-client" version = "0.20.0" @@ -7183,6 +7222,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"