llama 3.1

app.py

@@ -1,6 +1,4 @@
 """
-TODO: checkout langgraph
-TODO: clear screen between agent calls (see here https://github.com/langchain-ai/streamlit-agent/blob/main/streamlit_agent/clear_results.py)
 """
 
 from collections import defaultdict
@@ -11,7 +9,9 @@ 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.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
@@ -56,23 +56,22 @@ CONGRESS_GOV_TYPE_MAP = {
     "sjres": "senate-joint-resolution",
     "sres": "senate-resolution",
 }
-OPENAI_CHAT_MODELS = [
-    "gpt-4o-mini",
-    "gpt-4o",
-]
-ANTHROPIC_CHAT_MODELS = [
-    "claude-3-haiku-20240307",
-    "claude-3-5-sonnet-20240620",
-    "claude-3-opus-20240229",
-]
-TOGETHER_CHAT_MODELS = [
-    "meta-llama/Meta-Llama-3.1-8B-Instruct-Turbo",
-    "meta-llama/Meta-Llama-3.1-70B-Instruct-Turbo",
-    "meta-llama/Meta-Llama-3.1-405B-Instruct-Turbo",
-]
-    
-
-CHAT_MODELS = OPENAI_CHAT_MODELS + ANTHROPIC_CHAT_MODELS + TOGETHER_CHAT_MODELS
+OPENAI_CHAT_MODELS = {
+    "gpt-4o-mini": {"cost": {"pmi": 0.15, "pmo": 0.60}},
+    #    "gpt-4o": {"cost": {"pmi": 5.00, "pmo": 15.0}},
+}
+ANTHROPIC_CHAT_MODELS = {
+    "claude-3-haiku-20240307": {"cost": {"pmi": 0.25, "pmo": 1.25}},
+    #    "claude-3-5-sonnet-20240620": {"cost": {"pmi": 3.00, "pmo": 15.0}},
+    #    "claude-3-opus-20240229": {"cost": {"pmi": 15.0, "pmo": 75.0}},
+}
+TOGETHER_CHAT_MODELS = {
+    "meta-llama/Meta-Llama-3.1-8B-Instruct-Turbo": {"cost": {"pmi": 0.18, "pmo": 0.18}},
+    "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}},
+}
 
 PROVIDER_MODELS = {
     "OpenAI": OPENAI_CHAT_MODELS,
@@ -174,14 +173,20 @@ def escape_markdown(text):
     return text
 
 
-def get_vectorstore_filter():
+def get_vectorstore_filter(key_prefix: str):
     vs_filter = {}
-    if SS["filter_legis_id"] != "":
-        vs_filter["legis_id"] = SS["filter_legis_id"]
-    if SS["filter_bioguide_id"] != "":
-        vs_filter["sponsor_bioguide_id"] = SS["filter_bioguide_id"]
-    vs_filter = {**vs_filter, "congress_num": {"$in": SS["filter_congress_nums"]}}
-    vs_filter = {**vs_filter, "sponsor_party": {"$in": SS["filter_sponsor_parties"]}}
+    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"]
+    vs_filter = {
+        **vs_filter,
+        "congress_num": {"$in": SS[f"{key_prefix}|filter_congress_nums"]},
+    }
+    vs_filter = {
+        **vs_filter,
+        "sponsor_party": {"$in": SS[f"{key_prefix}|filter_sponsor_parties"]},
+    }
     return vs_filter
 
 
@@ -195,7 +200,6 @@ def render_doc_grp(legis_id: str, doc_grp: list[Document]):
     )
     congress_gov_link = f"[congress.gov]({congress_gov_url})"
 
-
     ref = "{} chunks from {}\n\n{}\n\n{}\n\n[{} ({}) ]({})".format(
         len(doc_grp),
         first_doc.metadata["legis_id"],
@@ -282,53 +286,118 @@ Suggest reforms that would benefit the Medicaid program.
         )
 
 
-def render_sidebar():
+def render_generative_config(key_prefix: str):
+    st.selectbox(
+        label="provider", options=PROVIDER_MODELS.keys(), key=f"{key_prefix}|provider"
+    )
+    st.selectbox(
+        label="model name",
+        options=PROVIDER_MODELS[SS[f"{key_prefix}|provider"]],
+        key=f"{key_prefix}|model_name",
+    )
+    st.slider(
+        "temperature",
+        min_value=0.0,
+        max_value=2.0,
+        value=0.01,
+        key=f"{key_prefix}|temperature",
+    )
+    st.slider(
+        "max_output_tokens",
+        min_value=512,
+        max_value=1024,
+        key=f"{key_prefix}|max_output_tokens",
+    )
+    st.slider(
+        "top_p", min_value=0.0, max_value=1.0, value=0.9, key=f"{key_prefix}|top_p"
+    )
+    st.checkbox(
+        "escape markdown in answer", key=f"{key_prefix}|response_escape_markdown"
+    )
+    st.checkbox(
+        "add legis urls in answer",
+        value=True,
+        key=f"{key_prefix}|response_add_legis_urls",
+    )
 
-    with st.container(border=True):
-        render_outreach_links()
 
-    st.checkbox("escape markdown in answer", key="response_escape_markdown")
-    st.checkbox("add legis urls in answer", value=True, key="response_add_legis_urls")
+def render_retrieval_config(key_prefix: str):
+    st.slider(
+        "Number of chunks to retrieve",
+        min_value=1,
+        max_value=32,
+        value=8,
+        key=f"{key_prefix}|n_ret_docs",
+    )
+    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(
+        "Congress Numbers",
+        CONGRESS_NUMBERS,
+        default=CONGRESS_NUMBERS,
+        key=f"{key_prefix}|filter_congress_nums",
+    )
+    st.multiselect(
+        "Sponsor Party",
+        SPONSOR_PARTIES,
+        default=SPONSOR_PARTIES,
+        key=f"{key_prefix}|filter_sponsor_parties",
+    )
+
 
-    with st.expander("Generative Config"):
-        st.selectbox(label="provider", options=PROVIDER_MODELS.keys(), key="provider")
-        st.selectbox(label="model name", options=PROVIDER_MODELS[SS["provider"]], key="model_name")
-        st.slider(
-            "temperature", min_value=0.0, max_value=2.0, value=0.01, key="temperature"
-        )
-        st.slider(
-            "max_output_tokens", min_value=512, max_value=1024, key="max_output_tokens"
-        )
-        st.slider("top_p", min_value=0.0, max_value=1.0, value=1.0, key="top_p")
-
-    with st.expander("Retrieval Config"):
-        st.slider(
-            "Number of chunks to retrieve",
-            min_value=1,
-            max_value=32,
-            value=8,
-            key="n_ret_docs",
+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"],
         )
-        st.text_input("Bill ID (e.g. 118-s-2293)", key="filter_legis_id")
-        st.text_input("Bioguide ID (e.g. R000595)", key="filter_bioguide_id")
-        st.multiselect(
-            "Congress Numbers",
-            CONGRESS_NUMBERS,
-            default=CONGRESS_NUMBERS,
-            key="filter_congress_nums",
+    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"],
         )
-        st.multiselect(
-            "Sponsor Party",
-            SPONSOR_PARTIES,
-            default=SPONSOR_PARTIES,
-            key="filter_sponsor_parties",
+    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 render_sidebar():
+
+    with st.container(border=True):
+        render_outreach_links()
 
 
 def render_query_rag_tab():
 
+    key_prefix = "query_rag"
     render_example_queries()
 
+    col1, col2 = st.columns(2)
+    with col1:
+        with st.expander("Generative Config"):
+            render_generative_config(key_prefix)
+    with col2:
+        with st.expander("Retrieval Config"):
+            render_retrieval_config(key_prefix)
+
     QUERY_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.
 
 ---
@@ -348,14 +417,18 @@ Query: {query}"""
     )
 
     with st.form("query_form"):
-        st.text_area("Enter a query that can be answered with congressional legislation:", key="query")
+        st.text_area(
+            "Enter a query that can be answered with congressional legislation:",
+            key=f"{key_prefix}|query",
+        )
         query_submitted = st.form_submit_button("Submit")
 
     if query_submitted:
 
-        vs_filter = get_vectorstore_filter()
+        llm = get_llm(key_prefix)
+        vs_filter = get_vectorstore_filter(key_prefix)
         retriever = vectorstore.as_retriever(
-            search_kwargs={"k": SS["n_ret_docs"], "filter": vs_filter},
+            search_kwargs={"k": SS[f"{key_prefix}|n_ret_docs"], "filter": vs_filter},
         )
 
         rag_chain = (
@@ -364,37 +437,41 @@ Query: {query}"""
                     "docs": retriever,  # list of docs
                     "query": RunnablePassthrough(),  # str
                 }
-            )
-            .assign(context=(lambda x: format_docs(x["docs"])))
-            .assign(output=prompt | llm | StrOutputParser())
+            ).assign(context=(lambda x: format_docs(x["docs"])))
+            #            .assign(output=prompt | llm | StrOutputParser())
+            .assign(output=prompt | llm)
         )
 
-        if SS["model_name"] in OPENAI_CHAT_MODELS:
-            with get_openai_callback() as cb:
-                SS["out"] = rag_chain.invoke(SS["query"])
-                SS["cb"] = cb
-        else:
-            SS.pop("cb", None)
-            SS["out"] = rag_chain.invoke(SS["query"])
+        SS[f"{key_prefix}|out"] = rag_chain.invoke(SS[f"{key_prefix}|query"])
 
-    if "out" in SS:
+    if f"{key_prefix}|out" in SS:
 
-        out_display = SS["out"]["output"]
-        if SS["response_escape_markdown"]:
+        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["response_add_legis_urls"]:
+        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)
 
-        if "cb" in SS:
-            with st.container(border=True):
-                st.write("API Usage")
-                st.warning(SS["cb"])
+        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["out"]["docs"])
+            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)"
             )
@@ -402,86 +479,68 @@ Query: {query}"""
                 render_doc_grp(legis_id, doc_grp)
 
         with st.expander("Debug"):
-            st.write(SS["out"])
-
-
-def render_query_agent_tab():
-
-    from retriever_tools import get_retriever_tool
-
-    from langchain_community.tools import WikipediaQueryRun
-    from langchain_community.utilities import WikipediaAPIWrapper
-#    from langchain.agents import load_tools
-    from langchain_community.agent_toolkits.load_tools import load_tools
-    from langchain.agents import create_react_agent
-    from langchain import hub
-
-    if SS["model_name"] not in OPENAI_CHAT_MODELS:
-        st.write("only supported with OpenAI for now")
-        return
-
-    vs_filter = get_vectorstore_filter()
-    retriever = vectorstore.as_retriever(
-        search_kwargs={"k": SS["n_ret_docs"], "filter": vs_filter},
+            st.write(SS[f"{key_prefix}|out"])
+
+
+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"]
     )
-    legis_retrieval_tool = get_retriever_tool(
-        retriever,
-        "search_legislation",
-        "Searches and returns excerpts from congressional legislation. Always call this tool first.",
-        format_docs,
+    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"]
     )
-
-    api_wrapper = WikipediaAPIWrapper(top_k_results=4, doc_content_chars_max=800)
-    wiki_search_tool = WikipediaQueryRun(api_wrapper=api_wrapper)
-
-    ddg_tool = load_tools(["ddg-search"])[0]
-
-    avatars = {"human": "user", "ai": "assistant"}
-    tools = [legis_retrieval_tool, wiki_search_tool, ddg_tool]
-    llm_with_tools = llm.bind_tools(tools)
-
-    agent_prompt = ChatPromptTemplate.from_messages(
-        [
-            ("system", "You are a helpful assistant."),
-            ("human", "{input}"),
-            MessagesPlaceholder(variable_name="agent_scratchpad"),
-        ]
+    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"]
     )
-    agent = (
-        {
-            "input": lambda x: x["input"],
-            "agent_scratchpad": lambda x: format_to_openai_tool_messages(
-                x["intermediate_steps"]
-            ),
-        }
-        | agent_prompt
-        | llm_with_tools
-        | OpenAIToolsAgentOutputParser()
-    )
-
-    prompt = hub.pull("hwchase17/react")
-    agent = create_react_agent(llm, tools, prompt)
-    agent_executor = AgentExecutor(
-        agent=agent,
-        tools=tools,
-        return_intermediate_steps=True,
-        handle_parsing_errors=True,
-        verbose=True,
-    )
-
-    if user_input := st.chat_input(key="single_query_agent_input"):
-        st.chat_message("user").write(user_input)
-        with st.chat_message("assistant"):
-            st_callback = StreamlitCallbackHandler(st.container())
-            response = agent_executor.invoke({"input": user_input}, {"callbacks": [st_callback]})
-            st.write(response["output"])
-
-
-def render_chat_agent_tab():
-    st.write("Coming Soon")
+    return {
+        "input_tokens": input_tokens,
+        "output_tokens": output_tokens,
+        "cost": cost,
+    }
 
 
+def render_query_rag_sbs_tab():
 
+    return
 
 
 ##################
@@ -495,53 +554,21 @@ with st.sidebar:
     render_sidebar()
 
 
-if SS["model_name"] in OPENAI_CHAT_MODELS:
-    llm = ChatOpenAI(
-        model=SS["model_name"],
-        temperature=SS["temperature"],
-        api_key=st.secrets["openai_api_key"],
-        top_p=SS["top_p"],
-        seed=SEED,
-        max_tokens=SS["max_output_tokens"],
-    )
-elif SS["model_name"] in ANTHROPIC_CHAT_MODELS:
-    llm = ChatAnthropic(
-        model_name=SS["model_name"],
-        temperature=SS["temperature"],
-        api_key=st.secrets["anthropic_api_key"],
-        top_p=SS["top_p"],
-        max_tokens_to_sample=SS["max_output_tokens"],
-    )
-elif SS["model_name"] in TOGETHER_CHAT_MODELS:
-    llm = ChatTogether(
-        model=SS["model_name"],
-        temperature=SS["temperature"],
-        max_tokens=SS["max_output_tokens"],
-        top_p=SS["top_p"],
-        seed=SEED,
-        api_key=st.secrets["together_api_key"],
-    )
-else:
-    raise ValueError()
-
-
 vectorstore = load_pinecone_vectorstore()
 
-query_rag_tab, query_agent_tab, chat_agent_tab, guide_tab = st.tabs([
-    "query_rag",
-    "query_agent",
-    "chat_agent",
-    "guide",
-])
+query_rag_tab, query_rag_sbs_tab, guide_tab = st.tabs(
+    [
+        "query_rag",
+        "query_rag_sbs",
+        "guide",
+    ]
+)
 
 with query_rag_tab:
     render_query_rag_tab()
 
-with query_agent_tab:
-    render_query_agent_tab()
-
-with chat_agent_tab:
-    render_chat_agent_tab()
+with query_rag_sbs_tab:
+    render_query_rag_sbs_tab()
 
 with guide_tab:
     render_guide()