daa4e9acc4
<!-- .github/pull_request_template.md --> ## Description <!-- Provide a clear description of the changes in this PR --> ## DCO Affirmation I affirm that all code in every commit of this pull request conforms to the terms of the Topoteretes Developer Certificate of Origin.
153 lines
4.2 KiB
Python
153 lines
4.2 KiB
Python
from enum import Enum
|
|
from typing import Optional
|
|
from pydantic import BaseModel
|
|
from cognee.infrastructure.databases.vector import get_vectordb_config
|
|
from cognee.infrastructure.llm import get_llm_config
|
|
|
|
|
|
class ConfigChoice(BaseModel):
|
|
value: str
|
|
label: str
|
|
|
|
|
|
class ModelName(Enum):
|
|
openai = "openai"
|
|
ollama = "ollama"
|
|
anthropic = "anthropic"
|
|
gemini = "gemini"
|
|
|
|
|
|
class LLMConfig(BaseModel):
|
|
api_key: str
|
|
model: str
|
|
provider: str
|
|
endpoint: Optional[str]
|
|
api_version: Optional[str]
|
|
models: dict[str, list[ConfigChoice]]
|
|
providers: list[ConfigChoice]
|
|
|
|
|
|
class VectorDBConfig(BaseModel):
|
|
api_key: str
|
|
url: str
|
|
provider: str
|
|
providers: list[ConfigChoice]
|
|
|
|
|
|
class SettingsDict(BaseModel):
|
|
llm: LLMConfig
|
|
vector_db: VectorDBConfig
|
|
|
|
|
|
def get_settings() -> SettingsDict:
|
|
llm_config = get_llm_config()
|
|
|
|
vector_dbs = [
|
|
{
|
|
"value": "qdrant",
|
|
"label": "Qdrant",
|
|
},
|
|
{
|
|
"value": "lancedb",
|
|
"label": "LanceDB",
|
|
},
|
|
{
|
|
"value": "pgvector",
|
|
"label": "PGVector",
|
|
},
|
|
]
|
|
|
|
vector_config = get_vectordb_config()
|
|
|
|
llm_providers = [
|
|
{
|
|
"value": "openai",
|
|
"label": "OpenAI",
|
|
},
|
|
{
|
|
"value": "ollama",
|
|
"label": "Ollama",
|
|
},
|
|
{
|
|
"value": "anthropic",
|
|
"label": "Anthropic",
|
|
},
|
|
{
|
|
"value": "gemini",
|
|
"label": "Gemini",
|
|
},
|
|
]
|
|
|
|
return SettingsDict.model_validate(
|
|
dict(
|
|
llm={
|
|
"provider": llm_config.llm_provider,
|
|
"model": llm_config.llm_model,
|
|
"endpoint": llm_config.llm_endpoint,
|
|
"api_version": llm_config.llm_api_version,
|
|
"api_key": (llm_config.llm_api_key[0:10] + "*" * (len(llm_config.llm_api_key) - 10))
|
|
if llm_config.llm_api_key
|
|
else None,
|
|
"providers": llm_providers,
|
|
"models": {
|
|
"openai": [
|
|
{
|
|
"value": "gpt-4o-mini",
|
|
"label": "gpt-4o-mini",
|
|
},
|
|
{
|
|
"value": "gpt-4o",
|
|
"label": "gpt-4o",
|
|
},
|
|
{
|
|
"value": "gpt-4-turbo",
|
|
"label": "gpt-4-turbo",
|
|
},
|
|
{
|
|
"value": "gpt-3.5-turbo",
|
|
"label": "gpt-3.5-turbo",
|
|
},
|
|
],
|
|
"ollama": [
|
|
{
|
|
"value": "llama3",
|
|
"label": "llama3",
|
|
},
|
|
{
|
|
"value": "mistral",
|
|
"label": "mistral",
|
|
},
|
|
],
|
|
"anthropic": [
|
|
{
|
|
"value": "Claude 3 Opus",
|
|
"label": "Claude 3 Opus",
|
|
},
|
|
{
|
|
"value": "Claude 3 Sonnet",
|
|
"label": "Claude 3 Sonnet",
|
|
},
|
|
{
|
|
"value": "Claude 3 Haiku",
|
|
"label": "Claude 3 Haiku",
|
|
},
|
|
],
|
|
"gemini": [
|
|
{
|
|
"value": "gemini-2.0-flash-exp",
|
|
"label": "Gemini 2.0 Flash",
|
|
},
|
|
],
|
|
},
|
|
},
|
|
vector_db={
|
|
"provider": vector_config.vector_db_provider,
|
|
"url": vector_config.vector_db_url,
|
|
"api_key": (
|
|
vector_config.vector_db_key[0:10]
|
|
+ "*" * (len(vector_config.vector_db_key) - 10)
|
|
),
|
|
"providers": vector_dbs,
|
|
},
|
|
)
|
|
)
|