camel-ai 0.1.1__py3-none-any.whl → 0.1.3__py3-none-any.whl

camel/storages/key_value_storages/json.py

@@ -13,29 +13,24 @@
 # =========== Copyright 2023 @ CAMEL-AI.org. All Rights Reserved. ===========
 
 import json
+from enum import EnumMeta
 from pathlib import Path
-from typing import Any, Dict, List, Optional
+from typing import Any, ClassVar, Dict, List, Optional
 
 from camel.storages.key_value_storages import BaseKeyValueStorage
-from camel.types import (
-    ModelType,
-    OpenAIBackendRole,
-    RoleType,
-    TaskType,
-    VectorDistance,
-)
+from camel.types import ModelType, OpenAIBackendRole, RoleType, TaskType
 
 
 class _CamelJSONEncoder(json.JSONEncoder):
     r"""A custom JSON encoder for serializing specifically enumerated types.
     Ensures enumerated types can be stored in and retrieved from JSON format.
     """
-    CAMEL_ENUMS = {
+
+    CAMEL_ENUMS: ClassVar[Dict[str, EnumMeta]] = {
         "RoleType": RoleType,
         "TaskType": TaskType,
         "ModelType": ModelType,
         "OpenAIBackendRole": OpenAIBackendRole,
-        "VectorDistance": VectorDistance,
     }
 
     def default(self, obj) -> Any:
@@ -75,7 +70,8 @@ class JsonStorage(BaseKeyValueStorage):
         """
         with self.json_path.open("a") as f:
             f.writelines(
-                [json.dumps(r, cls=_CamelJSONEncoder) + "\n" for r in records])
+                [json.dumps(r, cls=_CamelJSONEncoder) + "\n" for r in records]
+            )
 
     def load(self) -> List[Dict[str, Any]]:
         r"""Loads all stored records from the key-value storage system.
@@ -91,7 +87,6 @@ class JsonStorage(BaseKeyValueStorage):
             ]
 
     def clear(self) -> None:
-        r"""Removes all records from the key-value storage system.
-        """
+        r"""Removes all records from the key-value storage system."""
         with self.json_path.open("w"):
             pass

camel/storages/vectordb_storages/__init__.py

@@ -0,0 +1,33 @@
+# =========== Copyright 2023 @ 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 @ CAMEL-AI.org. All Rights Reserved. ===========
+
+from .base import (
+    BaseVectorStorage,
+    VectorDBQuery,
+    VectorDBQueryResult,
+    VectorDBStatus,
+    VectorRecord,
+)
+from .milvus import MilvusStorage
+from .qdrant import QdrantStorage
+
+__all__ = [
+    'BaseVectorStorage',
+    'VectorDBQuery',
+    'VectorDBQueryResult',
+    'QdrantStorage',
+    'MilvusStorage',
+    'VectorRecord',
+    'VectorDBStatus',
+]

camel/storages/vectordb_storages/base.py

@@ -0,0 +1,202 @@
+# =========== Copyright 2023 @ 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 @ CAMEL-AI.org. All Rights Reserved. ===========
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass, field
+from typing import Any, Dict, List, Optional
+from uuid import uuid4
+
+
+@dataclass
+class VectorRecord:
+    r"""Encapsulates information about a vector's unique identifier and its
+    payload, which is primarily used as a data transfer object when saving
+    to vector storage.
+
+    Attributes:
+        vector (List[float]): The numerical representation of the vector.
+        id (str, optional): A unique identifier for the vector. If not
+            provided, an random uuid will be assigned.
+        payload (Optional[Dict[str, Any]], optional): Any additional metadata
+            or information related to the vector. (default: :obj:`None`)
+    """
+
+    vector: List[float]
+    id: str = field(default_factory=lambda: str(uuid4()))
+    payload: Optional[Dict[str, Any]] = None
+
+
+@dataclass
+class VectorDBQuery:
+    r"""Represents a query to a vector database.
+
+    Attributes:
+        query_vector (List[float]): The numerical representation of the query
+            vector.
+        top_k (int, optional): The number of top similar vectors to retrieve
+            from the database. (default: :obj:`1`)
+    """
+
+    query_vector: List[float]
+    top_k: int = 1
+
+
+@dataclass
+class VectorDBQueryResult:
+    r"""Encapsulates the result of a query against a vector database.
+
+    Attributes:
+        record (VectorRecord): The target vector record.
+        similarity (float): The similarity score between the query vector and
+            the record.
+    """
+
+    record: VectorRecord
+    similarity: float
+
+    @classmethod
+    def construct(
+        cls,
+        similarity: float,
+        vector: List[float],
+        id: str,
+        payload: Optional[Dict[str, Any]] = None,
+    ) -> "VectorDBQueryResult":
+        r"""A class method to construct a `VectorDBQueryResult` instance."""
+        return cls(
+            record=VectorRecord(vector, id, payload),
+            similarity=similarity,
+        )
+
+
+@dataclass
+class VectorDBStatus:
+    r"""Vector database status.
+
+    Attributes:
+        vector_dim (int): The dimention of stored vectors.
+        vector_count (int): The number of stored vectors.
+
+    """
+
+    vector_dim: int
+    vector_count: int
+
+
+class BaseVectorStorage(ABC):
+    r"""An abstract base class for vector storage systems."""
+
+    @abstractmethod
+    def add(
+        self,
+        records: List[VectorRecord],
+        **kwargs: Any,
+    ) -> None:
+        r"""Saves a list of vector records to the storage.
+
+        Args:
+            records (List[VectorRecord]): List of vector records to be saved.
+            **kwargs (Any): Additional keyword arguments.
+
+        Raises:
+            RuntimeError: If there is an error during the saving process.
+        """
+        pass
+
+    @abstractmethod
+    def delete(
+        self,
+        ids: List[str],
+        **kwargs: Any,
+    ) -> None:
+        r"""Deletes a list of vectors identified by their IDs from the storage.
+
+        Args:
+            ids (List[str]): List of unique identifiers for the vectors to be
+                deleted.
+            **kwargs (Any): Additional keyword arguments.
+
+        Raises:
+            RuntimeError: If there is an error during the deletion process.
+        """
+        pass
+
+    @abstractmethod
+    def status(self) -> VectorDBStatus:
+        r"""Returns status of the vector database.
+
+        Returns:
+            VectorDBStatus: The vector database status.
+        """
+        pass
+
+    @abstractmethod
+    def query(
+        self,
+        query: VectorDBQuery,
+        **kwargs: Any,
+    ) -> List[VectorDBQueryResult]:
+        r"""Searches for similar vectors in the storage based on the provided query.
+
+        Args:
+            query (VectorDBQuery): The query object containing the search
+                vector and the number of top similar vectors to retrieve.
+            **kwargs (Any): Additional keyword arguments.
+
+        Returns:
+            List[VectorDBQueryResult]: A list of vectors retrieved from the
+                storage based on similarity to the query vector.
+        """
+        pass
+
+    @abstractmethod
+    def clear(self) -> None:
+        r"""Remove all vectors from the storage."""
+        pass
+
+    @abstractmethod
+    def load(self) -> None:
+        r"""Load the collection hosted on cloud service."""
+        pass
+
+    @property
+    @abstractmethod
+    def client(self) -> Any:
+        r"""Provides access to the underlying vector database client."""
+        pass
+
+    def get_payloads_by_vector(
+        self,
+        vector: List[float],
+        top_k: int,
+    ) -> List[Dict[str, Any]]:
+        r"""Returns payloads of top k vector records that closest to the given
+        vector.
+
+        This function is a wrapper of `BaseVectorStorage.query`.
+
+        Args:
+            vector (List[float]): The search vector.
+            top_k (int): The number of top similer vectors.
+
+        Returns:
+            List[List[Dict[str, Any]]]: A list of vector payloads retrieved
+                from the storage based on similarity to the query vector.
+        """
+        results = self.query(VectorDBQuery(vector, top_k))
+        return [
+            result.record.payload
+            for result in results
+            if result.record.payload is not None
+        ]

camel/storages/vectordb_storages/milvus.py

@@ -0,0 +1,396 @@
+# =========== Copyright 2023 @ 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 @ CAMEL-AI.org. All Rights Reserved. ===========
+import logging
+import re
+from datetime import datetime
+from typing import Any, Dict, List, Optional, Tuple
+
+from camel.storages.vectordb_storages import (
+    BaseVectorStorage,
+    VectorDBQuery,
+    VectorDBQueryResult,
+    VectorDBStatus,
+    VectorRecord,
+)
+
+logger = logging.getLogger(__name__)
+
+
+class MilvusStorage(BaseVectorStorage):
+    r"""An implementation of the `BaseVectorStorage` for interacting with
+    Milvus, a cloud-native vector search engine.
+
+    The detailed information about Milvus is available at:
+    `Milvus <https://milvus.io/docs/overview.md/>`_
+
+    Args:
+        vector_dim (int): The dimenstion of storing vectors.
+        url_and_api_key (Tuple[str, str]): Tuple containing
+           the URL and API key for connecting to a remote Milvus instance.
+           URL maps to Milvus uri concept, typically "endpoint:port".
+           API key maps to Milvus token concept, for self-hosted it's
+           "username:pwd", for Zilliz Cloud (fully-managed Milvus) it's API
+           Key.
+        collection_name (Optional[str], optional): Name for the collection in
+            the Milvus. If not provided, set it to the current time with iso
+            format. (default: :obj:`None`)
+        **kwargs (Any): Additional keyword arguments for initializing
+            `MilvusClient`.
+
+    Raises:
+        ImportError: If `pymilvus` package is not installed.
+    """
+
+    def __init__(
+        self,
+        vector_dim: int,
+        url_and_api_key: Tuple[str, str],
+        collection_name: Optional[str] = None,
+        **kwargs: Any,
+    ) -> None:
+        try:
+            from pymilvus import MilvusClient
+        except ImportError as exc:
+            raise ImportError(
+                "Please install `pymilvus` first. You can install it by "
+                "running `pip install pymilvus`."
+            ) from exc
+
+        self._client: MilvusClient
+        self._create_client(url_and_api_key, **kwargs)
+        self.vector_dim = vector_dim
+        self.collection_name = (
+            collection_name or self._generate_collection_name()
+        )
+        self._check_and_create_collection()
+
+    def _create_client(
+        self,
+        url_and_api_key: Tuple[str, str],
+        **kwargs: Any,
+    ) -> None:
+        r"""Initializes the Milvus client with the provided connection details.
+
+        Args:
+            url_and_api_key (Tuple[str, str]): The URL and API key for the
+                Milvus server.
+            **kwargs: Additional keyword arguments passed to the Milvus client.
+        """
+        from pymilvus import MilvusClient
+
+        self._client = MilvusClient(
+            uri=url_and_api_key[0],
+            token=url_and_api_key[1],
+            **kwargs,
+        )
+
+    def _check_and_create_collection(self) -> None:
+        r"""Checks if the specified collection exists in Milvus and creates it
+        if it doesn't, ensuring it matches the specified vector dimensionality.
+        """
+        if self._collection_exists(self.collection_name):
+            in_dim = self._get_collection_info(self.collection_name)[
+                "vector_dim"
+            ]
+            if in_dim != self.vector_dim:
+                # The name of collection has to be confirmed by the user
+                raise ValueError(
+                    "Vector dimension of the existing collection "
+                    f'"{self.collection_name}" ({in_dim}) is different from '
+                    f"the given embedding dim ({self.vector_dim})."
+                )
+        else:
+            self._create_collection(
+                collection_name=self.collection_name,
+            )
+
+    def _create_collection(
+        self,
+        collection_name: str,
+        **kwargs: Any,
+    ) -> None:
+        r"""Creates a new collection in the database.
+
+        Args:
+            collection_name (str): Name of the collection to be created.
+            **kwargs (Any): Additional keyword arguments pass to create
+                collection.
+        """
+
+        from pymilvus import DataType
+
+        # Set the schema
+        schema = self._client.create_schema(
+            auto_id=False,
+            enable_dynamic_field=True,
+            description='collection schema',
+        )
+
+        schema.add_field(
+            field_name="id",
+            datatype=DataType.VARCHAR,
+            descrition='A unique identifier for the vector',
+            is_primary=True,
+            max_length=65535,
+        )
+        # max_length reference: https://milvus.io/docs/limitations.md
+        schema.add_field(
+            field_name="vector",
+            datatype=DataType.FLOAT_VECTOR,
+            description='The numerical representation of the vector',
+            dim=self.vector_dim,
+        )
+        schema.add_field(
+            field_name="payload",
+            datatype=DataType.JSON,
+            description=(
+                'Any additional metadata or information related' 'to the vector'
+            ),
+        )
+
+        # Create the collection
+        self._client.create_collection(
+            collection_name=collection_name,
+            schema=schema,
+            **kwargs,
+        )
+
+        # Set the index of the parameters
+        index_params = self._client.prepare_index_params()
+
+        index_params.add_index(
+            field_name="vector",
+            metric_type="COSINE",
+            index_type="AUTOINDEX",
+            index_name="vector_index",
+        )
+
+        self._client.create_index(
+            collection_name=collection_name, index_params=index_params
+        )
+
+    def _delete_collection(
+        self,
+        collection_name: str,
+    ) -> None:
+        r"""Deletes an existing collection from the database.
+
+        Args:
+            collection (str): Name of the collection to be deleted.
+        """
+        self._client.drop_collection(collection_name=collection_name)
+
+    def _collection_exists(self, collection_name: str) -> bool:
+        r"""Checks whether a collection with the specified name exists in the
+        database.
+
+        Args:
+            collection_name (str): The name of the collection to check.
+
+        Returns:
+            bool: True if the collection exists, False otherwise.
+        """
+        return self._client.has_collection(collection_name)
+
+    def _generate_collection_name(self) -> str:
+        r"""Generates a unique name for a new collection based on the current
+        timestamp. Milvus collection names can only contain alphanumeric
+        characters and underscores.
+
+        Returns:
+            str: A unique, valid collection name.
+        """
+        timestamp = datetime.now().isoformat()
+        transformed_name = re.sub(r'[^a-zA-Z0-9_]', '_', timestamp)
+        valid_name = "Time" + transformed_name
+        return valid_name
+
+    def _get_collection_info(self, collection_name: str) -> Dict[str, Any]:
+        r"""Retrieves details of an existing collection.
+
+        Args:
+            collection_name (str): Name of the collection to be checked.
+
+        Returns:
+            Dict[str, Any]: A dictionary containing details about the
+                collection.
+        """
+        vector_count = self._client.get_collection_stats(collection_name)[
+            'row_count'
+        ]
+        collection_info = self._client.describe_collection(collection_name)
+        collection_id = collection_info['collection_id']
+
+        dim_value = next(
+            (
+                field['params']['dim']
+                for field in collection_info['fields']
+                if field['description']
+                == 'The numerical representation of the vector'
+            ),
+            None,
+        )
+
+        return {
+            "id": collection_id,  # the id of the collection
+            "vector_count": vector_count,  # the number of the vector
+            "vector_dim": dim_value,  # the dimension of the vector
+        }
+
+    def _validate_and_convert_vectors(
+        self, records: List[VectorRecord]
+    ) -> List[dict]:
+        r"""Validates and converts VectorRecord instances to the format
+        expected by Milvus.
+
+        Args:
+            records (List[VectorRecord]): List of vector records to validate
+            and convert.
+
+        Returns:
+            List[dict]: A list of dictionaries formatted for Milvus insertion.
+        """
+
+        validated_data = []
+
+        for record in records:
+            record_dict = {
+                "id": record.id,
+                "payload": record.payload if record.payload is not None else '',
+                "vector": record.vector,
+            }
+            validated_data.append(record_dict)
+
+        return validated_data
+
+    def add(
+        self,
+        records: List[VectorRecord],
+        **kwargs,
+    ) -> None:
+        r"""Adds a list of vectors to the specified collection.
+
+        Args:
+            records (List[VectorRecord]): List of vectors to be added.
+            **kwargs (Any): Additional keyword arguments pass to insert.
+
+        Raises:
+            RuntimeError: If there was an error in the addition process.
+        """
+        validated_records = self._validate_and_convert_vectors(records)
+
+        op_info = self._client.insert(
+            collection_name=self.collection_name,
+            data=validated_records,
+            **kwargs,
+        )
+        logger.debug(f"Successfully added vectors in Milvus: {op_info}")
+
+    def delete(
+        self,
+        ids: List[str],
+        **kwargs: Any,
+    ) -> None:
+        r"""Deletes a list of vectors identified by their IDs from the
+        storage. If unsure of ids you can first query the collection to grab
+        the corresponding data.
+
+        Args:
+            ids (List[str]): List of unique identifiers for the vectors to be
+                deleted.
+            **kwargs (Any): Additional keyword arguments passed to delete.
+
+        Raises:
+            RuntimeError: If there is an error during the deletion process.
+        """
+
+        op_info = self._client.delete(
+            collection_name=self.collection_name, pks=ids, **kwargs
+        )
+        logger.debug(f"Successfully deleted vectors in Milvus: {op_info}")
+
+    def status(self) -> VectorDBStatus:
+        r"""Retrieves the current status of the Milvus collection. This method
+        provides information about the collection, including its vector
+        dimensionality and the total number of vectors stored.
+
+        Returns:
+            VectorDBStatus: An object containing information about the
+                collection's status.
+        """
+        status = self._get_collection_info(self.collection_name)
+        return VectorDBStatus(
+            vector_dim=status["vector_dim"],
+            vector_count=status["vector_count"],
+        )
+
+    def query(
+        self,
+        query: VectorDBQuery,
+        **kwargs: Any,
+    ) -> List[VectorDBQueryResult]:
+        r"""Searches for similar vectors in the storage based on the provided
+        query.
+
+        Args:
+            query (VectorDBQuery): The query object containing the search
+                vector and the number of top similar vectors to retrieve.
+            **kwargs (Any): Additional keyword arguments passed to search.
+
+        Returns:
+            List[VectorDBQueryResult]: A list of vectors retrieved from the
+                storage based on similarity to the query vector.
+        """
+        search_result = self._client.search(
+            collection_name=self.collection_name,
+            data=[query.query_vector],
+            limit=query.top_k,
+            output_fields=['vector', 'payload'],
+            **kwargs,
+        )
+        query_results = []
+        for point in search_result:
+            query_results.append(
+                VectorDBQueryResult.construct(
+                    similarity=(point[0]['distance']),
+                    id=str(point[0]['id']),
+                    payload=(point[0]['entity'].get('payload')),
+                    vector=point[0]['entity'].get('vector'),
+                )
+            )
+
+        return query_results
+
+    def clear(self) -> None:
+        r"""Removes all vectors from the Milvus collection. This method
+        deletes the existing collection and then recreates it with the same
+        schema to effectively remove all stored vectors.
+        """
+        self._delete_collection(self.collection_name)
+        self._create_collection(collection_name=self.collection_name)
+
+    def load(self) -> None:
+        r"""Load the collection hosted on cloud service."""
+        self._client.load_collection(self.collection_name)
+
+    @property
+    def client(self) -> Any:
+        r"""Provides direct access to the Milvus client. This property allows
+        for direct interactions with the Milvus client for operations that are
+        not covered by the `MilvusStorage` class.
+
+        Returns:
+            Any: The Milvus client instance.
+        """
+        return self._client