mirror of
https://github.com/MODSetter/SurfSense.git
synced 2026-06-20 21:18:13 +02:00
148 lines
4.3 KiB
Python
148 lines
4.3 KiB
Python
import uuid
|
|
from datetime import datetime
|
|
from typing import Any
|
|
|
|
from pydantic import BaseModel, ConfigDict, Field
|
|
|
|
from app.db import ConnectionScope, ModelSource
|
|
|
|
|
|
class ModelRead(BaseModel):
|
|
id: int
|
|
connection_id: int
|
|
model_id: str
|
|
display_name: str | None = None
|
|
source: ModelSource | str
|
|
supports_chat: bool | None = None
|
|
max_input_tokens: int | None = None
|
|
supports_image_input: bool | None = None
|
|
supports_tools: bool | None = None
|
|
supports_image_generation: bool | None = None
|
|
capabilities_override: dict[str, Any] = Field(default_factory=dict)
|
|
enabled: bool
|
|
billing_tier: str | None = None
|
|
catalog: dict[str, Any] = Field(default_factory=dict)
|
|
created_at: datetime | None = None
|
|
|
|
model_config = ConfigDict(from_attributes=True)
|
|
|
|
|
|
class ConnectionRead(BaseModel):
|
|
id: int
|
|
provider: str
|
|
base_url: str | None = None
|
|
api_key: str | None = None
|
|
extra: dict[str, Any] = Field(default_factory=dict)
|
|
scope: ConnectionScope | str
|
|
search_space_id: int | None = None
|
|
user_id: uuid.UUID | None = None
|
|
enabled: bool
|
|
has_api_key: bool
|
|
models: list[ModelRead] = Field(default_factory=list)
|
|
created_at: datetime | None = None
|
|
|
|
model_config = ConfigDict(from_attributes=True)
|
|
|
|
|
|
class ModelSelection(BaseModel):
|
|
model_id: str = Field(..., max_length=255)
|
|
display_name: str | None = Field(None, max_length=255)
|
|
source: ModelSource | str = ModelSource.DISCOVERED
|
|
supports_chat: bool | None = None
|
|
max_input_tokens: int | None = None
|
|
supports_image_input: bool | None = None
|
|
supports_tools: bool | None = None
|
|
supports_image_generation: bool | None = None
|
|
enabled: bool = False
|
|
metadata: dict[str, Any] = Field(default_factory=dict)
|
|
|
|
|
|
class ModelPreviewRead(BaseModel):
|
|
model_id: str
|
|
display_name: str | None = None
|
|
source: ModelSource | str = ModelSource.DISCOVERED
|
|
supports_chat: bool | None = None
|
|
max_input_tokens: int | None = None
|
|
supports_image_input: bool | None = None
|
|
supports_tools: bool | None = None
|
|
supports_image_generation: bool | None = None
|
|
enabled: bool = False
|
|
metadata: dict[str, Any] = Field(default_factory=dict)
|
|
|
|
|
|
class ConnectionCreate(BaseModel):
|
|
provider: str = Field(..., max_length=100)
|
|
base_url: str | None = Field(None, max_length=500)
|
|
api_key: str | None = None
|
|
extra: dict[str, Any] = Field(default_factory=dict)
|
|
scope: ConnectionScope = ConnectionScope.SEARCH_SPACE
|
|
search_space_id: int | None = None
|
|
enabled: bool = True
|
|
models: list[ModelSelection] = Field(default_factory=list)
|
|
|
|
|
|
class ModelTestPreview(ConnectionCreate):
|
|
model_id: str = Field(..., max_length=255)
|
|
|
|
|
|
class ConnectionUpdate(BaseModel):
|
|
provider: str | None = Field(None, max_length=100)
|
|
base_url: str | None = Field(None, max_length=500)
|
|
api_key: str | None = None
|
|
extra: dict[str, Any] | None = None
|
|
enabled: bool | None = None
|
|
|
|
|
|
class ModelCreate(BaseModel):
|
|
"""Manually register a model id on a connection.
|
|
|
|
For providers without a usable ``/models`` endpoint (Perplexity, MiniMax,
|
|
Azure deployments, etc.) or to pin a single model from a noisy provider.
|
|
"""
|
|
|
|
model_id: str = Field(..., max_length=255)
|
|
display_name: str | None = Field(None, max_length=255)
|
|
|
|
|
|
class ModelUpdate(BaseModel):
|
|
display_name: str | None = Field(None, max_length=255)
|
|
enabled: bool | None = None
|
|
supports_chat: bool | None = None
|
|
max_input_tokens: int | None = None
|
|
supports_image_input: bool | None = None
|
|
supports_tools: bool | None = None
|
|
supports_image_generation: bool | None = None
|
|
capabilities_override: dict[str, Any] | None = None
|
|
|
|
|
|
class ModelsBulkUpdate(BaseModel):
|
|
model_ids: list[int] = Field(..., min_length=1, max_length=1000)
|
|
enabled: bool
|
|
|
|
|
|
class ModelProviderRead(BaseModel):
|
|
provider: str
|
|
transport: str
|
|
discovery: str
|
|
default_base_url: str | None = None
|
|
base_url_required: bool
|
|
auth_style: str
|
|
local_only: bool = False
|
|
|
|
|
|
class VerifyConnectionResponse(BaseModel):
|
|
status: str
|
|
ok: bool
|
|
message: str = ""
|
|
|
|
|
|
class ModelRolesRead(BaseModel):
|
|
chat_model_id: int | None = 0
|
|
vision_model_id: int | None = 0
|
|
image_gen_model_id: int | None = 0
|
|
|
|
|
|
class ModelRolesUpdate(BaseModel):
|
|
chat_model_id: int | None = None
|
|
vision_model_id: int | None = None
|
|
image_gen_model_id: int | None = None
|