23 lines
739 B
Python
23 lines
739 B
Python
from functools import lru_cache
|
|
from pydantic_settings import BaseSettings, SettingsConfigDict
|
|
from cognee.shared.data_models import DefaultContentPrediction, SummarizedContent
|
|
from typing import Optional
|
|
import os
|
|
|
|
|
|
class CognifyConfig(BaseSettings):
|
|
classification_model: object = DefaultContentPrediction
|
|
summarization_model: object = SummarizedContent
|
|
max_tokens: Optional[int] = os.getenv("MAX_TOKENS")
|
|
model_config = SettingsConfigDict(env_file=".env", extra="allow")
|
|
|
|
def to_dict(self) -> dict:
|
|
return {
|
|
"classification_model": self.classification_model,
|
|
"summarization_model": self.summarization_model,
|
|
}
|
|
|
|
|
|
@lru_cache
|
|
def get_cognify_config():
|
|
return CognifyConfig()
|