mirror of
https://github.com/MODSetter/SurfSense.git
synced 2026-05-21 18:55:16 +02:00
feat: migrated to surfsense deep agent
This commit is contained in:
DESKTOP-RTLN3BA\$punk
parent
b14283e300
commit
4a0c3e368a
90 changed files with 5337 additions and 6029 deletions
|
|
@ -1,14 +1,144 @@
|
|||
"""
|
||||
LLM configuration utilities for SurfSense agents.
|
||||
|
||||
This module provides functions for loading LLM configurations from YAML files
|
||||
and creating ChatLiteLLM instances from configuration dictionaries.
|
||||
This module provides functions for loading LLM configurations from:
|
||||
1. YAML files (global configs with negative IDs)
|
||||
2. Database NewLLMConfig table (user-created configs with positive IDs)
|
||||
|
||||
It also provides utilities for creating ChatLiteLLM instances and
|
||||
managing prompt configurations.
|
||||
"""
|
||||
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
|
||||
import yaml
|
||||
from langchain_litellm import ChatLiteLLM
|
||||
from sqlalchemy import select
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
|
||||
# Provider mapping for LiteLLM model string construction
|
||||
PROVIDER_MAP = {
|
||||
"OPENAI": "openai",
|
||||
"ANTHROPIC": "anthropic",
|
||||
"GROQ": "groq",
|
||||
"COHERE": "cohere",
|
||||
"GOOGLE": "gemini",
|
||||
"OLLAMA": "ollama",
|
||||
"MISTRAL": "mistral",
|
||||
"AZURE_OPENAI": "azure",
|
||||
"OPENROUTER": "openrouter",
|
||||
"XAI": "xai",
|
||||
"BEDROCK": "bedrock",
|
||||
"VERTEX_AI": "vertex_ai",
|
||||
"TOGETHER_AI": "together_ai",
|
||||
"FIREWORKS_AI": "fireworks_ai",
|
||||
"DEEPSEEK": "openai",
|
||||
"ALIBABA_QWEN": "openai",
|
||||
"MOONSHOT": "openai",
|
||||
"ZHIPU": "openai",
|
||||
"REPLICATE": "replicate",
|
||||
"PERPLEXITY": "perplexity",
|
||||
"ANYSCALE": "anyscale",
|
||||
"DEEPINFRA": "deepinfra",
|
||||
"CEREBRAS": "cerebras",
|
||||
"SAMBANOVA": "sambanova",
|
||||
"AI21": "ai21",
|
||||
"CLOUDFLARE": "cloudflare",
|
||||
"DATABRICKS": "databricks",
|
||||
"COMETAPI": "cometapi",
|
||||
"HUGGINGFACE": "huggingface",
|
||||
"CUSTOM": "custom",
|
||||
}
|
||||
|
||||
|
||||
@dataclass
|
||||
class AgentConfig:
|
||||
"""
|
||||
Complete configuration for the SurfSense agent.
|
||||
|
||||
This combines LLM settings with prompt configuration from NewLLMConfig.
|
||||
"""
|
||||
|
||||
# LLM Model Settings
|
||||
provider: str
|
||||
model_name: str
|
||||
api_key: str
|
||||
api_base: str | None = None
|
||||
custom_provider: str | None = None
|
||||
litellm_params: dict | None = None
|
||||
|
||||
# Prompt Configuration
|
||||
system_instructions: str | None = None
|
||||
use_default_system_instructions: bool = True
|
||||
citations_enabled: bool = True
|
||||
|
||||
# Metadata
|
||||
config_id: int | None = None
|
||||
config_name: str | None = None
|
||||
|
||||
@classmethod
|
||||
def from_new_llm_config(cls, config) -> "AgentConfig":
|
||||
"""
|
||||
Create an AgentConfig from a NewLLMConfig database model.
|
||||
|
||||
Args:
|
||||
config: NewLLMConfig database model instance
|
||||
|
||||
Returns:
|
||||
AgentConfig instance
|
||||
"""
|
||||
return cls(
|
||||
provider=config.provider.value
|
||||
if hasattr(config.provider, "value")
|
||||
else str(config.provider),
|
||||
model_name=config.model_name,
|
||||
api_key=config.api_key,
|
||||
api_base=config.api_base,
|
||||
custom_provider=config.custom_provider,
|
||||
litellm_params=config.litellm_params,
|
||||
system_instructions=config.system_instructions,
|
||||
use_default_system_instructions=config.use_default_system_instructions,
|
||||
citations_enabled=config.citations_enabled,
|
||||
config_id=config.id,
|
||||
config_name=config.name,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def from_yaml_config(cls, yaml_config: dict) -> "AgentConfig":
|
||||
"""
|
||||
Create an AgentConfig from a YAML configuration dictionary.
|
||||
|
||||
YAML configs now support the same prompt configuration fields as NewLLMConfig:
|
||||
- system_instructions: Custom system instructions (empty string uses defaults)
|
||||
- use_default_system_instructions: Whether to use default instructions
|
||||
- citations_enabled: Whether citations are enabled
|
||||
|
||||
Args:
|
||||
yaml_config: Configuration dictionary from YAML file
|
||||
|
||||
Returns:
|
||||
AgentConfig instance
|
||||
"""
|
||||
# Get system instructions from YAML, default to empty string
|
||||
system_instructions = yaml_config.get("system_instructions", "")
|
||||
|
||||
return cls(
|
||||
provider=yaml_config.get("provider", "").upper(),
|
||||
model_name=yaml_config.get("model_name", ""),
|
||||
api_key=yaml_config.get("api_key", ""),
|
||||
api_base=yaml_config.get("api_base"),
|
||||
custom_provider=yaml_config.get("custom_provider"),
|
||||
litellm_params=yaml_config.get("litellm_params"),
|
||||
# Prompt configuration from YAML (with defaults for backwards compatibility)
|
||||
system_instructions=system_instructions if system_instructions else None,
|
||||
use_default_system_instructions=yaml_config.get(
|
||||
"use_default_system_instructions", True
|
||||
),
|
||||
citations_enabled=yaml_config.get("citations_enabled", True),
|
||||
config_id=yaml_config.get("id"),
|
||||
config_name=yaml_config.get("name"),
|
||||
)
|
||||
|
||||
|
||||
def load_llm_config_from_yaml(llm_config_id: int = -1) -> dict | None:
|
||||
|
|
@ -47,9 +177,118 @@ def load_llm_config_from_yaml(llm_config_id: int = -1) -> dict | None:
|
|||
return None
|
||||
|
||||
|
||||
async def load_new_llm_config_from_db(
|
||||
session: AsyncSession,
|
||||
config_id: int,
|
||||
) -> "AgentConfig | None":
|
||||
"""
|
||||
Load a NewLLMConfig from the database by ID.
|
||||
|
||||
Args:
|
||||
session: AsyncSession for database access
|
||||
config_id: The ID of the NewLLMConfig to load
|
||||
|
||||
Returns:
|
||||
AgentConfig instance or None if not found
|
||||
"""
|
||||
# Import here to avoid circular imports
|
||||
from app.db import NewLLMConfig
|
||||
|
||||
try:
|
||||
result = await session.execute(
|
||||
select(NewLLMConfig).filter(NewLLMConfig.id == config_id)
|
||||
)
|
||||
config = result.scalars().first()
|
||||
|
||||
if not config:
|
||||
print(f"Error: NewLLMConfig with id {config_id} not found")
|
||||
return None
|
||||
|
||||
return AgentConfig.from_new_llm_config(config)
|
||||
except Exception as e:
|
||||
print(f"Error loading NewLLMConfig from database: {e}")
|
||||
return None
|
||||
|
||||
|
||||
async def load_agent_llm_config_for_search_space(
|
||||
session: AsyncSession,
|
||||
search_space_id: int,
|
||||
) -> "AgentConfig | None":
|
||||
"""
|
||||
Load the agent LLM configuration for a search space.
|
||||
|
||||
This loads the LLM config based on the search space's agent_llm_id setting:
|
||||
- Positive ID: Load from NewLLMConfig database table
|
||||
- Negative ID: Load from YAML global configs
|
||||
- None: Falls back to first global config (id=-1)
|
||||
|
||||
Args:
|
||||
session: AsyncSession for database access
|
||||
search_space_id: The search space ID
|
||||
|
||||
Returns:
|
||||
AgentConfig instance or None if not found
|
||||
"""
|
||||
# Import here to avoid circular imports
|
||||
from app.db import SearchSpace
|
||||
|
||||
try:
|
||||
# Get the search space to check its agent_llm_id preference
|
||||
result = await session.execute(
|
||||
select(SearchSpace).filter(SearchSpace.id == search_space_id)
|
||||
)
|
||||
search_space = result.scalars().first()
|
||||
|
||||
if not search_space:
|
||||
print(f"Error: SearchSpace with id {search_space_id} not found")
|
||||
return None
|
||||
|
||||
# Use agent_llm_id from search space, fallback to -1 (first global config)
|
||||
config_id = (
|
||||
search_space.agent_llm_id if search_space.agent_llm_id is not None else -1
|
||||
)
|
||||
|
||||
# Load the config using the unified loader
|
||||
return await load_agent_config(session, config_id, search_space_id)
|
||||
except Exception as e:
|
||||
print(f"Error loading agent LLM config for search space {search_space_id}: {e}")
|
||||
return None
|
||||
|
||||
|
||||
async def load_agent_config(
|
||||
session: AsyncSession,
|
||||
config_id: int,
|
||||
search_space_id: int | None = None,
|
||||
) -> "AgentConfig | None":
|
||||
"""
|
||||
Load an agent configuration, supporting both YAML (negative IDs) and database (positive IDs) configs.
|
||||
|
||||
This is the main entry point for loading configurations:
|
||||
- Negative IDs: Load from YAML file (global configs)
|
||||
- Positive IDs: Load from NewLLMConfig database table
|
||||
|
||||
Args:
|
||||
session: AsyncSession for database access
|
||||
config_id: The config ID (negative for YAML, positive for database)
|
||||
search_space_id: Optional search space ID for context
|
||||
|
||||
Returns:
|
||||
AgentConfig instance or None if not found
|
||||
"""
|
||||
if config_id < 0:
|
||||
# Load from YAML (global configs have negative IDs)
|
||||
yaml_config = load_llm_config_from_yaml(config_id)
|
||||
if yaml_config:
|
||||
return AgentConfig.from_yaml_config(yaml_config)
|
||||
return None
|
||||
else:
|
||||
# Load from database (NewLLMConfig)
|
||||
return await load_new_llm_config_from_db(session, config_id)
|
||||
|
||||
|
||||
def create_chat_litellm_from_config(llm_config: dict) -> ChatLiteLLM | None:
|
||||
"""
|
||||
Create a ChatLiteLLM instance from a global LLM config.
|
||||
Create a ChatLiteLLM instance from a global LLM config dictionary.
|
||||
|
||||
Args:
|
||||
llm_config: LLM configuration dictionary from YAML
|
||||
|
|
@ -57,34 +296,12 @@ def create_chat_litellm_from_config(llm_config: dict) -> ChatLiteLLM | None:
|
|||
Returns:
|
||||
ChatLiteLLM instance or None on error
|
||||
"""
|
||||
# Provider mapping (same as in llm_service.py)
|
||||
provider_map = {
|
||||
"OPENAI": "openai",
|
||||
"ANTHROPIC": "anthropic",
|
||||
"GROQ": "groq",
|
||||
"COHERE": "cohere",
|
||||
"GOOGLE": "gemini",
|
||||
"OLLAMA": "ollama",
|
||||
"MISTRAL": "mistral",
|
||||
"AZURE_OPENAI": "azure",
|
||||
"OPENROUTER": "openrouter",
|
||||
"XAI": "xai",
|
||||
"BEDROCK": "bedrock",
|
||||
"VERTEX_AI": "vertex_ai",
|
||||
"TOGETHER_AI": "together_ai",
|
||||
"FIREWORKS_AI": "fireworks_ai",
|
||||
"DEEPSEEK": "openai",
|
||||
"ALIBABA_QWEN": "openai",
|
||||
"MOONSHOT": "openai",
|
||||
"ZHIPU": "openai",
|
||||
}
|
||||
|
||||
# Build the model string
|
||||
if llm_config.get("custom_provider"):
|
||||
model_string = f"{llm_config['custom_provider']}/{llm_config['model_name']}"
|
||||
else:
|
||||
provider = llm_config.get("provider", "").upper()
|
||||
provider_prefix = provider_map.get(provider, provider.lower())
|
||||
provider_prefix = PROVIDER_MAP.get(provider, provider.lower())
|
||||
model_string = f"{provider_prefix}/{llm_config['model_name']}"
|
||||
|
||||
# Create ChatLiteLLM instance with streaming enabled
|
||||
|
|
@ -103,3 +320,42 @@ def create_chat_litellm_from_config(llm_config: dict) -> ChatLiteLLM | None:
|
|||
litellm_kwargs.update(llm_config["litellm_params"])
|
||||
|
||||
return ChatLiteLLM(**litellm_kwargs)
|
||||
|
||||
|
||||
def create_chat_litellm_from_agent_config(
|
||||
agent_config: AgentConfig,
|
||||
) -> ChatLiteLLM | None:
|
||||
"""
|
||||
Create a ChatLiteLLM instance from an AgentConfig.
|
||||
|
||||
Args:
|
||||
agent_config: AgentConfig instance
|
||||
|
||||
Returns:
|
||||
ChatLiteLLM instance or None on error
|
||||
"""
|
||||
# Build the model string
|
||||
if agent_config.custom_provider:
|
||||
model_string = f"{agent_config.custom_provider}/{agent_config.model_name}"
|
||||
else:
|
||||
provider_prefix = PROVIDER_MAP.get(
|
||||
agent_config.provider, agent_config.provider.lower()
|
||||
)
|
||||
model_string = f"{provider_prefix}/{agent_config.model_name}"
|
||||
|
||||
# Create ChatLiteLLM instance with streaming enabled
|
||||
litellm_kwargs = {
|
||||
"model": model_string,
|
||||
"api_key": agent_config.api_key,
|
||||
"streaming": True, # Enable streaming for real-time token streaming
|
||||
}
|
||||
|
||||
# Add optional parameters
|
||||
if agent_config.api_base:
|
||||
litellm_kwargs["api_base"] = agent_config.api_base
|
||||
|
||||
# Add any additional litellm parameters
|
||||
if agent_config.litellm_params:
|
||||
litellm_kwargs.update(agent_config.litellm_params)
|
||||
|
||||
return ChatLiteLLM(**litellm_kwargs)
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue