camel-ai 0.2.67__py3-none-any.whl → 0.2.80a2__py3-none-any.whl

camel/storages/vectordb_storages/oceanbase.py

@@ -74,7 +74,6 @@ class OceanBaseStorage(BaseVectorStorage):
             ObVecClient,
         )
         from pyobvector.client.index_param import (
-            IndexParam,
             IndexParams,
         )
         from pyobvector.schema import VECTOR
@@ -112,20 +111,21 @@ class OceanBaseStorage(BaseVectorStorage):
 
             # Create vector index
             index_params: IndexParams = IndexParams()
-            index_params.add_index_param(
-                IndexParam(
-                    index_name="embedding_idx",
-                    field_name="embedding",
-                    distance=self.distance,
-                    type="hnsw",
-                    m=16,
-                    ef_construction=256,
-                )
+            index_params.add_index(
+                field_name="embedding",
+                index_type="hnsw",
+                index_name="embedding_idx",
+                distance=self.distance,
+                m=16,
+                ef_construction=256,
             )
 
-            self._client.create_vidx_with_vec_index_param(
-                table_name=self.table_name, vidx_param=index_params.params[0]
-            )
+            # Get the first index parameter
+            first_index_param = next(iter(index_params), None)
+            if first_index_param is not None:
+                self._client.create_vidx_with_vec_index_param(
+                    table_name=self.table_name, vidx_param=first_index_param
+                )
 
             logger.info(f"Created table {self.table_name} with vector index")
         else:

camel/storages/vectordb_storages/pgvector.py

@@ -0,0 +1,349 @@
+# ========= Copyright 2023-2024 @ CAMEL-AI.org. All Rights Reserved. =========
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ========= Copyright 2023-2024 @ CAMEL-AI.org. All Rights Reserved. =========
+
+import json
+from typing import Any, Dict, List, Optional
+
+from camel.logger import get_logger
+from camel.storages.vectordb_storages import (
+    BaseVectorStorage,
+    VectorDBQuery,
+    VectorDBQueryResult,
+    VectorDBStatus,
+    VectorRecord,
+)
+from camel.types import VectorDistance
+from camel.utils import dependencies_required
+
+logger = get_logger(__name__)
+
+
+class PgVectorStorage(BaseVectorStorage):
+    r"""PgVectorStorage is an implementation of BaseVectorStorage for
+    PostgreSQL with pgvector extension.
+
+    This class provides methods to add, delete, query, and manage vector
+    records in a PostgreSQL database using the pgvector extension.
+    It supports different distance metrics for similarity search.
+
+    Args:
+        vector_dim (int): The dimension of the vectors to be stored.
+        conn_info (Dict[str, Any]): Connection information for
+            psycopg2.connect.
+        table_name (str, optional): Name of the table to store vectors.
+            (default: :obj:`None`)
+        distance (VectorDistance, optional): Distance metric for vector
+            comparison. (default: :obj:`VectorDistance.COSINE`)
+    """
+
+    @dependencies_required('psycopg', 'pgvector')
+    def __init__(
+        self,
+        vector_dim: int,
+        conn_info: Dict[str, Any],
+        table_name: Optional[str] = None,
+        distance: VectorDistance = VectorDistance.COSINE,
+        **kwargs: Any,
+    ) -> None:
+        r"""Initialize PgVectorStorage.
+
+        Args:
+            vector_dim (int): The dimension of the vectors.
+            conn_info (Dict[str, Any]): Connection info for psycopg2.connect.
+            table_name (str, optional): Table name. (default: :obj:`None`)
+            distance (VectorDistance, optional): Distance metric.
+                (default: :obj:`VectorDistance.COSINE`)
+        """
+        import psycopg
+        from pgvector.psycopg import register_vector
+
+        if vector_dim <= 0:
+            raise ValueError("vector_dim must be positive")
+
+        self.vector_dim = vector_dim
+        self.conn_info = conn_info
+        self.table_name = table_name or 'vectors'
+        self.distance = distance
+
+        try:
+            self._conn = psycopg.connect(**conn_info)
+            register_vector(self._conn)
+            self._ensure_table()
+            self._ensure_index()
+        except Exception as e:
+            logger.error(f"Failed to initialize PgVectorStorage: {e}")
+            raise
+
+    def _ensure_table(self) -> None:
+        r"""Ensure the vector table exists in the database.
+        Creates the table if it does not exist.
+        """
+        try:
+            from psycopg.sql import SQL, Identifier, Literal
+
+            with self._conn.cursor() as cur:
+                query = SQL("""
+                    CREATE TABLE IF NOT EXISTS {table} (
+                        id VARCHAR PRIMARY KEY,
+                        vector vector({dim}),
+                        payload JSONB
+                    )
+                """).format(
+                    table=Identifier(self.table_name),
+                    dim=Literal(self.vector_dim),
+                )
+                cur.execute(query)
+                self._conn.commit()
+        except Exception as e:
+            logger.error(f"Failed to create table {self.table_name}: {e}")
+            raise
+
+    def _ensure_index(self) -> None:
+        r"""Ensure vector similarity search index exists for better
+        performance.
+        """
+        try:
+            from psycopg.sql import SQL, Identifier
+
+            with self._conn.cursor() as cur:
+                index_name = f"{self.table_name}_vector_idx"
+                query = SQL("""
+                    CREATE INDEX IF NOT EXISTS {index_name} 
+                    ON {table} 
+                    USING hnsw (vector vector_cosine_ops)
+                """).format(
+                    index_name=Identifier(index_name),
+                    table=Identifier(self.table_name),
+                )
+                cur.execute(query)
+                self._conn.commit()
+        except Exception as e:
+            logger.warning(f"Failed to create vector index: {e}")
+
+    def add(self, records: List[VectorRecord], **kwargs: Any) -> None:
+        r"""Add or update vector records in the database.
+
+        Args:
+            records (List[VectorRecord]): List of vector records to
+                add or update.
+        """
+        if not records:
+            return
+
+        try:
+            with self._conn.cursor() as cur:
+                # Use batch insert for better performance
+                batch_data = []
+                for rec in records:
+                    if len(rec.vector) != self.vector_dim:
+                        raise ValueError(
+                            f"Vector dimension mismatch: expected "
+                            f"{self.vector_dim}, got {len(rec.vector)}"
+                        )
+
+                    batch_data.append(
+                        (
+                            rec.id,
+                            rec.vector,
+                            json.dumps(rec.payload)
+                            if rec.payload is not None
+                            else None,
+                        )
+                    )
+
+                # Use executemany for efficient batch insert
+                from psycopg.sql import SQL, Identifier
+
+                query = SQL("""
+                    INSERT INTO {table} (id, vector, payload)
+                    VALUES (%s, %s, %s)
+                    ON CONFLICT (id) DO UPDATE SET 
+                    vector=EXCLUDED.vector, 
+                    payload=EXCLUDED.payload
+                """).format(table=Identifier(self.table_name))
+
+                cur.executemany(query, batch_data)
+                self._conn.commit()
+        except Exception as e:
+            self._conn.rollback()
+            logger.error(f"Failed to add records: {e}")
+            raise
+
+    def delete(self, ids: List[str], **kwargs: Any) -> None:
+        r"""Delete vector records from the database by their IDs.
+
+        Args:
+            ids (List[str]): List of record IDs to delete.
+        """
+        from psycopg.sql import SQL, Identifier
+
+        if not ids:
+            return
+
+        try:
+            with self._conn.cursor() as cur:
+                query = SQL("DELETE FROM {table} WHERE id = ANY(%s)").format(
+                    table=Identifier(self.table_name)
+                )
+                cur.execute(query, (ids,))
+                self._conn.commit()
+        except Exception as e:
+            self._conn.rollback()
+            logger.error(f"Failed to delete records: {e}")
+            raise
+
+    def query(
+        self, query: VectorDBQuery, **kwargs: Any
+    ) -> List[VectorDBQueryResult]:
+        r"""Query the database for the most similar vectors to the given
+        query vector.
+
+        Args:
+            query (VectorDBQuery): Query object containing the query
+                vector and top_k.
+            **kwargs (Any): Additional keyword arguments for the query.
+
+        Returns:
+            List[VectorDBQueryResult]: List of query results sorted by
+                similarity.
+        """
+        if len(query.query_vector) != self.vector_dim:
+            raise ValueError(
+                f"Query vector dimension mismatch: "
+                f"expected {self.vector_dim}, got {len(query.query_vector)}"
+            )
+
+        try:
+            with self._conn.cursor() as cur:
+                # Fix distance metric mapping
+                metric_info = {
+                    VectorDistance.COSINE: ('<=>', 'ASC'),  # Cosine distance
+                    VectorDistance.EUCLIDEAN: (
+                        '<->',
+                        'ASC',
+                    ),  # Euclidean distance
+                    VectorDistance.DOT: (
+                        '<#>',
+                        'DESC',
+                    ),  # Negative dot product (higher is better)
+                }
+
+                if self.distance not in metric_info:
+                    raise ValueError(
+                        f"Unsupported distance metric: {self.distance}"
+                    )
+
+                metric, order = metric_info[self.distance]
+
+                from psycopg.sql import SQL, Identifier, Literal
+
+                query_sql = SQL("""
+                    SELECT id, vector, payload, (vector {} %s::vector) 
+                    AS similarity
+                    FROM {}
+                    ORDER BY similarity {}
+                    LIMIT %s
+                """).format(
+                    Literal(metric),
+                    Identifier(self.table_name),
+                    Literal(order),
+                )
+
+                cur.execute(query_sql, (query.query_vector, query.top_k))
+                results = []
+                for row in cur.fetchall():
+                    id, vector, payload, similarity = row
+                    results.append(
+                        VectorDBQueryResult.create(
+                            similarity=float(similarity),
+                            vector=list(vector),
+                            id=id,
+                            payload=payload,
+                        )
+                    )
+                return results
+        except Exception as e:
+            logger.error(f"Failed to query vectors: {e}")
+            raise
+
+    def status(self, **kwargs: Any) -> VectorDBStatus:
+        r"""Get the status of the vector database, including vector
+        dimension and count.
+
+        Args:
+            **kwargs (Any): Additional keyword arguments for the query.
+
+        Returns:
+            VectorDBStatus: Status object with vector dimension and count.
+        """
+        try:
+            with self._conn.cursor() as cur:
+                from psycopg.sql import SQL, Identifier
+
+                query = SQL('SELECT COUNT(*) FROM {}').format(
+                    Identifier(self.table_name)
+                )
+                cur.execute(query)
+                result = cur.fetchone()
+                count = result[0] if result else 0
+                return VectorDBStatus(
+                    vector_dim=self.vector_dim, vector_count=count
+                )
+        except Exception as e:
+            logger.error(f"Failed to get status: {e}")
+            raise
+
+    def clear(self) -> None:
+        r"""Remove all vectors from the storage by truncating the table."""
+        try:
+            with self._conn.cursor() as cur:
+                from psycopg.sql import SQL, Identifier
+
+                query = SQL("TRUNCATE TABLE {table}").format(
+                    table=Identifier(self.table_name)
+                )
+                cur.execute(query)
+                self._conn.commit()
+        except Exception as e:
+            self._conn.rollback()
+            logger.error(f"Failed to clear table: {e}")
+            raise
+
+    def load(self) -> None:
+        r"""Load the collection hosted on cloud service (no-op for pgvector).
+        This method is provided for interface compatibility.
+        """
+        # For PostgreSQL local/managed instances, no loading is required
+        pass
+
+    def close(self) -> None:
+        r"""Close the database connection."""
+        if hasattr(self, '_conn') and self._conn:
+            try:
+                self._conn.close()
+            except Exception as e:
+                logger.warning(f"Error closing connection: {e}")
+
+    def __del__(self) -> None:
+        r"""Ensure connection is closed when object is destroyed."""
+        self.close()
+
+    @property
+    def client(self) -> Any:
+        r"""Provides access to the underlying vector database client.
+
+        Returns:
+            Any: The underlying psycopg connection object.
+        """
+        return self._conn

camel/storages/vectordb_storages/qdrant.py

@@ -14,6 +14,7 @@
 import logging
 from datetime import datetime
 from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union, cast
+from uuid import UUID
 
 if TYPE_CHECKING:
     from qdrant_client import QdrantClient
@@ -251,7 +252,6 @@ class QdrantStorage(BaseVectorStorage):
             else None,
             "vector_count": collection_info.points_count,
             "status": collection_info.status,
-            "vectors_count": collection_info.vectors_count,
             "config": collection_info.config,
         }
 
@@ -304,7 +304,7 @@ class QdrantStorage(BaseVectorStorage):
         """
         from qdrant_client.http.models import PointIdsList, UpdateStatus
 
-        points = cast(List[Union[str, int]], ids)
+        points = cast(List[Union[int, str, UUID]], ids)
 
         op_info = self._client.set_payload(
             collection_name=self.collection_name,
@@ -376,7 +376,7 @@ class QdrantStorage(BaseVectorStorage):
             op_info = self._client.delete(
                 collection_name=self.collection_name,
                 points_selector=PointIdsList(
-                    points=cast(List[Union[int, str]], ids)
+                    points=cast(List[Union[int, str, UUID]], ids)
                 ),
                 **kwargs,
             )