camel-ai 0.1.5.3__tar.gz → 0.1.5.5__tar.gz

{camel_ai-0.1.5.3 → camel_ai-0.1.5.5}/PKG-INFO

@@ -1,12 +1,12 @@
 Metadata-Version: 2.1
 Name: camel-ai
-Version: 0.1.5.3
+Version: 0.1.5.5
 Summary: Communicative Agents for AI Society Study
 Home-page: https://www.camel-ai.org/
 License: Apache-2.0
 Keywords: communicative-ai,ai-societies,artificial-intelligence,deep-learning,multi-agent-systems,cooperative-ai,natural-language-processing,large-language-models
 Author: CAMEL-AI.org
-Requires-Python: >=3.8.1,<3.12
+Requires-Python: >=3.9.0,<3.12
 Classifier: License :: OSI Approved :: Apache Software License
 Classifier: Programming Language :: Python :: 3
 Classifier: Programming Language :: Python :: 3.9
@@ -16,6 +16,7 @@ Provides-Extra: all
 Provides-Extra: encoders
 Provides-Extra: graph-storages
 Provides-Extra: huggingface-agent
+Provides-Extra: kv-stroages
 Provides-Extra: model-platforms
 Provides-Extra: retrievers
 Provides-Extra: test
@@ -23,7 +24,7 @@ Provides-Extra: tools
 Provides-Extra: vector-databases
 Requires-Dist: PyMuPDF (>=1.22.5,<2.0.0) ; extra == "tools" or extra == "all"
 Requires-Dist: accelerate (>=0,<1) ; extra == "huggingface-agent" or extra == "all"
-Requires-Dist: anthropic (>=0.28.0,<0.29.0)
+Requires-Dist: anthropic (>=0.29.0,<0.30.0)
 Requires-Dist: beautifulsoup4 (>=4,<5) ; extra == "tools" or extra == "all"
 Requires-Dist: cohere (>=4.56,<5.0) ; extra == "retrievers" or extra == "all"
 Requires-Dist: colorama (>=0,<1)
@@ -31,13 +32,15 @@ Requires-Dist: curl_cffi (==0.6.2)
 Requires-Dist: datasets (>=2,<3) ; extra == "huggingface-agent" or extra == "all"
 Requires-Dist: diffusers (>=0,<1) ; extra == "huggingface-agent" or extra == "all"
 Requires-Dist: discord.py (>=2.3.2,<3.0.0) ; extra == "tools" or extra == "all"
+Requires-Dist: docker (>=7.1.0,<8.0.0) ; extra == "tools" or extra == "all"
 Requires-Dist: docstring-parser (>=0.15,<0.16)
 Requires-Dist: docx2txt (>=0.8,<0.9) ; extra == "tools" or extra == "all"
 Requires-Dist: duckduckgo-search (>=6.1.0,<7.0.0) ; extra == "tools" or extra == "all"
 Requires-Dist: googlemaps (>=4.10.0,<5.0.0) ; extra == "tools" or extra == "all"
-Requires-Dist: imageio (>=2.34.1,<3.0.0) ; extra == "tools" or extra == "all"
+Requires-Dist: imageio[pyav] (>=2.34.2,<3.0.0) ; extra == "tools" or extra == "all"
 Requires-Dist: jsonschema (>=4,<5)
 Requires-Dist: litellm (>=1.38.1,<2.0.0) ; extra == "model-platforms" or extra == "all"
+Requires-Dist: milvus-lite (>=2.4.0,<=2.4.7)
 Requires-Dist: mock (>=5,<6) ; extra == "test"
 Requires-Dist: neo4j (>=5.18.0,<6.0.0) ; extra == "graph-storages" or extra == "all"
 Requires-Dist: newspaper3k (>=0.2.8,<0.3.0) ; extra == "tools" or extra == "all"
@@ -59,13 +62,14 @@ Requires-Dist: pytest (>=7,<8) ; extra == "test"
 Requires-Dist: pytest-asyncio (>=0.23.0,<0.24.0) ; extra == "test"
 Requires-Dist: qdrant-client (>=1.9.0,<2.0.0) ; extra == "vector-databases" or extra == "all"
 Requires-Dist: rank-bm25 (>=0.2.2,<0.3.0) ; extra == "retrievers" or extra == "all"
+Requires-Dist: redis (>=5.0.6,<6.0.0) ; extra == "kv-stroages" or extra == "all"
 Requires-Dist: requests_oauthlib (>=1.3.1,<2.0.0) ; extra == "tools" or extra == "all"
-Requires-Dist: sentence-transformers (>=2.2.2,<3.0.0) ; extra == "encoders" or extra == "all"
+Requires-Dist: sentence-transformers (>=3.0.1,<4.0.0) ; extra == "encoders" or extra == "all"
 Requires-Dist: sentencepiece (>=0,<1) ; extra == "huggingface-agent" or extra == "all"
 Requires-Dist: slack-sdk (>=3.27.2,<4.0.0) ; extra == "tools" or extra == "all"
 Requires-Dist: soundfile (>=0,<1) ; extra == "huggingface-agent" or extra == "all"
 Requires-Dist: tiktoken (>=0.7.0,<0.8.0)
-Requires-Dist: torch (>=1,<2) ; extra == "huggingface-agent" or extra == "all"
+Requires-Dist: torch (>=2,<3) ; extra == "huggingface-agent" or extra == "all"
 Requires-Dist: transformers (>=4,<5) ; extra == "huggingface-agent" or extra == "all"
 Requires-Dist: unstructured[all-docs] (>=0.10.30,<0.11.0) ; extra == "tools" or extra == "all"
 Requires-Dist: wikipedia (>=1,<2) ; extra == "tools" or extra == "all"
@@ -180,13 +184,13 @@ exit
 Install `CAMEL` from source with conda and pip:
 ```sh
 # Create a conda virtual environment
-conda create --name camel python=3.10
+conda create --name camel python=3.9
 
 # Activate CAMEL conda environment
 conda activate camel
 
 # Clone github repo
-git clone -b v0.1.5.3 https://github.com/camel-ai/camel.git
+git clone -b v0.1.5.5 https://github.com/camel-ai/camel.git
 
 # Change directory into project directory
 cd camel
@@ -198,6 +202,10 @@ pip install -e .
 pip install -e .[all] # (Optional)
 ```
 
+### From Docker
+
+Detailed guidance can be find [here](https://github.com/camel-ai/camel/blob/master/.container/README.md)
+
 ## Documentation
 
 [CAMEL package documentation pages](https://camel-ai.github.io/camel/).

{camel_ai-0.1.5.3 → camel_ai-0.1.5.5}/README.md

@@ -104,13 +104,13 @@ exit
 Install `CAMEL` from source with conda and pip:
 ```sh
 # Create a conda virtual environment
-conda create --name camel python=3.10
+conda create --name camel python=3.9
 
 # Activate CAMEL conda environment
 conda activate camel
 
 # Clone github repo
-git clone -b v0.1.5.3 https://github.com/camel-ai/camel.git
+git clone -b v0.1.5.5 https://github.com/camel-ai/camel.git
 
 # Change directory into project directory
 cd camel
@@ -122,6 +122,10 @@ pip install -e .
 pip install -e .[all] # (Optional)
 ```
 
+### From Docker
+
+Detailed guidance can be find [here](https://github.com/camel-ai/camel/blob/master/.container/README.md)
+
 ## Documentation
 
 [CAMEL package documentation pages](https://camel-ai.github.io/camel/).

{camel_ai-0.1.5.3 → camel_ai-0.1.5.5}/camel/__init__.py

@@ -12,7 +12,7 @@
 # limitations under the License.
 # =========== Copyright 2023 @ CAMEL-AI.org. All Rights Reserved. ===========
 
-__version__ = '0.1.5'
+__version__ = '0.1.5.5'
 
 __all__ = [
     '__version__',

{camel_ai-0.1.5.3 → camel_ai-0.1.5.5}/camel/agents/knowledge_graph_agent.py

@@ -13,7 +13,10 @@
 # =========== Copyright 2023 @ CAMEL-AI.org. All Rights Reserved. ===========
 from typing import Optional, Union
 
-from unstructured.documents.elements import Element
+try:
+    from unstructured.documents.elements import Element
+except ImportError:
+    Element = None
 
 from camel.agents import ChatAgent
 from camel.messages import BaseMessage

{camel_ai-0.1.5.3 → camel_ai-0.1.5.5}/camel/configs/__init__.py

@@ -14,11 +14,13 @@
 from .anthropic_config import ANTHROPIC_API_PARAMS, AnthropicConfig
 from .base_config import BaseConfig
 from .litellm_config import LITELLM_API_PARAMS, LiteLLMConfig
+from .ollama_config import OLLAMA_API_PARAMS, OllamaConfig
 from .openai_config import (
     OPENAI_API_PARAMS,
     ChatGPTConfig,
     OpenSourceConfig,
 )
+from .zhipuai_config import ZHIPUAI_API_PARAMS, ZhipuAIConfig
 
 __all__ = [
     'BaseConfig',
@@ -29,4 +31,8 @@ __all__ = [
     'OpenSourceConfig',
     'LiteLLMConfig',
     'LITELLM_API_PARAMS',
+    'OllamaConfig',
+    'OLLAMA_API_PARAMS',
+    'ZhipuAIConfig',
+    'ZHIPUAI_API_PARAMS',
 ]

{camel_ai-0.1.5.3 → camel_ai-0.1.5.5}/camel/configs/litellm_config.py

@@ -13,11 +13,14 @@
 # =========== Copyright 2023 @ CAMEL-AI.org. All Rights Reserved. ===========
 from __future__ import annotations
 
-from dataclasses import asdict, dataclass, field
-from typing import List, Optional, Union
+from dataclasses import asdict, dataclass
+from typing import TYPE_CHECKING, List, Optional, Union
 
 from camel.configs.base_config import BaseConfig
 
+if TYPE_CHECKING:
+    from camel.functions import OpenAIFunction
+
 
 @dataclass(frozen=True)
 class LiteLLMConfig(BaseConfig):
@@ -25,9 +28,6 @@ class LiteLLMConfig(BaseConfig):
     LiteLLM API.
 
     Args:
-        model (str): The name of the language model to use for text completion.
-        messages (List): A list of message objects representing the
-            conversation context. (default: [])
         timeout (Optional[Union[float, str]], optional): Request timeout.
             (default: None)
         temperature (Optional[float], optional): Temperature parameter for
@@ -65,12 +65,7 @@ class LiteLLMConfig(BaseConfig):
         deployment_id (Optional[str], optional): Deployment ID. (default: None)
         extra_headers (Optional[dict], optional): Additional headers for the
             request. (default: None)
-        base_url (Optional[str], optional): Base URL for the API. (default:
-            None)
         api_version (Optional[str], optional): API version. (default: None)
-        api_key (Optional[str], optional): API key. (default: None)
-        model_list (Optional[list], optional): List of API base, version,
-            keys. (default: None)
         mock_response (Optional[str], optional): Mock completion response for
             testing or debugging. (default: None)
         custom_llm_provider (Optional[str], optional): Non-OpenAI LLM
@@ -79,8 +74,6 @@ class LiteLLMConfig(BaseConfig):
             (default: None)
     """
 
-    model: str = "gpt-3.5-turbo"
-    messages: List = field(default_factory=list)
     timeout: Optional[Union[float, str]] = None
     temperature: Optional[float] = None
     top_p: Optional[float] = None
@@ -91,20 +84,17 @@ class LiteLLMConfig(BaseConfig):
     max_tokens: Optional[int] = None
     presence_penalty: Optional[float] = None
     frequency_penalty: Optional[float] = None
-    logit_bias: Optional[dict] = field(default_factory=dict)
+    logit_bias: Optional[dict] = None
     user: Optional[str] = None
     response_format: Optional[dict] = None
     seed: Optional[int] = None
-    tools: Optional[List] = field(default_factory=list)
+    tools: Optional[list[OpenAIFunction]] = None
     tool_choice: Optional[Union[str, dict]] = None
     logprobs: Optional[bool] = None
     top_logprobs: Optional[int] = None
     deployment_id: Optional[str] = None
-    extra_headers: Optional[dict] = field(default_factory=dict)
-    base_url: Optional[str] = None
+    extra_headers: Optional[dict] = None
     api_version: Optional[str] = None
-    api_key: Optional[str] = None
-    model_list: Optional[list] = field(default_factory=list)
     mock_response: Optional[str] = None
     custom_llm_provider: Optional[str] = None
     max_retries: Optional[int] = None

camel_ai-0.1.5.5/camel/configs/ollama_config.py

@@ -0,0 +1,85 @@
+# =========== 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 __future__ import annotations
+
+from dataclasses import asdict, dataclass
+from typing import Sequence
+
+from openai._types import NOT_GIVEN, NotGiven
+
+from camel.configs.base_config import BaseConfig
+
+
+@dataclass(frozen=True)
+class OllamaConfig(BaseConfig):
+    r"""Defines the parameters for generating chat completions using OpenAI
+    compatibility
+
+    Reference: https://github.com/ollama/ollama/blob/main/docs/openai.md
+
+    Args:
+        temperature (float, optional): Sampling temperature to use, between
+            :obj:`0` and :obj:`2`. Higher values make the output more random,
+            while lower values make it more focused and deterministic.
+            (default: :obj:`0.2`)
+        top_p (float, optional): An alternative to sampling with temperature,
+            called nucleus sampling, where the model considers the results of
+            the tokens with top_p probability mass. So :obj:`0.1` means only
+            the tokens comprising the top 10% probability mass are considered.
+            (default: :obj:`1.0`)
+        response_format (object, optional): An object specifying the format
+            that the model must output. Compatible with GPT-4 Turbo and all
+            GPT-3.5 Turbo models newer than gpt-3.5-turbo-1106. Setting to
+            {"type": "json_object"} enables JSON mode, which guarantees the
+            message the model generates is valid JSON. Important: when using
+            JSON mode, you must also instruct the model to produce JSON
+            yourself via a system or user message. Without this, the model
+            may generate an unending stream of whitespace until the generation
+            reaches the token limit, resulting in a long-running and seemingly
+            "stuck" request. Also note that the message content may be
+            partially cut off if finish_reason="length", which indicates the
+            generation exceeded max_tokens or the conversation exceeded the
+            max context length.
+        stream (bool, optional): If True, partial message deltas will be sent
+            as data-only server-sent events as they become available.
+            (default: :obj:`False`)
+        stop (str or list, optional): Up to :obj:`4` sequences where the API
+            will stop generating further tokens. (default: :obj:`None`)
+        max_tokens (int, optional): The maximum number of tokens to generate
+            in the chat completion. The total length of input tokens and
+            generated tokens is limited by the model's context length.
+            (default: :obj:`None`)
+        presence_penalty (float, optional): Number between :obj:`-2.0` and
+            :obj:`2.0`. Positive values penalize new tokens based on whether
+            they appear in the text so far, increasing the model's likelihood
+            to talk about new topics. See more information about frequency and
+            presence penalties. (default: :obj:`0.0`)
+        frequency_penalty (float, optional): Number between :obj:`-2.0` and
+            :obj:`2.0`. Positive values penalize new tokens based on their
+            existing frequency in the text so far, decreasing the model's
+            likelihood to repeat the same line verbatim. See more information
+            about frequency and presence penalties. (default: :obj:`0.0`)
+    """
+
+    temperature: float = 0.2
+    top_p: float = 1.0
+    stream: bool = False
+    stop: str | Sequence[str] | NotGiven = NOT_GIVEN
+    max_tokens: int | NotGiven = NOT_GIVEN
+    presence_penalty: float = 0.0
+    response_format: dict | NotGiven = NOT_GIVEN
+    frequency_penalty: float = 0.0
+
+
+OLLAMA_API_PARAMS = {param for param in asdict(OllamaConfig()).keys()}

camel_ai-0.1.5.5/camel/configs/zhipuai_config.py

@@ -0,0 +1,78 @@
+# =========== 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 __future__ import annotations
+
+from dataclasses import asdict, dataclass
+from typing import TYPE_CHECKING, Optional, Sequence
+
+from openai._types import NOT_GIVEN, NotGiven
+
+from camel.configs.base_config import BaseConfig
+
+if TYPE_CHECKING:
+    from camel.functions import OpenAIFunction
+
+
+@dataclass(frozen=True)
+class ZhipuAIConfig(BaseConfig):
+    r"""Defines the parameters for generating chat completions using OpenAI
+    compatibility
+
+    Reference: https://open.bigmodel.cn/dev/api#glm-4v
+
+    Args:
+        temperature (float, optional): Sampling temperature to use, between
+            :obj:`0` and :obj:`2`. Higher values make the output more random,
+            while lower values make it more focused and deterministic.
+            (default: :obj:`0.2`)
+        top_p (float, optional): An alternative to sampling with temperature,
+            called nucleus sampling, where the model considers the results of
+            the tokens with top_p probability mass. So :obj:`0.1` means only
+            the tokens comprising the top 10% probability mass are considered.
+            (default: :obj:`0.6`)
+        stream (bool, optional): If True, partial message deltas will be sent
+            as data-only server-sent events as they become available.
+            (default: :obj:`False`)
+        stop (str or list, optional): Up to :obj:`4` sequences where the API
+            will stop generating further tokens. (default: :obj:`None`)
+        max_tokens (int, optional): The maximum number of tokens to generate
+            in the chat completion. The total length of input tokens and
+            generated tokens is limited by the model's context length.
+            (default: :obj:`None`)
+        tools (list[OpenAIFunction], optional): A list of tools the model may
+            call. Currently, only functions are supported as a tool. Use this
+            to provide a list of functions the model may generate JSON inputs
+            for. A max of 128 functions are supported.
+        tool_choice (Union[dict[str, str], str], optional): Controls which (if
+            any) tool is called by the model. :obj:`"none"` means the model
+            will not call any tool and instead generates a message.
+            :obj:`"auto"` means the model can pick between generating a
+            message or calling one or more tools.  :obj:`"required"` means the
+            model must call one or more tools. Specifying a particular tool
+            via {"type": "function", "function": {"name": "my_function"}}
+            forces the model to call that tool. :obj:`"none"` is the default
+            when no tools are present. :obj:`"auto"` is the default if tools
+            are present.
+    """
+
+    temperature: float = 0.2
+    top_p: float = 0.6
+    stream: bool = False
+    stop: str | Sequence[str] | NotGiven = NOT_GIVEN
+    max_tokens: int | NotGiven = NOT_GIVEN
+    tools: Optional[list[OpenAIFunction]] = None
+    tool_choice: Optional[dict[str, str] | str] = None
+
+
+ZHIPUAI_API_PARAMS = {param for param in asdict(ZhipuAIConfig()).keys()}

{camel_ai-0.1.5.3 → camel_ai-0.1.5.5}/camel/embeddings/base.py

@@ -11,8 +11,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # =========== Copyright 2023 @ CAMEL-AI.org. All Rights Reserved. ===========
+from __future__ import annotations
+
 from abc import ABC, abstractmethod
-from typing import Any, Generic, List, TypeVar
+from typing import Any, Generic, TypeVar
 
 T = TypeVar('T')
 
@@ -23,19 +25,18 @@ class BaseEmbedding(ABC, Generic[T]):
     @abstractmethod
     def embed_list(
         self,
-        objs: List[T],
+        objs: list[T],
         **kwargs: Any,
-    ) -> List[List[float]]:
+    ) -> list[list[float]]:
         r"""Generates embeddings for the given texts.
 
         Args:
-            objs (List[T]): The objects for which to generate the embeddings.
+            objs (list[T]): The objects for which to generate the embeddings.
             **kwargs (Any): Extra kwargs passed to the embedding API.
 
         Returns:
-            List[List[float]]: A list that represents the
-            generated embedding as a list of floating-point numbers or a
-            numpy matrix with embeddings.
+            list[list[float]]: A list that represents the
+                generated embedding as a list of floating-point numbers.
         """
         pass
 
@@ -43,7 +44,7 @@ class BaseEmbedding(ABC, Generic[T]):
         self,
         obj: T,
         **kwargs: Any,
-    ) -> List[float]:
+    ) -> list[float]:
         r"""Generates an embedding for the given text.
 
         Args:
@@ -51,7 +52,7 @@ class BaseEmbedding(ABC, Generic[T]):
             **kwargs (Any): Extra kwargs passed to the embedding API.
 
         Returns:
-            List[float]: A list of floating-point numbers representing the
+            list[float]: A list of floating-point numbers representing the
                 generated embedding.
         """
         return self.embed_list([obj], **kwargs)[0]

{camel_ai-0.1.5.3 → camel_ai-0.1.5.5}/camel/embeddings/openai_embedding.py

@@ -11,24 +11,29 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # =========== Copyright 2023 @ CAMEL-AI.org. All Rights Reserved. ===========
+from __future__ import annotations
+
 import os
-from typing import Any, List, Optional
+from typing import Any
 
-from openai import OpenAI
+from openai import NOT_GIVEN, NotGiven, OpenAI
 
 from camel.embeddings.base import BaseEmbedding
 from camel.types import EmbeddingModelType
-from camel.utils import model_api_key_required
+from camel.utils import api_keys_required
 
 
 class OpenAIEmbedding(BaseEmbedding[str]):
     r"""Provides text embedding functionalities using OpenAI's models.
 
     Args:
-        model (OpenAiEmbeddingModel, optional): The model type to be used for
-            generating embeddings. (default: :obj:`ModelType.ADA_2`)
-        api_key (Optional[str]): The API key for authenticating with the
+        model_type (EmbeddingModelType, optional): The model type to be
+            used for text embeddings.
+            (default: :obj:`TEXT_EMBEDDING_3_SMALL`)
+        api_key (str, optional): The API key for authenticating with the
             OpenAI service. (default: :obj:`None`)
+        dimensions (int, optional): The text embedding output dimensions.
+            (default: :obj:`NOT_GIVEN`)
 
     Raises:
         RuntimeError: If an unsupported model type is specified.
@@ -36,36 +41,44 @@ class OpenAIEmbedding(BaseEmbedding[str]):
 
     def __init__(
         self,
-        model_type: EmbeddingModelType = EmbeddingModelType.ADA_2,
-        api_key: Optional[str] = None,
+        model_type: EmbeddingModelType = (
+            EmbeddingModelType.TEXT_EMBEDDING_3_SMALL
+        ),
+        api_key: str | None = None,
+        dimensions: int | NotGiven = NOT_GIVEN,
     ) -> None:
         if not model_type.is_openai:
             raise ValueError("Invalid OpenAI embedding model type.")
         self.model_type = model_type
-        self.output_dim = model_type.output_dim
+        if dimensions == NOT_GIVEN:
+            self.output_dim = model_type.output_dim
+        else:
+            assert isinstance(dimensions, int)
+            self.output_dim = dimensions
         self._api_key = api_key or os.environ.get("OPENAI_API_KEY")
         self.client = OpenAI(timeout=60, max_retries=3, api_key=self._api_key)
 
-    @model_api_key_required
+    @api_keys_required("OPENAI_API_KEY")
     def embed_list(
         self,
-        objs: List[str],
+        objs: list[str],
         **kwargs: Any,
-    ) -> List[List[float]]:
+    ) -> list[list[float]]:
         r"""Generates embeddings for the given texts.
 
         Args:
-            objs (List[str]): The texts for which to generate the embeddings.
+            objs (list[str]): The texts for which to generate the embeddings.
             **kwargs (Any): Extra kwargs passed to the embedding API.
 
         Returns:
-            List[List[float]]: A list that represents the generated embedding
+            list[list[float]]: A list that represents the generated embedding
                 as a list of floating-point numbers.
         """
         # TODO: count tokens
         response = self.client.embeddings.create(
             input=objs,
             model=self.model_type.value,
+            dimensions=self.output_dim,
             **kwargs,
         )
         return [data.embedding for data in response.data]

{camel_ai-0.1.5.3 → camel_ai-0.1.5.5}/camel/embeddings/sentence_transformers_embeddings.py

@@ -11,51 +11,63 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # =========== Copyright 2023 @ CAMEL-AI.org. All Rights Reserved. ===========
-from typing import Any, List
+from __future__ import annotations
+
+from typing import Any
+
+from numpy import ndarray
 
 from camel.embeddings.base import BaseEmbedding
 
 
 class SentenceTransformerEncoder(BaseEmbedding[str]):
-    r"""This class provides functionalities to generate embeddings
-    using a specified model from `Sentence Transformers`.
+    r"""This class provides functionalities to generate text
+    embeddings using `Sentence Transformers`.
 
     References:
         https://www.sbert.net/
     """
 
-    def __init__(self, model_name: str = 'intfloat/e5-large-v2'):
+    def __init__(
+        self,
+        model_name: str = "intfloat/e5-large-v2",
+        **kwargs,
+    ):
         r"""Initializes the: obj: `SentenceTransformerEmbedding` class
         with the specified transformer model.
 
         Args:
             model_name (str, optional): The name of the model to use.
-                                        Defaults to `intfloat/e5-large-v2`.
+                (default: :obj:`intfloat/e5-large-v2`)
+            **kwargs (optional): Additional arguments of
+                :class:`SentenceTransformer`, such as :obj:`prompts` etc.
         """
         from sentence_transformers import SentenceTransformer
 
-        self.model = SentenceTransformer(model_name)
+        self.model = SentenceTransformer(model_name, **kwargs)
 
     def embed_list(
         self,
-        objs: List[str],
+        objs: list[str],
         **kwargs: Any,
-    ) -> List[List[float]]:
+    ) -> list[list[float]]:
         r"""Generates embeddings for the given texts using the model.
 
         Args:
-            objs (List[str]): The texts for which to generate the
-            embeddings.
+            objs (list[str]): The texts for which to generate the
+                embeddings.
 
         Returns:
-            List[List[float]]: A list that represents the generated embedding
+            list[list[float]]: A list that represents the generated embedding
                 as a list of floating-point numbers.
         """
         if not objs:
             raise ValueError("Input text list is empty")
-        return self.model.encode(
+        embeddings = self.model.encode(
             objs, normalize_embeddings=True, **kwargs
-        ).tolist()
+        )
+        assert isinstance(embeddings, ndarray)
+        return embeddings.tolist()
 
     def get_output_dim(self) -> int:
         r"""Returns the output dimension of the embeddings.
@@ -63,4 +75,6 @@ class SentenceTransformerEncoder(BaseEmbedding[str]):
         Returns:
             int: The dimensionality of the embeddings.
         """
-        return self.model.get_sentence_embedding_dimension()
+        output_dim = self.model.get_sentence_embedding_dimension()
+        assert isinstance(output_dim, int)
+        return output_dim