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

camel/storages/vectordb_storages/qdrant.py

@@ -0,0 +1,371 @@
+# =========== 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 dataclasses import asdict
+from datetime import datetime
+from typing import Any, Dict, List, Optional, Tuple, Union, cast
+
+from camel.storages.vectordb_storages import (
+    BaseVectorStorage,
+    VectorDBQuery,
+    VectorDBQueryResult,
+    VectorDBStatus,
+    VectorRecord,
+)
+from camel.types import VectorDistance
+
+_qdrant_local_client_map: Dict[str, Tuple[Any, int]] = {}
+
+
+class QdrantStorage(BaseVectorStorage):
+    r"""An implementation of the `BaseVectorStorage` for interacting with
+    Qdrant, a vector search engine.
+
+    The detailed information about Qdrant is available at:
+    `Qdrant <https://qdrant.tech/>`_
+
+    Args:
+        vector_dim (int): The dimenstion of storing vectors.
+        collection_name (Optional[str], optional): Name for the collection in
+            the Qdrant. If not provided, set it to the current time with iso
+            format. (default: :obj:`None`)
+        url_and_api_key (Optional[Tuple[str, str]], optional): Tuple containing
+            the URL and API key for connecting to a remote Qdrant instance.
+            (default: :obj:`None`)
+        path (Optional[str], optional): Path to a directory for initializing a
+            local Qdrant client. (default: :obj:`None`)
+        distance (VectorDistance, optional): The distance metric for vector
+            comparison (default: :obj:`VectorDistance.COSINE`)
+        delete_collection_on_del (bool, optional): Flag to determine if the
+            collection should be deleted upon object destruction.
+            (default: :obj:`False`)
+        **kwargs (Any): Additional keyword arguments for initializing
+            `QdrantClient`.
+
+    Notes:
+        - If `url_and_api_key` is provided, it takes priority and the client
+          will attempt to connect to the remote Qdrant instance using the URL
+          endpoint.
+        - If `url_and_api_key` is not provided and `path` is given, the client
+          will use the local path to initialize Qdrant.
+        - If neither `url_and_api_key` nor `path` is provided, the client will
+          be initialized with an in-memory storage (`":memory:"`).
+    """
+
+    def __init__(
+        self,
+        vector_dim: int,
+        collection_name: Optional[str] = None,
+        url_and_api_key: Optional[Tuple[str, str]] = None,
+        path: Optional[str] = None,
+        distance: VectorDistance = VectorDistance.COSINE,
+        delete_collection_on_del: bool = False,
+        **kwargs: Any,
+    ) -> None:
+        try:
+            from qdrant_client import QdrantClient
+        except ImportError as exc:
+            raise ImportError(
+                "Please install `qdrant-client` first. You can install it by "
+                "running `pip install qdrant-client`."
+            ) from exc
+
+        self._client: QdrantClient
+        self._local_path: Optional[str] = None
+        self._create_client(url_and_api_key, path, **kwargs)
+
+        self.vector_dim = vector_dim
+        self.distance = distance
+        self.collection_name = (
+            collection_name or self._generate_collection_name()
+        )
+
+        self._check_and_create_collection()
+
+        self.delete_collection_on_del = delete_collection_on_del
+
+    def __del__(self):
+        r"""Deletes the collection if :obj:`del_collection` is set to
+        :obj:`True`.
+        """
+        # If the client is a local client, decrease count by 1
+        if self._local_path is not None:
+            # if count decrease to 0, remove it from the map
+            _client, _count = _qdrant_local_client_map.pop(self._local_path)
+            if _count > 1:
+                _qdrant_local_client_map[self._local_path] = (
+                    _client,
+                    _count - 1,
+                )
+
+        if (
+            hasattr(self, "delete_collection_on_del")
+            and self.delete_collection_on_del
+        ):
+            self._delete_collection(self.collection_name)
+
+    def _create_client(
+        self,
+        url_and_api_key: Optional[Tuple[str, str]],
+        path: Optional[str],
+        **kwargs: Any,
+    ) -> None:
+        from qdrant_client import QdrantClient
+
+        if url_and_api_key is not None:
+            self._client = QdrantClient(
+                url=url_and_api_key[0],
+                api_key=url_and_api_key[1],
+                **kwargs,
+            )
+        elif path is not None:
+            # Avoid creating a local client multiple times,
+            # which is prohibited by Qdrant
+            self._local_path = path
+            if path in _qdrant_local_client_map:
+                # Store client instance in the map and maintain counts
+                self._client, count = _qdrant_local_client_map[path]
+                _qdrant_local_client_map[path] = (self._client, count + 1)
+            else:
+                self._client = QdrantClient(path=path, **kwargs)
+                _qdrant_local_client_map[path] = (self._client, 1)
+        else:
+            self._client = QdrantClient(":memory:", **kwargs)
+
+    def _check_and_create_collection(self) -> None:
+        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,
+                size=self.vector_dim,
+                distance=self.distance,
+            )
+
+    def _create_collection(
+        self,
+        collection_name: str,
+        size: int,
+        distance: VectorDistance = VectorDistance.COSINE,
+        **kwargs: Any,
+    ) -> None:
+        r"""Creates a new collection in the database.
+
+        Args:
+            collection_name (str): Name of the collection to be created.
+            size (int): Dimensionality of vectors to be stored in this
+                collection.
+            distance (VectorDistance, optional): The distance metric to be used
+                for vector similarity. (default: :obj:`VectorDistance.COSINE`)
+            **kwargs (Any): Additional keyword arguments.
+        """
+        from qdrant_client.http.models import Distance, VectorParams
+
+        distance_map = {
+            VectorDistance.DOT: Distance.DOT,
+            VectorDistance.COSINE: Distance.COSINE,
+            VectorDistance.EUCLIDEAN: Distance.EUCLID,
+        }
+        self._client.recreate_collection(
+            collection_name=collection_name,
+            vectors_config=VectorParams(
+                size=size,
+                distance=distance_map[distance],
+            ),
+            **kwargs,
+        )
+
+    def _delete_collection(
+        self,
+        collection_name: str,
+        **kwargs: Any,
+    ) -> None:
+        r"""Deletes an existing collection from the database.
+
+        Args:
+            collection (str): Name of the collection to be deleted.
+            **kwargs (Any): Additional keyword arguments.
+        """
+        self._client.delete_collection(
+            collection_name=collection_name, **kwargs
+        )
+
+    def _collection_exists(self, collection_name: str) -> bool:
+        r"""Returns wether the collection exists in the database"""
+        for c in self._client.get_collections().collections:
+            if collection_name == c.name:
+                return True
+        return False
+
+    def _generate_collection_name(self) -> str:
+        r"""Generates a collection name if user doesn't provide"""
+        return datetime.now().isoformat()
+
+    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.
+        """
+        from qdrant_client.http.models import VectorParams
+
+        # TODO: check more information
+        collection_info = self._client.get_collection(
+            collection_name=collection_name
+        )
+        vector_config = collection_info.config.params.vectors
+        return {
+            "vector_dim": vector_config.size
+            if isinstance(vector_config, VectorParams)
+            else None,
+            "vector_count": collection_info.points_count,
+            "status": collection_info.status,
+            "vectors_count": collection_info.vectors_count,
+            "config": collection_info.config,
+        }
+
+    def add(
+        self,
+        records: List[VectorRecord],
+        **kwargs,
+    ) -> None:
+        r"""Adds a list of vectors to the specified collection.
+
+        Args:
+            vectors (List[VectorRecord]): List of vectors to be added.
+            **kwargs (Any): Additional keyword arguments.
+
+        Raises:
+            RuntimeError: If there was an error in the addition process.
+        """
+        from qdrant_client.http.models import PointStruct, UpdateStatus
+
+        qdrant_points = [PointStruct(**asdict(p)) for p in records]
+        op_info = self._client.upsert(
+            collection_name=self.collection_name,
+            points=qdrant_points,
+            wait=True,
+            **kwargs,
+        )
+        if op_info.status != UpdateStatus.COMPLETED:
+            raise RuntimeError(
+                "Failed to add vectors in Qdrant, operation info: "
+                f"{op_info}."
+            )
+
+    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.
+        """
+        from qdrant_client.http.models import PointIdsList, UpdateStatus
+
+        points = cast(List[Union[str, int]], ids)
+        op_info = self._client.delete(
+            collection_name=self.collection_name,
+            points_selector=PointIdsList(points=points),
+            wait=True,
+            **kwargs,
+        )
+        if op_info.status != UpdateStatus.COMPLETED:
+            raise RuntimeError(
+                "Failed to delete vectors in Qdrant, operation info: "
+                f"{op_info}"
+            )
+
+    def status(self) -> VectorDBStatus:
+        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.
+
+        Returns:
+            List[VectorDBQueryResult]: A list of vectors retrieved from the
+                storage based on similarity to the query vector.
+        """
+        # TODO: filter
+        search_result = self._client.search(
+            collection_name=self.collection_name,
+            query_vector=query.query_vector,
+            with_payload=True,
+            with_vectors=True,
+            limit=query.top_k,
+            **kwargs,
+        )
+        query_results = []
+        for point in search_result:
+            query_results.append(
+                VectorDBQueryResult.construct(
+                    similarity=point.score,
+                    id=str(point.id),
+                    payload=point.payload,
+                    vector=point.vector,  # type: ignore[arg-type]
+                )
+            )
+
+        return query_results
+
+    def clear(self) -> None:
+        r"""Remove all vectors from the storage."""
+        self._delete_collection(self.collection_name)
+        self._create_collection(
+            collection_name=self.collection_name,
+            size=self.vector_dim,
+            distance=self.distance,
+        )
+
+    def load(self) -> None:
+        r"""Load the collection hosted on cloud service."""
+        pass
+
+    @property
+    def client(self) -> Any:
+        r"""Provides access to the underlying vector database client."""
+        return self._client

camel/terminators/__init__.py

@@ -12,7 +12,7 @@
 # limitations under the License.
 # =========== Copyright 2023 @ CAMEL-AI.org. All Rights Reserved. ===========
 from .base import BaseTerminator
-from .response_terminator import ResponseWordsTerminator, ResponseTerminator
+from .response_terminator import ResponseTerminator, ResponseWordsTerminator
 from .token_limit_terminator import TokenLimitTerminator
 
 __all__ = [

camel/terminators/base.py

@@ -18,7 +18,6 @@ from camel.messages import BaseMessage
 
 
 class BaseTerminator(ABC):
-
     def __init__(self, *args, **kwargs) -> None:
         self._terminated: bool = False
         self._termination_reason: Optional[str] = None
@@ -33,10 +32,10 @@ class BaseTerminator(ABC):
 
 
 class ResponseTerminator(BaseTerminator):
-
     @abstractmethod
     def is_terminated(
-            self, messages: List[BaseMessage]) -> Tuple[bool, Optional[str]]:
+        self, messages: List[BaseMessage]
+    ) -> Tuple[bool, Optional[str]]:
         pass
 
     @abstractmethod

camel/terminators/response_terminator.py

@@ -34,9 +34,12 @@ class ResponseWordsTerminator(ResponseTerminator):
             (default: :obj:`TerminationMode.ANY`)
     """
 
-    def __init__(self, words_dict: Dict[str,
-                                        int], case_sensitive: bool = False,
-                 mode: TerminationMode = TerminationMode.ANY):
+    def __init__(
+        self,
+        words_dict: Dict[str, int],
+        case_sensitive: bool = False,
+        mode: TerminationMode = TerminationMode.ANY,
+    ):
         super().__init__()
         self.words_dict = words_dict
         self.case_sensitive = case_sensitive
@@ -50,11 +53,14 @@ class ResponseWordsTerminator(ResponseTerminator):
         for word in self.words_dict:
             threshold = self.words_dict[word]
             if threshold <= 0:
-                raise ValueError(f"Threshold for word `{word}` should "
-                                 f"be larger than 0, got `{threshold}`")
+                raise ValueError(
+                    f"Threshold for word `{word}` should "
+                    f"be larger than 0, got `{threshold}`"
+                )
 
     def is_terminated(
-            self, messages: List[BaseMessage]) -> Tuple[bool, Optional[str]]:
+        self, messages: List[BaseMessage]
+    ) -> Tuple[bool, Optional[str]]:
         r"""Whether terminate the agent by checking the occurrence
         of specified words reached to preset thresholds.
 
@@ -90,10 +96,12 @@ class ResponseWordsTerminator(ResponseTerminator):
             for word, value in self._word_count_dict[i].items():
                 if value >= self.words_dict[word]:
                     reached += 1
-                    reason = (f"Word `{word}` appears {value} times in the "
-                              f"{i + 1} message of the response which has "
-                              f"reached termination threshold "
-                              f"{self.words_dict[word]}.")
+                    reason = (
+                        f"Word `{word}` appears {value} times in the "
+                        f"{i + 1} message of the response which has "
+                        f"reached termination threshold "
+                        f"{self.words_dict[word]}."
+                    )
                     reasons.append(reason)
             all_reasons.append(reasons)
             num_reached.append(reached)
@@ -108,8 +116,9 @@ class ResponseWordsTerminator(ResponseTerminator):
                     self._terminated = True
                     self._termination_reason = "\n".join(all_reasons[i])
             else:
-                raise ValueError(f"Unsupported termination mode "
-                                 f"`{self.mode}`")
+                raise ValueError(
+                    f"Unsupported termination mode " f"`{self.mode}`"
+                )
         return self._terminated, self._termination_reason
 
     def reset(self):

camel/terminators/token_limit_terminator.py

@@ -13,7 +13,7 @@
 # =========== Copyright 2023 @ CAMEL-AI.org. All Rights Reserved. ===========
 from typing import Optional, Tuple
 
-from camel.terminators import BaseTerminator
+from camel.terminators.base import BaseTerminator
 
 
 class TokenLimitTerminator(BaseTerminator):
@@ -29,8 +29,10 @@ class TokenLimitTerminator(BaseTerminator):
 
     def _validate(self):
         if self.token_limit <= 0:
-            raise ValueError(f"`token_limit` should be a "
-                             f"value larger than 0, got {self.token_limit}.")
+            raise ValueError(
+                f"`token_limit` should be a "
+                f"value larger than 0, got {self.token_limit}."
+            )
 
     def is_terminated(self, num_tokens: int) -> Tuple[bool, Optional[str]]:
         r"""Whether terminate the agent by checking number of

camel/types/__init__.py

@@ -12,24 +12,27 @@
 # limitations under the License.
 # =========== Copyright 2023 @ CAMEL-AI.org. All Rights Reserved. ===========
 from .enums import (
-    RoleType,
+    EmbeddingModelType,
     ModelType,
+    OpenAIBackendRole,
+    OpenAIImageDetailType,
+    OpenAIImageType,
+    RoleType,
+    StorageType,
     TaskType,
     TerminationMode,
-    OpenAIBackendRole,
-    EmbeddingModelType,
     VectorDistance,
 )
 from .openai_types import (
-    Choice,
     ChatCompletion,
+    ChatCompletionAssistantMessageParam,
     ChatCompletionChunk,
+    ChatCompletionFunctionMessageParam,
     ChatCompletionMessage,
     ChatCompletionMessageParam,
     ChatCompletionSystemMessageParam,
     ChatCompletionUserMessageParam,
-    ChatCompletionAssistantMessageParam,
-    ChatCompletionFunctionMessageParam,
+    Choice,
     CompletionUsage,
 )
 
@@ -41,6 +44,7 @@ __all__ = [
     'OpenAIBackendRole',
     'EmbeddingModelType',
     'VectorDistance',
+    'StorageType',
     'Choice',
     'ChatCompletion',
     'ChatCompletionChunk',
@@ -51,4 +55,6 @@ __all__ = [
     'ChatCompletionAssistantMessageParam',
     'ChatCompletionFunctionMessageParam',
     'CompletionUsage',
+    'OpenAIImageType',
+    'OpenAIImageDetailType',
 ]