fixes in function calling args (#135)

premAI-io · Jun 10, 2024 · b98134a · b98134a
1 parent 7c62b2d
commit b98134a
Show file tree

Hide file tree

Showing 3 changed files with 23 additions and 3 deletions.
diff --git a/prem_utils/connectors/groq.py b/prem_utils/connectors/groq.py
@@ -1,3 +1,4 @@
+import json
 import logging
 
 from groq import AsyncGroq, Groq
@@ -63,6 +64,12 @@ def parse_chunk(self, chunk):
             ],
         }
 
+    def _get_arguments(self, arguments):
+        try:
+            return json.loads(arguments)
+        except json.JSONDecodeError:
+            return None
+
     async def chat_completion(
         self,
         model: str,
@@ -144,7 +151,7 @@ async def chat_completion(
                             {
                                 "id": tool_call.id,
                                 "function": {
-                                    "arguments": tool_call.function.arguments,
+                                    "arguments": self._get_arguments(tool_call.function.arguments),
                                     "name": tool_call.function.name,
                                 },
                                 "type": tool_call.type,

diff --git a/prem_utils/connectors/mistral.py b/prem_utils/connectors/mistral.py
@@ -1,3 +1,4 @@
+import json
 from collections.abc import Sequence
 
 from mistralai.async_client import MistralAsyncClient
@@ -46,6 +47,12 @@ def build_messages(self, messages):
                 chat_messages.append(chat_message)
         return chat_messages
 
+    def _get_arguments(self, arguments):
+        try:
+            return json.loads(arguments)
+        except json.JSONDecodeError:
+            return None
+
     async def chat_completion(
         self,
         model: str,
@@ -100,7 +107,7 @@ async def chat_completion(
                                 {
                                     "id": tool_call.id,
                                     "function": {
-                                        "arguments": tool_call.function.arguments,
+                                        "arguments": self._get_arguments(tool_call.function.arguments),
                                         "name": tool_call.function.name,
                                     },
                                     "type": tool_call.type.value if tool_call.type else None,

diff --git a/prem_utils/connectors/openai.py b/prem_utils/connectors/openai.py
@@ -70,6 +70,12 @@ def parse_chunk(self, chunk):
             ],
         }
 
+    def _get_arguments(self, arguments):
+        try:
+            return json.loads(arguments)
+        except json.JSONDecodeError:
+            return None
+
     async def chat_completion(
         self,
         model: str,
@@ -148,7 +154,7 @@ async def chat_completion(
                             {
                                 "id": tool_call.id,
                                 "function": {
-                                    "arguments": tool_call.function.arguments,
+                                    "arguments": self._get_arguments(tool_call.function.arguments),
                                     "name": tool_call.function.name,
                                 },
                                 "type": tool_call.type,