base.py 1.0 KB

123456789101112131415161718192021222324252627
  1. from embedchain import CustomApp
  2. from embedchain.config import AddConfig, CustomAppConfig, LlmConfig
  3. from embedchain.embedder.openai_embedder import OpenAiEmbedder
  4. from embedchain.helper_classes.json_serializable import (
  5. JSONSerializable, register_deserializable)
  6. from embedchain.llm.openai_llm import OpenAiLlm
  7. from embedchain.vectordb.chroma_db import ChromaDB
  8. @register_deserializable
  9. class BaseBot(JSONSerializable):
  10. def __init__(self):
  11. self.app = CustomApp(config=CustomAppConfig(), llm=OpenAiLlm(), db=ChromaDB(), embedder=OpenAiEmbedder())
  12. def add(self, data, config: AddConfig = None):
  13. """Add data to the bot"""
  14. config = config if config else AddConfig()
  15. self.app.add(data, config=config)
  16. def query(self, query, config: LlmConfig = None):
  17. """Query bot"""
  18. config = config
  19. return self.app.query(query, config=config)
  20. def start(self):
  21. """Start the bot's functionality."""
  22. raise NotImplementedError("Subclasses must implement the start method.")