Currently TikToken is used for tokenizing by default which is only supported by OpenAI, this is an initial commit in an attempt to add Cognee tokenizing support for multiple LLMs
22 lines
683 B
Python
22 lines
683 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
|
|
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()
|