b297289060
<!-- .github/pull_request_template.md --> ## Description 1. Cleans up VectorDB adapters that have been migrated to `cognee-community` repo 2. Adds async lock protection create_collection method in remaining VectorDB - ChromaDB See #1222 ## 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. --------- Co-authored-by: Igor Ilic <30923996+dexters1@users.noreply.github.com>
149 lines
4.1 KiB
Python
149 lines
4.1 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": "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,
|
|
},
|
|
)
|
|
)
|