Merge pull request #22 from small-thinking/convert-param-to-json

Convert param to json

.flake8

@@ -0,0 +1,5 @@
+[flake8]
+max-line-length = 120
+exclude = __init__.py
+per-file-ignores =
+    __init__.py:F401

.github/workflows/test.yaml

@@ -33,6 +33,12 @@ jobs:
     - name: Install pytest
       run: |
         pip install pytest
+
+    - name: Run Flake8
+      run: poetry run flake8 polymind/
+
+    - name: Run isort
+      run: poetry run isort --check-only .
 
     - name: Run Pytest
-      run: poetry run pytest
+      run: poetry run pytest

polymind/__init__.py

@@ -1,8 +1,9 @@
+# noqa: D104
 # Expose the core classes
-from .core.tool import BaseTool
+from .core.agent import Agent, ThoughtProcess
 from .core.message import Message
 from .core.task import BaseTask, CompositeTask, SequentialTask
-from .core.agent import Agent, ThoughtProcess
+from .core.tool import BaseTool
 
 # Expose the tools
 from .tools.oai_tools import OpenAIChatTool

polymind/core/agent.py

@@ -1,8 +1,10 @@
-from pydantic import BaseModel, Field
+from abc import ABC, abstractmethod
+from typing import Dict, Optional
+
+from pydantic import BaseModel
+
 from polymind.core.message import Message
 from polymind.core.tool import BaseTool
-from abc import ABC, abstractmethod
-from typing import Dict, Optional, TYPE_CHECKING
 
 
 class ThoughtProcess(BaseModel, ABC):
@@ -20,7 +22,8 @@ def __str__(self):
 
     async def __call__(self, input: Message, agent: "Agent") -> Message:
         """Makes the instance callable, delegating to the execute method.
-        This allows the instance to be used as a callable object, simplifying the syntax for executing the thought process.
+        This allows the instance to be used as a callable object, simplifying
+        the syntax for executing the thought process.
 
         Args:
             input (Message): The input message to the thought process.

polymind/core/message.py

@@ -1,6 +1,7 @@
-from pydantic import BaseModel, field_validator
 from typing import Any, Dict
 
+from pydantic import BaseModel, field_validator
+
 
 class Message(BaseModel):
     """Message is a class that represents a message that can carry any information."""

polymind/core/task.py

@@ -1,9 +1,11 @@
-from pydantic import BaseModel, Field
 from abc import ABC, abstractmethod
+from typing import List
+
+from dotenv import load_dotenv
+from pydantic import BaseModel, Field
+
 from polymind.core.message import Message
 from polymind.core.tool import BaseTool
-from typing import Dict, List
-from dotenv import load_dotenv
 
 
 class BaseTask(BaseModel, ABC):

polymind/core/tool.py

@@ -1,17 +1,34 @@
+import json
+import re
 from abc import ABC, abstractmethod
+from typing import Dict, List
+
 from dotenv import load_dotenv
+from pydantic import BaseModel, Field, validator
+
 from polymind.core.message import Message
-from pydantic import BaseModel, validator, Field
-import re
-from typing import List, Tuple
 
 
 class Param(BaseModel):
     """Param is used to describe the specification of a parameter for a tool."""
 
-    name: str
-    type: str = Field(...)
-    description: str
+    name: str = Field(description="The name of the parameter.")
+    type: str = Field(
+        description="The type of the parameter: str, int, float, Dict[KeyType, ValueType], or List[ElementType]."
+    )
+    description: str = Field(description="A description of the parameter.")
+    example: str = Field(default="", description="An example value for the parameter.")
+
+    def to_json_obj(self) -> Dict[str, str]:
+        return {
+            "name": self.name,
+            "type": self.type,
+            "description": self.description,
+            "example": self.example,
+        }
+
+    def __str__(self) -> str:
+        return json.dumps(self.to_json_obj(), indent=4)
 
     @validator("type")
     def check_type(cls, v: str) -> str:
@@ -58,13 +75,23 @@ async def __call__(self, input: Message) -> Message:
         """
         return await self._execute(input)
 
-    def get_spec(self) -> Tuple[List[Param], List[Param]]:
+    def get_spec(self) -> str:
         """Return the input and output specification of the tool.
 
         Returns:
             Tuple[List[Param], List[Param]]: The input and output specification of the tool.
         """
-        return self.input_spec(), self.output_spec()
+        input_json_obj = []
+        for param in self.input_spec():
+            input_json_obj.append(param.to_json_obj())
+        output_json_obj = []
+        for param in self.output_spec():
+            output_json_obj.append(param.to_json_obj())
+        spec_json_obj = {
+            "input_message": input_json_obj,
+            "output_message": output_json_obj,
+        }
+        return json.dumps(spec_json_obj, indent=4)
 
     @abstractmethod
     def input_spec(self) -> List[Param]:

polymind/tools/oai_tools.py

@@ -2,13 +2,14 @@
 This file contains the necessary tools of using OpenAI models.
 """
 
+import os
 from typing import List
+
+from openai import AsyncOpenAI
 from pydantic import Field
-from polymind.core.tool import BaseTool, Param
+
 from polymind.core.message import Message
-from openai import AsyncOpenAI
-import os
-from dotenv import load_dotenv
+from polymind.core.tool import BaseTool, Param
 
 
 class OpenAIChatTool(BaseTool):
@@ -41,12 +42,18 @@ def input_spec(self) -> List[Param]:
             List[Param]: The input specification of the tool.
         """
         return [
-            Param(name="prompt", type="str", description="The prompt for the chat."),
             Param(
                 name="system_prompt",
                 type="str",
+                example="You are a helpful AI assistant.",
                 description="The system prompt for the chat.",
             ),
+            Param(
+                name="prompt",
+                type="str",
+                example="hello, how are you?",
+                description="The prompt for the chat.",
+            ),
         ]
 
     def output_spec(self) -> List[Param]:
@@ -58,7 +65,10 @@ def output_spec(self) -> List[Param]:
         """
         return [
             Param(
-                name="response", type="str", description="The response from the chat."
+                name="response",
+                type="str",
+                example="I'm good, how are you?",
+                description="The response from the chat.",
             ),
         ]
 

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "polymind"
-version = "0.0.14"  # Update this version before publishing to PyPI
+version = "0.0.15"  # Update this version before publishing to PyPI
 description = "PolyMind is a customizable collaborative multi-agent framework for collective intelligence and distributed problem solving."
 authors = ["Yx Jiang <[email protected]>"]
 license = "MIT License"
@@ -20,6 +20,12 @@ isort = "^5.13.2"
 flake8 = "^7.0.0"
 pytest-asyncio = "^0.23.5.post1"
 
+[tool.black]
+line-length = 120
+
+[tool.isort]
+skip = ["__init__.py"]
+
 [build-system]
 requires = ["poetry-core"]
 build-backend = "poetry.core.masonry.api"

tests/polymind/core/test_agent.py

@@ -3,10 +3,10 @@
     poetry run pytest tests/polymind/core/test_thought_process.py
 """
 
-from pydantic import BaseModel
 import pytest
-from polymind.core.message import Message
+
 from polymind.core.agent import Agent, ThoughtProcess
+from polymind.core.message import Message
 
 
 class MockThoughtProcess(ThoughtProcess):

tests/polymind/core/test_task.py

@@ -3,9 +3,11 @@
 """
 
 import os
+
 import pytest
-from polymind.core.task import BaseTask, SequentialTask
+
 from polymind.core.message import Message
+from polymind.core.task import BaseTask, SequentialTask
 from polymind.core.tool import BaseTool, Param