AgentOps Integration

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')