Feature (OpenMemory): Add support for LLM and Embedding Providers in OpenMemory (#2794)
This commit is contained in:
Saket Aryan
240
openmemory/api/app/routers/config.py
Normal file
240
openmemory/api/app/routers/config.py
Normal file
@@ -0,0 +1,240 @@
|
||||
import os
|
||||
import json
|
||||
from typing import Dict, Any, Optional
|
||||
from fastapi import APIRouter, HTTPException, Depends
|
||||
from pydantic import BaseModel, Field
|
||||
from sqlalchemy.orm import Session
|
||||
from app.database import get_db
|
||||
from app.models import Config as ConfigModel
|
||||
from app.utils.memory import reset_memory_client
|
||||
|
||||
router = APIRouter(prefix="/api/v1/config", tags=["config"])
|
||||
|
||||
class LLMConfig(BaseModel):
|
||||
model: str = Field(..., description="LLM model name")
|
||||
temperature: float = Field(..., description="Temperature setting for the model")
|
||||
max_tokens: int = Field(..., description="Maximum tokens to generate")
|
||||
api_key: Optional[str] = Field(None, description="API key or 'env:API_KEY' to use environment variable")
|
||||
ollama_base_url: Optional[str] = Field(None, description="Base URL for Ollama server (e.g., http://host.docker.internal:11434)")
|
||||
|
||||
class LLMProvider(BaseModel):
|
||||
provider: str = Field(..., description="LLM provider name")
|
||||
config: LLMConfig
|
||||
|
||||
class EmbedderConfig(BaseModel):
|
||||
model: str = Field(..., description="Embedder model name")
|
||||
api_key: Optional[str] = Field(None, description="API key or 'env:API_KEY' to use environment variable")
|
||||
ollama_base_url: Optional[str] = Field(None, description="Base URL for Ollama server (e.g., http://host.docker.internal:11434)")
|
||||
|
||||
class EmbedderProvider(BaseModel):
|
||||
provider: str = Field(..., description="Embedder provider name")
|
||||
config: EmbedderConfig
|
||||
|
||||
class OpenMemoryConfig(BaseModel):
|
||||
custom_instructions: Optional[str] = Field(None, description="Custom instructions for memory management and fact extraction")
|
||||
|
||||
class Mem0Config(BaseModel):
|
||||
llm: Optional[LLMProvider] = None
|
||||
embedder: Optional[EmbedderProvider] = None
|
||||
|
||||
class ConfigSchema(BaseModel):
|
||||
openmemory: Optional[OpenMemoryConfig] = None
|
||||
mem0: Mem0Config
|
||||
|
||||
def get_default_configuration():
|
||||
"""Get the default configuration with sensible defaults for LLM and embedder."""
|
||||
return {
|
||||
"openmemory": {
|
||||
"custom_instructions": None
|
||||
},
|
||||
"mem0": {
|
||||
"llm": {
|
||||
"provider": "openai",
|
||||
"config": {
|
||||
"model": "gpt-4o-mini",
|
||||
"temperature": 0.1,
|
||||
"max_tokens": 2000,
|
||||
"api_key": "env:OPENAI_API_KEY"
|
||||
}
|
||||
},
|
||||
"embedder": {
|
||||
"provider": "openai",
|
||||
"config": {
|
||||
"model": "text-embedding-3-small",
|
||||
"api_key": "env:OPENAI_API_KEY"
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
def get_config_from_db(db: Session, key: str = "main"):
|
||||
"""Get configuration from database."""
|
||||
config = db.query(ConfigModel).filter(ConfigModel.key == key).first()
|
||||
|
||||
if not config:
|
||||
# Create default config with proper provider configurations
|
||||
default_config = get_default_configuration()
|
||||
db_config = ConfigModel(key=key, value=default_config)
|
||||
db.add(db_config)
|
||||
db.commit()
|
||||
db.refresh(db_config)
|
||||
return default_config
|
||||
|
||||
# Ensure the config has all required sections with defaults
|
||||
config_value = config.value
|
||||
default_config = get_default_configuration()
|
||||
|
||||
# Merge with defaults to ensure all required fields exist
|
||||
if "openmemory" not in config_value:
|
||||
config_value["openmemory"] = default_config["openmemory"]
|
||||
|
||||
if "mem0" not in config_value:
|
||||
config_value["mem0"] = default_config["mem0"]
|
||||
else:
|
||||
# Ensure LLM config exists with defaults
|
||||
if "llm" not in config_value["mem0"] or config_value["mem0"]["llm"] is None:
|
||||
config_value["mem0"]["llm"] = default_config["mem0"]["llm"]
|
||||
|
||||
# Ensure embedder config exists with defaults
|
||||
if "embedder" not in config_value["mem0"] or config_value["mem0"]["embedder"] is None:
|
||||
config_value["mem0"]["embedder"] = default_config["mem0"]["embedder"]
|
||||
|
||||
# Save the updated config back to database if it was modified
|
||||
if config_value != config.value:
|
||||
config.value = config_value
|
||||
db.commit()
|
||||
db.refresh(config)
|
||||
|
||||
return config_value
|
||||
|
||||
def save_config_to_db(db: Session, config: Dict[str, Any], key: str = "main"):
|
||||
"""Save configuration to database."""
|
||||
db_config = db.query(ConfigModel).filter(ConfigModel.key == key).first()
|
||||
|
||||
if db_config:
|
||||
db_config.value = config
|
||||
db_config.updated_at = None # Will trigger the onupdate to set current time
|
||||
else:
|
||||
db_config = ConfigModel(key=key, value=config)
|
||||
db.add(db_config)
|
||||
|
||||
db.commit()
|
||||
db.refresh(db_config)
|
||||
return db_config.value
|
||||
|
||||
@router.get("/", response_model=ConfigSchema)
|
||||
async def get_configuration(db: Session = Depends(get_db)):
|
||||
"""Get the current configuration."""
|
||||
config = get_config_from_db(db)
|
||||
return config
|
||||
|
||||
@router.put("/", response_model=ConfigSchema)
|
||||
async def update_configuration(config: ConfigSchema, db: Session = Depends(get_db)):
|
||||
"""Update the configuration."""
|
||||
current_config = get_config_from_db(db)
|
||||
|
||||
# Convert to dict for processing
|
||||
updated_config = current_config.copy()
|
||||
|
||||
# Update openmemory settings if provided
|
||||
if config.openmemory is not None:
|
||||
if "openmemory" not in updated_config:
|
||||
updated_config["openmemory"] = {}
|
||||
updated_config["openmemory"].update(config.openmemory.dict(exclude_none=True))
|
||||
|
||||
# Update mem0 settings
|
||||
updated_config["mem0"] = config.mem0.dict(exclude_none=True)
|
||||
|
||||
# Save the configuration to database
|
||||
save_config_to_db(db, updated_config)
|
||||
reset_memory_client()
|
||||
return updated_config
|
||||
|
||||
@router.post("/reset", response_model=ConfigSchema)
|
||||
async def reset_configuration(db: Session = Depends(get_db)):
|
||||
"""Reset the configuration to default values."""
|
||||
try:
|
||||
# Get the default configuration with proper provider setups
|
||||
default_config = get_default_configuration()
|
||||
|
||||
# Save it as the current configuration in the database
|
||||
save_config_to_db(db, default_config)
|
||||
reset_memory_client()
|
||||
return default_config
|
||||
except Exception as e:
|
||||
raise HTTPException(
|
||||
status_code=500,
|
||||
detail=f"Failed to reset configuration: {str(e)}"
|
||||
)
|
||||
|
||||
@router.get("/mem0/llm", response_model=LLMProvider)
|
||||
async def get_llm_configuration(db: Session = Depends(get_db)):
|
||||
"""Get only the LLM configuration."""
|
||||
config = get_config_from_db(db)
|
||||
llm_config = config.get("mem0", {}).get("llm", {})
|
||||
return llm_config
|
||||
|
||||
@router.put("/mem0/llm", response_model=LLMProvider)
|
||||
async def update_llm_configuration(llm_config: LLMProvider, db: Session = Depends(get_db)):
|
||||
"""Update only the LLM configuration."""
|
||||
current_config = get_config_from_db(db)
|
||||
|
||||
# Ensure mem0 key exists
|
||||
if "mem0" not in current_config:
|
||||
current_config["mem0"] = {}
|
||||
|
||||
# Update the LLM configuration
|
||||
current_config["mem0"]["llm"] = llm_config.dict(exclude_none=True)
|
||||
|
||||
# Save the configuration to database
|
||||
save_config_to_db(db, current_config)
|
||||
reset_memory_client()
|
||||
return current_config["mem0"]["llm"]
|
||||
|
||||
@router.get("/mem0/embedder", response_model=EmbedderProvider)
|
||||
async def get_embedder_configuration(db: Session = Depends(get_db)):
|
||||
"""Get only the Embedder configuration."""
|
||||
config = get_config_from_db(db)
|
||||
embedder_config = config.get("mem0", {}).get("embedder", {})
|
||||
return embedder_config
|
||||
|
||||
@router.put("/mem0/embedder", response_model=EmbedderProvider)
|
||||
async def update_embedder_configuration(embedder_config: EmbedderProvider, db: Session = Depends(get_db)):
|
||||
"""Update only the Embedder configuration."""
|
||||
current_config = get_config_from_db(db)
|
||||
|
||||
# Ensure mem0 key exists
|
||||
if "mem0" not in current_config:
|
||||
current_config["mem0"] = {}
|
||||
|
||||
# Update the Embedder configuration
|
||||
current_config["mem0"]["embedder"] = embedder_config.dict(exclude_none=True)
|
||||
|
||||
# Save the configuration to database
|
||||
save_config_to_db(db, current_config)
|
||||
reset_memory_client()
|
||||
return current_config["mem0"]["embedder"]
|
||||
|
||||
@router.get("/openmemory", response_model=OpenMemoryConfig)
|
||||
async def get_openmemory_configuration(db: Session = Depends(get_db)):
|
||||
"""Get only the OpenMemory configuration."""
|
||||
config = get_config_from_db(db)
|
||||
openmemory_config = config.get("openmemory", {})
|
||||
return openmemory_config
|
||||
|
||||
@router.put("/openmemory", response_model=OpenMemoryConfig)
|
||||
async def update_openmemory_configuration(openmemory_config: OpenMemoryConfig, db: Session = Depends(get_db)):
|
||||
"""Update only the OpenMemory configuration."""
|
||||
current_config = get_config_from_db(db)
|
||||
|
||||
# Ensure openmemory key exists
|
||||
if "openmemory" not in current_config:
|
||||
current_config["openmemory"] = {}
|
||||
|
||||
# Update the OpenMemory configuration
|
||||
current_config["openmemory"].update(openmemory_config.dict(exclude_none=True))
|
||||
|
||||
# Save the configuration to database
|
||||
save_config_to_db(db, current_config)
|
||||
reset_memory_client()
|
||||
return current_config["openmemory"]
|
||||
Reference in New Issue
Block a user