mirror of
https://github.com/MODSetter/SurfSense.git
synced 2026-06-30 21:59:46 +02:00
feat(llm): expand LLM provider options and improve model selection UI
- Added new LLM providers including Google, Azure OpenAI, Bedrock, and others to the backend. - Updated the model selection UI to dynamically display available models based on the selected provider. - Enhanced the provider change handling to reset the model selection when the provider is changed. - Improved the overall user experience by providing contextual information for model selection.
This commit is contained in:
DESKTOP-RTLN3BA\$punk
parent
c1e8753567
commit
38dffaffa3
10 changed files with 2197 additions and 126 deletions
|
|
@ -0,0 +1,199 @@
|
|||
"""Update LiteLLMProvider enum with comprehensive provider support
|
||||
|
||||
This migration adds support for all major LiteLLM providers including:
|
||||
- Fast inference platforms (XAI, Cerebras, SambaNova, Fireworks AI)
|
||||
- Cloud platforms (Cloudflare, Databricks)
|
||||
- Renames AWS_BEDROCK to BEDROCK for consistency
|
||||
- Adds CUSTOM for custom OpenAI-compatible endpoints
|
||||
|
||||
Revision ID: 35
|
||||
Revises: 34
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
from alembic import op
|
||||
|
||||
# revision identifiers, used by Alembic.
|
||||
revision: str = "35"
|
||||
down_revision: str | None = "34"
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
"""
|
||||
Add new LLM providers to LiteLLMProvider enum and migrate existing data.
|
||||
|
||||
New providers added:
|
||||
- XAI: xAI's Grok models
|
||||
- FIREWORKS_AI: Fireworks AI platform
|
||||
- CEREBRAS: Cerebras inference platform (fastest)
|
||||
- SAMBANOVA: SambaNova inference platform
|
||||
- CLOUDFLARE: Cloudflare Workers AI
|
||||
- DATABRICKS: Databricks Model Serving
|
||||
- BEDROCK: AWS Bedrock (replaces AWS_BEDROCK)
|
||||
- CUSTOM: Custom OpenAI-compatible endpoints
|
||||
"""
|
||||
|
||||
# Add XAI (xAI Grok models)
|
||||
op.execute(
|
||||
"""
|
||||
DO $$
|
||||
BEGIN
|
||||
IF NOT EXISTS (
|
||||
SELECT 1 FROM pg_enum
|
||||
WHERE enumtypid = 'litellmprovider'::regtype
|
||||
AND enumlabel = 'XAI'
|
||||
) THEN
|
||||
ALTER TYPE litellmprovider ADD VALUE 'XAI';
|
||||
END IF;
|
||||
END$$;
|
||||
"""
|
||||
)
|
||||
|
||||
# Add FIREWORKS_AI
|
||||
op.execute(
|
||||
"""
|
||||
DO $$
|
||||
BEGIN
|
||||
IF NOT EXISTS (
|
||||
SELECT 1 FROM pg_enum
|
||||
WHERE enumtypid = 'litellmprovider'::regtype
|
||||
AND enumlabel = 'FIREWORKS_AI'
|
||||
) THEN
|
||||
ALTER TYPE litellmprovider ADD VALUE 'FIREWORKS_AI';
|
||||
END IF;
|
||||
END$$;
|
||||
"""
|
||||
)
|
||||
|
||||
# Add CEREBRAS (fastest inference platform)
|
||||
op.execute(
|
||||
"""
|
||||
DO $$
|
||||
BEGIN
|
||||
IF NOT EXISTS (
|
||||
SELECT 1 FROM pg_enum
|
||||
WHERE enumtypid = 'litellmprovider'::regtype
|
||||
AND enumlabel = 'CEREBRAS'
|
||||
) THEN
|
||||
ALTER TYPE litellmprovider ADD VALUE 'CEREBRAS';
|
||||
END IF;
|
||||
END$$;
|
||||
"""
|
||||
)
|
||||
|
||||
# Add SAMBANOVA
|
||||
op.execute(
|
||||
"""
|
||||
DO $$
|
||||
BEGIN
|
||||
IF NOT EXISTS (
|
||||
SELECT 1 FROM pg_enum
|
||||
WHERE enumtypid = 'litellmprovider'::regtype
|
||||
AND enumlabel = 'SAMBANOVA'
|
||||
) THEN
|
||||
ALTER TYPE litellmprovider ADD VALUE 'SAMBANOVA';
|
||||
END IF;
|
||||
END$$;
|
||||
"""
|
||||
)
|
||||
|
||||
# Add CLOUDFLARE (Cloudflare Workers AI)
|
||||
op.execute(
|
||||
"""
|
||||
DO $$
|
||||
BEGIN
|
||||
IF NOT EXISTS (
|
||||
SELECT 1 FROM pg_enum
|
||||
WHERE enumtypid = 'litellmprovider'::regtype
|
||||
AND enumlabel = 'CLOUDFLARE'
|
||||
) THEN
|
||||
ALTER TYPE litellmprovider ADD VALUE 'CLOUDFLARE';
|
||||
END IF;
|
||||
END$$;
|
||||
"""
|
||||
)
|
||||
|
||||
# Add DATABRICKS
|
||||
op.execute(
|
||||
"""
|
||||
DO $$
|
||||
BEGIN
|
||||
IF NOT EXISTS (
|
||||
SELECT 1 FROM pg_enum
|
||||
WHERE enumtypid = 'litellmprovider'::regtype
|
||||
AND enumlabel = 'DATABRICKS'
|
||||
) THEN
|
||||
ALTER TYPE litellmprovider ADD VALUE 'DATABRICKS';
|
||||
END IF;
|
||||
END$$;
|
||||
"""
|
||||
)
|
||||
|
||||
# Add BEDROCK (new standardized name)
|
||||
op.execute(
|
||||
"""
|
||||
DO $$
|
||||
BEGIN
|
||||
IF NOT EXISTS (
|
||||
SELECT 1 FROM pg_enum
|
||||
WHERE enumtypid = 'litellmprovider'::regtype
|
||||
AND enumlabel = 'BEDROCK'
|
||||
) THEN
|
||||
ALTER TYPE litellmprovider ADD VALUE 'BEDROCK';
|
||||
END IF;
|
||||
END$$;
|
||||
"""
|
||||
)
|
||||
|
||||
# Add CUSTOM for custom OpenAI-compatible endpoints
|
||||
op.execute(
|
||||
"""
|
||||
DO $$
|
||||
BEGIN
|
||||
IF NOT EXISTS (
|
||||
SELECT 1 FROM pg_enum
|
||||
WHERE enumtypid = 'litellmprovider'::regtype
|
||||
AND enumlabel = 'CUSTOM'
|
||||
) THEN
|
||||
ALTER TYPE litellmprovider ADD VALUE 'CUSTOM';
|
||||
END IF;
|
||||
END$$;
|
||||
"""
|
||||
)
|
||||
|
||||
# Note: Both AWS_BEDROCK and BEDROCK will coexist in the enum for backward compatibility.
|
||||
# - New configurations will use BEDROCK (via the frontend)
|
||||
# - Existing AWS_BEDROCK configurations will continue to work
|
||||
# - The backend service handles both values correctly
|
||||
#
|
||||
# Legacy enum values (PALM, NLPCLOUD, ALEPH_ALPHA, PETALS, etc.) also remain in the enum.
|
||||
# PostgreSQL doesn't support removing enum values without recreating the entire type.
|
||||
#
|
||||
# Data migration from AWS_BEDROCK -> BEDROCK is NOT performed because:
|
||||
# 1. PostgreSQL doesn't allow using new enum values in the same transaction
|
||||
# 2. Both values work correctly with the backend service
|
||||
# 3. Users can manually update old configs if desired (not required)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
"""
|
||||
Downgrade migration.
|
||||
|
||||
Note: PostgreSQL doesn't support removing enum values directly.
|
||||
This would require recreating the entire enum type and updating all dependent objects.
|
||||
For safety and data preservation, this downgrade is a no-op.
|
||||
|
||||
If you need to downgrade, you should:
|
||||
1. Ensure no llm_configs are using the new providers
|
||||
2. Manually remove enum values (requires enum recreation - complex operation)
|
||||
|
||||
This is not automated to prevent accidental data loss.
|
||||
New enum values (XAI, FIREWORKS_AI, CEREBRAS, BEDROCK, etc.) will remain
|
||||
in the database but won't be selectable in the application after downgrade.
|
||||
"""
|
||||
# PostgreSQL doesn't support removing enum values directly
|
||||
# This is a no-op for safety
|
||||
pass
|
||||
|
|
@ -86,32 +86,33 @@ class LiteLLMProvider(str, Enum):
|
|||
|
||||
OPENAI = "OPENAI"
|
||||
ANTHROPIC = "ANTHROPIC"
|
||||
GOOGLE = "GOOGLE"
|
||||
AZURE_OPENAI = "AZURE_OPENAI"
|
||||
BEDROCK = "BEDROCK"
|
||||
VERTEX_AI = "VERTEX_AI"
|
||||
GROQ = "GROQ"
|
||||
COHERE = "COHERE"
|
||||
HUGGINGFACE = "HUGGINGFACE"
|
||||
AZURE_OPENAI = "AZURE_OPENAI"
|
||||
GOOGLE = "GOOGLE"
|
||||
AWS_BEDROCK = "AWS_BEDROCK"
|
||||
OLLAMA = "OLLAMA"
|
||||
MISTRAL = "MISTRAL"
|
||||
TOGETHER_AI = "TOGETHER_AI"
|
||||
OPENROUTER = "OPENROUTER"
|
||||
REPLICATE = "REPLICATE"
|
||||
PALM = "PALM"
|
||||
VERTEX_AI = "VERTEX_AI"
|
||||
ANYSCALE = "ANYSCALE"
|
||||
PERPLEXITY = "PERPLEXITY"
|
||||
DEEPINFRA = "DEEPINFRA"
|
||||
AI21 = "AI21"
|
||||
NLPCLOUD = "NLPCLOUD"
|
||||
ALEPH_ALPHA = "ALEPH_ALPHA"
|
||||
PETALS = "PETALS"
|
||||
COMETAPI = "COMETAPI"
|
||||
# Chinese LLM Providers (OpenAI-compatible)
|
||||
DEEPSEEK = "DEEPSEEK"
|
||||
XAI = "XAI"
|
||||
OPENROUTER = "OPENROUTER"
|
||||
TOGETHER_AI = "TOGETHER_AI"
|
||||
FIREWORKS_AI = "FIREWORKS_AI"
|
||||
REPLICATE = "REPLICATE"
|
||||
PERPLEXITY = "PERPLEXITY"
|
||||
OLLAMA = "OLLAMA"
|
||||
ALIBABA_QWEN = "ALIBABA_QWEN"
|
||||
MOONSHOT = "MOONSHOT"
|
||||
ZHIPU = "ZHIPU"
|
||||
ANYSCALE = "ANYSCALE"
|
||||
DEEPINFRA = "DEEPINFRA"
|
||||
CEREBRAS = "CEREBRAS"
|
||||
SAMBANOVA = "SAMBANOVA"
|
||||
AI21 = "AI21"
|
||||
CLOUDFLARE = "CLOUDFLARE"
|
||||
DATABRICKS = "DATABRICKS"
|
||||
COMETAPI = "COMETAPI"
|
||||
HUGGINGFACE = "HUGGINGFACE"
|
||||
CUSTOM = "CUSTOM"
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -61,11 +61,26 @@ async def validate_llm_config(
|
|||
"AZURE_OPENAI": "azure",
|
||||
"OPENROUTER": "openrouter",
|
||||
"COMETAPI": "cometapi",
|
||||
# Chinese LLM providers (OpenAI-compatible)
|
||||
"XAI": "xai",
|
||||
"BEDROCK": "bedrock",
|
||||
"AWS_BEDROCK": "bedrock", # Legacy support (backward compatibility)
|
||||
"VERTEX_AI": "vertex_ai",
|
||||
"TOGETHER_AI": "together_ai",
|
||||
"FIREWORKS_AI": "fireworks_ai",
|
||||
"REPLICATE": "replicate",
|
||||
"PERPLEXITY": "perplexity",
|
||||
"ANYSCALE": "anyscale",
|
||||
"DEEPINFRA": "deepinfra",
|
||||
"CEREBRAS": "cerebras",
|
||||
"SAMBANOVA": "sambanova",
|
||||
"AI21": "ai21",
|
||||
"CLOUDFLARE": "cloudflare",
|
||||
"DATABRICKS": "databricks",
|
||||
# Chinese LLM providers
|
||||
"DEEPSEEK": "openai",
|
||||
"ALIBABA_QWEN": "openai",
|
||||
"MOONSHOT": "openai",
|
||||
"ZHIPU": "openai",
|
||||
"ZHIPU": "openai", # GLM needs special handling
|
||||
}
|
||||
provider_prefix = provider_map.get(provider, provider.lower())
|
||||
model_string = f"{provider_prefix}/{model_name}"
|
||||
|
|
@ -187,12 +202,26 @@ async def get_user_llm_instance(
|
|||
"AZURE_OPENAI": "azure",
|
||||
"OPENROUTER": "openrouter",
|
||||
"COMETAPI": "cometapi",
|
||||
# Chinese LLM providers (OpenAI-compatible)
|
||||
"DEEPSEEK": "openai", # DeepSeek uses OpenAI-compatible API
|
||||
"ALIBABA_QWEN": "openai", # Qwen uses OpenAI-compatible API
|
||||
"MOONSHOT": "openai", # Moonshot (Kimi) uses OpenAI-compatible API
|
||||
"ZHIPU": "openai", # Zhipu (GLM) uses OpenAI-compatible API
|
||||
# Add more mappings as needed
|
||||
"XAI": "xai",
|
||||
"BEDROCK": "bedrock",
|
||||
"AWS_BEDROCK": "bedrock", # Legacy support (backward compatibility)
|
||||
"VERTEX_AI": "vertex_ai",
|
||||
"TOGETHER_AI": "together_ai",
|
||||
"FIREWORKS_AI": "fireworks_ai",
|
||||
"REPLICATE": "replicate",
|
||||
"PERPLEXITY": "perplexity",
|
||||
"ANYSCALE": "anyscale",
|
||||
"DEEPINFRA": "deepinfra",
|
||||
"CEREBRAS": "cerebras",
|
||||
"SAMBANOVA": "sambanova",
|
||||
"AI21": "ai21",
|
||||
"CLOUDFLARE": "cloudflare",
|
||||
"DATABRICKS": "databricks",
|
||||
# Chinese LLM providers
|
||||
"DEEPSEEK": "openai",
|
||||
"ALIBABA_QWEN": "openai",
|
||||
"MOONSHOT": "openai",
|
||||
"ZHIPU": "openai",
|
||||
}
|
||||
provider_prefix = provider_map.get(
|
||||
llm_config.provider.value, llm_config.provider.value.lower()
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue