Fixed Ollama tool_query after tests

antoninoLorenzo · Aug 12, 2024 · ecc6f16 · ecc6f16
1 parent a63f0d4
commit ecc6f16
Showing 1 changed file with 9 additions and 3 deletions.
diff --git a/src/agent/llm/llm.py b/src/agent/llm/llm.py
@@ -14,7 +14,6 @@
 
 from src.agent.memory.base import Role
 
-
 AVAILABLE_MODELS = {
     'llama3': {
         'options': {
@@ -140,16 +139,23 @@ def tool_query(self, messages: list, tools: list | None = None):
         if not AVAILABLE_MODELS[self.model]['tools']:
             raise NotImplementedError(f'Model {self.model} do not implement tool calling')
 
+        try:
+            self.verify_messages_format(messages)
+        except (TypeError, ValueError) as input_err:
+            raise input_err from input_err
+
         if not tools:
-            # TODO: should add validation for tools
             raise ValueError('Empty tool list')
 
-        return self.client.chat(
+        tool_response = self.client.chat(
             model=self.model,
             messages=messages,
             tools=tools
         )
 
+        return tool_response if tool_response['message'].get('tool_calls') \
+            else None
+
 
 @dataclass
 class OpenRouter(Provider):