upates

app.py

@@ -6,36 +6,22 @@ import json
 import os
 import re
 
-from langchain.tools.retriever import create_retriever_tool
-from langchain.agents import AgentExecutor
-from langchain.agents import create_openai_tools_agent
-from langchain.agents.format_scratchpad.openai_tools import (
-    format_to_openai_tool_messages,
-)
-from langchain.agents.output_parsers.openai_tools import OpenAIToolsAgentOutputParser
 from langchain_core.documents import Document
-from langchain_core.prompts import PromptTemplate
 from langchain_core.prompts import ChatPromptTemplate
-from langchain_core.prompts import MessagesPlaceholder
-from langchain_core.messages import AIMessage
-from langchain_core.messages import HumanMessage
 from langchain_core.runnables import RunnableParallel
 from langchain_core.runnables import RunnablePassthrough
-from langchain_core.output_parsers import StrOutputParser
-from langchain_community.callbacks import get_openai_callback
-from langchain_community.callbacks import StreamlitCallbackHandler
 from langchain_community.embeddings import HuggingFaceBgeEmbeddings
 from langchain_community.vectorstores.utils import DistanceStrategy
 from langchain_openai import ChatOpenAI
 from langchain_anthropic import ChatAnthropic
 from langchain_together import ChatTogether
 from langchain_pinecone import PineconeVectorStore
-from pinecone import Pinecone
 import streamlit as st
 
+import usage
 
-st.set_page_config(layout="wide", page_title="LegisQA")
 
+st.set_page_config(layout="wide", page_title="LegisQA")
 os.environ["LANGCHAIN_API_KEY"] = st.secrets["langchain_api_key"]
 os.environ["LANGCHAIN_TRACING_V2"] = "true"
 os.environ["LANGCHAIN_PROJECT"] = st.secrets["langchain_project"]
@@ -70,7 +56,9 @@ TOGETHER_CHAT_MODELS = {
     "meta-llama/Meta-Llama-3.1-70B-Instruct-Turbo": {
         "cost": {"pmi": 0.88, "pmo": 0.88}
     },
-    "meta-llama/Meta-Llama-3.1-405B-Instruct-Turbo": {"cost": {"pmi": 5.00, "pmo": 5.00}},
+    "meta-llama/Meta-Llama-3.1-405B-Instruct-Turbo": {
+        "cost": {"pmi": 5.00, "pmo": 5.00}
+    },
 }
 
 PROVIDER_MODELS = {
@@ -128,6 +116,12 @@ def render_outreach_links():
     st.subheader(f":pancakes: Inference [together.ai]({together_url})")
 
 
+def render_sidebar():
+
+    with st.container(border=True):
+        render_outreach_links()
+
+
 def group_docs(docs) -> list[tuple[str, list[Document]]]:
     doc_grps = defaultdict(list)
 
@@ -151,7 +145,7 @@ def group_docs(docs) -> list[tuple[str, list[Document]]]:
     return doc_grps
 
 
-def format_docs(docs):
+def format_docs(docs: list[Document]) -> str:
     """JSON grouped"""
 
     doc_grps = group_docs(docs)
@@ -168,26 +162,26 @@ def format_docs(docs):
     return json.dumps(out, indent=4)
 
 
-def escape_markdown(text):
+def escape_markdown(text: str) -> str:
     MD_SPECIAL_CHARS = r"\`*_{}[]()#+-.!$"
     for char in MD_SPECIAL_CHARS:
         text = text.replace(char, "\\" + char)
     return text
 
 
-def get_vectorstore_filter(key_prefix: str):
+def get_vectorstore_filter(ret_config: dict) -> dict:
     vs_filter = {}
-    if SS[f"{key_prefix}|filter_legis_id"] != "":
-        vs_filter["legis_id"] = SS[f"{key_prefix}|filter_legis_id"]
-    if SS[f"{key_prefix}|filter_bioguide_id"] != "":
-        vs_filter["sponsor_bioguide_id"] = SS[f"{key_prefix}|filter_bioguide_id"]
+    if ret_config["filter_legis_id"] != "":
+        vs_filter["legis_id"] = ret_config["filter_legis_id"]
+    if ret_config["filter_bioguide_id"] != "":
+        vs_filter["sponsor_bioguide_id"] = ret_config["filter_bioguide_id"]
     vs_filter = {
         **vs_filter,
-        "congress_num": {"$in": SS[f"{key_prefix}|filter_congress_nums"]},
+        "congress_num": {"$in": ret_config["filter_congress_nums"]},
     }
     vs_filter = {
         **vs_filter,
-        "sponsor_party": {"$in": SS[f"{key_prefix}|filter_sponsor_parties"]},
+        "sponsor_party": {"$in": ret_config["filter_sponsor_parties"]},
     }
     return vs_filter
 
@@ -288,163 +282,137 @@ Suggest reforms that would benefit the Medicaid program.
         )
 
 
-def render_generative_config(key_prefix: str):
-    st.selectbox(
-        label="provider", options=PROVIDER_MODELS.keys(), key=f"{key_prefix}|provider"
+def get_generative_config(key_prefix: str) -> dict:
+    output = {}
+
+    key = "provider"
+    output[key] = st.selectbox(
+        label=key, options=PROVIDER_MODELS.keys(), key=f"{key_prefix}|{key}"
     )
-    st.selectbox(
-        label="model name",
-        options=PROVIDER_MODELS[SS[f"{key_prefix}|provider"]],
-        key=f"{key_prefix}|model_name",
+
+    key = "model_name"
+    output[key] = st.selectbox(
+        label=key,
+        options=PROVIDER_MODELS[output["provider"]],
+        key=f"{key_prefix}|{key}",
     )
-    st.slider(
-        "temperature",
+
+    key = "temperature"
+    output[key] = st.slider(
+        key,
         min_value=0.0,
         max_value=2.0,
-        value=0.01,
-        key=f"{key_prefix}|temperature",
+        value=0.0,
+        key=f"{key_prefix}|{key}",
     )
-    st.slider(
-        "max_output_tokens",
+
+    key = "max_output_tokens"
+    output[key] = st.slider(
+        key,
         min_value=1024,
         max_value=2048,
-        key=f"{key_prefix}|max_output_tokens",
+        key=f"{key_prefix}|{key}",
     )
-    st.slider(
-        "top_p", min_value=0.0, max_value=1.0, value=0.9, key=f"{key_prefix}|top_p"
+
+    key = "top_p"
+    output[key] = st.slider(
+        key, min_value=0.0, max_value=1.0, value=0.9, key=f"{key_prefix}|{key}"
     )
-    st.checkbox(
-        "escape markdown in answer", key=f"{key_prefix}|response_escape_markdown"
+
+    key = "should_escape_markdown"
+    output[key] = st.checkbox(
+        key,
+        value=False,
+        key=f"{key_prefix}|{key}",
     )
-    st.checkbox(
-        "add legis urls in answer",
+
+    key = "should_add_legis_urls"
+    output[key] = st.checkbox(
+        key,
         value=True,
-        key=f"{key_prefix}|response_add_legis_urls",
+        key=f"{key_prefix}|{key}",
     )
 
+    return output
+
+
+def get_retrieval_config(key_prefix: str) -> dict:
+    output = {}
 
-def render_retrieval_config(key_prefix: str):
-    st.slider(
+    key = "n_ret_docs"
+    output[key] = st.slider(
         "Number of chunks to retrieve",
         min_value=1,
         max_value=32,
         value=8,
-        key=f"{key_prefix}|n_ret_docs",
+        key=f"{key_prefix}|{key}",
     )
-    st.text_input("Bill ID (e.g. 118-s-2293)", key=f"{key_prefix}|filter_legis_id")
-    st.text_input("Bioguide ID (e.g. R000595)", key=f"{key_prefix}|filter_bioguide_id")
-    st.multiselect(
+
+    key = "filter_legis_id"
+    output[key] = st.text_input("Bill ID (e.g. 118-s-2293)", key=f"{key_prefix}|{key}")
+
+    key = "filter_bioguide_id"
+    output[key] = st.text_input("Bioguide ID (e.g. R000595)", key=f"{key_prefix}|{key}")
+
+    key = "filter_congress_nums"
+    output[key] = st.multiselect(
         "Congress Numbers",
         CONGRESS_NUMBERS,
         default=CONGRESS_NUMBERS,
-        key=f"{key_prefix}|filter_congress_nums",
+        key=f"{key_prefix}|{key}",
     )
-    st.multiselect(
+
+    key = "filter_sponsor_parties"
+    output[key] = st.multiselect(
         "Sponsor Party",
         SPONSOR_PARTIES,
         default=SPONSOR_PARTIES,
-        key=f"{key_prefix}|filter_sponsor_parties",
+        key=f"{key_prefix}|{key}",
     )
 
+    return output
 
-def get_llm(key_prefix: str):
-
-    if SS[f"{key_prefix}|model_name"] in OPENAI_CHAT_MODELS:
-        llm = ChatOpenAI(
-            model=SS[f"{key_prefix}|model_name"],
-            temperature=SS[f"{key_prefix}|temperature"],
-            api_key=st.secrets["openai_api_key"],
-            top_p=SS[f"{key_prefix}|top_p"],
-            seed=SEED,
-            max_tokens=SS[f"{key_prefix}|max_output_tokens"],
-        )
-    elif SS[f"{key_prefix}|model_name"] in ANTHROPIC_CHAT_MODELS:
-        llm = ChatAnthropic(
-            model_name=SS[f"{key_prefix}|model_name"],
-            temperature=SS[f"{key_prefix}|temperature"],
-            api_key=st.secrets["anthropic_api_key"],
-            top_p=SS[f"{key_prefix}|top_p"],
-            max_tokens_to_sample=SS[f"{key_prefix}|max_output_tokens"],
-        )
-    elif SS[f"{key_prefix}|model_name"] in TOGETHER_CHAT_MODELS:
-        llm = ChatTogether(
-            model=SS[f"{key_prefix}|model_name"],
-            temperature=SS[f"{key_prefix}|temperature"],
-            max_tokens=SS[f"{key_prefix}|max_output_tokens"],
-            top_p=SS[f"{key_prefix}|top_p"],
-            seed=SEED,
-            api_key=st.secrets["together_api_key"],
-        )
-    else:
-        raise ValueError()
 
-    return llm
+def get_llm(gen_config: dict):
 
+    match gen_config["provider"]:
 
-def get_token_usage(key_prefix: str, metadata: dict):
-    if SS[f"{key_prefix}|model_name"] in OPENAI_CHAT_MODELS:
-        model_info = PROVIDER_MODELS["OpenAI"][SS[f"{key_prefix}|model_name"]]
-        return get_openai_token_usage(metadata, model_info)
-    elif SS[f"{key_prefix}|model_name"] in ANTHROPIC_CHAT_MODELS:
-        model_info = PROVIDER_MODELS["Anthropic"][SS[f"{key_prefix}|model_name"]]
-        return get_anthropic_token_usage(metadata, model_info)
-    elif SS[f"{key_prefix}|model_name"] in TOGETHER_CHAT_MODELS:
-        model_info = PROVIDER_MODELS["Together"][SS[f"{key_prefix}|model_name"]]
-        return get_together_token_usage(metadata, model_info)
-    else:
-        raise ValueError()
-
-
-def get_openai_token_usage(metadata: dict, model_info: dict):
-    input_tokens = metadata["token_usage"]["prompt_tokens"]
-    output_tokens = metadata["token_usage"]["completion_tokens"]
-    cost = (
-        input_tokens * 1e-6 * model_info["cost"]["pmi"]
-        + output_tokens * 1e-6 * model_info["cost"]["pmo"]
-    )
-    return {
-        "input_tokens": input_tokens,
-        "output_tokens": output_tokens,
-        "cost": cost,
-    }
-
-
-def get_anthropic_token_usage(metadata: dict, model_info: dict):
-    input_tokens = metadata["usage"]["input_tokens"]
-    output_tokens = metadata["usage"]["output_tokens"]
-    cost = (
-        input_tokens * 1e-6 * model_info["cost"]["pmi"]
-        + output_tokens * 1e-6 * model_info["cost"]["pmo"]
-    )
-    return {
-        "input_tokens": input_tokens,
-        "output_tokens": output_tokens,
-        "cost": cost,
-    }
-
-
-def get_together_token_usage(metadata: dict, model_info: dict):
-    input_tokens = metadata["token_usage"]["prompt_tokens"]
-    output_tokens = metadata["token_usage"]["completion_tokens"]
-    cost = (
-        input_tokens * 1e-6 * model_info["cost"]["pmi"]
-        + output_tokens * 1e-6 * model_info["cost"]["pmo"]
-    )
-    return {
-        "input_tokens": input_tokens,
-        "output_tokens": output_tokens,
-        "cost": cost,
-    }
+        case "OpenAI":
+            llm = ChatOpenAI(
+                model=gen_config["model_name"],
+                temperature=gen_config["temperature"],
+                api_key=st.secrets["openai_api_key"],
+                top_p=gen_config["top_p"],
+                seed=SEED,
+                max_tokens=gen_config["max_output_tokens"],
+            )
 
+        case "Anthropic":
+            llm = ChatAnthropic(
+                model_name=gen_config["model_name"],
+                temperature=gen_config["temperature"],
+                api_key=st.secrets["anthropic_api_key"],
+                top_p=gen_config["top_p"],
+                max_tokens_to_sample=gen_config["max_output_tokens"],
+            )
 
-def render_sidebar():
+        case "Together":
+            llm = ChatTogether(
+                model=gen_config["model_name"],
+                temperature=gen_config["temperature"],
+                max_tokens=gen_config["max_output_tokens"],
+                top_p=gen_config["top_p"],
+                seed=SEED,
+                api_key=st.secrets["together_api_key"],
+            )
 
-    with st.container(border=True):
-        render_outreach_links()
+        case _:
+            raise ValueError()
 
+    return llm
 
-def render_query_rag_tab():
 
+def create_rag_chain(llm, retriever):
     QUERY_RAG_TEMPLATE = """You are an expert legislative analyst. Use the following excerpts from US congressional legislation to respond to the user's query. The excerpts are formatted as a JSON list. Each JSON object has "legis_id", "title", "introduced_date", "sponsor", and "snippets" keys. If a snippet is useful in writing part of your response, then cite the "legis_id", "title", "introduced_date", and "sponsor" in the response. If you don't know how to respond, just tell the user.
 
 ---
@@ -463,219 +431,191 @@ Query: {query}"""
         ]
     )
 
+    rag_chain = (
+        RunnableParallel(
+            {
+                "docs": retriever,
+                "query": RunnablePassthrough(),
+            }
+        )
+        .assign(context=lambda x: format_docs(x["docs"]))
+        .assign(aimessage=prompt | llm)
+    )
+
+    return rag_chain
+
+
+def process_query(gen_config: dict, ret_config: dict, query: str):
+    vectorstore = load_pinecone_vectorstore()
+    llm = get_llm(gen_config)
+    vs_filter = get_vectorstore_filter(ret_config)
+    retriever = vectorstore.as_retriever(
+        search_kwargs={"k": ret_config["n_ret_docs"], "filter": vs_filter},
+    )
+    rag_chain = create_rag_chain(llm, retriever)
+    response = rag_chain.invoke(query)
+    return response
+
+
+def display_retrieved_chunks(response):
+    with st.container(border=True):
+        doc_grps = group_docs(response["docs"])
+        st.write(
+            "Retrieved Chunks (note that you may need to 'right click' on links in the expanders to follow them)"
+        )
+        for legis_id, doc_grp in doc_grps:
+            render_doc_grp(legis_id, doc_grp)
+
+
+def display_response(
+    response, model_info, provider, should_escape_markdown, should_add_legis_urls
+):
+    out_display = response["aimessage"].content
+    if should_escape_markdown:
+        out_display = escape_markdown(out_display)
+    if should_add_legis_urls:
+        out_display = replace_legis_ids_with_urls(out_display)
+
+    with st.container(border=True):
+        st.write("Response")
+        st.info(out_display)
+
+    usage.display_api_usage(response, model_info, provider)
+    display_retrieved_chunks(response)
+
+
+def render_query_rag_tab():
     key_prefix = "query_rag"
     render_example_queries()
 
     with st.form(f"{key_prefix}|query_form"):
-        st.text_area(
-            "Enter a query that can be answered with congressional legislation:",
-            key=f"{key_prefix}|query",
+        query = st.text_area(
+            "Enter a query that can be answered with congressional legislation:"
         )
-        query_submitted = st.form_submit_button("Submit")
+        cols = st.columns(2)
+        with cols[0]:
+            query_submitted = st.form_submit_button("Submit")
+        with cols[1]:
+            status_placeholder = st.empty()
 
     col1, col2 = st.columns(2)
     with col1:
         with st.expander("Generative Config"):
-            render_generative_config(key_prefix)
+            gen_config = get_generative_config(key_prefix)
     with col2:
         with st.expander("Retrieval Config"):
-            render_retrieval_config(key_prefix)
+            ret_config = get_retrieval_config(key_prefix)
 
+    rkey = f"{key_prefix}|response"
     if query_submitted:
-
-        llm = get_llm(key_prefix)
-        vs_filter = get_vectorstore_filter(key_prefix)
-        retriever = vectorstore.as_retriever(
-            search_kwargs={"k": SS[f"{key_prefix}|n_ret_docs"], "filter": vs_filter},
+        with status_placeholder:
+            with st.spinner("generating response"):
+                SS[rkey] = process_query(gen_config, ret_config, query)
+
+    if response := SS.get(rkey):
+        model_info = PROVIDER_MODELS[gen_config["provider"]][gen_config["model_name"]]
+        display_response(
+            response,
+            model_info,
+            gen_config["provider"],
+            gen_config["should_escape_markdown"],
+            gen_config["should_add_legis_urls"],
         )
 
-        rag_chain = (
-            RunnableParallel(
-                {
-                    "docs": retriever,  # list of docs
-                    "query": RunnablePassthrough(),  # str
-                }
-            )
-            .assign(context=(lambda x: format_docs(x["docs"])))
-            .assign(output=prompt | llm)
-        )
-
-        SS[f"{key_prefix}|out"] = rag_chain.invoke(SS[f"{key_prefix}|query"])
-
-    if f"{key_prefix}|out" in SS:
-
-        out_display = SS[f"{key_prefix}|out"]["output"].content
-        if SS[f"{key_prefix}|response_escape_markdown"]:
-            out_display = escape_markdown(out_display)
-        if SS[f"{key_prefix}|response_add_legis_urls"]:
-            out_display = replace_legis_ids_with_urls(out_display)
-        with st.container(border=True):
-            st.write("Response")
-            st.info(out_display)
-
-        with st.container(border=True):
-            st.write("API Usage")
-            token_usage = get_token_usage(
-                key_prefix, SS[f"{key_prefix}|out"]["output"].response_metadata
-            )
-            col1, col2, col3 = st.columns(3)
-            with col1:
-                st.metric("Input Tokens", token_usage["input_tokens"])
-            with col2:
-                st.metric("Output Tokens", token_usage["output_tokens"])
-            with col3:
-                st.metric("Cost", f"${token_usage['cost']:.4f}")
-            with st.expander("Response Metadata"):
-                st.warning(SS[f"{key_prefix}|out"]["output"].response_metadata)
-
-        with st.container(border=True):
-            doc_grps = group_docs(SS[f"{key_prefix}|out"]["docs"])
-            st.write(
-                "Retrieved Chunks (note that you may need to 'right click' on links in the expanders to follow them)"
-            )
-            for legis_id, doc_grp in doc_grps:
-                render_doc_grp(legis_id, doc_grp)
-
         with st.expander("Debug"):
-            st.write(SS[f"{key_prefix}|out"])
+            st.write(response)
 
 
 def render_query_rag_sbs_tab():
-
-    QUERY_RAG_TEMPLATE = """You are an expert legislative analyst. Use the following excerpts from US congressional legislation to respond to the user's query. The excerpts are formatted as a JSON list. Each JSON object has "legis_id", "title", "introduced_date", "sponsor", and "snippets" keys. If a snippet is useful in writing part of your response, then cite the "legis_id", "title", "introduced_date", and "sponsor" in the response. If you don't know how to respond, just tell the user.
-
----
-
-Congressional Legislation Excerpts:
-
-{context}
-
----
-
-Query: {query}"""
-
     base_key_prefix = "query_rag_sbs"
 
-    prompt = ChatPromptTemplate.from_messages(
-        [
-            ("human", QUERY_RAG_TEMPLATE),
-        ]
-    )
-
     with st.form(f"{base_key_prefix}|query_form"):
-        st.text_area(
-            "Enter a query that can be answered with congressional legislation:",
-            key=f"{base_key_prefix}|query",
+        query = st.text_area(
+            "Enter a query that can be answered with congressional legislation:"
         )
-        query_submitted = st.form_submit_button("Submit")
+        cols = st.columns(2)
+        with cols[0]:
+            query_submitted = st.form_submit_button("Submit")
+        with cols[1]:
+            status_placeholder = st.empty()
 
     grp1a, grp2a = st.columns(2)
 
+    gen_configs = {}
+    ret_configs = {}
     with grp1a:
         st.header("Group 1")
         key_prefix = f"{base_key_prefix}|grp1"
         with st.expander("Generative Config"):
-            render_generative_config(key_prefix)
+            gen_configs["grp1"] = get_generative_config(key_prefix)
         with st.expander("Retrieval Config"):
-            render_retrieval_config(key_prefix)
+            ret_configs["grp1"] = get_retrieval_config(key_prefix)
 
     with grp2a:
         st.header("Group 2")
         key_prefix = f"{base_key_prefix}|grp2"
         with st.expander("Generative Config"):
-            render_generative_config(key_prefix)
+            gen_configs["grp2"] = get_generative_config(key_prefix)
         with st.expander("Retrieval Config"):
-            render_retrieval_config(key_prefix)
+            ret_configs["grp2"] = get_retrieval_config(key_prefix)
 
     grp1b, grp2b = st.columns(2)
     sbs_cols = {"grp1": grp1b, "grp2": grp2b}
+    grp_names = {"grp1": "Group 1", "grp2": "Group 2"}
 
     for post_key_prefix in ["grp1", "grp2"]:
+        with sbs_cols[post_key_prefix]:
+            key_prefix = f"{base_key_prefix}|{post_key_prefix}"
+            rkey = f"{key_prefix}|response"
+            if query_submitted:
+                with status_placeholder:
+                    with st.spinner(
+                        "generating response for {}".format(grp_names[post_key_prefix])
+                    ):
+                        SS[rkey] = process_query(
+                            gen_configs[post_key_prefix],
+                            ret_configs[post_key_prefix],
+                            query,
+                        )
+
+            if response := SS.get(rkey):
+                model_info = PROVIDER_MODELS[gen_configs[post_key_prefix]["provider"]][
+                    gen_configs[post_key_prefix]["model_name"]
+                ]
+                display_response(
+                    response,
+                    model_info,
+                    gen_configs[post_key_prefix]["provider"],
+                    gen_configs[post_key_prefix]["should_escape_markdown"],
+                    gen_configs[post_key_prefix]["should_add_legis_urls"],
+                )
 
-        key_prefix = f"{base_key_prefix}|{post_key_prefix}"
 
-        if query_submitted:
-            llm = get_llm(key_prefix)
-            vs_filter = get_vectorstore_filter(key_prefix)
-            retriever = vectorstore.as_retriever(
-                search_kwargs={
-                    "k": SS[f"{key_prefix}|n_ret_docs"],
-                    "filter": vs_filter,
-                },
-            )
-            rag_chain = (
-                RunnableParallel(
-                    {
-                        "docs": retriever,  # list of docs
-                        "query": RunnablePassthrough(),  # str
-                    }
-                )
-                .assign(context=(lambda x: format_docs(x["docs"])))
-                .assign(output=prompt | llm)
-            )
-            SS[f"{key_prefix}|out"] = rag_chain.invoke(SS[f"{base_key_prefix}|query"])
-
-        if f"{key_prefix}|out" in SS:
-            with sbs_cols[post_key_prefix]:
-                out_display = SS[f"{key_prefix}|out"]["output"].content
-                if SS[f"{key_prefix}|response_escape_markdown"]:
-                    out_display = escape_markdown(out_display)
-                if SS[f"{key_prefix}|response_add_legis_urls"]:
-                    out_display = replace_legis_ids_with_urls(out_display)
-                with st.container(border=True):
-                    st.write("Response")
-                    st.info(out_display)
-
-                with st.container(border=True):
-                    st.write("API Usage")
-                    token_usage = get_token_usage(
-                        key_prefix, SS[f"{key_prefix}|out"]["output"].response_metadata
-                    )
-                    col1, col2, col3 = st.columns(3)
-                    with col1:
-                        st.metric("Input Tokens", token_usage["input_tokens"])
-                    with col2:
-                        st.metric("Output Tokens", token_usage["output_tokens"])
-                    with col3:
-                        st.metric("Cost", f"${token_usage['cost']:.4f}")
-                    with st.expander("Response Metadata"):
-                        st.warning(SS[f"{key_prefix}|out"]["output"].response_metadata)
-
-                with st.container(border=True):
-                    doc_grps = group_docs(SS[f"{key_prefix}|out"]["docs"])
-                    st.write(
-                        "Retrieved Chunks (note that you may need to 'right click' on links in the expanders to follow them)"
-                    )
-                    for legis_id, doc_grp in doc_grps:
-                        render_doc_grp(legis_id, doc_grp)
-
-
-##################
-
-
-st.title(":classical_building: LegisQA :classical_building:")
-st.header("Chat With Congressional Bills")
-
-
-with st.sidebar:
-    render_sidebar()
-
-
-vectorstore = load_pinecone_vectorstore()
-
-query_rag_tab, query_rag_sbs_tab, guide_tab = st.tabs(
-    [
-        "RAG",
-        "RAG (side-by-side)",
-        "Guide",
-    ]
-)
+def main():
+
+    st.title(":classical_building: LegisQA :classical_building:")
+    st.header("Query Congressional Bills")
+
+    with st.sidebar:
+        render_sidebar()
+
+    query_rag_tab, query_rag_sbs_tab, guide_tab = st.tabs(
+        [
+            "RAG",
+            "RAG (side-by-side)",
+            "Guide",
+        ]
+    )
+
+    with query_rag_tab:
+        render_query_rag_tab()
+
+    with query_rag_sbs_tab:
+        render_query_rag_sbs_tab()
 
-with query_rag_tab:
-    render_query_rag_tab()
+    with guide_tab:
+        render_guide()
 
-with query_rag_sbs_tab:
-    render_query_rag_sbs_tab()
 
-with guide_tab:
-    render_guide()
+if __name__ == "__main__":
+    main()

usage.py

@@ -0,0 +1,72 @@
+import streamlit as st
+
+
+def get_openai_token_usage(metadata: dict, model_info: dict):
+    input_tokens = metadata["token_usage"]["prompt_tokens"]
+    output_tokens = metadata["token_usage"]["completion_tokens"]
+    cost = (
+        input_tokens * 1e-6 * model_info["cost"]["pmi"]
+        + output_tokens * 1e-6 * model_info["cost"]["pmo"]
+    )
+    return {
+        "input_tokens": input_tokens,
+        "output_tokens": output_tokens,
+        "cost": cost,
+    }
+
+
+def get_anthropic_token_usage(metadata: dict, model_info: dict):
+    input_tokens = metadata["usage"]["input_tokens"]
+    output_tokens = metadata["usage"]["output_tokens"]
+    cost = (
+        input_tokens * 1e-6 * model_info["cost"]["pmi"]
+        + output_tokens * 1e-6 * model_info["cost"]["pmo"]
+    )
+    return {
+        "input_tokens": input_tokens,
+        "output_tokens": output_tokens,
+        "cost": cost,
+    }
+
+
+def get_together_token_usage(metadata: dict, model_info: dict):
+    input_tokens = metadata["token_usage"]["prompt_tokens"]
+    output_tokens = metadata["token_usage"]["completion_tokens"]
+    cost = (
+        input_tokens * 1e-6 * model_info["cost"]["pmi"]
+        + output_tokens * 1e-6 * model_info["cost"]["pmo"]
+    )
+    return {
+        "input_tokens": input_tokens,
+        "output_tokens": output_tokens,
+        "cost": cost,
+    }
+
+
+def get_token_usage(metadata: dict, model_info: dict, provider: str):
+    match provider:
+        case "OpenAI":
+            return get_openai_token_usage(metadata, model_info)
+        case "Anthropic":
+            return get_anthropic_token_usage(metadata, model_info)
+        case "Together":
+            return get_together_token_usage(metadata, model_info)
+        case _:
+            raise ValueError()
+
+
+def display_api_usage(response, model_info, provider: str):
+    with st.container(border=True):
+        st.write("API Usage")
+        token_usage = get_token_usage(
+            response["aimessage"].response_metadata, model_info, provider
+        )
+        col1, col2, col3 = st.columns(3)
+        with col1:
+            st.metric("Input Tokens", token_usage["input_tokens"])
+        with col2:
+            st.metric("Output Tokens", token_usage["output_tokens"])
+        with col3:
+            st.metric("Cost", f"${token_usage['cost']:.4f}")
+        with st.expander("Response Metadata"):
+            st.warning(response["aimessage"].response_metadata)