File size: 10,636 Bytes
b5e0c7e
 
 
 
25b67f4
 
a75f490
a6a85af
 
 
 
25b67f4
a75f490
b5e0c7e
 
e1452a4
b5e0c7e
 
 
 
 
 
 
 
 
 
 
 
ece9872
 
b5e0c7e
 
 
 
 
 
 
 
be67d35
b5e0c7e
 
 
 
 
 
 
 
be67d35
b5e0c7e
 
 
25b67f4
 
 
 
 
 
 
e1452a4
25b67f4
 
 
 
 
 
 
 
 
 
 
e1452a4
25b67f4
 
 
 
 
 
5b2b247
72e1546
ed18c6a
b5e0c7e
 
 
 
 
5b2b247
 
b5e0c7e
308fc11
6960aa1
b5e0c7e
5b2b247
72e1546
ed18c6a
72e1546
b5e0c7e
ed18c6a
b5e0c7e
 
 
 
 
 
25b67f4
b5e0c7e
 
 
 
a441318
5b2b247
b5e0c7e
 
e1452a4
e7ab0c3
 
 
ed18c6a
 
 
e1452a4
be67d35
ed18c6a
9f650ed
91ec79e
 
 
 
d015953
 
 
91ec79e
 
 
 
 
 
 
 
 
ed18c6a
 
 
 
a6a85af
 
 
 
 
 
 
 
 
e1452a4
b5e0c7e
 
 
25b67f4
ed18c6a
25b67f4
a6a85af
b5e0c7e
91ec79e
b5e0c7e
 
 
 
 
a6a85af
b5e0c7e
 
a6a85af
 
 
b5e0c7e
ed18c6a
b5e0c7e
ed18c6a
 
b5e0c7e
 
 
 
308fc11
b5e0c7e
 
 
dea99b8
b5e0c7e
 
 
ed18c6a
25b67f4
 
 
b5e0c7e
 
 
 
 
25b67f4
b5e0c7e
 
 
 
 
ed18c6a
b5e0c7e
 
 
 
 
a6a85af
 
 
 
 
 
b5e0c7e
a6a85af
 
 
 
 
b5e0c7e
ed18c6a
 
 
b5e0c7e
 
 
a6a85af
ed18c6a
b5e0c7e
ed18c6a
b5e0c7e
 
ed18c6a
 
 
b5e0c7e
ed18c6a
a6a85af
 
 
 
ed18c6a
 
 
 
25b67f4
d015953
ed18c6a
 
 
25b67f4
b5e0c7e
 
 
e1452a4
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
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260

import os
from PIL import Image
import sys
import pandas as pd
import requests

from omegaconf import OmegaConf
import streamlit as st
from streamlit_pills import pills

from dotenv import load_dotenv
load_dotenv(override=True)

from pydantic import Field, BaseModel
from vectara_agent.agent import Agent, 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",
    "INTC": "Intel",
    "NFLX": "Netflix",
}
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. Always check this before using any other tool.
        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.
        Always check this before using any other tool.
        """
        return years
    
    # Tool to get the income statement for a given company and year using the FMP API
    def get_income_statement(
        ticker=Field(description="the ticker symbol of the company."),
        year=Field(description="the year for which to get the income statement."),
    ) -> str:
        """
        Get the income statement for a given company and year using the FMP (https://financialmodelingprep.com) API.
        Returns a dictionary with the income statement data. All data is in USD, but you can convert it to more compact form like K, M, B.
        """
        fmp_api_key = os.environ.get("FMP_API_KEY", None)
        if fmp_api_key is None:
            return "FMP_API_KEY environment variable not set. This tool does not work."
        url = f"https://financialmodelingprep.com/api/v3/income-statement/{ticker}?apikey={fmp_api_key}"
        response = requests.get(url)
        if response.status_code == 200:
            data = response.json()
            income_statement = pd.DataFrame(data)
            income_statement["date"] = pd.to_datetime(income_statement["date"])
            income_statement_specific_year = income_statement[
                income_statement["date"].dt.year == int(year)
            ]
            values_dict = income_statement_specific_year.to_dict(orient="records")[0]
            return f"Financial results: {', '.join([f'{key}: {value}' for key, value in values_dict.items() if key not in ['date', 'cik', 'link', 'finalLink']])}"
        else:
            return "FMP API returned error. This tool does not work."

    class QueryTranscriptsArgs(BaseModel):
        query: str = Field(..., description="The user query.")
        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)
    ask_transcripts = tools_factory.create_rag_tool(
        tool_name = "ask_transcripts",
        tool_description = """
        Given a company name and year, response to a user question about the company, based on analyst call transcripts about the company's financial reports for that year.
        You can ask this tool any question about the compaany including risks, opportunities, financial performance, competitors and more.
        """,
        tool_args_schema = QueryTranscriptsArgs,
        reranker = "multilingual_reranker_v1", rerank_k = 100, 
        n_sentences_before = 2, n_sentences_after = 2, lambda_val = 0.005,
        summary_num_results = 10,
        vectara_summarizer = 'vectara-summary-ext-24-05-med-omni',
        include_citations = False,
    )

    return (tools_factory.get_tools(
                [
                    get_company_info, 
                    get_valid_years,
                    get_income_statement,
                ]
            ) +
            tools_factory.standard_tools() + 
            tools_factory.financial_tools() + 
            tools_factory.guardrail_tools() +
            [ask_transcripts]
    )

def initialize_agent(_cfg):
    if 'agent' in st.session_state:
        return st.session_state.agent

    financial_bot_instructions = """
    - You are a helpful financial assistant, with expertise in financial reporting, in conversation with a user. 
    - Never discuss politics, and always respond politely.
    - Respond in a compact format by using appropriate units of measure (e.g., K for thousands, M for millions, B for billions). 
      Do not report the same number twice (e.g. $100K and 100,000 USD).
    - Always check the get_company_info and get_valid_years tools to validate company and year are valid.
    - When querying a tool for a numeric value or KPI, use a concise and non-ambiguous description of what you are looking for. 
    - If you calculate a metric, make sure you have all the necessary information to complete the calculation. Don't guess.
    """

    def update_func(status_type: AgentStatusType, msg: str):
        if status_type != AgentStatusType.AGENT_UPDATE:
            output = f"{status_type.value} - {msg}"
            st.session_state.log_messages.append(output)

    agent = Agent(
        tools=create_tools(_cfg),
        topic="10-K annual financial reports",
        custom_instructions=financial_bot_instructions,
        update_func=update_func
    )
    return agent


def toggle_logs():
    st.session_state.show_logs = not st.session_state.show_logs

def show_example_questions():        
    if len(st.session_state.example_messages) > 0 and st.session_state.first_turn:            
        selected_example = pills("Queries to Try:", st.session_state.example_messages, index=None)
        if selected_example:
            st.session_state.ex_prompt = selected_example
            st.session_state.first_turn = False
            return True
    return False

def launch_bot():
    def reset():
        st.session_state.messages = [{"role": "assistant", "content": initial_prompt, "avatar": "πŸ¦–"}]
        st.session_state.thinking_message = "Agent at work..."
        st.session_state.log_messages = []
        st.session_state.prompt = None
        st.session_state.show_logs = False
        st.session_state.first_turn = True

    st.set_page_config(page_title="Financial Assistant", layout="wide")
    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']),
            'examples': os.environ.get('QUERY_EXAMPLES', None)
        })
        st.session_state.cfg = cfg
        st.session_state.ex_prompt = None
        example_messages = [example.strip() for example in cfg.examples.split(",")] if cfg.examples else []
        st.session_state.example_messages = [em for em in example_messages if len(em)>0]
        reset()

    cfg = st.session_state.cfg
    if 'agent' not in st.session_state:
        st.session_state.agent = initialize_agent(cfg)

    # left side content
    with st.sidebar:
        image = Image.open('Vectara-logo.png')
        st.image(image, width=175)
        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 several companies:\n\n **{companies}**.\n"
        )

        st.markdown("\n\n")
        bc1, _ = st.columns([1, 1])
        with bc1:
            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 RAG 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"])

    example_container = st.empty()
    with example_container:
        if show_example_questions():
            example_container.empty()
            st.rerun()

    # User-provided prompt
    if st.session_state.ex_prompt:
        prompt = st.session_state.ex_prompt
    else:
        prompt = st.chat_input()
    if prompt:
        st.session_state.messages.append({"role": "user", "content": prompt, "avatar": 'πŸ§‘β€πŸ’»'})
        st.session_state.prompt = prompt  # Save the prompt in session state
        st.session_state.log_messages = []
        st.session_state.show_logs = False
        with st.chat_message("user", avatar='πŸ§‘β€πŸ’»'):
            print(f"Starting new question: {prompt}\n")
            st.write(prompt)
        st.session_state.ex_prompt = None
        
    # Generate a new response if last message is not from assistant
    if st.session_state.prompt:
        with st.chat_message("assistant", avatar='πŸ€–'):
            with st.spinner(st.session_state.thinking_message):
                res = st.session_state.agent.chat(st.session_state.prompt)
                res = res.replace('$', '\\$')  # escape dollar sign for markdown
            message = {"role": "assistant", "content": res, "avatar": 'πŸ€–'}
            st.session_state.messages.append(message)
            st.markdown(res)
        st.session_state.ex_prompt = None
        st.session_state.prompt = None
        st.rerun()
    
    log_placeholder = st.empty()
    with log_placeholder.container():
        if st.session_state.show_logs:
            st.button("Hide Logs", on_click=toggle_logs)
            for msg in st.session_state.log_messages:
                st.text(msg)
        else:
            if len(st.session_state.log_messages) > 0:
                st.button("Show Logs", on_click=toggle_logs)

    sys.stdout.flush()

if __name__ == "__main__":
    launch_bot()