refactor agent response handling

- Introduce LLMResponse class to separate text and tool calls from LLM responses
- Modify Agent.process_response() to Agent.call_tools() to handle tool calls
- Update LLM classes to use LLMResponse
- Remove redundant code and improve error handling

Author: theonlyamos

cognitrix/agents/base.py

@@ -4,15 +4,14 @@
 import asyncio
 import logging
 import aiofiles
-from pathlib import Path
 from threading import Thread
 from datetime import datetime
 from typing import Dict, List, Optional, Self, TypeAlias, Union, Type, Any
 
 from pydantic import BaseModel, Field
 
 from cognitrix.tasks import Task
-from cognitrix.llms.base import LLM
+from cognitrix.llms.base import LLM, LLMResponse
 from cognitrix.tools.base import Tool
 from cognitrix.utils import extract_json, json_return_format
 from cognitrix.agents.templates import AUTONOMOUSE_AGENT_2
@@ -184,27 +183,12 @@ def get_sub_agent_by_name(self, name: str) -> Optional['Agent']:
     def get_tool_by_name(self, name: str) -> Optional[Tool]:
         return next((tool for tool in self.tools if tool.name.lower() == name.lower()), None)
 
-    async def process_response(self, response: str|dict) -> Union[dict, str]:
-        # print(response)
-        # response = response.replace("'", '"')
-        response_data = response
-        if isinstance(response, str):
-            response = response.replace('\\n', '')
-            # response = response.replace("'", "\""
-            # response = response.replace('"', '\\"')
-            response_data = extract_json(response)
-        
-
+    async def call_tools(self, tool_calls: list) -> Union[dict, str]:
+ 
         try:
-            if isinstance(response_data, dict):
-                # final_result_keys = ['final_answer', 'tool_calls_result', 'response']
-                
+            if tool_calls:
                 tool_calls_result = []
-                
-                if response_data['type'].replace('\\', '') != 'tool_calls':
-                    return response_data['result']
-
-                for t in response_data['tool_calls']:
+                for t in tool_calls:
                     tool = self.get_tool_by_name(t['name'])
 
                     if not tool:
@@ -235,8 +219,8 @@ async def process_response(self, response: str|dict) -> Union[dict, str]:
             else:
                 raise Exception('Not a json object')
         except Exception as e:
-            # logger.exception(e)
-            return response_data
+            logger.exception(e)
+            return str(e)
 
     def add_tool(self, tool: Tool):
         self.tools.append(tool)
@@ -282,18 +266,22 @@ def initialize(self, session_id: Optional[str] = None):
                 response: Any = self.llm(full_prompt)
 
                 self.llm.chat_history.append(full_prompt)
-                self.llm.chat_history.append({'role': self.name, 'type': 'text', 'message': response})
-
-                if self.verbose:
-                    print(response)
-
-                result: dict[Any, Any] | str = asyncio.run(self.process_response(response))
+                
+                if response.text:
+                    self.llm.chat_history.append({'role': self.name, 'type': 'text', 'message': response.text})
+                    print(f"\n{self.name}: {response.text}")
+                
+                if response.tool_calls:
+                    result: dict[Any, Any] | str = asyncio.run(self.call_tools(response.tool_calls))
 
-                if isinstance(result, dict) and result['type'] == 'tool_calls_result':
-                    query = result
+                    if isinstance(result, dict) and result['type'] == 'tool_calls_result':
+                        query = result
+                    else:
+                        print(result)
                 else:
-                    print(f"\n{self.name}: {result}")
                     query = input("\nUser (q to quit): ")
+
+                # query = input("\nUser (q to quit): ")
 
                 self.save_session(session)
 
@@ -322,7 +310,7 @@ def handle_transcription(self, sentence: str, transcriber: Transcriber):
             if self.verbose:
                 print(response)
 
-            processsed_response: dict[Any, Any] | str = asyncio.run(self.process_response(response))
+            processsed_response: dict[Any, Any] | str = asyncio.run(self.call_tools(response))
 
             if isinstance(processsed_response, dict) and processsed_response['type'] == 'tool_calls_result':
                 query = processsed_response
@@ -361,7 +349,7 @@ def run_task(self, parent: Self):
             if parent.verbose:
                 print(response)
 
-            agent_result = asyncio.run(self.process_response(response))
+            agent_result = asyncio.run(self.call_tools(response))
             if isinstance(agent_result, dict) and agent_result['type'] == 'tool_calls_result':
                 query = agent_result
             else:
@@ -371,7 +359,7 @@ def run_task(self, parent: Self):
                 parent_response: Any = parent.llm(parent_prompt)
                 parent.llm.chat_history.append(parent_prompt)
                 parent.llm.chat_history.append({'role': 'assistant', 'type': 'text', 'message': parent_response})
-                parent_result = asyncio.run(parent.process_response(parent_response))
+                parent_result = asyncio.run(parent.call_tools(parent_response))
                 print(f"\n\n{parent.name}: {parent_result}")
                 query = ""
 

cognitrix/llms/__init__.py

@@ -1,4 +1,5 @@
 from cognitrix.llms.base import LLM
+from cognitrix.llms.base import LLMResponse
 from cognitrix.llms.groq_llm import Groq
 from cognitrix.llms.local_llm import Local
 from cognitrix.llms.cohere_llm import Cohere

cognitrix/llms/base.py

@@ -1,10 +1,10 @@
 import json
 from pydantic import BaseModel, Field
-from typing import Any, List, Dict
+from typing import Any, List, Dict, Optional, TypedDict
 import logging
 import inspect
 
-from cognitrix.utils import image_to_base64
+from cognitrix.utils import extract_json, image_to_base64
 from cognitrix.tools.base import Tool
 
 logging.basicConfig(
@@ -14,6 +14,35 @@
 )
 logger = logging.getLogger('cognitrix.log')
 
+class LLMResponse:
+    """Class to handle and separate LLM responses into text and tool calls."""
+    
+    def __init__(self, llm_response: Optional[str]=None):
+        self.llm_response = llm_response
+        self.text: Optional[str] = None
+        self.tool_calls: Optional[List[Dict[str, Any]]] = None
+        self.parse_llm_response()
+    
+    def parse_llm_response(self):
+        """Parse the LLM response into text and tool calls."""
+        
+        if not self.llm_response: return 
+        
+        response_data = extract_json(self.llm_response)
+
+        try:
+            if isinstance(response_data, dict):
+                if 'result' in response_data.keys():
+                    self.text = response_data['result']
+                else:
+                    self.tool_calls = response_data['tool_calls']
+            else:
+                self.text = str(response_data)
+        except Exception as e:
+            logger.exception(e)
+            self.text = str(response_data)
+            
+
 class LLM(BaseModel):
     """
     A class for representing a large language model.
@@ -162,4 +191,4 @@ def load_llm(model_name: str):
             return None
 
     def __call__(*args, **kwargs):
-        pass
+        return LLMResponse()

cognitrix/llms/clarifai_llm.py

@@ -1,5 +1,5 @@
 from clarifai.client.model import Model
-from cognitrix.llms.base import LLM
+from cognitrix.llms.base import LLM, LLMResponse
 from typing import Any, Optional
 from dotenv import load_dotenv
 import logging
@@ -36,7 +36,7 @@ class Clarifai(LLM):
     api_key: str = os.getenv('CLARIFAI_API_KEY', '')
     """Clarifai Personal Access Token""" 
 
-    def __call__(self, query, **kwds: Any)->str:
+    def __call__(self, query, **kwds: Any):
         """Generates a response to a query using the Clarifai API.
 
         Args:
@@ -54,5 +54,5 @@ def __call__(self, query, **kwds: Any)->str:
         query = f"{self.system_prompt}\n {json.dumps(formatted_messages)}"
         result = self.client.predict_by_bytes(query.encode(), input_type="text")
 
-        return result.outputs[0].data.text.raw
+        return LLMResponse(result.outputs[0].data.text.raw)
 

cognitrix/llms/cohere_llm.py

@@ -1,5 +1,5 @@
 import cohere
-from cognitrix.llms.base import LLM
+from cognitrix.llms.base import LLM, LLMResponse
 from typing import Any, Optional
 from dotenv import load_dotenv
 import logging
@@ -54,7 +54,7 @@ def format_tools(self, tools: list[dict[str, Any]]):
 
             self.tools.append(f_tool)
 
-    def __call__(self, query, **kwds: Any)->str:
+    def __call__(self, query, **kwds: Any):
         """Generates a response to a query using the Cohere API.
 
         Args:
@@ -79,7 +79,7 @@ def __call__(self, query, **kwds: Any)->str:
             connectors=[{"id": "web-search"}]
         )
 
-        return response.text
+        return LLMResponse(response.text)
 
 if __name__ == "__main__":
     try:

cognitrix/llms/google_llm.py

@@ -1,4 +1,4 @@
-from cognitrix.llms.base import LLM
+from cognitrix.llms.base import LLM, LLMResponse
 from typing import Any
 import google.generativeai as genai
 from dotenv import load_dotenv
@@ -72,7 +72,7 @@ def format_query(self, message: dict[str, str]) -> list:
 
         return messages
 
-    def __call__(self, query, **kwds: Any)->str|None:
+    def __call__(self, query, **kwds: Any):
         """Generates a response to a query using the Gemini API.
 
         Args:
@@ -104,15 +104,5 @@ def __call__(self, query, **kwds: Any)->str|None:
 
         response.resolve()
 
-        return response.text
-    
-if __name__ == "__main__":
-    try:
-        assistant = Google()
-        # assistant.add_tool(calculator)
-        while True:
-            message = input("\nEnter Query$ ")
-            result = assistant(message)
-            print(result)
-    except KeyboardInterrupt:
-        sys.exit(1)
+        return LLMResponse(response.text)
+

cognitrix/llms/groq_llm.py

@@ -1,5 +1,5 @@
 from groq import Groq as GroqLLM
-from cognitrix.llms.base import LLM
+from cognitrix.llms.base import LLM, LLMResponse
 from typing import Any, Optional
 from dotenv import load_dotenv
 import logging
@@ -74,7 +74,7 @@ def format_query(self, message: dict[str, str]) -> list:
         return messages
 
 
-    def __call__(self, query, **kwds: Any)->str|None:
+    def __call__(self, query, **kwds: Any):
         """Generates a response to a query using the Groq API.
 
         Args:
@@ -107,7 +107,7 @@ def __call__(self, query, **kwds: Any)->str|None:
         # tool_calls = response_message.tool_calls
         # print(tool_calls)
 
-        return response_message.content
+        return LLMResponse(response_message.content)
 
 if __name__ == "__main__":
     try:

cognitrix/llms/local_llm.py

@@ -1,5 +1,5 @@
 from openai import OpenAI as OpenAILLM
-from cognitrix.llms.base import LLM
+from cognitrix.llms.base import LLM, LLMResponse
 from cognitrix.utils import image_to_base64
 from typing import Any, Optional
 from dotenv import load_dotenv
@@ -95,7 +95,7 @@ def format_query(self, message: dict[str, str]) -> list:
 
         return messages
 
-    def __call__(self, query: dict, **kwds: dict)->Optional[str]:
+    def __call__(self, query: dict, **kwds: dict):
         """Generates a response to a query using the OpenAI API.
 
         Args:
@@ -121,4 +121,4 @@ def __call__(self, query: dict, **kwds: dict)->Optional[str]:
             max_tokens=self.max_tokens
         )
 
-        return response.choices[0].message.content                          #type: ignore
+        return LLMResponse(response.choices[0].message.content)                   #type: ignore

cognitrix/llms/mindsdb_llm.py

@@ -3,6 +3,8 @@
 from typing import Optional
 import os
 
+from cognitrix.llms.base import LLMResponse
+
 class MindsDB(OpenAI):
     """A class for interacting with the MindsDB API."""
 
@@ -12,7 +14,7 @@ class MindsDB(OpenAI):
 
     base_url: str = 'https://llm.mdb.ai'
 
-    def __call__(self, query: dict, **kwds: dict)->Optional[str]:
+    def __call__(self, query: dict, **kwds: dict):
         """Generates a response to a query using the OpenAI API.
 
         Args:
@@ -38,4 +40,4 @@ def __call__(self, query: dict, **kwds: dict)->Optional[str]:
             max_tokens=self.max_tokens
         )
 
-        return response.choices[0].message.content
+        return LLMResponse(response.choices[0].message.content)

cognitrix/llms/ollama_llm.py

@@ -1,4 +1,4 @@
-from cognitrix.llms.base import LLM
+from cognitrix.llms.base import LLM, LLMResponse
 from cognitrix.utils import image_to_base64
 from typing import Any, Optional
 from dotenv import load_dotenv
@@ -88,7 +88,7 @@ def format_query(self, message: dict[str, str]) -> list:
 
         return messages
 
-    def __call__(self, query: dict, **kwds: dict)->Optional[str]:
+    def __call__(self, query: dict, **kwds: dict):
         """Generates a response to a query using the OpenAI API.
 
         Args:
@@ -112,5 +112,5 @@ def __call__(self, query: dict, **kwds: dict)->Optional[str]:
             ],
             stream=False,
         )
-        print(response)
-        return response['message']['content']                        #type: ignore
+        
+        return LLMResponse(response['message']['content'])                        #type: ignore