app.py

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)