| import os |
| import sys |
| import unittest |
| from unittest import mock |
| from uuid import uuid4 |
|
|
| from weaviate import Client |
| from weaviate.util import get_valid_uuid |
|
|
| from autogpt.config import Config |
| from autogpt.memory.base import get_ada_embedding |
| from autogpt.memory.weaviate import WeaviateMemory |
|
|
|
|
| class TestWeaviateMemory(unittest.TestCase): |
| cfg = None |
| client = None |
| index = None |
|
|
| @classmethod |
| def setUpClass(cls): |
| |
| cls.cfg = Config() |
|
|
| if cls.cfg.use_weaviate_embedded: |
| from weaviate.embedded import EmbeddedOptions |
|
|
| cls.client = Client( |
| embedded_options=EmbeddedOptions( |
| hostname=cls.cfg.weaviate_host, |
| port=int(cls.cfg.weaviate_port), |
| persistence_data_path=cls.cfg.weaviate_embedded_path, |
| ) |
| ) |
| else: |
| cls.client = Client( |
| f"{cls.cfg.weaviate_protocol}://{cls.cfg.weaviate_host}:{self.cfg.weaviate_port}" |
| ) |
|
|
| cls.index = WeaviateMemory.format_classname(cls.cfg.memory_index) |
|
|
| """ |
| In order to run these tests you will need a local instance of |
| Weaviate running. Refer to https://weaviate.io/developers/weaviate/installation/docker-compose |
| for creating local instances using docker. |
| Alternatively in your .env file set the following environmental variables to run Weaviate embedded (see: https://weaviate.io/developers/weaviate/installation/embedded): |
| |
| USE_WEAVIATE_EMBEDDED=True |
| WEAVIATE_EMBEDDED_PATH="/home/me/.local/share/weaviate" |
| """ |
|
|
| def setUp(self): |
| try: |
| self.client.schema.delete_class(self.index) |
| except: |
| pass |
|
|
| self.memory = WeaviateMemory(self.cfg) |
|
|
| def test_add(self): |
| doc = "You are a Titan name Thanos and you are looking for the Infinity Stones" |
| self.memory.add(doc) |
| result = self.client.query.get(self.index, ["raw_text"]).do() |
| actual = result["data"]["Get"][self.index] |
|
|
| self.assertEqual(len(actual), 1) |
| self.assertEqual(actual[0]["raw_text"], doc) |
|
|
| def test_get(self): |
| doc = "You are an Avenger and swore to defend the Galaxy from a menace called Thanos" |
|
|
| with self.client.batch as batch: |
| batch.add_data_object( |
| uuid=get_valid_uuid(uuid4()), |
| data_object={"raw_text": doc}, |
| class_name=self.index, |
| vector=get_ada_embedding(doc), |
| ) |
|
|
| batch.flush() |
|
|
| actual = self.memory.get(doc) |
|
|
| self.assertEqual(len(actual), 1) |
| self.assertEqual(actual[0], doc) |
|
|
| def test_get_stats(self): |
| docs = [ |
| "You are now about to count the number of docs in this index", |
| "And then you about to find out if you can count correctly", |
| ] |
|
|
| [self.memory.add(doc) for doc in docs] |
|
|
| stats = self.memory.get_stats() |
|
|
| self.assertTrue(stats) |
| self.assertTrue("count" in stats) |
| self.assertEqual(stats["count"], 2) |
|
|
| def test_clear(self): |
| docs = [ |
| "Shame this is the last test for this class", |
| "Testing is fun when someone else is doing it", |
| ] |
|
|
| [self.memory.add(doc) for doc in docs] |
|
|
| self.assertEqual(self.memory.get_stats()["count"], 2) |
|
|
| self.memory.clear() |
|
|
| self.assertEqual(self.memory.get_stats()["count"], 0) |
|
|
|
|
| if __name__ == "__main__": |
| unittest.main() |
|
|
