982501f6fa
--------- Co-authored-by: Ari Webb <arijwebb@gmail.com> Co-authored-by: Ari Webb <ari@letta.com>
76 lines
3.3 KiB
Python
76 lines
3.3 KiB
Python
from typing import TYPE_CHECKING, Optional
|
|
|
|
from sqlalchemy import Boolean, ForeignKey, String, UniqueConstraint
|
|
from sqlalchemy.orm import Mapped, mapped_column, relationship
|
|
|
|
from letta.orm.sqlalchemy_base import SqlalchemyBase
|
|
from letta.schemas.provider_model import ProviderModel as PydanticProviderModel
|
|
|
|
if TYPE_CHECKING:
|
|
from letta.orm.organization import Organization
|
|
from letta.orm.provider import Provider
|
|
|
|
|
|
class ProviderModel(SqlalchemyBase):
|
|
"""ProviderModel ORM class - represents individual models available from providers"""
|
|
|
|
__tablename__ = "provider_models"
|
|
__pydantic_model__ = PydanticProviderModel
|
|
__table_args__ = (
|
|
UniqueConstraint(
|
|
"handle",
|
|
"organization_id",
|
|
"model_type",
|
|
name="unique_handle_per_org_and_type",
|
|
),
|
|
UniqueConstraint(
|
|
"name",
|
|
"provider_id",
|
|
"model_type",
|
|
name="unique_model_per_provider_and_type",
|
|
),
|
|
)
|
|
|
|
# The unique handle used in the API (e.g., "openai/gpt-4o-mini", "anthropic/claude-3-5-sonnet")
|
|
# Format: {provider_name}/{display_name}
|
|
handle: Mapped[str] = mapped_column(String, nullable=False, index=True, doc="Unique handle for API reference")
|
|
|
|
# Display name shown in the UI for the model
|
|
display_name: Mapped[str] = mapped_column(String, nullable=False, doc="Display name for the model")
|
|
|
|
# The actual model name used by the provider (e.g., "gpt-4o-mini", "openai/gpt-4" for OpenRouter)
|
|
name: Mapped[str] = mapped_column(String, nullable=False, doc="The actual model name used by the provider")
|
|
|
|
# Foreign key to the provider
|
|
provider_id: Mapped[str] = mapped_column(
|
|
String, ForeignKey("providers.id", ondelete="CASCADE"), nullable=False, index=True, doc="Provider ID reference"
|
|
)
|
|
|
|
# Optional organization ID - NULL for global models, set for org-scoped models
|
|
organization_id: Mapped[Optional[str]] = mapped_column(
|
|
String,
|
|
ForeignKey("organizations.id", ondelete="CASCADE"),
|
|
nullable=True,
|
|
index=True,
|
|
doc="Organization ID if org-scoped, NULL if global",
|
|
)
|
|
|
|
# Model type: llm or embedding
|
|
model_type: Mapped[str] = mapped_column(String, nullable=False, index=True, doc="Type of model (llm or embedding)")
|
|
|
|
# Whether the model is enabled (default True)
|
|
enabled: Mapped[bool] = mapped_column(Boolean, nullable=False, default=True, server_default="TRUE", doc="Whether the model is enabled")
|
|
|
|
# Model endpoint type (e.g., "openai", "anthropic", etc.)
|
|
model_endpoint_type: Mapped[str] = mapped_column(String, nullable=False, doc="The endpoint type for the model")
|
|
|
|
# Additional metadata fields
|
|
max_context_window: Mapped[int] = mapped_column(nullable=True, doc="Context window size for the model")
|
|
supports_token_streaming: Mapped[bool] = mapped_column(Boolean, nullable=True, doc="Whether streaming is supported")
|
|
supports_tool_calling: Mapped[bool] = mapped_column(Boolean, nullable=True, doc="Whether tool calling is supported")
|
|
embedding_dim: Mapped[Optional[int]] = mapped_column(nullable=True, doc="Embedding dimension for embedding models")
|
|
|
|
# relationships
|
|
provider: Mapped["Provider"] = relationship("Provider", back_populates="models")
|
|
organization: Mapped[Optional["Organization"]] = relationship("Organization", back_populates="provider_models")
|