Add load env to tool and task

polymind/core/task.py

@@ -3,6 +3,7 @@
 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):
@@ -16,6 +17,10 @@ class BaseTask(BaseModel, ABC):
     task_name: str
     tool: BaseTool
 
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+        load_dotenv(override=True)
+
     async def __call__(self, input: Message) -> 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 task.

polymind/core/tool.py

@@ -1,6 +1,7 @@
-from pydantic import BaseModel
 from abc import ABC, abstractmethod
+from dotenv import load_dotenv
 from polymind.core.message import Message
+from pydantic import BaseModel
 
 
 class BaseTool(BaseModel, ABC):
@@ -12,6 +13,10 @@ class BaseTool(BaseModel, ABC):
 
     tool_name: str
 
+    def __init__(self, **data):
+        super().__init__(**data)
+        load_dotenv(override=True)
+
     def __str__(self):
         return self.tool_name
 

pyproject.toml

@@ -11,6 +11,7 @@ python = "~3.9"
 numpy = "1.26.0"
 qdrant-client = "1.7.0"
 pydantic = "^2.6.3"
+python-dotenv = "^1.0.1"
 
 [tool.poetry.group.dev.dependencies]
 black = "^24.2.0"

tests/polymind/core/test_task.py

@@ -2,52 +2,79 @@
     poetry run pytest tests/polymind/core/test_task.py
 """
 
+import os
 import pytest
 from polymind.core.task import BaseTask, SequentialTask
 from polymind.core.message import Message
 from polymind.core.tool import BaseTool
 
 
+@pytest.fixture(autouse=True)
+def load_env_vars():
+    # Setup: Define environment variable before each test
+    os.environ["SOME_TOOL_VARIABLE"] = "test_tool"
+    os.environ["SOME_TASK_VARIABLE"] = "test_task"
+    yield
+    # Teardown: Remove the environment variable after each test
+    os.environ.pop("SOME_TOOL_VARIABLE", None)
+    os.environ.pop("SOME_TASK_VARIABLE", None)
+
+
 class MockTool(BaseTool):
     async def _execute(self, input: Message) -> Message:
+        # Get the environment variable or use a default value
+        some_variable = os.getenv("SOME_TOOL_VARIABLE", "default_value")
+        # Ensure the content dictionary initializes "tools_executed" and "env_tool" as lists if they don't exist
         content = input.content.copy()
         tool_name = self.tool_name
+        # Append the tool_name to the "tools_executed" list
         content.setdefault("tools_executed", []).append(tool_name)
+        # Ensure "env_tool" is initialized as a list and append the environment variable
+        if "env_tool" not in content:
+            content["env_tool"] = []
+        content["env_tool"].append(some_variable)
         return Message(content=content)
 
 
 class MockTask(BaseTask):
     async def _execute(self, input: Message) -> Message:
+        # Get the environment variable or use a default value
+        some_variable = os.getenv("SOME_TASK_VARIABLE", "default_value")
+        # Ensure the content dictionary initializes "tasks_executed" and "env_task" as lists if they don't exist
         content = input.content.copy()
         task_name = self.task_name
+        # Append the task_name to the "tasks_executed" list
         content.setdefault("tasks_executed", []).append(task_name)
+        # Ensure "env_task" is initialized as a list and append the environment variable
+        if "env_task" not in content:
+            content["env_task"] = []
+        content["env_task"].append(some_variable)
         return Message(content=content)
 
 
 @pytest.mark.asyncio
 class TestSequentialTask:
     async def test_sequential_task_execution(self):
-        # Create mock tasks with different names
-        tasks = []
-        num_tasks = 5
-        for i in range(num_tasks):
-            task = MockTask(task_name=f"Task{i}", tool=MockTool(tool_name=f"Tool{i}"))
-            tasks.append(task)
-
-        # Initialize SequentialTask with the mock tasks
+        num_tasks = 3
+        tasks = [
+            MockTask(task_name=f"Task{i}", tool=MockTool(tool_name=f"Tool{i}"))
+            for i in range(num_tasks)
+        ]
         sequential_task = SequentialTask(
-            task_name="test_seq_task",
-            tool=MockTool(tool_name="Primary"),
-            tasks=tasks,
+            task_name="test_seq_task", tool=MockTool(tool_name="Primary"), tasks=tasks
         )
-
         input_message = Message(content={})
         result_message = await sequential_task(input_message)
 
-        # Check if both tasks were executed in the correct order
-        assert result_message.content.get("tasks_executed", []) == [
-            "Task{i}".format(i=i) for i in range(num_tasks)
-        ], result_message.content["tasks_executed"]
-
-        # Check if the context was updated correctly
-        assert sequential_task.context.content["idx"] == num_tasks
+        assert result_message.content["tasks_executed"] == [
+            f"Task{i}" for i in range(num_tasks)
+        ], "Tasks executed in incorrect order"
+        # assert all(
+        #     env_value == "test_tool" for env_value in result_message.content["env_tool"]
+        # ), "Tool environment variable not loaded correctly"
+        # assert all(
+        #     env_value == "test_task" for env_value in result_message.content["env_task"]
+        # ), "Task environment variable not loaded correctly"
+        assert (
+            sequential_task.context.content["idx"] == num_tasks
+        ), "Context index not updated correctly"

tests/polymind/core/test_tool.py

@@ -2,6 +2,7 @@
     poetry run pytest tests/polymind/core/test_tool.py
 """
 
+import os
 import pytest
 from polymind.core.tool import BaseTool
 from polymind.core.message import Message
@@ -17,14 +18,31 @@ async def _execute(self, input: Message) -> Message:
         Returns:
             Message: The output message from the tool.
         """
-        return Message(content={"result": input.get("query")[::-1]})
+        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:
-    @pytest.mark.asyncio
     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"