rework agent decorator and agent tests

backend/src/agents/agent.py

@@ -1,7 +1,7 @@
 from abc import ABC
 import json
 import logging
-from typing import List, Type
+from typing import List, Type, Union
 from src.llm import LLM, get_llm
 from src.utils.log_publisher import LogPrefix, publish_log_info
 
@@ -64,11 +64,12 @@ async def invoke(self, utterance: str) -> str:
         return result_of_action
 
 
-def agent(name: str, description: str, tools: List[Tool]):
-    def decorator(agent: Type[ChatAgent]):
+def agent(name: str, description: str, tools: List[Tool] = None):
+    def decorator(agent: Type[Union[Agent, ChatAgent]]):
         agent.name = name
         agent.description = description
-        agent.tools = tools
+        if tools is not None:
+            agent.tools = tools
         return agent
 
     return decorator

backend/src/agents/materiality_agent.py

@@ -2,6 +2,7 @@
 from pathlib import Path
 import logging
 
+from src.agents import agent
 from src.llm import LLMFileFromPath
 from src.agents import Agent
 from src.prompts import PromptEngine
@@ -10,6 +11,10 @@
 logger = logging.getLogger(__name__)
 
 
+@agent(
+    name="MaterialityAgent",
+    description="This agent is responsible for answering questions about materiality using it's library of documents"
+)
 class MaterialityAgent(Agent):
     async def list_material_topics(self, company_name: str) -> dict[str, str]:
         with open('./library/catalogue.json') as file:

backend/src/agents/report_agent.py

@@ -1,13 +1,18 @@
 import json
 import logging
 
+from src.agents import agent
 from src.agents import Agent
 from src.prompts import PromptEngine
 
 logger = logging.getLogger(__name__)
 engine = PromptEngine()
 
 
+@agent(
+    name="ReportAgent",
+    description="This agent is responsible for generating an ESG focused report on a narrative document"
+)
 class ReportAgent(Agent):
     async def create_report(self, file_content: str, materiality_topics: dict[str, str]) -> str:
         user_prompt = engine.load_prompt(

backend/tests/agents/agent_test.py

@@ -0,0 +1,55 @@
+import json
+
+import pytest
+from pytest import raises
+from tests.agents import MockChatAgent
+from src.llm.factory import get_llm
+
+
+mock_model = "mockmodel"
+mock_llm = get_llm("mockllm")
+mock_agent_instance = MockChatAgent("mockllm", mock_model)
+
+
+def mock_response(tool_name: str, tool_parameters: dict[str,str]) -> str:
+    return json.dumps({"tool_name": tool_name, "tool_parameters": tool_parameters, "reasoning": "Mock Reasoning"})
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("mocked_response, expect_success, expected", [
+        (
+            mock_response("Mock Tool A", {"input": "string for tool to output"}),
+            True,
+            "string for tool to output"
+        ),
+        (
+            mock_response("Undefined Tool", {}),
+            False,
+            "Unable to extract chosen tool and parameters from {'tool_name': 'None', 'tool_parameters': {},"
+            " 'reasoning': 'No tool was appropriate for the task'}"
+        ),
+        (
+            mock_response("None", {}),
+            False,
+            "Unable to extract chosen tool and parameters from {'tool_name': 'None', 'tool_parameters': {},"
+            " 'reasoning': 'No tool was appropriate for the task'}"
+        )
+    ],
+    ids=[
+        "When appropriate tool selected, Test Chat Agent invoke func will call tool with parameters",
+        "When 'Undefined Tool' selected, Test Chat Agent invoke func will explain no tool was appropriate for task",
+        "When 'None' selected, Test Chat Agent invoke func will explain no tool was appropriate for task"
+    ]
+)
+async def test_chat_agent_invoke_uses_tool(mocker, mocked_response: str, expect_success: bool, expected: str):
+    mock_llm.chat = mocker.AsyncMock(return_value=str(mocked_response))
+
+    if expect_success:
+        response = await mock_agent_instance.invoke("Mock task to solve")
+
+        assert response == expected
+    else:
+        with raises(Exception) as error:
+            await mock_agent_instance.invoke("Mock task to solve")
+
+            assert str(error.value) == expected