File size: 7,429 Bytes
b5e0c7e
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
72b3bfc
b5e0c7e
 
 
0ea5146
b5e0c7e
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178

from omegaconf import OmegaConf
import streamlit as st
import os
from PIL import Image
import re
import sys

from pydantic import Field, BaseModel
from vectara_agent.agent import Agent, AgentType, AgentStatusType
from vectara_agent.tools import ToolsFactory


tickers = {
    "AAPL": "Apple Computer", 
    "GOOG": "Google", 
    "AMZN": "Amazon",
    "SNOW": "Snowflake",
    "TEAM": "Atlassian",
    "TSLA": "Tesla",
    "NVDA": "Nvidia",
    "MSFT": "Microsoft",
    "AMD": "Advanced Micro Devices",
}
years = [2020, 2021, 2022, 2023, 2024]
initial_prompt = "How can I help you today?"

def create_tools(cfg):    

    def get_company_info() -> list[str]:
        """
        Returns a dictionary of companies you can query about their financial reports.
        The output is a dictionary of valid ticker symbols mapped to company names.
        You can use this to identify the companies you can query about, and their ticker information.
        """
        return tickers

    def get_valid_years() -> list[str]:
        """
        Returns a list of the years for which financial reports are available.
        """
        return years
    
    class QueryFinancialReportsArgs(BaseModel):
        query: str = Field(..., description="The user query. Must be a question about the company's financials, and should not include the company name, ticker or year.")
        year: int = Field(..., description=f"The year. an integer between {min(years)} and {max(years)}.")
        ticker: str = Field(..., description=f"The company ticker. Must be a valid ticket symbol from the list {tickers.keys()}.")

    tools_factory = ToolsFactory(vectara_api_key=cfg.api_key, 
                                 vectara_customer_id=cfg.customer_id, 
                                 vectara_corpus_id=cfg.corpus_id)
    query_financial_reports = tools_factory.create_rag_tool(
        tool_name = "query_financial_reports",
        tool_description = """
        Given a company name and year, 
        returns a response (str) to a user query about the company's financials for that year.
        When using this tool, make sure to provide the a valid company ticker and a year. 
        Use this tool to get financial information one metric at a time.
        """,
        tool_args_schema = QueryFinancialReportsArgs,
        tool_filter_template = "doc.year = {year} and doc.ticker = '{ticker}'",
        reranker = "slingshot", rerank_k = 100, 
        n_sentences_before = 2, n_sentences_after = 2, lambda_val = 0.0,
        summary_num_results = 15,
        vectara_summarizer = 'vectara-summary-ext-24-05-med-omni',
    )

    return (tools_factory.get_tools(
                [
                    get_company_info, 
                    get_valid_years,
                ]
            ) +
            tools_factory.standard_tools() + 
            tools_factory.financial_tools() + 
            [query_financial_reports]
    )

def launch_bot(agent_type: AgentType):
    def reset():
        cfg = st.session_state.cfg
        st.session_state.messages = [{"role": "assistant", "content": initial_prompt, "avatar": "πŸ¦–"}]
        st.session_state.thinking_message = "Agent at work..."

        # Create the agent
        print("Creating agent...")

        def update_func(status_type: AgentStatusType, msg: str):
            output = f"{status_type.value} - {msg}"
            st.session_state.thinking_placeholder.text(output)

        financial_bot_instructions = """
        - You are a helpful financial assistant in conversation with a user. Use your financial expertise when crafting a query to the tool, to ensure you get the most accurate information.
        - You can answer questions, provide insights, or summarize any information from financial reports.
        - A user may refer to a company's ticker instead of its full name - consider those the same when a user is asking about a company.
        - When calculating a financial metric, make sure you have all the information from tools to complete the calculation.
        - In many cases you may need to query tools on each sub-metric separately before computing the final metric.
        - When using a tool to obtain financial data, consider the fact that information for a certain year may be reported in the the following year's report.
        - Report financial data in a consistent manner. For example if you report revenue in thousands, always report revenue in thousands.
        """

        st.session_state.agent = Agent(
            agent_type = agent_type,
            tools = create_tools(cfg),
            topic = "10-K financial reports",
            custom_instructions = financial_bot_instructions,
            update_func = update_func
        )

    if 'cfg' not in st.session_state:
        cfg = OmegaConf.create({
            'customer_id': str(os.environ['VECTARA_CUSTOMER_ID']),
            'corpus_id': str(os.environ['VECTARA_CORPUS_ID']),
            'api_key': str(os.environ['VECTARA_API_KEY']),
        })
        st.session_state.cfg = cfg
        reset()

    cfg = st.session_state.cfg
    st.set_page_config(page_title="Financial Assistant", layout="wide")

    # left side content
    with st.sidebar:
        image = Image.open('Vectara-logo.png')
        st.image(image, width=250)
        st.markdown("## Welcome to the financial assistant demo.\n\n\n")
        companies = ", ".join(tickers.values())
        st.markdown(
            f"This assistant can help you with any questions about the financials of the following companies:\n\n **{companies}**.\n\n"
            "You can ask questions, analyze data, provide insights, or summarize any information from financial reports."
        )

        st.markdown("\n\n")
        if st.button('Start Over'):
            reset()

        st.markdown("---")
        st.markdown(
            "## How this works?\n"
            "This app was built with [Vectara](https://vectara.com).\n\n"
            "It demonstrates the use of Agentic Chat functionality with Vectara"
        )
        st.markdown("---")


    if "messages" not in st.session_state.keys():
        reset()

    # Display chat messages
    for message in st.session_state.messages:
        with st.chat_message(message["role"], avatar=message["avatar"]):
            st.write(message["content"])

    # User-provided prompt
    if prompt := st.chat_input():
        st.session_state.messages.append({"role": "user", "content": prompt, "avatar": 'πŸ§‘β€πŸ’»'})
        with st.chat_message("user", avatar='πŸ§‘β€πŸ’»'):
            print(f"Starting new question: {prompt}\n")
            st.write(prompt)
    
    # Generate a new response if last message is not from assistant
    if st.session_state.messages[-1]["role"] != "assistant":
        with st.chat_message("assistant", avatar='πŸ€–'):
            with st.spinner(st.session_state.thinking_message):
                st.session_state.thinking_placeholder = st.empty()
                res = st.session_state.agent.chat(prompt)
                cleaned = re.sub(r'\[\d+\]', '', res).replace('$', '\\$')
            message = {"role": "assistant", "content": cleaned, "avatar": 'πŸ€–'}
            st.session_state.messages.append(message)
            st.session_state.thinking_placeholder.empty()
            st.rerun()

    sys.stdout.flush()

if __name__ == "__main__":
    print("Starting up...")
    launch_bot(agent_type = AgentType.REACT)