refactor: classes and configs (#528)
This commit is contained in:
@@ -1,26 +1,25 @@
|
||||
from embedchain import CustomApp
|
||||
from embedchain.config import AddConfig, CustomAppConfig, QueryConfig
|
||||
from embedchain.config import AddConfig, CustomAppConfig, LlmConfig
|
||||
from embedchain.embedder.openai_embedder import OpenAiEmbedder
|
||||
from embedchain.helper_classes.json_serializable import (
|
||||
JSONSerializable, register_deserializable)
|
||||
from embedchain.models import EmbeddingFunctions, Providers
|
||||
from embedchain.llm.openai_llm import OpenAiLlm
|
||||
from embedchain.vectordb.chroma_db import ChromaDB
|
||||
|
||||
|
||||
@register_deserializable
|
||||
class BaseBot(JSONSerializable):
|
||||
def __init__(self, app_config=None):
|
||||
if app_config is None:
|
||||
app_config = CustomAppConfig(embedding_fn=EmbeddingFunctions.OPENAI, provider=Providers.OPENAI)
|
||||
self.app_config = app_config
|
||||
self.app = CustomApp(config=self.app_config)
|
||||
def __init__(self):
|
||||
self.app = CustomApp(config=CustomAppConfig(), llm=OpenAiLlm(), db=ChromaDB(), embedder=OpenAiEmbedder())
|
||||
|
||||
def add(self, data, config: AddConfig = None):
|
||||
"""Add data to the bot"""
|
||||
config = config if config else AddConfig()
|
||||
self.app.add(data, config=config)
|
||||
|
||||
def query(self, query, config: QueryConfig = None):
|
||||
def query(self, query, config: LlmConfig = None):
|
||||
"""Query bot"""
|
||||
config = config if config else QueryConfig()
|
||||
config = config
|
||||
return self.app.query(query, config=config)
|
||||
|
||||
def start(self):
|
||||
|
||||
Reference in New Issue
Block a user