From 190613f05ce8d0b64dd38caea34e7a2617c61a05 Mon Sep 17 00:00:00 2001
From: sarath2496 <sarathc225@gmail.com>
Date: Tue, 23 Jul 2024 20:58:18 -0500
Subject: [PATCH] AgentOps Integration

---
 autogenui/manager.py | 43 ++++++++++++++++++++++++++++++++-----------
 1 file changed, 32 insertions(+), 11 deletions(-)

diff --git a/autogenui/manager.py b/autogenui/manager.py
index 7667a92..ef73b2f 100644
--- a/autogenui/manager.py
+++ b/autogenui/manager.py
@@ -1,18 +1,23 @@
-# a manager class that can
-# load an autogen flow run an autogen flow and return the response to the client
-
-
 from typing import Dict
 import autogen
 from .utils import parse_token_usage
 import time
+import agentops
+from dotenv import load_dotenv
+import os
 
+# Load environment variables
+load_dotenv()
+
+# Initialize AgentOps
+agentops.init(os.getenv('AGENTOPS_API_KEY'))
 
 class Manager(object):
+    @agentops.record_function('Manager_init')
     def __init__(self) -> None:
-
         pass
 
+    @agentops.record_function('Manager_run_flow')
     def run_flow(self, prompt: str, flow: str = "default") -> None:
         #autogen.ChatCompletion.start_logging(compact=False)
         config_list = autogen.config_list_openai_aoai()
@@ -25,7 +30,9 @@ def run_flow(self, prompt: str, flow: str = "default") -> None:
 
         assistant = autogen.AssistantAgent(
             name="assistant",
-            max_consecutive_auto_reply=3, llm_config=llm_config,)
+            max_consecutive_auto_reply=3, 
+            llm_config=llm_config,
+        )
 
         # create a UserProxyAgent instance named "user_proxy"
         user_proxy = autogen.UserProxyAgent(
@@ -39,18 +46,32 @@ def run_flow(self, prompt: str, flow: str = "default") -> None:
                 "use_docker": False
             },
         )
+        
         start_time = time.time()
-        user_proxy.initiate_chat(
-            assistant,
-            message=prompt,
-        )
+        
+        with agentops.record_span('user_proxy_chat'):
+            user_proxy.initiate_chat(
+                assistant,
+                message=prompt,
+            )
 
         messages = user_proxy.chat_messages[assistant]
         #logged_history = autogen.ChatCompletion.logged_history
         autogen.ChatCompletion.stop_logging()
+        
+        duration = time.time() - start_time
+        
         response = {
             "messages": messages[1:],
             "usage": "", #parse_token_usage(logged_history),
-            "duration": time.time() - start_time,
+            "duration": duration,
         }
+        
+        # Record metrics
+        agentops.record_metric('chat_duration', duration)
+        agentops.record_metric('message_count', len(messages) - 1)
+        
         return response
+
+# End of program
+agentops.end_session('Success')
\ No newline at end of file