| | 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() |
| |
|
