From 02903c0a644a7f38d66068e8b0cdc086d3cd5b5c Mon Sep 17 00:00:00 2001
From: f1delius <tagazeem@gmail.com>
Date: Sun, 28 Apr 2024 23:14:40 +0530
Subject: [PATCH] agent thread update

---
 src/bot.py      | 12 +++++-------
 src/workflow.py |  7 ++++---
 2 files changed, 9 insertions(+), 10 deletions(-)

diff --git a/src/bot.py b/src/bot.py
index cb030e6..b1f7dd1 100644
--- a/src/bot.py
+++ b/src/bot.py
@@ -1,6 +1,4 @@
 import asyncio
-import os
-from pathlib import Path
 import re
 import sys
 import traceback
@@ -177,6 +175,10 @@ async def message_callback(self, room: MatrixRoom, event: RoomMessageText) -> No
         if event.formatted_body:
             if self.bot_username in event.formatted_body:
                 tagged = True
+        if thread_id:
+            thread_event_id = thread_id
+        else:
+            thread_event_id = reply_to_event_id
 
         if self.user_id != event.sender and (tagged or room.is_group):
             if self.owner_id != sender_id and self.msg_limit[sender_id] > 10:
@@ -197,10 +199,6 @@ async def message_callback(self, room: MatrixRoom, event: RoomMessageText) -> No
                 if self.workflow:
                     get_steps = await workflow_steps(self.superagent_url, self.workflow_id, self.api_key, self.httpx_client)
                     api_url = f"{self.superagent_url}/api/v1/workflows/{self.workflow_id}/invoke"
-                    if thread_id:
-                        thread_event_id = thread_id
-                    else:
-                        thread_event_id = reply_to_event_id
                     await stream_json_response_with_auth(api_url, self.api_key, content_body, get_steps, thread_event_id, reply_to_event_id, room_id, self.httpx_client)
                     self.msg_limit[sender_id] += len(get_steps)
                     return
@@ -235,7 +233,7 @@ async def message_callback(self, room: MatrixRoom, event: RoomMessageText) -> No
                 )
                 self.msg_limit[sender_id] += 1
             except Exception as e:
-                print(e)
+                logger.error(e)
 
     # message_callback decryption_failure event
 
diff --git a/src/workflow.py b/src/workflow.py
index 660f692..215f98f 100644
--- a/src/workflow.py
+++ b/src/workflow.py
@@ -1,5 +1,6 @@
 import httpx
 import aiohttp
+import logging
 
 from api import send_message_as_tool
 
@@ -45,7 +46,7 @@ async def stream_json_response_with_auth(api_url, api_key, msg_data, agent, thre
                         # Print the complete message for the previous event
                         prev_data = prev_data.replace("````","`\n```")
                         prev_data = prev_data.replace("```", "\n```")
-                        print(
+                        logging.info(
                             f'Event: {prev_event}, Data: {prev_data}')
                         await send_agent_message(agent[prev_event], thread_id, reply_id, prev_data, room_id)
                         # Reset the previous data
@@ -67,10 +68,10 @@ async def stream_json_response_with_auth(api_url, api_key, msg_data, agent, thre
     if prev_event is not None:
         prev_data = prev_data.replace("````","`\n```")
         prev_data = prev_data.replace("```", "\n```")
-        print(f'Event: {prev_event}, Data: {prev_data}')
+        logging.info(f'Event: {prev_event}, Data: {prev_data}')
         await send_agent_message(agent[prev_event], thread_id, reply_id, prev_data, room_id)
     else:
-        print('Failed to fetch streaming data')
+        logging.info('Failed to fetch streaming data')
 
 
 async def send_agent_message(agent, thread_event_id, reply_id, data, room_id):