coding-alt/AutoGPT

1""" Milvus memory storage provider."""2from pymilvus import Collection, CollectionSchema, DataType, FieldSchema, connections3 4from autogpt.memory.base import MemoryProviderSingleton, get_ada_embedding5 6 7class MilvusMemory(MemoryProviderSingleton):8    """Milvus memory storage provider."""9 10    def __init__(self, cfg) -> None:11        """Construct a milvus memory storage connection.12 13        Args:14            cfg (Config): Auto-GPT global config.15        """16        # connect to milvus server.17        connections.connect(address=cfg.milvus_addr)18        fields = [19            FieldSchema(name="pk", dtype=DataType.INT64, is_primary=True, auto_id=True),20            FieldSchema(name="embeddings", dtype=DataType.FLOAT_VECTOR, dim=1536),21            FieldSchema(name="raw_text", dtype=DataType.VARCHAR, max_length=65535),22        ]23 24        # create collection if not exist and load it.25        self.milvus_collection = cfg.milvus_collection26        self.schema = CollectionSchema(fields, "auto-gpt memory storage")27        self.collection = Collection(self.milvus_collection, self.schema)28        # create index if not exist.29        if not self.collection.has_index():30            self.collection.release()31            self.collection.create_index(32                "embeddings",33                {34                    "metric_type": "IP",35                    "index_type": "HNSW",36                    "params": {"M": 8, "efConstruction": 64},37                },38                index_name="embeddings",39            )40        self.collection.load()41 42    def add(self, data) -> str:43        """Add an embedding of data into memory.44 45        Args:46            data (str): The raw text to construct embedding index.47 48        Returns:49            str: log.50        """51        embedding = get_ada_embedding(data)52        result = self.collection.insert([[embedding], [data]])53        _text = (54            "Inserting data into memory at primary key: "55            f"{result.primary_keys[0]}:\n data: {data}"56        )57        return _text58 59    def get(self, data):60        """Return the most relevant data in memory.61        Args:62            data: The data to compare to.63        """64        return self.get_relevant(data, 1)65 66    def clear(self) -> str:67        """Drop the index in memory.68 69        Returns:70            str: log.71        """72        self.collection.drop()73        self.collection = Collection(self.milvus_collection, self.schema)74        self.collection.create_index(75            "embeddings",76            {77                "metric_type": "IP",78                "index_type": "HNSW",79                "params": {"M": 8, "efConstruction": 64},80            },81            index_name="embeddings",82        )83        self.collection.load()84        return "Obliviated"85 86    def get_relevant(self, data: str, num_relevant: int = 5):87        """Return the top-k relevant data in memory.88        Args:89            data: The data to compare to.90            num_relevant (int, optional): The max number of relevant data.91                Defaults to 5.92 93        Returns:94            list: The top-k relevant data.95        """96        # search the embedding and return the most relevant text.97        embedding = get_ada_embedding(data)98        search_params = {99            "metrics_type": "IP",100            "params": {"nprobe": 8},101        }102        result = self.collection.search(103            [embedding],104            "embeddings",105            search_params,106            num_relevant,107            output_fields=["raw_text"],108        )109        return [item.entity.value_of_field("raw_text") for item in result[0]]110 111    def get_stats(self) -> str:112        """113        Returns: The stats of the milvus cache.114        """115        return f"Entities num: {self.collection.num_entities}"116