diff --git a/agent_builder_demo.py b/agent_builder_demo.py index ee5a89d..45cefed 100644 --- 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('') 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 index 7f834fa..969a79f 100644 --- 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,13 +23,6 @@ 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"]) @@ -37,7 +30,12 @@ ] ) -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,7 +67,7 @@ ) 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, @@ -75,13 +75,10 @@ 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') \ No newline at end of file diff --git a/experiments/portfolio_optimization.py b/experiments/portfolio_optimization.py index 50e64c0..92b1a28 100644 --- a/experiments/portfolio_optimization.py +++ b/experiments/portfolio_optimization.py @@ -5,29 +5,38 @@ 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", @@ -35,15 +44,19 @@ }, ) -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, @@ -51,24 +64,27 @@ "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') \ No newline at end of file diff --git a/finrobot/agents/workflow.py b/finrobot/agents/workflow.py index 1bbe652..5ad2f23 100644 --- 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('') +@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') diff --git a/finrobot/functional/analyzer.py b/finrobot/functional/analyzer.py index 47737ae..8d4df84 100644 --- a/finrobot/functional/analyzer.py +++ b/finrobot/functional/analyzer.py @@ -3,8 +3,12 @@ from typing import Annotated from datetime import timedelta, datetime from ..data_source import YFinanceUtils, SECUtils, FMPUtils +import agentops +# Initialize AgentOps +agentops.init('') +@agentops.record_function('combine_prompt') def combine_prompt(instruction, resource, table_str=None): if table_str: prompt = f"{table_str}\n\nResource: {resource}\n\nInstruction: {instruction}" @@ -12,15 +16,16 @@ def combine_prompt(instruction, resource, table_str=None): prompt = f"Resource: {resource}\n\nInstruction: {instruction}" return prompt - +@agentops.record_function('save_to_file') def save_to_file(data: str, file_path: str): os.makedirs(os.path.dirname(file_path), exist_ok=True) with open(file_path, "w") as f: f.write(data) - class ReportAnalysisUtils: + @staticmethod + @agentops.record_function('analyze_income_stmt') def analyze_income_stmt( ticker_symbol: Annotated[str, "ticker symbol"], fyear: Annotated[str, "fiscal year of the 10-K report"], @@ -30,11 +35,9 @@ def analyze_income_stmt( Retrieve the income statement for the given ticker symbol with the related section of its 10-K report. Then return with an instruction on how to analyze the income statement. """ - # Retrieve the income statement income_stmt = YFinanceUtils.get_income_stmt(ticker_symbol) df_string = "Income statement:\n" + income_stmt.to_string().strip() - # Analysis instruction instruction = dedent( """ Conduct a comprehensive analysis of the company's income statement for the current fiscal year. @@ -44,21 +47,19 @@ def analyze_income_stmt( and net profit margins to evaluate cost efficiency, operational effectiveness, and overall profitability. Analyze Earnings Per Share to understand investor perspectives. Compare these metrics with historical data and industry or competitor benchmarks to identify growth patterns, profitability trends, and - operational challenges. The output should be a strategic overview of the company’s financial health + operational challenges. The output should be a strategic overview of the company's financial health in a single paragraph, less than 130 words, summarizing the previous analysis into 4-5 key points under respective subheadings with specific discussion and strong data support. """ ) - # Retrieve the related section from the 10-K report section_text = SECUtils.get_10k_section(ticker_symbol, fyear, 7) - - # Combine the instruction, section text, and income statement prompt = combine_prompt(instruction, section_text, df_string) - save_to_file(prompt, save_path) return f"instruction & resources saved to {save_path}" + @staticmethod + @agentops.record_function('analyze_balance_sheet') def analyze_balance_sheet( ticker_symbol: Annotated[str, "ticker symbol"], fyear: Annotated[str, "fiscal year of the 10-K report"], @@ -88,6 +89,8 @@ def analyze_balance_sheet( save_to_file(prompt, save_path) return f"instruction & resources saved to {save_path}" + @staticmethod + @agentops.record_function('analyze_cash_flow') def analyze_cash_flow( ticker_symbol: Annotated[str, "ticker symbol"], fyear: Annotated[str, "fiscal year of the 10-K report"], @@ -117,6 +120,8 @@ def analyze_cash_flow( save_to_file(prompt, save_path) return f"instruction & resources saved to {save_path}" + @staticmethod + @agentops.record_function('analyze_segment_stmt') def analyze_segment_stmt( ticker_symbol: Annotated[str, "ticker symbol"], fyear: Annotated[str, "fiscal year of the 10-K report"], @@ -149,6 +154,8 @@ def analyze_segment_stmt( save_to_file(prompt, save_path) return f"instruction & resources saved to {save_path}" + @staticmethod + @agentops.record_function('income_summarization') def income_summarization( ticker_symbol: Annotated[str, "ticker symbol"], fyear: Annotated[str, "fiscal year of the 10-K report"], @@ -160,9 +167,6 @@ def income_summarization( With the income statement and segment analysis for the given ticker symbol. Then return with an instruction on how to synthesize these analyses into a single coherent paragraph. """ - # income_stmt_analysis = analyze_income_stmt(ticker_symbol) - # segment_analysis = analyze_segment_stmt(ticker_symbol) - instruction = dedent( f""" Income statement analysis: {income_stmt_analysis}, @@ -183,6 +187,8 @@ def income_summarization( save_to_file(prompt, save_path) return f"instruction & resources saved to {save_path}" + @staticmethod + @agentops.record_function('get_risk_assessment') def get_risk_assessment( ticker_symbol: Annotated[str, "ticker symbol"], fyear: Annotated[str, "fiscal year of the 10-K report"], @@ -207,6 +213,8 @@ def get_risk_assessment( save_to_file(prompt, save_path) return f"instruction & resources saved to {save_path}" + @staticmethod + @agentops.record_function('analyze_business_highlights') def analyze_business_highlights( ticker_symbol: Annotated[str, "ticker symbol"], fyear: Annotated[str, "fiscal year of the 10-K report"], @@ -236,6 +244,8 @@ def analyze_business_highlights( save_to_file(prompt, save_path) return f"instruction & resources saved to {save_path}" + @staticmethod + @agentops.record_function('analyze_company_description') def analyze_company_description( ticker_symbol: Annotated[str, "ticker symbol"], fyear: Annotated[str, "fiscal year of the 10-K report"], @@ -263,9 +273,9 @@ def analyze_company_description( instruction = dedent( """ According to the given information, - 1. Briefly describe the company’s industry, + 1. Briefly describe the company's industry, 2. Highlight core strengths and competitive advantages key products or services, - 3. Identify current industry trends, opportunities, and challenges that influence the company’s strategy, + 3. Identify current industry trends, opportunities, and challenges that influence the company's strategy, 4. Outline recent strategic initiatives such as product launches, acquisitions, or new partnerships, and describe the company's response to market conditions. Less than 400 words. """ @@ -276,6 +286,8 @@ def analyze_company_description( save_to_file(prompt, save_path) return f"instruction & resources saved to {save_path}" + @staticmethod + @agentops.record_function('get_key_data') def get_key_data( ticker_symbol: Annotated[str, "ticker symbol"], filing_date: Annotated[ @@ -315,13 +327,9 @@ def get_key_data( fiftyTwoWeekLow = hist["High"].min() fiftyTwoWeekHigh = hist["Low"].max() - # avg_daily_volume_6m = hist['Volume'].mean() - # convert back to str for function calling filing_date = filing_date.strftime("%Y-%m-%d") - # Print the result - # print(f"Over the past 6 months, the average daily trading volume for {ticker_symbol} was: {avg_daily_volume_6m:.2f}") rating, _ = YFinanceUtils.get_analyst_recommendations(ticker_symbol) target_price = FMPUtils.get_target_price(ticker_symbol, filing_date) result = { @@ -342,3 +350,8 @@ def get_key_data( ), } return result + +# End of ReportAnalysisUtils class + +# Add this at the end of your script or in your main execution point +agentops.end_session('Success') \ No newline at end of file diff --git a/finrobot/functional/rag.py b/finrobot/functional/rag.py index bb98441..4bd4cfe 100644 --- a/finrobot/functional/rag.py +++ b/finrobot/functional/rag.py @@ -1,6 +1,9 @@ from autogen.agentchat.contrib.retrieve_user_proxy_agent import RetrieveUserProxyAgent from typing import Annotated +import agentops +# Initialize AgentOps +agentops.init('') PROMPT_RAG_FUNC = """Below is the context retrieved from the required file based on your query. If you can't answer the question with or without the current context, you should try using a more refined search query according to your requirements, or ask for more contexts. @@ -10,7 +13,7 @@ Retrieved context is: {input_context} """ - +@agentops.record_function('get_rag_function') def get_rag_function(retrieve_config, description=""): def termination_msg(x): @@ -22,7 +25,7 @@ def termination_msg(x): if "customized_prompt" not in retrieve_config: retrieve_config["customized_prompt"] = PROMPT_RAG_FUNC - rag_assitant = RetrieveUserProxyAgent( + rag_assistant = RetrieveUserProxyAgent( name="RAG_Assistant", is_termination_msg=termination_msg, human_input_mode="NEVER", @@ -33,6 +36,7 @@ def termination_msg(x): description="Assistant who has extra content retrieval power for solving difficult problems.", ) + @agentops.record_function('retrieve_content') def retrieve_content( message: Annotated[ str, @@ -42,23 +46,23 @@ def retrieve_content( n_results: Annotated[int, "Number of results to retrieve, default to 3"] = 3, ) -> str: - rag_assitant.n_results = n_results # Set the number of results to be retrieved. + rag_assistant.n_results = n_results # Set the number of results to be retrieved. # Check if we need to update the context. - update_context_case1, update_context_case2 = rag_assitant._check_update_context( + update_context_case1, update_context_case2 = rag_assistant._check_update_context( message ) if ( update_context_case1 or update_context_case2 - ) and rag_assitant.update_context: - rag_assitant.problem = ( + ) and rag_assistant.update_context: + rag_assistant.problem = ( message - if not hasattr(rag_assitant, "problem") - else rag_assitant.problem + if not hasattr(rag_assistant, "problem") + else rag_assistant.problem ) - _, ret_msg = rag_assitant._generate_retrieve_user_reply(message) + _, ret_msg = rag_assistant._generate_retrieve_user_reply(message) else: _context = {"problem": message, "n_results": n_results} - ret_msg = rag_assitant.message_generator(rag_assitant, None, _context) + ret_msg = rag_assistant.message_generator(rag_assistant, None, _context) return ret_msg if ret_msg else message if description: @@ -68,6 +72,9 @@ def retrieve_content( docs = retrieve_config.get("docs_path", []) if docs: docs_str = "\n".join(docs if isinstance(docs, list) else [docs]) - retrieve_content.__doc__ += f"Availale Documents:\n{docs_str}" + retrieve_content.__doc__ += f"Available Documents:\n{docs_str}" + + return retrieve_content, rag_assistant # for debug use - return retrieve_content, rag_assitant # for debug use +# Add this at the end of your script or in your main execution point +agentops.end_session('Success') \ No newline at end of file diff --git a/finrobot/functional/ragquery.py b/finrobot/functional/ragquery.py index d797071..9822a2b 100644 --- a/finrobot/functional/ragquery.py +++ b/finrobot/functional/ragquery.py @@ -11,104 +11,103 @@ from finrobot.data_source.finance_data import get_data from typing import List, Optional import os -SAVE_DIR = "output/SEC_EDGAR_FILINGS_MD" +import agentops + +# Initialize AgentOps with API key from environment variable +agentops.init(os.getenv('AGENTOPS_API_KEY')) +SAVE_DIR = "output/SEC_EDGAR_FILINGS_MD" +@agentops.record_function('rag_database_earnings_call') def rag_database_earnings_call( ticker: str, - year: str)->str: - - #assert quarter in earnings_call_quarter_vals, "The quarter should be from Q1, Q2, Q3, Q4" - earnings_docs, earnings_call_quarter_vals, speakers_list_1, speakers_list_2, speakers_list_3, speakers_list_4 = get_data(ticker=ticker,year=year,data_source='earnings_calls') - - emb_fn = SentenceTransformerEmbeddings(model_name="all-MiniLM-L6-v2") - - text_splitter = RecursiveCharacterTextSplitter( - chunk_size=1024, - chunk_overlap=100, - length_function=len,) - earnings_calls_split_docs = text_splitter.split_documents(earnings_docs) - - earnings_call_db = Chroma.from_documents(earnings_calls_split_docs, emb_fn, persist_directory="./earnings-call-db",collection_name="earnings_call") - - - quarter_speaker_dict = { - "Q1":speakers_list_1, - "Q2":speakers_list_2, - "Q3":speakers_list_3, - "Q4":speakers_list_4} + year: str) -> str: - def query_database_earnings_call( - question: str, - quarter: str)->str: - """This tool will query the earnings call transcripts database for a given question and quarter and it will retrieve - the relevant text along from the earnings call and the speaker who addressed the relevant documents. This tool helps in answering questions - from the earnings call transcripts. - - Args: - question (str): _description_. Question to query the database for relevant documents. - quarter (str): _description_. the financial quarter that is discussed in the question and possible options are Q1, Q2, Q3, Q4 - - Returns: - str: relevant text along from the earnings call and the speaker who addressed the relevant documents - """ - assert quarter in earnings_call_quarter_vals, "The quarter should be from Q1, Q2, Q3, Q4" - - req_speaker_list = [] - quarter_speaker_list = quarter_speaker_dict[quarter] - - for sl in quarter_speaker_list: - if sl in question or sl.lower() in question: - req_speaker_list.append(sl) - if len(req_speaker_list) == 0: - req_speaker_list = quarter_speaker_list - - relevant_docs = earnings_call_db.similarity_search( - question, - k=5, - filter={ - "$and":[ - { - "quarter":{"$eq":quarter} - }, - { - "speaker":{"$in":req_speaker_list} - } - ] - } - ) - - speaker_releavnt_dict = {} - for doc in relevant_docs: - speaker = doc.metadata['speaker'] - speaker_text = doc.page_content - if speaker not in speaker_releavnt_dict: - speaker_releavnt_dict[speaker] = speaker_text - else: - speaker_releavnt_dict[speaker] += " "+speaker_text - - relevant_speaker_text = "" - for speaker, text in speaker_releavnt_dict.items(): - relevant_speaker_text += speaker + ": " - relevant_speaker_text += text + "\n\n" - - return relevant_speaker_text - - return query_database_earnings_call, earnings_call_quarter_vals, quarter_speaker_dict - - - - - - - + earnings_docs, earnings_call_quarter_vals, speakers_list_1, speakers_list_2, speakers_list_3, speakers_list_4 = get_data(ticker=ticker, year=year, data_source='earnings_calls') + + emb_fn = SentenceTransformerEmbeddings(model_name="all-MiniLM-L6-v2") + + text_splitter = RecursiveCharacterTextSplitter( + chunk_size=1024, + chunk_overlap=100, + length_function=len,) + earnings_calls_split_docs = text_splitter.split_documents(earnings_docs) + + earnings_call_db = Chroma.from_documents(earnings_calls_split_docs, emb_fn, persist_directory="./earnings-call-db", collection_name="earnings_call") + + quarter_speaker_dict = { + "Q1": speakers_list_1, + "Q2": speakers_list_2, + "Q3": speakers_list_3, + "Q4": speakers_list_4} + + @agentops.record_function('query_database_earnings_call') + def query_database_earnings_call( + question: str, + quarter: str) -> str: + """This tool will query the earnings call transcripts database for a given question and quarter and it will retrieve + the relevant text along from the earnings call and the speaker who addressed the relevant documents. This tool helps in answering questions + from the earnings call transcripts. + + Args: + question (str): _description_. Question to query the database for relevant documents. + quarter (str): _description_. the financial quarter that is discussed in the question and possible options are Q1, Q2, Q3, Q4 + + Returns: + str: relevant text along from the earnings call and the speaker who addressed the relevant documents + """ + assert quarter in earnings_call_quarter_vals, "The quarter should be from Q1, Q2, Q3, Q4" + + req_speaker_list = [] + quarter_speaker_list = quarter_speaker_dict[quarter] + + for sl in quarter_speaker_list: + if sl in question or sl.lower() in question: + req_speaker_list.append(sl) + if len(req_speaker_list) == 0: + req_speaker_list = quarter_speaker_list + + relevant_docs = earnings_call_db.similarity_search( + question, + k=5, + filter={ + "$and":[ + { + "quarter":{"$eq":quarter} + }, + { + "speaker":{"$in":req_speaker_list} + } + ] + } + ) + + speaker_releavnt_dict = {} + for doc in relevant_docs: + speaker = doc.metadata['speaker'] + speaker_text = doc.page_content + if speaker not in speaker_releavnt_dict: + speaker_releavnt_dict[speaker] = speaker_text + else: + speaker_releavnt_dict[speaker] += " "+speaker_text + + relevant_speaker_text = "" + for speaker, text in speaker_releavnt_dict.items(): + relevant_speaker_text += speaker + ": " + relevant_speaker_text += text + "\n\n" + + return relevant_speaker_text + + return query_database_earnings_call, earnings_call_quarter_vals, quarter_speaker_dict + +@agentops.record_function('rag_database_sec') def rag_database_sec( ticker: str, year: str, FROM_MARKDOWN = False, - filing_types = ['10-K','10-Q'])->str: + filing_types = ['10-K','10-Q']) -> str: if not FROM_MARKDOWN: - sec_data,sec_form_names = get_data(ticker=ticker, year=year,data_source='unstructured',include_amends=True,filing_types=filing_types) + sec_data, sec_form_names = get_data(ticker=ticker, year=year, data_source='unstructured', include_amends=True, filing_types=filing_types) emb_fn = SentenceTransformerEmbeddings(model_name="all-MiniLM-L6-v2") text_splitter = RecursiveCharacterTextSplitter( chunk_size=1024, @@ -116,9 +115,10 @@ def rag_database_sec( length_function=len,) sec_filings_split_docs = text_splitter.split_documents(sec_data) - sec_filings_unstructured_db = Chroma.from_documents(sec_filings_split_docs, emb_fn, persist_directory="./sec-filings-db",collection_name="sec_filings") + sec_filings_unstructured_db = Chroma.from_documents(sec_filings_split_docs, emb_fn, persist_directory="./sec-filings-db", collection_name="sec_filings") - def query_database_unstructured_sec(question: str,sec_form_name: str)->str: + @agentops.record_function('query_database_unstructured_sec') + def query_database_unstructured_sec(question: str, sec_form_name: str) -> str: """This tool will query the SEC Filings database for a given question and form name, and it will retrieve the relevant text along from the SEC filings and the section names. This tool helps in answering questions from the sec filings. @@ -127,8 +127,6 @@ def query_database_unstructured_sec(question: str,sec_form_name: str)->str: question (str): _description_. Question to query the database for relevant documents sec_form_name (str): _description_. SEC FORM NAME that the question is talking about. It can be 10-K for yearly data and 10-Q for quarterly data. For quarterly data, it can be 10-Q2 to represent Quarter 2 and similarly for other quarters. - - Returns: str: Relevant context for the question from the sec filings """ @@ -157,8 +155,8 @@ def query_database_unstructured_sec(question: str,sec_form_name: str)->str: return query_database_unstructured_sec, sec_form_names elif FROM_MARKDOWN: - sec_data,sec_form_names = get_data(ticker=ticker, year=year,data_source='unstructured',include_amends=True,filing_types=filing_types) - get_data(ticker=ticker,year=year,data_source='marker_pdf',batch_processing=False,batch_multiplier=1) + sec_data, sec_form_names = get_data(ticker=ticker, year=year, data_source='unstructured', include_amends=True, filing_types=filing_types) + get_data(ticker=ticker, year=year, data_source='marker_pdf', batch_processing=False, batch_multiplier=1) headers_to_split_on = [ ("#", "Header 1"), @@ -168,11 +166,11 @@ def query_database_unstructured_sec(question: str,sec_form_name: str)->str: markdown_splitter = MarkdownHeaderTextSplitter(headers_to_split_on=headers_to_split_on) markdown_dir = "output/SEC_EDGAR_FILINGS_MD" md_content_list = [] - for md_dirs in os.listdir(os.path.join(markdown_dir,f"{ticker}-{year}")): - md_file_path = os.path.join(markdown_dir,f"{ticker}-{year}",md_dirs,f"{md_dirs}.md") + for md_dirs in os.listdir(os.path.join(markdown_dir, f"{ticker}-{year}")): + md_file_path = os.path.join(markdown_dir, f"{ticker}-{year}", md_dirs, f"{md_dirs}.md") with open(md_file_path, 'r') as file: content = file.read() - md_content_list.append([content,'-'.join(md_dirs.split('-')[-2:])]) + md_content_list.append([content, '-'.join(md_dirs.split('-')[-2:])]) sec_markdown_docs = [] @@ -182,11 +180,12 @@ def query_database_unstructured_sec(question: str,sec_form_name: str)->str: md_header_docs.metadata.update({"filing_type":md_content[1]}) sec_markdown_docs.extend(md_header_splits) - sec_filings_md_db = Chroma.from_documents(sec_markdown_docs, emb_fn, persist_directory="./sec-filings-md-db",collection_name="sec_filings_md") + sec_filings_md_db = Chroma.from_documents(sec_markdown_docs, emb_fn, persist_directory="./sec-filings-md-db", collection_name="sec_filings_md") + @agentops.record_function('query_database_markdown_sec') def query_database_markdown_sec( question: str, - sec_form_name: str)->str: + sec_form_name: str) -> str: """This tool will query the SEC Filings database for a given question and form name, and it will retrieve the relevant text along from the SEC filings and the section names. This tool helps in answering questions from the sec filings. @@ -195,8 +194,6 @@ def query_database_markdown_sec( question (str): _description_. Question to query the database for relevant documents sec_form_name (str): _description_. SEC FORM NAME that the question is talking about. It can be 10-K for yearly data and 10-Q for quarterly data. For quarterly data, it can be 10-Q2 to represent Quarter 2 and similarly for other quarters. - - Returns: str: Relevant context for the question from the sec filings """ @@ -216,3 +213,6 @@ def query_database_markdown_sec( return relevant_section_text return query_database_markdown_sec, sec_form_names + +# Add this at the end of your script or in your main execution point +agentops.end_session('Success') \ No newline at end of file diff --git a/finrobot/toolkits.py b/finrobot/toolkits.py index c3bdf50..373c5fd 100644 --- a/finrobot/toolkits.py +++ b/finrobot/toolkits.py @@ -1,12 +1,15 @@ from autogen import register_function, ConversableAgent from .data_source import * from .functional.coding import CodingUtils - from typing import List, Callable from functools import wraps from pandas import DataFrame +import agentops +# Initialize AgentOps +agentops.init('') +@agentops.record_function('stringify_output') def stringify_output(func): @wraps(func) def wrapper(*args, **kwargs): @@ -15,10 +18,9 @@ def wrapper(*args, **kwargs): return result.to_string() else: return str(result) - return wrapper - +@agentops.record_function('register_toolkits') def register_toolkits( config: List[dict | Callable | type], caller: ConversableAgent, @@ -26,19 +28,15 @@ def register_toolkits( **kwargs ): """Register tools from a configuration list.""" - for tool in config: - if isinstance(tool, type): - register_tookits_from_cls(caller, executor, tool, **kwargs) + register_toolkits_from_cls(caller, executor, tool, **kwargs) continue - tool_dict = {"function": tool} if callable(tool) else tool if "function" not in tool_dict or not callable(tool_dict["function"]): raise ValueError( "Function not found in tool configuration or not callable." ) - tool_function = tool_dict["function"] name = tool_dict.get("name", tool_function.__name__) description = tool_dict.get("description", tool_function.__doc__) @@ -50,10 +48,9 @@ def register_toolkits( description=description, ) - +@agentops.record_function('register_code_writing') def register_code_writing(caller: ConversableAgent, executor: ConversableAgent): """Register code writing tools.""" - register_toolkits( [ { @@ -81,8 +78,8 @@ def register_code_writing(caller: ConversableAgent, executor: ConversableAgent): executor, ) - -def register_tookits_from_cls( +@agentops.record_function('register_toolkits_from_cls') +def register_toolkits_from_cls( caller: ConversableAgent, executor: ConversableAgent, cls: type, @@ -104,3 +101,16 @@ def register_tookits_from_cls( and not func.startswith("_") ] register_toolkits([getattr(cls, func) for func in funcs], caller, executor) + +# If you have specific agent classes, you can decorate them like this: +# @agentops.track_agent(name='my-expert-agent') +# class MyExpertAgent(ConversableAgent): +# ... + +# Assuming you have a main function or execution point, you should add this at the end: +# def main(): +# ... +# agentops.end_session('Success') + +# if __name__ == "__main__": +# main() \ No newline at end of file