From f294dc5c99fd194470df4a3a2def38fd1bfc25c5 Mon Sep 17 00:00:00 2001
From: Yx Jiang <2237303+yxjiang@users.noreply.github.com>
Date: Fri, 15 Mar 2024 13:21:44 -0700
Subject: [PATCH 1/4] convert param to json

---
 polymind/core/tool.py            |  23 ++++-
 tests/polymind/core/test_tool.py | 158 ++++++++++++++++++++-----------
 2 files changed, 120 insertions(+), 61 deletions(-)

diff --git a/polymind/core/tool.py b/polymind/core/tool.py
index c02e30c..b05b089 100644
--- a/polymind/core/tool.py
+++ b/polymind/core/tool.py
@@ -9,9 +9,15 @@
 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 __str__(self) -> str:
+        return f"{self.name}: {self.type} - {self.description}"
 
     @validator("type")
     def check_type(cls, v: str) -> str:
@@ -58,13 +64,20 @@ 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_str = ""
+        for param in self.input_spec():
+            input_json_str += f"{param}\n"
+        output_json_str = ""
+        for param in self.output_spec():
+            output_json_str += f"{param}\n"
+        spec_json_str = f"Input:\n{input_json_str}\nOutput:\n{output_json_str}"
+        return spec_json_str
 
     @abstractmethod
     def input_spec(self) -> List[Param]:
diff --git a/tests/polymind/core/test_tool.py b/tests/polymind/core/test_tool.py
index 30a7905..1713248 100644
--- a/tests/polymind/core/test_tool.py
+++ b/tests/polymind/core/test_tool.py
@@ -9,67 +9,37 @@
 from pydantic import ValidationError
 
 
-class ToolForTest(BaseTool):
-
-    def input_spec(self) -> list[Param]:
-        return [Param(name="query", type="str", description="The query to reverse")]
-
-    def output_spec(self) -> list[Param]:
-        return [Param(name="result", type="str", description="The reversed query")]
-
-    async def _execute(self, input: Message) -> Message:
-        """Reverse the prompt and return the result.
-
-        Args:
-            input (Message): The input message to the tool.
-
-        Returns:
-            Message: The output message from the tool.
-        """
-        some_value = os.getenv("SOME_VARIABLE", "default")
-        return Message(content={"result": input.get("query")[::-1], "env": some_value})
-
-
-@pytest.fixture(autouse=True)
-def load_env_vars():
-    # Setup: Define environment variable before each test
-    os.environ["SOME_VARIABLE"] = "test_value"
-    yield
-    # Teardown: Remove the environment variable after each test
-    os.environ.pop("SOME_VARIABLE", None)
-
-
-@pytest.mark.asyncio
-class TestBaseTool:
-    async def test_tool_execute(self):
-        tool = ToolForTest(tool_name="test_tool")
-        input_message = Message(content={"query": "test"})
-        result_message = await tool(input_message)
-        assert result_message.get("result") == "tset"
-
-    async def test_tool_execute_with_env(self):
-        tool = ToolForTest(tool_name="test_tool")
-        input_message = Message(content={"query": "test"})
-        result_message = await tool(input_message)
-        # Assert both the tool's execution result and the loaded environment variable
-        assert result_message.get("result") == "tset"
-        assert result_message.get("env") == "test_value"
-
-
 class TestParam:
     @pytest.mark.parametrize("type_str", ["str", "int", "float"])
     def test_valid_simple_types(self, type_str):
         """Test that Param accepts valid simple type strings."""
-        param = Param(name="test_param", type=type_str, description="A test parameter")
-        assert param.type == type_str, f"Param type should be {type_str}"
+        param = Param(
+            name="test_param",
+            type=type_str,
+            description="A test parameter",
+            example="example value",
+        )
+        assert param.type == type_str, "Param type should be {}".format(type_str)
+        assert (
+            param.example == "example value"
+        ), "Param example should match the provided example"
 
-    @pytest.mark.parametrize("type_str", ["Dict[str, int]", "List[int]"])
-    def test_valid_complex_types(self, type_str):
-        """Test that Param accepts valid complex type strings with appropriate element types."""
+    @pytest.mark.parametrize(
+        "type_str, example",
+        [("Dict[str, int]", "{'key': 123}"), ("List[int]", "[1, 2, 3]")],
+    )
+    def test_valid_complex_types_with_example(self, type_str, example):
+        """Test that Param accepts valid complex type strings with appropriate element types and examples."""
         param = Param(
-            name="complex_param", type=type_str, description="A complex parameter"
+            name="complex_param",
+            type=type_str,
+            description="A complex parameter",
+            example=example,
         )
-        assert param.type == type_str, f"Param type should be {type_str}"
+        assert param.type == type_str, "Param type should be {}".format(type_str)
+        assert (
+            param.example == example
+        ), "Param example should match the provided example"
 
     @pytest.mark.parametrize(
         "type_str",
@@ -91,10 +61,86 @@ def test_invalid_types(self, type_str):
         with pytest.raises(ValidationError):
             Param(name="test_param", type=type_str, description="A test parameter")
 
-    def test_param_with_description(self):
-        """Test that Param correctly stores a description."""
+    def test_param_with_description_and_default_example(self):
+        """Test that Param correctly stores a description and uses the default example if not specified."""
         description = "This parameter is for testing."
         param = Param(name="test_param", type="str", description=description)
         assert (
             param.description == description
         ), "Param description should match the input description"
+        assert (
+            param.example == ""
+        ), "Param example should use the default empty string if not specified"
+
+
+# class ToolForTest(BaseTool):
+
+#     def input_spec(self) -> list[Param]:
+#         return [Param(name="query", type="str", description="The query to reverse")]
+
+#     def output_spec(self) -> list[Param]:
+#         return [Param(name="result", type="str", description="The reversed query")]
+
+#     async def _execute(self, input: Message) -> Message:
+#         """Reverse the prompt and return the result.
+
+#         Args:
+#             input (Message): The input message to the tool.
+
+#         Returns:
+#             Message: The output message from the tool.
+#         """
+#         some_value = os.getenv("SOME_VARIABLE", "default")
+#         return Message(content={"result": input.get("query")[::-1], "env": some_value})
+
+
+# @pytest.fixture(autouse=True)
+# def load_env_vars():
+#     # Setup: Define environment variable before each test
+#     os.environ["SOME_VARIABLE"] = "test_value"
+#     yield
+#     # Teardown: Remove the environment variable after each test
+#     os.environ.pop("SOME_VARIABLE", None)
+
+
+# @pytest.mark.asyncio
+# class TestBaseTool:
+#     async def test_tool_execute(self):
+#         tool = ToolForTest(tool_name="test_tool")
+#         input_message = Message(content={"query": "test"})
+#         result_message = await tool(input_message)
+#         assert (
+#             result_message.get("result") == "tset"
+#         ), "The result should be the reverse of the input query"
+
+#     async def test_tool_execute_with_env(self):
+#         tool = ToolForTest(tool_name="test_tool")
+#         input_message = Message(content={"query": "test"})
+#         result_message = await tool(input_message)
+#         assert (
+#             result_message.get("result") == "tset"
+#         ), "The result should be the reverse of the input query"
+#         assert (
+#             result_message.get("env") == "test_value"
+#         ), "The environment variable should be loaded correctly"
+
+# def test_get_spec(self):
+#     tool = ToolForTest()
+#     input_spec, output_spec = tool.get_spec()
+
+#     expected_input_spec = [
+#         Param(name="query", type="str", description="The query to reverse")
+#     ]
+#     expected_output_spec = [
+#         Param(name="result", type="str", description="The reversed query")
+#     ]
+
+#     # Since Param is a Pydantic model, comparing the dicts is a straightforward way to assert equality
+#     assert all(
+#         param.dict() == expected.dict()
+#         for param, expected in zip(input_spec, expected_input_spec)
+#     ), "Input specs do not match"
+#     assert all(
+#         param.dict() == expected.dict()
+#         for param, expected in zip(output_spec, expected_output_spec)
+#     ), "Output specs do not match"

From 8696a950f7927dc1afd5c4e99ff7eea2f2b99b43 Mon Sep 17 00:00:00 2001
From: Yx Jiang <2237303+yxjiang@users.noreply.github.com>
Date: Fri, 15 Mar 2024 13:42:12 -0700
Subject: [PATCH 2/4] Convert param to json

---
 polymind/core/tool.py            |  28 +++--
 pyproject.toml                   |   2 +-
 tests/polymind/core/test_tool.py | 197 ++++++++++++++++++++-----------
 3 files changed, 147 insertions(+), 80 deletions(-)

diff --git a/polymind/core/tool.py b/polymind/core/tool.py
index b05b089..2dee4e3 100644
--- a/polymind/core/tool.py
+++ b/polymind/core/tool.py
@@ -3,7 +3,8 @@
 from polymind.core.message import Message
 from pydantic import BaseModel, validator, Field
 import re
-from typing import List, Tuple
+from typing import Dict, List
+import json
 
 
 class Param(BaseModel):
@@ -16,8 +17,16 @@ class Param(BaseModel):
     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 f"{self.name}: {self.type} - {self.description}"
+        return json.dumps(self.to_json_obj(), indent=4)
 
     @validator("type")
     def check_type(cls, v: str) -> str:
@@ -70,14 +79,17 @@ def get_spec(self) -> str:
         Returns:
             Tuple[List[Param], List[Param]]: The input and output specification of the tool.
         """
-        input_json_str = ""
+        input_json_obj = []
         for param in self.input_spec():
-            input_json_str += f"{param}\n"
-        output_json_str = ""
+            input_json_obj.append(param.to_json_obj())
+        output_json_obj = []
         for param in self.output_spec():
-            output_json_str += f"{param}\n"
-        spec_json_str = f"Input:\n{input_json_str}\nOutput:\n{output_json_str}"
-        return spec_json_str
+            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]:
diff --git a/pyproject.toml b/pyproject.toml
index 3d3bf3a..f2fef95 100644
--- a/pyproject.toml
+++ b/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 <2237303+yxjiang@users.noreply.github.com>"]
 license = "MIT License"
diff --git a/tests/polymind/core/test_tool.py b/tests/polymind/core/test_tool.py
index 1713248..82f4fb2 100644
--- a/tests/polymind/core/test_tool.py
+++ b/tests/polymind/core/test_tool.py
@@ -7,6 +7,7 @@
 from polymind.core.tool import BaseTool, Param
 from polymind.core.message import Message
 from pydantic import ValidationError
+import json
 
 
 class TestParam:
@@ -73,74 +74,128 @@ def test_param_with_description_and_default_example(self):
         ), "Param example should use the default empty string if not specified"
 
 
-# class ToolForTest(BaseTool):
-
-#     def input_spec(self) -> list[Param]:
-#         return [Param(name="query", type="str", description="The query to reverse")]
-
-#     def output_spec(self) -> list[Param]:
-#         return [Param(name="result", type="str", description="The reversed query")]
-
-#     async def _execute(self, input: Message) -> Message:
-#         """Reverse the prompt and return the result.
-
-#         Args:
-#             input (Message): The input message to the tool.
-
-#         Returns:
-#             Message: The output message from the tool.
-#         """
-#         some_value = os.getenv("SOME_VARIABLE", "default")
-#         return Message(content={"result": input.get("query")[::-1], "env": some_value})
-
-
-# @pytest.fixture(autouse=True)
-# def load_env_vars():
-#     # Setup: Define environment variable before each test
-#     os.environ["SOME_VARIABLE"] = "test_value"
-#     yield
-#     # Teardown: Remove the environment variable after each test
-#     os.environ.pop("SOME_VARIABLE", None)
-
-
-# @pytest.mark.asyncio
-# class TestBaseTool:
-#     async def test_tool_execute(self):
-#         tool = ToolForTest(tool_name="test_tool")
-#         input_message = Message(content={"query": "test"})
-#         result_message = await tool(input_message)
-#         assert (
-#             result_message.get("result") == "tset"
-#         ), "The result should be the reverse of the input query"
-
-#     async def test_tool_execute_with_env(self):
-#         tool = ToolForTest(tool_name="test_tool")
-#         input_message = Message(content={"query": "test"})
-#         result_message = await tool(input_message)
-#         assert (
-#             result_message.get("result") == "tset"
-#         ), "The result should be the reverse of the input query"
-#         assert (
-#             result_message.get("env") == "test_value"
-#         ), "The environment variable should be loaded correctly"
-
-# def test_get_spec(self):
-#     tool = ToolForTest()
-#     input_spec, output_spec = tool.get_spec()
-
-#     expected_input_spec = [
-#         Param(name="query", type="str", description="The query to reverse")
-#     ]
-#     expected_output_spec = [
-#         Param(name="result", type="str", description="The reversed query")
-#     ]
-
-#     # Since Param is a Pydantic model, comparing the dicts is a straightforward way to assert equality
-#     assert all(
-#         param.dict() == expected.dict()
-#         for param, expected in zip(input_spec, expected_input_spec)
-#     ), "Input specs do not match"
-#     assert all(
-#         param.dict() == expected.dict()
-#         for param, expected in zip(output_spec, expected_output_spec)
-#     ), "Output specs do not match"
+class ToolForTest(BaseTool):
+
+    def input_spec(self) -> list[Param]:
+        return [
+            Param(
+                name="query",
+                type="str",
+                example="example-str",
+                description="The query to reverse",
+            ),
+            Param(
+                name="query2",
+                type="str",
+                example="example-str",
+                description="The query2 to reverse",
+            ),
+        ]
+
+    def output_spec(self) -> list[Param]:
+        return [
+            Param(
+                name="result",
+                type="str",
+                example="example-output-str",
+                description="The reversed query",
+            ),
+            Param(
+                name="result2",
+                type="str",
+                example="example-output-str",
+                description="The reversed query2",
+            ),
+        ]
+
+    async def _execute(self, input: Message) -> Message:
+        """Reverse the prompt and return the result.
+
+        Args:
+            input (Message): The input message to the tool.
+
+        Returns:
+            Message: The output message from the tool.
+        """
+        some_value = os.getenv("SOME_VARIABLE", "default")
+        return Message(
+            content={
+                "result": input.get("query")[::-1],
+                "result2": input.get("query2")[::-1],
+                "env": some_value,
+            }
+        )
+
+
+@pytest.fixture(autouse=True)
+def load_env_vars():
+    # Setup: Define environment variable before each test
+    os.environ["SOME_VARIABLE"] = "test_value"
+    yield
+    # Teardown: Remove the environment variable after each test
+    os.environ.pop("SOME_VARIABLE", None)
+
+
+@pytest.mark.asyncio
+class TestBaseTool:
+    async def test_tool_execute(self):
+        tool = ToolForTest(tool_name="test_tool")
+        input_message = Message(content={"query": "test", "query2": "hello"})
+        result_message = await tool(input_message)
+        assert (
+            result_message.get("result") == "tset"
+        ), "The result should be the reverse of the input query"
+        assert (
+            result_message.get("result2") == "olleh"
+        ), "The result should be the reverse of the input query2"
+
+    async def test_tool_execute_with_env(self):
+        tool = ToolForTest(tool_name="test_tool")
+        input_message = Message(content={"query": "test", "query2": "hello"})
+        result_message = await tool(input_message)
+        assert (
+            result_message.get("result") == "tset"
+        ), "The result should be the reverse of the input query"
+        assert (
+            result_message.get("result2") == "olleh"
+        ), "The result should be the reverse of the input query2"
+        assert (
+            result_message.get("env") == "test_value"
+        ), "The environment variable should be loaded correctly"
+
+    def test_get_spec(self):
+        tool = ToolForTest(tool_name="test_tool")
+        spec_str = tool.get_spec()
+        expected_json_str = """{
+        "input_message": [
+            {
+                "name": "query",
+                "type": "str",
+                "description": "The query to reverse",
+                "example": "example-str"
+            },
+            {
+                "name": "query2",
+                "type": "str",
+                "description": "The query2 to reverse",
+                "example": "example-str"
+            }
+        ],
+        "output_message": [
+            {
+                "name": "result",
+                "type": "str",
+                "description": "The reversed query",
+                "example": "example-output-str"
+            },
+            {
+                "name": "result2",
+                "type": "str",
+                "description": "The reversed query2",
+                "example": "example-output-str"
+            }
+        ]
+        }"""
+        assert json.loads(spec_str) == json.loads(
+            expected_json_str
+        ), "The spec string should match the expected JSON string"

From e6acb4ea8b8da3ef10f8b65a2ac1408f7049d8f3 Mon Sep 17 00:00:00 2001
From: Yx Jiang <2237303+yxjiang@users.noreply.github.com>
Date: Fri, 15 Mar 2024 13:58:57 -0700
Subject: [PATCH 3/4] update

---
 .flake8                                |  5 +++
 .github/workflows/test.yaml            |  8 ++++-
 polymind/__init__.py                   |  4 +--
 polymind/core/agent.py                 | 11 +++---
 polymind/core/message.py               |  3 +-
 polymind/core/task.py                  |  8 +++--
 polymind/core/tool.py                  | 10 +++---
 polymind/tools/oai_tools.py            | 22 ++++++++----
 pyproject.toml                         |  3 ++
 tests/polymind/core/test_agent.py      |  5 +--
 tests/polymind/core/test_task.py       |  4 ++-
 tests/polymind/core/test_tool.py       |  8 +++--
 tests/polymind/tools/test_oai_tools.py | 50 ++++++++++++++++----------
 13 files changed, 96 insertions(+), 45 deletions(-)
 create mode 100644 .flake8

diff --git a/.flake8 b/.flake8
new file mode 100644
index 0000000..cde5cdc
--- /dev/null
+++ b/.flake8
@@ -0,0 +1,5 @@
+[flake8]
+max-line-length = 120
+exclude = __init__.py
+per-file-ignores =
+    __init__.py:F401
\ No newline at end of file
diff --git a/.github/workflows/test.yaml b/.github/workflows/test.yaml
index 19a1ec3..d20a8ec 100644
--- a/.github/workflows/test.yaml
+++ b/.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
\ No newline at end of file
+      run: poetry run pytest
diff --git a/polymind/__init__.py b/polymind/__init__.py
index 2562c55..67a1829 100644
--- a/polymind/__init__.py
+++ b/polymind/__init__.py
@@ -1,8 +1,8 @@
 # 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
diff --git a/polymind/core/agent.py b/polymind/core/agent.py
index 1139207..236d4f4 100644
--- a/polymind/core/agent.py
+++ b/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.
diff --git a/polymind/core/message.py b/polymind/core/message.py
index 355e066..5196a6d 100644
--- a/polymind/core/message.py
+++ b/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."""
diff --git a/polymind/core/task.py b/polymind/core/task.py
index 2113a56..0b5f5cf 100644
--- a/polymind/core/task.py
+++ b/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):
diff --git a/polymind/core/tool.py b/polymind/core/tool.py
index 2dee4e3..09fbe92 100644
--- a/polymind/core/tool.py
+++ b/polymind/core/tool.py
@@ -1,10 +1,12 @@
+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 Dict, List
-import json
 
 
 class Param(BaseModel):
diff --git a/polymind/tools/oai_tools.py b/polymind/tools/oai_tools.py
index e9e3bbf..b8bf63a 100644
--- a/polymind/tools/oai_tools.py
+++ b/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.",
             ),
         ]
 
diff --git a/pyproject.toml b/pyproject.toml
index f2fef95..b21c9d6 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -20,6 +20,9 @@ isort = "^5.13.2"
 flake8 = "^7.0.0"
 pytest-asyncio = "^0.23.5.post1"
 
+[tool.black]
+line-length = 120
+
 [build-system]
 requires = ["poetry-core"]
 build-backend = "poetry.core.masonry.api"
diff --git a/tests/polymind/core/test_agent.py b/tests/polymind/core/test_agent.py
index 33a910c..d172d18 100644
--- a/tests/polymind/core/test_agent.py
+++ b/tests/polymind/core/test_agent.py
@@ -3,10 +3,11 @@
     poetry run pytest tests/polymind/core/test_thought_process.py
 """
 
-from pydantic import BaseModel
 import pytest
-from polymind.core.message import Message
+from pydantic import BaseModel
+
 from polymind.core.agent import Agent, ThoughtProcess
+from polymind.core.message import Message
 
 
 class MockThoughtProcess(ThoughtProcess):
diff --git a/tests/polymind/core/test_task.py b/tests/polymind/core/test_task.py
index ff374a5..1bec0b5 100644
--- a/tests/polymind/core/test_task.py
+++ b/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
 
 
diff --git a/tests/polymind/core/test_tool.py b/tests/polymind/core/test_tool.py
index 82f4fb2..ec251c1 100644
--- a/tests/polymind/core/test_tool.py
+++ b/tests/polymind/core/test_tool.py
@@ -2,12 +2,14 @@
     poetry run pytest tests/polymind/core/test_tool.py
 """
 
+import json
 import os
+
 import pytest
-from polymind.core.tool import BaseTool, Param
-from polymind.core.message import Message
 from pydantic import ValidationError
-import json
+
+from polymind.core.message import Message
+from polymind.core.tool import BaseTool, Param
 
 
 class TestParam:
diff --git a/tests/polymind/tools/test_oai_tools.py b/tests/polymind/tools/test_oai_tools.py
index f9c863f..97a8cde 100644
--- a/tests/polymind/tools/test_oai_tools.py
+++ b/tests/polymind/tools/test_oai_tools.py
@@ -3,8 +3,11 @@
     poetry run pytest tests/polymind/tools/test_oai_tools.py
 """
 
-import pytest
+import json
 from unittest.mock import AsyncMock, patch
+
+import pytest
+
 from polymind.core.message import Message
 from polymind.tools.oai_tools import OpenAIChatTool
 
@@ -55,20 +58,31 @@ async def test_execute_failure_empty_prompt(self, tool):
 
     def test_get_spec(self, tool):
         """Test get_spec method of OpenAIChatTool."""
-        input_spec = tool.get_spec()[0]
-        output_spec = tool.get_spec()[1]
-
-        assert len(input_spec) == 2
-        assert input_spec[0].name == "prompt"
-        assert input_spec[0].type == "str"
-        assert input_spec[0].description == "The prompt for the chat."
-
-        assert len(input_spec) == 2
-        assert input_spec[1].name == "system_prompt"
-        assert input_spec[1].type == "str"
-        assert input_spec[1].description == "The system prompt for the chat."
-
-        assert len(output_spec) == 1
-        assert output_spec[0].name == "response"
-        assert output_spec[0].type == "str"
-        assert output_spec[0].description == "The response from the chat."
+        spec_str = tool.get_spec()
+        expected_json_str = """{
+        "input_message": [
+            {
+                "name": "system_prompt",
+                "type": "str",
+                "description": "The system prompt for the chat.",
+                "example": "You are a helpful AI assistant."
+            },
+            {
+                "name": "prompt",
+                "type": "str",
+                "description": "The prompt for the chat.",
+                "example": "hello, how are you?"
+            }
+        ],
+        "output_message": [
+            {
+                "name": "response",
+                "type": "str",
+                "description": "The response from the chat.",
+                "example": "I'm good, how are you?"
+            }
+        ]
+        }"""
+        assert json.loads(spec_str) == json.loads(
+            expected_json_str
+        ), "The spec string should match the expected JSON string"

From f2292e3c64dd4acdbecdb31754747c2b73a9ce29 Mon Sep 17 00:00:00 2001
From: Yx Jiang <2237303+yxjiang@users.noreply.github.com>
Date: Fri, 15 Mar 2024 14:06:51 -0700
Subject: [PATCH 4/4] lint

---
 polymind/__init__.py              | 1 +
 pyproject.toml                    | 3 +++
 tests/polymind/core/test_agent.py | 1 -
 3 files changed, 4 insertions(+), 1 deletion(-)

diff --git a/polymind/__init__.py b/polymind/__init__.py
index 67a1829..40a53f1 100644
--- a/polymind/__init__.py
+++ b/polymind/__init__.py
@@ -1,3 +1,4 @@
+# noqa: D104
 # Expose the core classes
 from .core.agent import Agent, ThoughtProcess
 from .core.message import Message
diff --git a/pyproject.toml b/pyproject.toml
index b21c9d6..4a5c9e9 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -23,6 +23,9 @@ 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"
diff --git a/tests/polymind/core/test_agent.py b/tests/polymind/core/test_agent.py
index d172d18..4260884 100644
--- a/tests/polymind/core/test_agent.py
+++ b/tests/polymind/core/test_agent.py
@@ -4,7 +4,6 @@
 """
 
 import pytest
-from pydantic import BaseModel
 
 from polymind.core.agent import Agent, ThoughtProcess
 from polymind.core.message import Message