From 0bb26ae50477ca70ada0824e3d75aa668c187906 Mon Sep 17 00:00:00 2001
From: "Timothy J. Baek" <timothyjrbeck@gmail.com>
Date: Tue, 11 Jun 2024 11:31:14 -0700
Subject: [PATCH] refac: tools

---
 backend/apps/socket/main.py | 12 ------------
 backend/main.py             | 20 ++++++++++++--------
 2 files changed, 12 insertions(+), 20 deletions(-)

diff --git a/backend/apps/socket/main.py b/backend/apps/socket/main.py
index e70812867..123ff31cd 100644
--- a/backend/apps/socket/main.py
+++ b/backend/apps/socket/main.py
@@ -19,8 +19,6 @@ TIMEOUT_DURATION = 3
 
 @sio.event
 async def connect(sid, environ, auth):
-    print("connect ", sid)
-
     user = None
     if auth and "token" in auth:
         data = decode_token(auth["token"])
@@ -37,7 +35,6 @@ async def connect(sid, environ, auth):
 
             print(f"user {user.name}({user.id}) connected with session ID {sid}")
 
-            print(len(set(USER_POOL)))
             await sio.emit("user-count", {"count": len(set(USER_POOL))})
             await sio.emit("usage", {"models": get_models_in_use()})
 
@@ -64,13 +61,11 @@ async def user_join(sid, data):
 
             print(f"user {user.name}({user.id}) connected with session ID {sid}")
 
-            print(len(set(USER_POOL)))
             await sio.emit("user-count", {"count": len(set(USER_POOL))})
 
 
 @sio.on("user-count")
 async def user_count(sid):
-    print("user-count", sid)
     await sio.emit("user-count", {"count": len(set(USER_POOL))})
 
 
@@ -79,14 +74,12 @@ def get_models_in_use():
     models_in_use = []
     for model_id, data in USAGE_POOL.items():
         models_in_use.append(model_id)
-    print(f"Models in use: {models_in_use}")
 
     return models_in_use
 
 
 @sio.on("usage")
 async def usage(sid, data):
-    print(f'Received "usage" event from {sid}: {data}')
 
     model_id = data["model"]
 
@@ -114,7 +107,6 @@ async def usage(sid, data):
 
 async def remove_after_timeout(sid, model_id):
     try:
-        print("remove_after_timeout", sid, model_id)
         await asyncio.sleep(TIMEOUT_DURATION)
         if model_id in USAGE_POOL:
             print(USAGE_POOL[model_id]["sids"])
@@ -124,7 +116,6 @@ async def remove_after_timeout(sid, model_id):
             if len(USAGE_POOL[model_id]["sids"]) == 0:
                 del USAGE_POOL[model_id]
 
-            print(f"Removed usage data for {model_id} due to timeout")
             # Broadcast the usage data to all clients
             await sio.emit("usage", {"models": get_models_in_use()})
     except asyncio.CancelledError:
@@ -143,9 +134,6 @@ async def disconnect(sid):
         if len(USER_POOL[user_id]) == 0:
             del USER_POOL[user_id]
 
-        print(f"user {user_id} disconnected with session ID {sid}")
-        print(USER_POOL)
-
         await sio.emit("user-count", {"count": len(USER_POOL)})
     else:
         print(f"Unknown session ID {sid} disconnected")
diff --git a/backend/main.py b/backend/main.py
index 187a0b720..04b4ebffc 100644
--- a/backend/main.py
+++ b/backend/main.py
@@ -178,7 +178,7 @@ async def get_function_call_response(messages, tool_id, template, task_model_id,
         "History:\n"
         + "\n".join(
             [
-                f"{message['role']}: {message['content']}"
+                f"{message['role'].upper()}: \"\"\"{message['content']}\"\"\""
                 for message in messages[::-1][:4]
             ]
         )
@@ -209,17 +209,21 @@ async def get_function_call_response(messages, tool_id, template, task_model_id,
             response = await generate_openai_chat_completion(payload, user=user)
 
         content = None
-        async for chunk in response.body_iterator:
-            data = json.loads(chunk.decode("utf-8"))
-            content = data["choices"][0]["message"]["content"]
 
-        # Cleanup any remaining background tasks if necessary
-        if response.background is not None:
-            await response.background()
+        if hasattr(response, "body_iterator"):
+            async for chunk in response.body_iterator:
+                data = json.loads(chunk.decode("utf-8"))
+                content = data["choices"][0]["message"]["content"]
+
+            # Cleanup any remaining background tasks if necessary
+            if response.background is not None:
+                await response.background()
+        else:
+            content = response["choices"][0]["message"]["content"]
 
         # Parse the function response
         if content is not None:
-            print(content)
+            print(f"content: {content}")
             result = json.loads(content)
             print(result)