Changed class order in settings router to allow proper work of settings Refactor #COG-334
52 lines
No EOL
1.7 KiB
Python
52 lines
No EOL
1.7 KiB
Python
from fastapi import APIRouter
|
|
from cognee.api.DTO import InDTO, OutDTO
|
|
from typing import Union, Optional, Literal
|
|
from cognee.modules.users.methods import get_authenticated_user
|
|
from fastapi import Depends
|
|
from cognee.modules.users.models import User
|
|
|
|
def get_settings_router():
|
|
router = APIRouter()
|
|
|
|
from cognee.modules.settings.get_settings import LLMConfig, VectorDBConfig
|
|
|
|
class LLMConfigDTO(OutDTO, LLMConfig):
|
|
pass
|
|
|
|
class VectorDBConfigDTO(OutDTO, VectorDBConfig):
|
|
pass
|
|
|
|
class SettingsDTO(OutDTO):
|
|
llm: LLMConfigDTO
|
|
vector_db: VectorDBConfigDTO
|
|
|
|
@router.get("/", response_model=SettingsDTO)
|
|
async def get_settings(user: User = Depends(get_authenticated_user)):
|
|
from cognee.modules.settings import get_settings as get_cognee_settings
|
|
return get_cognee_settings()
|
|
|
|
class LLMConfigDTO(InDTO):
|
|
provider: Union[Literal["openai"], Literal["ollama"], Literal["anthropic"]]
|
|
model: str
|
|
api_key: str
|
|
|
|
class VectorDBConfigDTO(InDTO):
|
|
provider: Union[Literal["lancedb"], Literal["qdrant"], Literal["weaviate"], Literal["pgvector"]]
|
|
url: str
|
|
api_key: str
|
|
|
|
class SettingsPayloadDTO(InDTO):
|
|
llm: Optional[LLMConfigDTO] = None
|
|
vector_db: Optional[VectorDBConfigDTO] = None
|
|
|
|
@router.post("/", response_model=None)
|
|
async def save_settings(new_settings: SettingsPayloadDTO, user: User = Depends(get_authenticated_user)):
|
|
from cognee.modules.settings import save_llm_config, save_vector_db_config
|
|
|
|
if new_settings.llm is not None:
|
|
await save_llm_config(new_settings.llm)
|
|
|
|
if new_settings.vector_db is not None:
|
|
await save_vector_db_config(new_settings.vector_db)
|
|
|
|
return router |