coding-alt/AutoGPT

1import uuid2 3import weaviate4from weaviate import Client5from weaviate.embedded import EmbeddedOptions6from weaviate.util import generate_uuid57 8from autogpt.config import Config9from autogpt.memory.base import MemoryProviderSingleton, get_ada_embedding10 11 12def default_schema(weaviate_index):13    return {14        "class": weaviate_index,15        "properties": [16            {17                "name": "raw_text",18                "dataType": ["text"],19                "description": "original text for the embedding",20            }21        ],22    }23 24 25class WeaviateMemory(MemoryProviderSingleton):26    def __init__(self, cfg):27        auth_credentials = self._build_auth_credentials(cfg)28 29        url = f"{cfg.weaviate_protocol}://{cfg.weaviate_host}:{cfg.weaviate_port}"30 31        if cfg.use_weaviate_embedded:32            self.client = Client(33                embedded_options=EmbeddedOptions(34                    hostname=cfg.weaviate_host,35                    port=int(cfg.weaviate_port),36                    persistence_data_path=cfg.weaviate_embedded_path,37                )38            )39 40            print(41                f"Weaviate Embedded running on: {url} with persistence path: {cfg.weaviate_embedded_path}"42            )43        else:44            self.client = Client(url, auth_client_secret=auth_credentials)45 46        self.index = WeaviateMemory.format_classname(cfg.memory_index)47        self._create_schema()48 49    @staticmethod50    def format_classname(index):51        # weaviate uses capitalised index names52        # The python client uses the following code to format53        # index names before the corresponding class is created54        if len(index) == 1:55            return index.capitalize()56        return index[0].capitalize() + index[1:]57 58    def _create_schema(self):59        schema = default_schema(self.index)60        if not self.client.schema.contains(schema):61            self.client.schema.create_class(schema)62 63    def _build_auth_credentials(self, cfg):64        if cfg.weaviate_username and cfg.weaviate_password:65            return weaviate.AuthClientPassword(66                cfg.weaviate_username, cfg.weaviate_password67            )68        if cfg.weaviate_api_key:69            return weaviate.AuthApiKey(api_key=cfg.weaviate_api_key)70        else:71            return None72 73    def add(self, data):74        vector = get_ada_embedding(data)75 76        doc_uuid = generate_uuid5(data, self.index)77        data_object = {"raw_text": data}78 79        with self.client.batch as batch:80            batch.add_data_object(81                uuid=doc_uuid,82                data_object=data_object,83                class_name=self.index,84                vector=vector,85            )86 87        return f"Inserting data into memory at uuid: {doc_uuid}:\n data: {data}"88 89    def get(self, data):90        return self.get_relevant(data, 1)91 92    def clear(self):93        self.client.schema.delete_all()94 95        # weaviate does not yet have a neat way to just remove the items in an index96        # without removing the entire schema, therefore we need to re-create it97        # after a call to delete_all98        self._create_schema()99 100        return "Obliterated"101 102    def get_relevant(self, data, num_relevant=5):103        query_embedding = get_ada_embedding(data)104        try:105            results = (106                self.client.query.get(self.index, ["raw_text"])107                .with_near_vector({"vector": query_embedding, "certainty": 0.7})108                .with_limit(num_relevant)109                .do()110            )111 112            if len(results["data"]["Get"][self.index]) > 0:113                return [114                    str(item["raw_text"]) for item in results["data"]["Get"][self.index]115                ]116            else:117                return []118 119        except Exception as err:120            print(f"Unexpected error {err=}, {type(err)=}")121            return []122 123    def get_stats(self):124        result = self.client.query.aggregate(self.index).with_meta_count().do()125        class_data = result["data"]["Aggregate"][self.index]126 127        return class_data[0]["meta"] if class_data else {}128