camel-ai 0.1.5.6__py3-none-any.whl → 0.1.6.1__py3-none-any.whl

camel/utils/token_counting.py

@@ -26,6 +26,8 @@ from PIL import Image
 from camel.types import ModelType, OpenAIImageType, OpenAIVisionDetailType
 
 if TYPE_CHECKING:
+    from mistral_common.protocol.instruct.request import ChatCompletionRequest
+
     from camel.messages import OpenAIMessage
 
 LOW_DETAIL_TOKENS = 85
@@ -37,7 +39,7 @@ EXTRA_TOKENS = 85
 
 
 def messages_to_prompt(messages: List[OpenAIMessage], model: ModelType) -> str:
-    r"""Parse the message list into a single prompt following model-specifc
+    r"""Parse the message list into a single prompt following model-specific
     formats.
 
     Args:
@@ -51,7 +53,12 @@ def messages_to_prompt(messages: List[OpenAIMessage], model: ModelType) -> str:
     system_message = messages[0]["content"]
 
     ret: str
-    if model == ModelType.LLAMA_2 or model == ModelType.LLAMA_3:
+    if model in [
+        ModelType.LLAMA_2,
+        ModelType.LLAMA_3,
+        ModelType.GROQ_LLAMA_3_8B,
+        ModelType.GROQ_LLAMA_3_70B,
+    ]:
         # reference: https://github.com/facebookresearch/llama/blob/cfc3fc8c1968d390eb830e65c63865e980873a06/llama/generation.py#L212
         seps = [" ", " </s><s>"]
         role_map = {"user": "[INST]", "assistant": "[/INST]"}
@@ -74,7 +81,7 @@ def messages_to_prompt(messages: List[OpenAIMessage], model: ModelType) -> str:
             else:
                 ret += role
         return ret
-    elif model == ModelType.VICUNA or model == ModelType.VICUNA_16K:
+    elif model in [ModelType.VICUNA, ModelType.VICUNA_16K]:
         seps = [" ", "</s>"]
         role_map = {"user": "USER", "assistant": "ASSISTANT"}
 
@@ -132,6 +139,40 @@ def messages_to_prompt(messages: List[OpenAIMessage], model: ModelType) -> str:
             else:
                 ret += '<|im_start|>' + role + '\n'
         return ret
+    elif model == ModelType.GROQ_MIXTRAL_8_7B:
+        # Mistral/Mixtral format
+        system_prompt = f"<s>[INST] {system_message} [/INST]\n"
+        ret = system_prompt
+
+        for msg in messages[1:]:
+            if msg["role"] == "user":
+                ret += f"[INST] {msg['content']} [/INST]\n"
+            elif msg["role"] == "assistant":
+                ret += f"{msg['content']}</s>\n"
+
+            if not isinstance(msg['content'], str):
+                raise ValueError(
+                    "Currently multimodal context is not "
+                    "supported by the token counter."
+                )
+
+        return ret.strip()
+    elif model in [ModelType.GROQ_GEMMA_7B_IT, ModelType.GROQ_GEMMA_2_9B_IT]:
+        # Gemma format
+        ret = f"<bos>{system_message}\n"
+        for msg in messages:
+            if msg["role"] == "user":
+                ret += f"Human: {msg['content']}\n"
+            elif msg["role"] == "assistant":
+                ret += f"Assistant: {msg['content']}\n"
+
+            if not isinstance(msg['content'], str):
+                raise ValueError(
+                    "Currently multimodal context is not supported by the token counter."
+                )
+
+        ret += "<eos>"
+        return ret
     else:
         raise ValueError(f"Invalid model type: {model}")
 
@@ -232,6 +273,7 @@ class OpenAITokenCounter(BaseTokenCounter):
             model (ModelType): Model type for which tokens will be counted.
         """
         self.model: str = model.value_for_tiktoken
+        self.model_type = model
 
         self.tokens_per_message: int
         self.tokens_per_name: int
@@ -300,7 +342,7 @@ class OpenAITokenCounter(BaseTokenCounter):
                                 base64.b64decode(encoded_image)
                             )
                             image = Image.open(image_bytes)
-                            num_tokens += count_tokens_from_image(
+                            num_tokens += self._count_tokens_from_image(
                                 image, OpenAIVisionDetailType(detail)
                             )
                 if key == "name":
@@ -310,6 +352,45 @@ class OpenAITokenCounter(BaseTokenCounter):
         num_tokens += 3
         return num_tokens
 
+    def _count_tokens_from_image(
+        self, image: Image.Image, detail: OpenAIVisionDetailType
+    ) -> int:
+        r"""Count image tokens for OpenAI vision model. An :obj:`"auto"`
+        resolution model will be treated as :obj:`"high"`. All images with
+        :obj:`"low"` detail cost 85 tokens each. Images with :obj:`"high"` detail
+        are first scaled to fit within a 2048 x 2048 square, maintaining their
+        aspect ratio. Then, they are scaled such that the shortest side of the
+        image is 768px long. Finally, we count how many 512px squares the image
+        consists of. Each of those squares costs 170 tokens. Another 85 tokens are
+        always added to the final total. For more details please refer to `OpenAI
+        vision docs <https://platform.openai.com/docs/guides/vision>`_
+
+        Args:
+            image (PIL.Image.Image): Image to count number of tokens.
+            detail (OpenAIVisionDetailType): Image detail type to count
+                number of tokens.
+
+        Returns:
+            int: Number of tokens for the image given a detail type.
+        """
+        if detail == OpenAIVisionDetailType.LOW:
+            return LOW_DETAIL_TOKENS
+
+        width, height = image.size
+        if width > FIT_SQUARE_PIXELS or height > FIT_SQUARE_PIXELS:
+            scaling_factor = max(width, height) / FIT_SQUARE_PIXELS
+            width = int(width / scaling_factor)
+            height = int(height / scaling_factor)
+
+        scaling_factor = min(width, height) / SHORTEST_SIDE_PIXELS
+        scaled_width = int(width / scaling_factor)
+        scaled_height = int(height / scaling_factor)
+
+        h = ceil(scaled_height / SQUARE_PIXELS)
+        w = ceil(scaled_width / SQUARE_PIXELS)
+        total = EXTRA_TOKENS + SQUARE_TOKENS * h * w
+        return total
+
 
 class AnthropicTokenCounter(BaseTokenCounter):
     def __init__(self, model_type: ModelType):
@@ -428,41 +509,68 @@ class LiteLLMTokenCounter:
         return self.completion_cost(completion_response=response)
 
 
-def count_tokens_from_image(
-    image: Image.Image, detail: OpenAIVisionDetailType
-) -> int:
-    r"""Count image tokens for OpenAI vision model. An :obj:`"auto"`
-    resolution model will be treated as :obj:`"high"`. All images with
-    :obj:`"low"` detail cost 85 tokens each. Images with :obj:`"high"` detail
-    are first scaled to fit within a 2048 x 2048 square, maintaining their
-    aspect ratio. Then, they are scaled such that the shortest side of the
-    image is 768px long. Finally, we count how many 512px squares the image
-    consists of. Each of those squares costs 170 tokens. Another 85 tokens are
-    always added to the final total. For more details please refer to `OpenAI
-    vision docs <https://platform.openai.com/docs/guides/vision>`_
+class MistralTokenCounter(BaseTokenCounter):
+    def __init__(self, model_type: ModelType):
+        r"""Constructor for the token counter for Mistral models.
 
-    Args:
-        image (PIL.Image.Image): Image to count number of tokens.
-        detail (OpenAIVisionDetailType): Image detail type to count
-            number of tokens.
+        Args:
+            model_type (ModelType): Model type for which tokens will be
+                counted.
+        """
+        from mistral_common.tokens.tokenizers.mistral import MistralTokenizer
 
-    Returns:
-        int: Number of tokens for the image given a detail type.
-    """
-    if detail == OpenAIVisionDetailType.LOW:
-        return LOW_DETAIL_TOKENS
-
-    width, height = image.size
-    if width > FIT_SQUARE_PIXELS or height > FIT_SQUARE_PIXELS:
-        scaling_factor = max(width, height) / FIT_SQUARE_PIXELS
-        width = int(width / scaling_factor)
-        height = int(height / scaling_factor)
-
-    scaling_factor = min(width, height) / SHORTEST_SIDE_PIXELS
-    scaled_width = int(width / scaling_factor)
-    scaled_height = int(height / scaling_factor)
-
-    h = ceil(scaled_height / SQUARE_PIXELS)
-    w = ceil(scaled_width / SQUARE_PIXELS)
-    total = EXTRA_TOKENS + SQUARE_TOKENS * h * w
-    return total
+        self.model_type = model_type
+
+        # Determine the model type and set the tokenizer accordingly
+        model_name = (
+            "codestral-22b"
+            if self.model_type
+            in {ModelType.MISTRAL_CODESTRAL, ModelType.MISTRAL_CODESTRAL_MAMBA}
+            else self.model_type.value
+        )
+
+        self.tokenizer = MistralTokenizer.from_model(model_name)
+
+    def count_tokens_from_messages(self, messages: List[OpenAIMessage]) -> int:
+        r"""Count number of tokens in the provided message list using
+        loaded tokenizer specific for this type of model.
+
+        Args:
+            messages (List[OpenAIMessage]): Message list with the chat history
+                in OpenAI API format.
+
+        Returns:
+            int: Total number of tokens in the messages.
+        """
+        total_tokens = 0
+        for msg in messages:
+            tokens = self.tokenizer.encode_chat_completion(
+                self._convert_response_from_openai_to_mistral(msg)
+            ).tokens
+            total_tokens += len(tokens)
+        return total_tokens
+
+    def _convert_response_from_openai_to_mistral(
+        self, openai_msg: OpenAIMessage
+    ) -> ChatCompletionRequest:
+        r"""Convert an OpenAI message to a Mistral ChatCompletionRequest.
+
+        Args:
+            openai_msg (OpenAIMessage): An individual message with OpenAI
+                format.
+
+        Returns:
+            ChatCompletionRequest: The converted message in Mistral's request
+                format.
+        """
+
+        from mistral_common.protocol.instruct.request import (
+            ChatCompletionRequest,
+        )
+
+        mistral_request = ChatCompletionRequest(  # type: ignore[type-var]
+            model=self.model_type.value,
+            messages=[openai_msg],
+        )
+
+        return mistral_request

camel/workforce/__init__.py

@@ -0,0 +1,23 @@
+# =========== 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 BaseNode
+from .manager_node import ManagerNode
+from .worker_node import WorkerNode
+
+__all__ = [
+    "BaseNode",
+    "WorkerNode",
+    "ManagerNode",
+]

camel/workforce/base.py

@@ -0,0 +1,50 @@
+# =========== 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 typing import Any
+
+from camel.workforce.task_channel import TaskChannel
+
+
+class BaseNode(ABC):
+    def __init__(self, description: str) -> None:
+        self.node_id = str(id(self))
+        self.description = description
+        # every node is initialized to use its own channel
+        self._channel: TaskChannel = TaskChannel()
+        self._running = False
+
+    def reset(self, *args: Any, **kwargs: Any) -> Any:
+        """Resets the node to its initial state."""
+        raise NotImplementedError()
+
+    @abstractmethod
+    def set_channel(self, channel: TaskChannel):
+        r"""Sets the channel for the node."""
+
+    @abstractmethod
+    async def _listen_to_channel(self):
+        r"""Listens to the channel and handle tasks. This method should be
+        the main loop for the node.
+        """
+
+    @abstractmethod
+    async def start(self):
+        r"""Start the node."""
+
+    @abstractmethod
+    def stop(self):
+        r"""
+        Stop the node.
+        """

camel/workforce/manager_node.py

@@ -0,0 +1,299 @@
+# =========== 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
+
+import asyncio
+from collections import deque
+from typing import Deque, Dict, List, Optional
+
+from colorama import Fore
+
+from camel.agents import ChatAgent
+from camel.messages.base import BaseMessage
+from camel.tasks.task import Task, TaskState
+from camel.workforce.base import BaseNode
+from camel.workforce.single_agent_node import SingleAgentNode
+from camel.workforce.task_channel import TaskChannel
+from camel.workforce.utils import (
+    check_if_running,
+    parse_assign_task_resp,
+    parse_create_node_resp,
+)
+from camel.workforce.worker_node import WorkerNode
+from camel.workforce.workforce_prompt import (
+    ASSIGN_TASK_PROMPT,
+    CREATE_NODE_PROMPT,
+)
+
+
+class ManagerNode(BaseNode):
+    r"""A node that manages multiple nodes. It will split the task it
+    receives into subtasks and assign them to the child nodes under
+    it, and also handles the situation when the task fails.
+
+    Args:
+        description (str): Description of the node.
+        coordinator_agent_kwargs (Optional[Dict]): Keyword arguments for the
+            coordinator agent, e.g. `model`, `api_key`, `tools`, etc.
+        task_agent_kwargs (Optional[Dict]): Keyword arguments for the task
+            agent, e.g. `model`, `api_key`, `tools`, etc.
+    """
+
+    def __init__(
+        self,
+        description: str,
+        children: List[BaseNode],
+        coordinator_agent_kwargs: Optional[Dict] = None,
+        task_agent_kwargs: Optional[Dict] = None,
+    ) -> None:
+        super().__init__(description)
+        self._child_listening_tasks: Deque[asyncio.Task] = deque()
+        self._children = children
+
+        coord_agent_sysmsg = BaseMessage.make_assistant_message(
+            role_name="Workforce Manager",
+            content="You are coordinating a group of workers. A worker can be "
+            "a group of agents or a single agent. Each worker is created to"
+            " solve a specific kind of task. Your job includes assigning "
+            "tasks to a existing worker, creating a new worker for a task, "
+            "etc.",
+        )
+        self.coordinator_agent = ChatAgent(
+            coord_agent_sysmsg, **(coordinator_agent_kwargs or {})
+        )
+
+        task_sys_msg = BaseMessage.make_assistant_message(
+            role_name="Task Planner",
+            content="You are going to compose and decompose tasks.",
+        )
+        self.task_agent = ChatAgent(task_sys_msg, **(task_agent_kwargs or {}))
+
+        # if there is one, will set by the workforce class wrapping this
+        self._task: Optional[Task] = None
+        self._pending_tasks: Deque[Task] = deque()
+
+    @check_if_running(False)
+    def set_main_task(self, task: Task) -> None:
+        r"""Set the main task for the node."""
+        self._task = task
+
+    def _get_child_nodes_info(self) -> str:
+        r"""Get the information of all the child nodes under this node."""
+        return '\n'.join(
+            f'{child.node_id}: {child.description}' for child in self._children
+        )
+
+    def _find_assignee(
+        self,
+        task: Task,
+        failed_log: Optional[str] = None,
+    ) -> str:
+        r"""Assigns a task to a child node if capable, otherwise create a
+        new worker node.
+
+        Parameters:
+            task (Task): The task to be assigned.
+            failed_log (Optional[str]): Optional log of a previous failed
+                attempt.
+
+        Returns:
+            str: ID of the assigned node.
+        """
+        prompt = ASSIGN_TASK_PROMPT.format(
+            content=task.content,
+            child_nodes_info=self._get_child_nodes_info(),
+        )
+        req = BaseMessage.make_user_message(
+            role_name="User",
+            content=prompt,
+        )
+        response = self.coordinator_agent.step(req)
+        try:
+            print(f"{Fore.YELLOW}{response.msg.content}{Fore.RESET}")
+            assignee_id = parse_assign_task_resp(response.msg.content)
+        except ValueError:
+            assignee_id = self._create_worker_node_for_task(task).node_id
+        return assignee_id
+
+    async def _post_task(self, task: Task, assignee_id: str) -> None:
+        await self._channel.post_task(task, self.node_id, assignee_id)
+
+    async def _post_dependency(self, dependency: Task) -> None:
+        await self._channel.post_dependency(dependency, self.node_id)
+
+    def _create_worker_node_for_task(self, task: Task) -> WorkerNode:
+        r"""Creates a new worker node for a given task and add it to the
+        children list of this node. This is one of the actions that
+        the coordinator can take when a task has failed.
+
+        Args:
+            task (Task): The task for which the worker node is created.
+
+        Returns:
+            WorkerNode: The created worker node.
+        """
+        prompt = CREATE_NODE_PROMPT.format(
+            content=task.content,
+            child_nodes_info=self._get_child_nodes_info(),
+        )
+        req = BaseMessage.make_user_message(
+            role_name="User",
+            content=prompt,
+        )
+        response = self.coordinator_agent.step(req)
+        new_node_conf = parse_create_node_resp(response.msg.content)
+
+        worker_sysmsg = BaseMessage.make_assistant_message(
+            role_name=new_node_conf.role,
+            content=new_node_conf.system,
+        )
+
+        # TODO: add a default selection of tools for the worker
+        worker = ChatAgent(worker_sysmsg)
+
+        new_node = SingleAgentNode(
+            description=new_node_conf.description,
+            worker=worker,
+        )
+        new_node.set_channel(self._channel)
+
+        print(
+            f"{Fore.GREEN}New worker node {new_node.node_id} created."
+            f"{Fore.RESET}"
+        )
+
+        self._children.append(new_node)
+        self._child_listening_tasks.append(
+            asyncio.create_task(new_node.start())
+        )
+        return new_node
+
+    async def _get_returned_task(self) -> Task:
+        r"""Get the task that's published by this node and just get returned
+        from the assignee.
+        """
+        return await self._channel.get_returned_task_by_publisher(self.node_id)
+
+    async def _post_ready_tasks(self) -> None:
+        r"""Send all the pending tasks that have all the dependencies met to
+        the channel, or directly return if there is none. For now, we will
+        directly send the first task in the pending list because all the tasks
+        are linearly dependent."""
+
+        if not self._pending_tasks:
+            return
+
+        ready_task = self._pending_tasks[0]
+
+        # if the task has failed previously, just compose and send the task
+        # to the channel as a dependency
+        if ready_task.state == TaskState.FAILED:
+            # TODO: the composing of tasks seems not work very well
+            ready_task.compose(self.task_agent)
+            # remove the subtasks from the channel
+            for subtask in ready_task.subtasks:
+                await self._channel.remove_task(subtask.id)
+            # send the task to the channel as a dependency
+            await self._post_dependency(ready_task)
+            self._pending_tasks.popleft()
+            # try to send the next task in the pending list
+            await self._post_ready_tasks()
+        else:
+            # directly post the task to the channel if it's a new one
+            # find a node to assign the task
+            assignee_id = self._find_assignee(task=ready_task)
+            await self._post_task(ready_task, assignee_id)
+
+    async def _handle_failed_task(self, task: Task) -> None:
+        # remove the failed task from the channel
+        await self._channel.remove_task(task.id)
+        if task.get_depth() >= 3:
+            # create a new WF and reassign
+            # TODO: add a state for reassign?
+            assignee = self._create_worker_node_for_task(task)
+            # print('create_new_assignee:', assignee)
+            await self._post_task(task, assignee.node_id)
+        else:
+            subtasks = task.decompose(self.task_agent)
+            # Insert packets at the head of the queue
+            self._pending_tasks.extendleft(reversed(subtasks))
+            await self._post_ready_tasks()
+
+    async def _handle_completed_task(self, task: Task) -> None:
+        # archive the packet, making it into a dependency
+        self._pending_tasks.popleft()
+        await self._channel.archive_task(task.id)
+        await self._post_ready_tasks()
+
+    @check_if_running(False)
+    def set_channel(self, channel: TaskChannel):
+        r"""Set the channel for the node and all the child nodes under it."""
+        self._channel = channel
+        for child in self._children:
+            child.set_channel(channel)
+
+    @check_if_running(False)
+    async def _listen_to_channel(self) -> None:
+        r"""Continuously listen to the channel, post task to the channel and
+        track the status of posted tasks.
+        """
+
+        self._running = True
+        print(f"{Fore.GREEN}Manager node {self.node_id} started.{Fore.RESET}")
+
+        # if this node is at the top level, it will have an initial task
+        # the initial task must be decomposed into subtasks first
+        if self._task is not None:
+            subtasks = self._task.decompose(self.task_agent)
+            self._pending_tasks.extend(subtasks)
+            self._task.state = TaskState.FAILED
+            self._pending_tasks.append(self._task)
+
+        # before starting the loop, send ready pending tasks to the channel
+        await self._post_ready_tasks()
+
+        await self._channel.print_channel()
+
+        while self._task is None or self._pending_tasks:
+            returned_task = await self._get_returned_task()
+            if returned_task.state == TaskState.DONE:
+                await self._handle_completed_task(returned_task)
+            elif returned_task.state == TaskState.FAILED:
+                await self._handle_failed_task(returned_task)
+            else:
+                raise ValueError(
+                    f"Task {returned_task.id} has an unexpected state."
+                )
+
+        # shut down the whole workforce tree
+        self.stop()
+
+    @check_if_running(False)
+    async def start(self) -> None:
+        r"""Start itself and all the child nodes under it."""
+        for child in self._children:
+            child_listening_task = asyncio.create_task(child.start())
+            self._child_listening_tasks.append(child_listening_task)
+        await self._listen_to_channel()
+
+    @check_if_running(True)
+    def stop(self) -> None:
+        r"""Stop all the child nodes under it. The node itself will be stopped
+        by its parent node.
+        """
+        for child in self._children:
+            child.stop()
+        for child_task in self._child_listening_tasks:
+            child_task.cancel()
+        self._running = False