AI4Finance-Foundation · sarath59 · Jul 12, 2024
diff --git a/agent_builder_demo.py b/agent_builder_demo.py
@@ -2,12 +2,19 @@
 import autogen
 from autogen.agentchat.contrib.agent_builder import AgentBuilder
 from finrobot.utils import get_current_date
+import agentops
 
+# Initialize AgentOps
+agentops.init('<INSERT YOUR API KEY HERE>')
 
 config_file_or_env = "OAI_CONFIG_LIST"
 llm_config = {"temperature": 0}
 
-builder = AgentBuilder(
+@agentops.track_agent(name='AgentBuilder')
+class TrackedAgentBuilder(AgentBuilder):
+    pass
+
+builder = TrackedAgentBuilder(
     config_file_or_env=config_file_or_env,
     builder_model="gpt-4-0125-preview",
     agent_model="gpt-4-0125-preview",
@@ -34,11 +41,31 @@
     )
     builder.save(config_path)
 
-group_chat = autogen.GroupChat(agents=agent_list, messages=[], max_round=20)
-manager = autogen.GroupChatManager(
+@agentops.track_agent(name='GroupChat')
+class TrackedGroupChat(autogen.GroupChat):
+    pass
+
+@agentops.track_agent(name='GroupChatManager')
+class TrackedGroupChatManager(autogen.GroupChatManager):
+    pass
+
+group_chat = TrackedGroupChat(agents=agent_list, messages=[], max_round=20)
+manager = TrackedGroupChatManager(
     groupchat=group_chat, llm_config={"config_list": config_list, **llm_config}
 )
-agent_list[0].initiate_chat(
+
+@agentops.record_function('initiate_chat')
+def initiate_chat(agent, manager, message):
+    agent.initiate_chat(
+        manager,
+        message=message,
+    )
+
+initiate_chat(
+    agent_list[0],
     manager,
-    message=f"Today is {get_current_date()}, predict next week's stock price for Nvidia with its recent market news and stock price movements.",
+    message=f"Today is {get_current_date()}, predict next week's stock price for Nvidia with its recent market news and stock price movements."
 )
+
+# End of program
+agentops.end_session('Success')
diff --git a/experiments/multi_factor_agents.py b/experiments/multi_factor_agents.py
@@ -2,11 +2,11 @@
 import json
 import autogen
 from autogen.cache import Cache
-
-# from finrobot.utils import create_inner_assistant
-
 from functools import partial
+import agentops
 
+# Initialize AgentOps
+agentops.init("YOUR_AGENTOPS_API_KEY_HERE")
 
 config_list_gpt4 = autogen.config_list_from_json(
     "OAI_CONFIG_LIST",
@@ -23,21 +23,19 @@
 
 quant_group_config = json.load(open("quantitative_investment_group_config.json"))
 
-# user_proxy = autogen.UserProxyAgent(
-#     name="User",
-#     # human_input_mode="ALWAYS",
-#     human_input_mode="NEVER",
-#     code_execution_config=False
-# )
-
 group_descs = "\n\n".join(
     [
         "Name: {} \nResponsibility: {}".format(c["name"], c["profile"])
         for c in quant_group_config
     ]
 )
 
-group_leader = autogen.AssistantAgent(
+@agentops.track_agent(name="Group_Leader")
+class GroupLeader(autogen.AssistantAgent):
+    def __init__(self, name, system_message, llm_config):
+        super().__init__(name=name, system_message=system_message, llm_config=llm_config)
+
+group_leader = GroupLeader(
     name="Group_Leader",
     system_message="""
     As a group leader, you are responsible for coordinating the team's efforts to achieve the project's objectives. 
@@ -52,13 +50,15 @@
     llm_config=llm_config,
 )
 
-executor = autogen.UserProxyAgent(
+class EnhancedExecutor(autogen.UserProxyAgent):
+    @agentops.record_function("execute_code")
+    def execute_code(self, code, **kwargs):
+        return super().execute_code(code, **kwargs)
+
+executor = EnhancedExecutor(
     name="Executor",
     human_input_mode="NEVER",
-    # human_input_mode="ALWAYS",
-    is_termination_msg=lambda x: x.get("content", "")
-    and "TERMINATE" in x.get("content", ""),
-    # max_consecutive_auto_reply=3,
+    is_termination_msg=lambda x: x.get("content", "") and "TERMINATE" in x.get("content", ""),
     code_execution_config={
         "last_n_messages": 3,
         "work_dir": "quant",
@@ -67,21 +67,18 @@
 )
 
 quant_group = {
-    c["name"]: autogen.agentchat.AssistantAgent(
+    c["name"]: agentops.track_agent(name=c["name"])(autogen.agentchat.AssistantAgent)(
         name=c["name"],
         system_message=c["profile"],
         llm_config=llm_config,
     )
     for c in quant_group_config
 }
 
-
 def order_trigger(pattern, sender):
-    # print(pattern)
-    # print(sender.last_message()['content'])
     return pattern in sender.last_message()["content"]
 
-
+@agentops.record_function("process_order")
 def order_message(pattern, recipient, messages, sender, config):
     full_order = recipient.chat_messages_for_summary(sender)[-1]["content"]
     pattern = rf"\[{pattern}\](?::)?\s*(.+?)(?=\n\[|$)"
@@ -97,8 +94,6 @@ def order_message(pattern, recipient, messages, sender, config):
     DO NOT include TERMINATE in your response until you have received the results from the execution of the Python scripts.
     If the task cannot be done currently or need assistance from other members, report the reasons or requirements to group leader ended with TERMINATE. 
     """
-    # For coding tasks, only use the functions you have been provided with.
-
 
 for name, agent in quant_group.items():
     executor.register_nested_chats(
@@ -117,5 +112,11 @@ def order_message(pattern, recipient, messages, sender, config):
 
 quant_task = "Develop and test the feasibility of a quantitative investment strategy focusing on the Dow Jones 30 stocks, utilizing your multi-factor analysis expertise to identify potential investment opportunities and optimize the portfolio's performance. Ensure the strategy is robust, data-driven, and aligns with our risk management principles."
 
-with Cache.disk() as cache:
-    executor.initiate_chat(group_leader, message=quant_task, cache=cache)
+@agentops.record_function("run_quant_strategy")
+def run_quant_strategy():
+    with Cache.disk() as cache:
+        executor.initiate_chat(group_leader, message=quant_task, cache=cache)
+
+if __name__ == "__main__":
+    run_quant_strategy()
+    agentops.end_session('Success')
diff --git a/experiments/portfolio_optimization.py b/experiments/portfolio_optimization.py
@@ -5,70 +5,86 @@
 from textwrap import dedent
 from autogen import register_function
 from investment_group import group_config
-
-
-llm_config = {
-    "config_list": autogen.config_list_from_json(
-        "../OAI_CONFIG_LIST",
-        filter_dict={
-            "model": ["gpt-4-0125-preview"],
-        },
-    ),
-    "cache_seed": 42,
-    "temperature": 0,
-}
+import agentops
+
+# Initialize AgentOps
+agentops.init("YOUR_AGENTOPS_API_KEY_HERE")
+
+@agentops.record_function("setup_llm_config")
+def setup_llm_config():
+    return {
+        "config_list": autogen.config_list_from_json(
+            "../OAI_CONFIG_LIST",
+            filter_dict={
+                "model": ["gpt-4-0125-preview"],
+            },
+        ),
+        "cache_seed": 42,
+        "temperature": 0,
+    }
+
+llm_config = setup_llm_config()
 
 register_keys_from_json("../config_api_keys")
 
-# group_config = json.load(open("investment_group.json"))
+@agentops.track_agent(name="UserProxy")
+class EnhancedUserProxy(autogen.UserProxyAgent):
+    @agentops.record_function("execute_code")
+    def execute_code(self, code, **kwargs):
+        return super().execute_code(code, **kwargs)
 
-user_proxy = autogen.UserProxyAgent(
+user_proxy = EnhancedUserProxy(
     name="User",
-    # human_input_mode="ALWAYS",
     human_input_mode="NEVER",
-    is_termination_msg=lambda x: x.get("content", "")
-    and "TERMINATE" in x.get("content", ""),
+    is_termination_msg=lambda x: x.get("content", "") and "TERMINATE" in x.get("content", ""),
     code_execution_config={
         "last_n_messages": 3,
         "work_dir": "quant",
         "use_docker": False,
     },
 )
 
-rag_func = get_rag_function(
-    retrieve_config={
-        "task": "qa",
-        "docs_path": "https://www.sec.gov/Archives/edgar/data/1737806/000110465923049927/pdd-20221231x20f.htm",
-        "chunk_token_size": 1000,
-        "collection_name": "pdd2022",
-        "get_or_create": True,
-    },
-)
+@agentops.record_function("setup_rag_function")
+def setup_rag_function():
+    return get_rag_function(
+        retrieve_config={
+            "task": "qa",
+            "docs_path": "https://www.sec.gov/Archives/edgar/data/1737806/000110465923049927/pdd-20221231x20f.htm",
+            "chunk_token_size": 1000,
+            "collection_name": "pdd2022",
+            "get_or_create": True,
+        },
+    )
+
+rag_func = setup_rag_function()
 
 with_leader_config = {
     "Market Sentiment Analysts": True,
     "Risk Assessment Analysts": True,
     "Fundamental Analysts": True,
 }
 
-representatives = []
-
-for group_name, single_group_config in group_config["groups"].items():
-
-    with_leader = with_leader_config.get(group_name)
+@agentops.record_function("create_group")
+def create_group(group_name, single_group_config, with_leader):
     if with_leader:
         group_members = single_group_config["with_leader"]
         group_members["agents"] = group_members.pop("employees")
-        group = MultiAssistantWithLeader(
+        return MultiAssistantWithLeader(
             group_members, llm_config=llm_config, user_proxy=user_proxy
         )
     else:
         group_members = single_group_config["without_leader"]
         group_members["agents"] = group_members.pop("employees")
-        group = MultiAssistant(
+        return MultiAssistant(
             group_members, llm_config=llm_config, user_proxy=user_proxy
         )
 
+representatives = []
+
+for group_name, single_group_config in group_config["groups"].items():
+    with_leader = with_leader_config.get(group_name)
+    group = create_group(group_name, single_group_config, with_leader)
+
     for agent in group.agents:
         register_function(
             rag_func,
@@ -79,12 +95,15 @@
 
     representatives.append(group.representative)
 
+@agentops.record_function("create_main_group")
+def create_main_group():
+    cio_config = group_config["CIO"]
+    main_group_config = {"leader": cio_config, "agents": representatives}
+    return MultiAssistantWithLeader(
+        main_group_config, llm_config=llm_config, user_proxy=user_proxy
+    )
 
-cio_config = group_config["CIO"]
-main_group_config = {"leader": cio_config, "agents": representatives}
-main_group = MultiAssistantWithLeader(
-    main_group_config, llm_config=llm_config, user_proxy=user_proxy
-)
+main_group = create_main_group()
 
 task = dedent(
     """
@@ -124,46 +143,10 @@
     """
 )
 
-# task = dedent(
-#     """
-#     As the Chief Investment Officer, your task is to evaluate the potential investment in Company ABC based on the provided data. You will need to coordinate with the Market Sentiment Analysts, Risk Assessment Analysts, and Fundamental Analysts to gather and analyze the relevant information. Your final deliverable should include a comprehensive evaluation and a recommendation on whether to invest in Company ABC.
-
-#     Specific Instructions:
-
-#     Coordinate with Market Sentiment Analysts:
-
-#     Task: Calculate the sentiment score based on the provided market sentiment data.
-#     Data: Positive mentions (80), Negative mentions (20)
-#     Formula: Sentiment Score = (Positive Mentions - Negative Mentions) / Total Mentions
-#     Expected Output: Sentiment Score (percentage)
-
-#     Coordinate with Risk Assessment Analysts:
-
-#     Task: Calculate the risk score using the provided financial ratios.
-#     Data:
-#     Debt-to-Equity Ratio: 1.5
-#     Current Ratio: 2.0
-#     Return on Equity (ROE): 0.1 (10%)
-#     Weights: Debt-to-Equity (0.5), Current Ratio (0.3), ROE (0.2)
-#     Formula: Risk Score = 0.5 * Debt-to-Equity + 0.3 * (1 / Current Ratio) - 0.2 * ROE
-#     Expected Output: Risk Score
-
-#     Coordinate with Fundamental Analysts:
-
-#     Task: Calculate the Profit Margin and Return on Assets (ROA) based on the provided financial data.
-#     Data:
-#     Revenue: $1,000,000
-#     Net Income: $100,000
-#     Total Assets: $500,000
-#     Formulas:
-#     Profit Margin = (Net Income / Revenue) * 100
-#     ROA = (Net Income / Total Assets) * 100
-#     Expected Outputs: Profit Margin (percentage) and ROA (percentage)
-
-#     Final Deliverable:
-#     Integrate Findings: Compile the insights from all three groups to get a holistic view of Company ABC's potential.
-#     Evaluation and Recommendation: Based on the integrated analysis, provide a recommendation on whether to invest in Company ABC, including the rationale behind your decision.
-# """
-# )
-
-main_group.chat(message=task, use_cache=True)
+@agentops.record_function("run_investment_analysis")
+def run_investment_analysis():
+    main_group.chat(message=task, use_cache=True)
+
+if __name__ == "__main__":
+    run_investment_analysis()
+    agentops.end_session('Success')
diff --git a/finrobot/agents/workflow.py b/finrobot/agents/workflow.py
@@ -1,3 +1,6 @@
+# Install AgentOps SDK
+# pip install agentops
+
 from .agent_library import library
 from typing import Any, Callable, Dict, List, Optional, Annotated
 import autogen
@@ -17,8 +20,12 @@
 from ..functional.rag import get_rag_function
 from .utils import *
 from .prompts import leader_system_message, role_system_message
+import agentops
 
+# Initialize AgentOps
+agentops.init('<INSERT YOUR API KEY HERE>')
 
+@agentops.track_agent(name='FinRobot')
 class FinRobot(AssistantAgent):
 
     def __init__(
@@ -467,3 +474,7 @@ def _get_representative(self):
                 ),
             )
         return leader
+
+
+# End of program
+agentops.end_session('Success')