Finchat

FinanceChatPack

Bases: BaseLlamaPack

Source code in llama-index-packs/llama-index-packs-finchat/llama_index/packs/finchat/base.py

class FinanceChatPack(BaseLlamaPack):
    def __init__(
        self,
        polygon_api_key: str,
        finnhub_api_key: str,
        alpha_vantage_api_key: str,
        newsapi_api_key: str,
        openai_api_key: str,
        postgres_db_uri: str,
        gpt_model_name: str = "gpt-4-0613",
    ):
        llm = OpenAI(temperature=0, model=gpt_model_name, api_key=openai_api_key)
        self.db_tool_spec = SQLDatabaseToolSpec(uri=postgres_db_uri)
        self.fin_tool_spec = FinanceAgentToolSpec(
            polygon_api_key, finnhub_api_key, alpha_vantage_api_key, newsapi_api_key
        )

        self.db_table_info = self.db_tool_spec.get_table_info()
        prefix_messages = self.construct_prefix_db_message(self.db_table_info)
        # add some role play in the system .
        database_agent = OpenAIAgent.from_tools(
            [
                tool
                for tool in self.db_tool_spec.to_tool_list()
                if tool.metadata.name == "run_sql_query"
            ],
            prefix_messages=prefix_messages,
            llm=llm,
            verbose=True,
        )
        database_agent_tool = QueryEngineTool.from_defaults(
            database_agent,
            name="database_agent",
            description=""""
                This agent analyzes a text query and add further explanations and thoughts to help a data scientist who has access to following tables:

                {table_info}

                Be concise and do not lose any information about original query while passing to the data scientist.
                """,
        )

        fin_api_agent = OpenAIAgent.from_tools(
            self.fin_tool_spec.to_tool_list(),
            system_prompt=f"""
                You are a helpful AI financial assistant designed to understand the intent of the user query and then use relevant tools/apis to help answer it.
                You can use more than one tool/api only if needed, but final response should be concise and relevant. If you are not able to find
                relevant tool/api, respond respectfully suggesting that you don't know. Think step by step""",
            llm=llm,
            verbose=True,
        )

        fin_api_agent_tool = QueryEngineTool.from_defaults(
            fin_api_agent,
            name="fin_api_agent",
            description=f"""
                This agent has access to another agent which can access certain open APIs to provide information based on user query.
                Analyze the query and add any information if needed which can help to decide which API to call.
                Be concise and do not lose any information about original query.
                """,
        )

        self.fin_hierarchical_agent = OpenAIAgent.from_tools(
            [database_agent_tool, fin_api_agent_tool],
            system_prompt="""
                You are a specialized financial assistant with access to certain tools which can access open APIs and SP500 companies database containing information on
                daily opening price, closing price, high, low, volume, reported earnings, estimated earnings since 2010 to 2023. Before answering query you should check
                if the question can be answered via querying the database or using specific open APIs. If you try to find answer via querying database first and it did
                not work out, think if you can use other tool APIs available before replying gracefully.
                """,
            llm=llm,
            verbose=True,
        )

    def construct_prefix_db_message(self, table_info: str) -> str:
        system_prompt = f"""
        You are a smart data scientist working in a reputed trading firm like Jump Trading developing automated trading algorithms. Take a deep breathe and think
        step by step to design queries over a SQL database.

        Here is a complete description of tables in SQL database you have access to:

        {table_info}

        Use responses to past questions also to guide you.


        """

        prefix_messages = []
        prefix_messages.append(ChatMessage(role="system", content=system_prompt))

        prefix_messages.append(
            ChatMessage(
                role="user",
                content="What is the average price of Google in the month of July in 2023",
            )
        )
        prefix_messages.append(
            ChatMessage(
                role="assistant",
                content="""
        SELECT AVG(close) AS AvgPrice
        FROM stock_data
        WHERE stock = 'GOOG'
            AND date >= '2023-07-01'
            AND date <= '2023-07-31';
        """,
            )
        )

        prefix_messages.append(
            ChatMessage(
                role="user",
                content="Which stock has the maximum % change in any month in 2023",
            )
        )
        # prefix_messages.append(ChatMessage(role="user", content="Which stocks gave more than 2% return constantly in month of July from past 5 years"))
        prefix_messages.append(
            ChatMessage(
                role="assistant",
                content="""
        WITH MonthlyPrices AS (
            SELECT
                stock,
                EXTRACT(YEAR FROM date) AS year,
                EXTRACT(MONTH FROM date) AS month,
                FIRST_VALUE(close) OVER (PARTITION BY stock, EXTRACT(YEAR FROM date), EXTRACT(MONTH FROM date) ORDER BY date ASC) AS opening_price,
                LAST_VALUE(close) OVER (PARTITION BY stock, EXTRACT(YEAR FROM date), EXTRACT(MONTH FROM date) ORDER BY date ASC ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING) AS closing_price
            FROM
                stock_data
            WHERE
                EXTRACT(YEAR FROM date) = 2023
        ),
        PercentageChanges AS (
            SELECT
                stock,
                year,
                month,
                CASE
                    WHEN opening_rice IS NULL OR closing_price IS NULL THEN NULL
                    WHEN opening_price = 0 THEN NULL
                    ELSE ((closing_price - opening_price) / opening_price) * 100
                END AS pct
            FROM
                MonthlyPrices
        )
        SELECT *
        FROM
            PercentageChanges
        WHERE pct IS NOT NULL
        ORDER BY
            pct DESC
        LIMIT 1;
        """,
            )
        )

        prefix_messages.append(
            ChatMessage(
                role="user",
                content="How many times Microsoft beat earnings estimates in 2022",
            )
        )
        prefix_messages.append(
            ChatMessage(
                role="assistant",
                content="""
        SELECT
            COUNT(*)
        FROM
            earnings
        WHERE
            stock = 'MSFT' AND reported > estimated and EXTRACT(YEAR FROM date) = 2022
        """,
            )
        )

        prefix_messages.append(
            ChatMessage(
                role="user",
                content="Which stocks have beaten earnings estimate by more than 1$ consecutively from last 4 reportings?",
            )
        )
        prefix_messages.append(
            ChatMessage(
                role="assistant",
                content="""
        WITH RankedEarnings AS(
            SELECT
                stock,
                date,
                reported,
                estimated,
                RANK() OVER (PARTITION BY stock ORDER BY date DESC) as ranking
            FROM
                earnings
        )
        SELECT
            stock
        FROM
            RankedEarnings
        WHERE
            ranking <= 4 AND reported - estimated > 1
        GROUP BY
            stock
        HAVING COUNT(*) = 4
        """,
            )
        )

        return prefix_messages

    def run(self, query: str):
        return self.fin_hierarchical_agent.chat(query)