test_embedchain.py 1.3 KB

123456789101112131415161718192021222324252627282930313233343536373839
  1. import os
  2. import unittest
  3. from unittest.mock import patch
  4. from embedchain import App
  5. from embedchain.config import AppConfig
  6. class TestChromaDbHostsLoglevel(unittest.TestCase):
  7. os.environ["OPENAI_API_KEY"] = "test_key"
  8. @patch("chromadb.api.models.Collection.Collection.add")
  9. @patch("chromadb.api.models.Collection.Collection.get")
  10. @patch("embedchain.embedchain.EmbedChain.retrieve_from_database")
  11. @patch("embedchain.embedchain.EmbedChain.get_answer_from_llm")
  12. @patch("embedchain.embedchain.EmbedChain.get_llm_model_answer")
  13. def test_whole_app(
  14. self,
  15. _mock_get,
  16. _mock_add,
  17. _mock_ec_retrieve_from_database,
  18. _mock_get_answer_from_llm,
  19. mock_ec_get_llm_model_answer,
  20. ):
  21. """
  22. Test if the `App` instance is initialized without a config that does not contain default hosts and ports.
  23. """
  24. config = AppConfig(log_level="DEBUG", collect_metrics=False)
  25. app = App(config)
  26. knowledge = "lorem ipsum dolor sit amet, consectetur adipiscing"
  27. app.add_local("text", knowledge)
  28. app.query("What text did I give you?")
  29. app.chat("What text did I give you?")
  30. self.assertEqual(mock_ec_get_llm_model_answer.call_args[1]["documents"], [knowledge])