Repository: mem0ai/mem0
File: mem0/vector_stores/milvus.py (Python)

Code before:
    def create_col(
        self, collection_name : str, vector_size : str, metric_type : MetricType = MetricType.COSINE
    ) -> None:
        """Create a new collection with index_type IVF_FLAT.

        Args:
            collection_name (str): Name of the collection (defaults to mem0).
            vector_size (str): Dimensions of the embedding model (defaults to 1536).
            metric_type (MetricType, optional): etric type for similarity search. Defaults to MetricType.COSINE.
        """

        if self.client.has_collection(collection_name):
            print(f"Collection {collection_name} already exists. Skipping creation.")
        else:
            fields = [
                FieldSchema(name="id", dtype=DataType.VARCHAR, is_primary=True, max_length=512),
                FieldSchema(name="vectors", dtype=DataType.FLOAT_VECTOR, dim=vector_size),
                FieldSchema(name="metadata", dtype=DataType.JSON),
            ]
            
            schema = CollectionSchema(fields, enable_dynamic_field=True)

            index = self.client.prepare_index_params(
                field_name="vectors",
                metric_type=metric_type,
                index_type="IVF_FLAT",
                index_name="vector_index",
                params={ "nlist": 128 }
            )
            
            self.client.create_collection(collection_name=collection_name, schema=schema, index_params=index)
            
            
    def insert(self, ids, vectors, payloads, **kwargs: Optional[dict[str, any]]):
        """Insert vectors into a collection.

        Args:
            vectors (List[List[float]]): List of vectors to insert.
            payloads (List[Dict], optional): List of payloads corresponding to vectors.
            ids (List[str], optional): List of IDs corresponding to vectors.
        """
        for idx, embedding, metadata in zip(ids, vectors, payloads):
            data = {"id": idx, "vectors": embedding, "metadata": metadata}
            self.client.insert(collection_name=self.collection_name, data=data, **kwargs)


    def _create_filter(self, filters: dict):
        """Prepare filters for efficient query. 

        Args:
            filters (dict): filters [user_id, agent_id, run_id]

Review comment (suggestion):
maybe "AUTOINDEX" is better , because milvus lite only support 'AUTOINDEX', 'HNSW', 'FLAT' so far. If users are using milvus lite with `IVF_FLAT` index and the `./milvus.db` as uri, it will raise an error.

Code after:
        
        self.client = MilvusClient(uri=url,token=token)
        
        self.create_col(
            collection_name=self.collection_name,
            vector_size=self.embedding_model_dims,
            metric_type=self.metric_type
        )
       
        
    def create_col(
        self, collection_name : str, vector_size : str, metric_type : MetricType = MetricType.COSINE
    ) -> None:
        """Create a new collection with index_type AUTOINDEX.

        Args:
            collection_name (str): Name of the collection (defaults to mem0).
            vector_size (str): Dimensions of the embedding model (defaults to 1536).
            metric_type (MetricType, optional): etric type for similarity search. Defaults to MetricType.COSINE.
        """

        if self.client.has_collection(collection_name):
            logger.info(f"Collection {collection_name} already exists. Skipping creation.")
        else:
            fields = [
                FieldSchema(name="id", dtype=DataType.VARCHAR, is_primary=True, max_length=512),
                FieldSchema(name="vectors", dtype=DataType.FLOAT_VECTOR, dim=vector_size),
                FieldSchema(name="metadata", dtype=DataType.JSON),
            ]
            
            schema = CollectionSchema(fields, enable_dynamic_field=True)

            index = self.client.prepare_index_params(
                field_name="vectors",
                metric_type=metric_type,
                index_type="AUTOINDEX",
                index_name="vector_index",
                params={ "nlist": 128 }
            )
            
            self.client.create_collection(collection_name=collection_name, schema=schema, index_params=index)
            
            
    def insert(self, ids, vectors, payloads, **kwargs: Optional[dict[str, any]]):
        """Insert vectors into a collection.

        Args:
            vectors (List[List[float]]): List of vectors to insert.
            payloads (List[Dict], optional): List of payloads corresponding to vectors.
            ids (List[str], optional): List of IDs corresponding to vectors.
        """