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

Code before:
            port=port
        )
        self.cur = self.conn.cursor()

        collections = self.list_cols()
        if collection_name not in collections:
            self.create_col(collection_name, embedding_model_dims)
        
    def create_col(self, name, embedding_model_dims):
        """
        Create a new collection (table in PostgreSQL).

        Args:
            name (str): Name of the collection.
            embedding_model_dims (int, optional): Dimension of the embedding vector.
        """
        self.cur.execute(f"""
            CREATE TABLE IF NOT EXISTS {name} (
                id UUID PRIMARY KEY,
                vector vector({embedding_model_dims}),
                payload JSONB
            );
        """)
        self.conn.commit()

    def insert(self, name, vectors, payloads = None, ids = None):
        """
        Insert vectors into a collection.

        Args:
            name (str): Name of the collection.
            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.
        """
        json_payloads = [json.dumps(payload) for payload in payloads]

        data = [(id, vector, payload) for id, vector, payload in zip(ids, vectors, json_payloads)]
        execute_values(self.cur, f"INSERT INTO {name} (id, vector, payload) VALUES %s", data)
        self.conn.commit()

    def search(self, name, query, limit = 5, filters = None):
        """
        Search for similar vectors.

        Args:
            name (str): Name of the collection.
            query (List[float]): Query vector.
            limit (int, optional): Number of results to return. Defaults to 5.
            filters (Dict, optional): Filters to apply to the search. Defaults to None.


Review comment (suggestion):
Don't need to pass `name` again and again to all methods I believe. Please correct me if I am missing something. 

Code after:
            password=password,
            host=host,
            port=port
        )
        self.cur = self.conn.cursor()

        collections = self.list_cols()
        if collection_name not in collections:
            self.create_col(embedding_model_dims)
        
    def create_col(self, embedding_model_dims):
        """
        Create a new collection (table in PostgreSQL).

        Args:
            name (str): Name of the collection.
            embedding_model_dims (int, optional): Dimension of the embedding vector.
        """
        self.cur.execute(f"""
            CREATE TABLE IF NOT EXISTS {self.collection_name} (
                id UUID PRIMARY KEY,
                vector vector({embedding_model_dims}),
                payload JSONB
            );
        """)
        self.conn.commit()

    def insert(self, vectors, payloads = None, ids = None):
        """
        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.
        """
        json_payloads = [json.dumps(payload) for payload in payloads]

        data = [(id, vector, payload) for id, vector, payload in zip(ids, vectors, json_payloads)]
        execute_values(self.cur, f"INSERT INTO {self.collection_name} (id, vector, payload) VALUES %s", data)
        self.conn.commit()

    def search(self, query, limit = 5, filters = None):
        """
        Search for similar vectors.

        Args:
            query (List[float]): Query vector.
            limit (int, optional): Number of results to return. Defaults to 5.
            filters (Dict, optional): Filters to apply to the search. Defaults to None.