chore: linting (#230)

This commit is contained in:
cachho
2023-07-11 14:04:57 +02:00
committed by GitHub
parent 1c4f629fc0
commit 9409e5605a
3 changed files with 25 additions and 23 deletions

View File

@@ -46,25 +46,31 @@ class InitConfig(BaseConfig):
def _set_embedding_function(self, ef): def _set_embedding_function(self, ef):
self.ef = ef self.ef = ef
return return
def _set_embedding_function_to_default(self): def _set_embedding_function_to_default(self):
""" """
Sets embedding function to default (`text-embedding-ada-002`). Sets embedding function to default (`text-embedding-ada-002`).
:raises ValueError: If the template is not valid as template should contain $context and $query :raises ValueError: If the template is not valid as template should contain
$context and $query
""" """
if os.getenv("OPENAI_API_KEY") is None or os.getenv("OPENAI_ORGANIZATION") is None: if (
raise ValueError("OPENAI_API_KEY or OPENAI_ORGANIZATION environment variables not provided") os.getenv("OPENAI_API_KEY") is None
self.ef = embedding_functions.OpenAIEmbeddingFunction( or os.getenv("OPENAI_ORGANIZATION") is None
api_key=os.getenv("OPENAI_API_KEY"), ):
organization_id=os.getenv("OPENAI_ORGANIZATION"), raise ValueError(
model_name="text-embedding-ada-002" "OPENAI_API_KEY or OPENAI_ORGANIZATION environment variables not provided" # noqa:E501
) )
self.ef = embedding_functions.OpenAIEmbeddingFunction(
api_key=os.getenv("OPENAI_API_KEY"),
organization_id=os.getenv("OPENAI_ORGANIZATION"),
model_name="text-embedding-ada-002",
)
return return
def _set_db(self, db): def _set_db(self, db):
if db: if db:
self.db = db self.db = db
return return
def _set_db_to_default(self): def _set_db_to_default(self):
@@ -72,6 +78,7 @@ class InitConfig(BaseConfig):
Sets database to default (`ChromaDb`). Sets database to default (`ChromaDb`).
""" """
from embedchain.vectordb.chroma_db import ChromaDB from embedchain.vectordb.chroma_db import ChromaDB
self.db = ChromaDB(ef=self.ef) self.db = ChromaDB(ef=self.ef)
def _setup_logging(self, debug_level): def _setup_logging(self, debug_level):

View File

@@ -301,13 +301,13 @@ class App(EmbedChain):
""" """
if config is None: if config is None:
config = InitConfig() config = InitConfig()
if not config.ef: if not config.ef:
config._set_embedding_function_to_default() config._set_embedding_function_to_default()
if not config.db: if not config.db:
config._set_db_to_default() config._set_db_to_default()
super().__init__(config) super().__init__(config)
def get_llm_model_answer(self, prompt, config: ChatConfig): def get_llm_model_answer(self, prompt, config: ChatConfig):
@@ -357,12 +357,13 @@ class OpenSourceApp(EmbedChain):
) # noqa:E501 ) # noqa:E501
if not config: if not config:
config = InitConfig() config = InitConfig()
if not config.ef: if not config.ef:
config._set_embedding_function( config._set_embedding_function(
embedding_functions.SentenceTransformerEmbeddingFunction( embedding_functions.SentenceTransformerEmbeddingFunction(
model_name="all-MiniLM-L6-v2" model_name="all-MiniLM-L6-v2"
)) )
)
if not config.db: if not config.db:
config._set_db_to_default() config._set_db_to_default()

View File

@@ -34,11 +34,5 @@ setuptools.setup(
"docx2txt", "docx2txt",
"pydantic==1.10.8", "pydantic==1.10.8",
], ],
extras_require={ extras_require={"dev": ["black", "ruff", "isort"]},
"dev": [
"black",
"ruff",
"isort"
]
}
) )