Spaces:
Running
Running
| 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() | |