adding unit tests

docker-compose.yaml

@@ -37,8 +37,8 @@ services:
     networks:
       - super_network
 #    uncomment to expose redis port to host
-    ports:
-      - "6379:6379"
+    # ports:
+    #   - "6379:6379"
     volumes:
       - redis_data:/data
 

superagi/tools/thinking/tools.py

@@ -60,8 +60,6 @@ def _execute(self, task_description: str):
             metadata = {"agent_execution_id":self.agent_execution_id}
             relevant_tool_response = self.tool_response_manager.get_relevant_response(query=task_description,metadata=metadata)
             prompt = prompt.replace("{relevant_tool_response}",relevant_tool_response)
-            print("Here is the relevant tool response: ",relevant_tool_response,"ENDD")
-            print("Here is the prompt response: ",prompt,"ENDD")
             messages = [{"role": "system", "content": prompt}]
             result = self.llm.chat_completion(messages, max_tokens=self.max_token_limit)
             return result["content"]

tests/unit_tests/agent/test_output_handler.py

@@ -9,6 +9,8 @@
 from superagi.helper.json_cleaner import JsonCleaner
 from superagi.models.agent import Agent
 from superagi.models.agent_execution_permission import AgentExecutionPermission
+import numpy as np
+from superagi.agent.output_handler import ToolOutputHandler
 
 
 # Test for ToolOutputHandler
@@ -21,6 +23,7 @@ def test_tool_output_handle(parse_mock, execute_mock, get_completed_tasks_mock,
     agent_execution_id = 11
     agent_config = {"agent_id": 22, "permission_type": "unrestricted"}
     assistant_reply = '{"tool": {"name": "someAction", "args": ["arg1", "arg2"]}}'
+
     parse_mock.return_value = AgentGPTAction(name="someAction", args=["arg1", "arg2"])
 
     # Define what the mock response status should be
@@ -40,6 +43,33 @@ def test_tool_output_handle(parse_mock, execute_mock, get_completed_tasks_mock,
     assert response.status == "PENDING"
     parse_mock.assert_called_with(assistant_reply)
     assert session_mock.add.call_count == 2
+
+@patch('superagi.agent.output_handler.TokenTextSplitter')
+def test_add_text_to_memory(TokenTextSplitter_mock):
+    # Arrange
+    agent_execution_id = 1
+    agent_config = {"agent_id": 2}
+    tool_output_handler = ToolOutputHandler(agent_execution_id, agent_config,[], None)
+
+    assistant_reply = '{"thoughts": {"text": "This is a task."}}'
+    tool_response_result = '["Task completed."]'
+
+    text_splitter_mock = MagicMock()
+    TokenTextSplitter_mock.return_value = text_splitter_mock
+    text_splitter_mock.split_text.return_value = ["This is a task.", "Task completed."]
+
+    # Mock the VectorStore memory
+    memory_mock = MagicMock()
+    tool_output_handler.memory = memory_mock
+
+    # Act
+    tool_output_handler.add_text_to_memory(assistant_reply, tool_response_result)
+
+    # Assert
+    TokenTextSplitter_mock.assert_called_once_with(chunk_size=1024, chunk_overlap=10)
+    text_splitter_mock.split_text.assert_called_once_with('This is a task.["Task completed."]')
+    memory_mock.add_texts.assert_called_once_with(["This is a task.", "Task completed."], [{"agent_execution_id": agent_execution_id}, {"agent_execution_id": agent_execution_id}])    
+
 
 @patch('superagi.models.agent_execution_permission.AgentExecutionPermission')
 def test_tool_handler_check_permission_in_restricted_mode(op_mock):

tests/unit_tests/vector_store/test_redis.py

@@ -0,0 +1,42 @@
+from unittest.mock import MagicMock, patch
+import numpy as np
+from superagi.vector_store.document import Document
+from superagi.vector_store.redis import Redis
+
+
+def test_escape_token():
+    redis_object = Redis(None, None)
+    escaped_token = redis_object.escape_token("An,example.<> string!")
+    assert escaped_token == "An\\,example\\.\\<\\>\\ string\\!"
+
+@patch('redis.Redis')
+def test_add_texts(redis_mock):
+    # Arrange
+    mock_index = "mock_index"
+    mock_embedding_model = MagicMock()
+    redis_object = Redis(mock_index, mock_embedding_model)
+    redis_object.build_redis_key = MagicMock(return_value="mock_key")
+    texts = ["Hello", "World"]
+    metadatas = [{"data": 1}, {"data": 2}]
+
+    # Act
+    redis_object.add_texts(texts, metadatas)
+
+    # Assert
+    assert redis_object.redis_client.pipeline().hset.call_count == len(texts)
+
+@patch('redis.Redis')
+def test_get_matching_text(redis_mock):
+    # Arrange
+    mock_index = "mock_index"
+    redis_object = Redis(mock_index, None)
+    redis_object.embedding_model = MagicMock()
+    redis_object.embedding_model.get_embedding.return_value = np.array([0.1, 0.2, 0.3])
+    query = "mock_query"
+
+    # Act
+    result = redis_object.get_matching_text(query, metadata={})
+
+    # Assert
+    redis_object.embedding_model.get_embedding.assert_called_once_with(query)
+    assert "documents" in result