Pass tool call and app function response back in metadata (#193)

model_server/app/function_calling/model_utils.py

@@ -13,62 +13,38 @@ logger = get_model_server_logger()
 
 class Message(BaseModel):
     role: str
-    content: str
+    content: str = ""
+    tool_calls: List[Dict[str, Any]] = []
+    tool_call_id: str = ""
 
 
 class ChatMessage(BaseModel):
     messages: list[Message]
     tools: List[Dict[str, Any]]
 
-    # TODO: make it default none
-    metadata: Dict[str, str] = {}
 
-
-def process_state(arch_state, history: list[Message]):
-    logger.info("state: {}".format(arch_state))
-    state_json = json.loads(arch_state)
-
-    state_map = {}
-    if state_json:
-        for tools_state in state_json:
-            for tool_state in tools_state:
-                state_map[tool_state["key"]] = tool_state
-
-    logger.info(f"state_map: {json.dumps(state_map)}")
-
-    sha_history = []
+def process_messages(history: list[Message]):
     updated_history = []
     for hist in history:
-        updated_history.append({"role": hist.role, "content": hist.content})
-        if hist.role == "user":
-            sha_history.append(hist.content)
-            sha256_hash = hashlib.sha256()
-            joined_key_str = ("#.#").join(sha_history)
-            sha256_hash.update(joined_key_str.encode())
-            sha_key = sha256_hash.hexdigest()
-            logger.info(f"sha_key: {sha_key}")
-            if sha_key in state_map:
-                tool_call_state = state_map[sha_key]
-                if "tool_call" in tool_call_state:
-                    tool_call_str = json.dumps(tool_call_state["tool_call"])
-                    updated_history.append(
-                        {
-                            "role": "assistant",
-                            "content": f"<tool_call>\n{tool_call_str}\n</tool_call>",
-                        }
-                    )
-                if "tool_response" in tool_call_state:
-                    tool_resp = tool_call_state["tool_response"]
-                    # TODO: try with role = user as well
-                    updated_history.append(
-                        {
-                            "role": "user",
-                            "content": f"<tool_response>\n{tool_resp}\n</tool_response>",
-                        }
-                    )
-                # we dont want to match this state with any other messages
-                del state_map[sha_key]
-
+        if hist.tool_calls:
+            if len(hist.tool_calls) > 1:
+                raise ValueError("Only one tool call is supported")
+            tool_call_str = json.dumps(hist.tool_calls[0]["function"])
+            updated_history.append(
+                {
+                    "role": "assistant",
+                    "content": f"<tool_call>\n{tool_call_str}\n</tool_call>",
+                }
+            )
+        elif hist.role == "tool":
+            updated_history.append(
+                {
+                    "role": "user",
+                    "content": f"<tool_response>\n{hist.content}\n</tool_response>",
+                }
+            )
+        else:
+            updated_history.append({"role": hist.role, "content": hist.content})
     return updated_history
 
 
@@ -79,10 +55,7 @@ async def chat_completion(req: ChatMessage, res: Response):
 
     messages = [{"role": "system", "content": tools_encoded}]
 
-    metadata = req.metadata
-    arch_state = metadata.get("x-arch-state", "[]")
-
-    updated_history = process_state(arch_state, req.messages)
+    updated_history = process_messages(req.messages)
     for message in updated_history:
         messages.append({"role": message["role"], "content": message["content"]})
 

model_server/app/tests/test_state.py

@@ -0,0 +1,66 @@
+from typing import List
+import pytest
+import json
+from app.function_calling.model_utils import Message, process_messages
+
+test_input_history = """
+[
+    {
+      "role": "user",
+      "content": "how is the weather in chicago for next 5 days?"
+    },
+    {
+      "role": "assistant",
+      "model": "Arch-Function-1.5B",
+      "tool_calls": [
+        {
+          "id": "call_3394",
+          "type": "function",
+          "function": {
+            "name": "weather_forecast",
+            "arguments": { "city": "Chicago", "days": 5 }
+          }
+        }
+      ]
+    },
+    {
+      "role": "tool",
+      "content": "--",
+      "tool_call_id": "call_3394"
+    },
+    {
+      "role": "assistant",
+      "content": "--",
+      "model": "gpt-3.5-turbo-0125"
+    },
+    {
+      "role": "user",
+      "content": "how is the weather in chicago for next 5 days?"
+    },
+    {
+      "role": "assistant",
+      "tool_calls": [
+        {
+          "id": "call_5306",
+          "type": "function",
+          "function": {
+            "name": "weather_forecast",
+            "arguments": { "city": "Chicago", "days": 5 }
+          }
+        }
+      ]
+    }
+  ]
+  """
+
+
+def test_update_fc_history():
+    history = json.loads(test_input_history)
+    message_history = []
+    for h in history:
+        message_history.append(Message(**h))
+
+    updated_history = process_messages(message_history)
+    assert len(updated_history) == 6
+    # ensure that tool role does not exist anymore
+    assert all([h["role"] != "tool" for h in updated_history])