lets go

app.py

@@ -1,13 +1,23 @@
+from collections import defaultdict
+import json
+
+from langchain_core.documents import Document
+from langchain_core.prompts import PromptTemplate
+from langchain_core.runnables import RunnableParallel
+from langchain_core.runnables import RunnablePassthrough
+from langchain_core.output_parsers import StrOutputParser
 from langchain_community.embeddings import HuggingFaceBgeEmbeddings
-from langchain_community.vectorstores import Pinecone
 from langchain_community.vectorstores.utils import DistanceStrategy
 from langchain_openai import ChatOpenAI
-from pinecone import Pinecone as PineconeClient
+from langchain_pinecone import PineconeVectorStore
+from pinecone import Pinecone
 import streamlit as st
 
-st.set_page_config(layout="wide", page_title="LegisQA")
 
+st.set_page_config(layout="wide", page_title="LegisQA")
+SS = st.session_state
 
+SEED = 292764
 CONGRESS_GOV_TYPE_MAP = {
     "hconres": "house-concurrent-resolution",
     "hjres": "house-joint-resolution",
@@ -25,6 +35,48 @@ OPENAI_CHAT_MODELS = [
 ]
 
 
+PREAMBLE = "You are an expert analyst. Use the following excerpts from US congressional legislation to respond to the user's query."
+PROMPT_TEMPLATES = {
+    "v1": PREAMBLE
+    + """ If you don't know how to respond, just tell the user.
+
+{context}
+
+Question: {question}""",
+    "v2": PREAMBLE
+    + """ Each snippet starts with a header that includes a unique snippet number (snippet_num), a legis_id, and a title. Your response should reference particular snippets using legis_id and title. If you don't know how to respond, just tell the user.
+
+{context}
+
+Question: {question}""",
+    "v3": PREAMBLE
+    + """ Each excerpt starts with a header that includes a legis_id, and a title followed by one or more text snippets. When using text snippets in your response, you should mention the legis_id and title. If you don't know how to respond, just tell the user.
+
+{context}
+
+Question: {question}""",
+    "v4": PREAMBLE
+    + """ The excerpts are formatted as a JSON list. Each JSON object has "legis_id", "title", and "snippets" keys. If a snippet is useful in writing part of your response, then mention the "title" and "legis_id" inline as you write. If you don't know how to respond, just tell the user.
+
+{context}
+
+Query: {question}""",
+}
+
+
+def get_sponsor_url(bioguide_id: str) -> str:
+    return f"https://bioguide.congress.gov/search/bio/{bioguide_id}"
+
+
+def get_congress_gov_url(congress_num: int, legis_type: str, legis_num: int) -> str:
+    lt = CONGRESS_GOV_TYPE_MAP[legis_type]
+    return f"https://www.congress.gov/bill/{int(congress_num)}th-congress/{lt}/{int(legis_num)}"
+
+
+def get_govtrack_url(congress_num: int, legis_type: str, legis_num: int) -> str:
+    return f"https://www.govtrack.us/congress/bills/{int(congress_num)}/{legis_type}{int(legis_num)}"
+
+
 def load_bge_embeddings():
     model_name = "BAAI/bge-small-en-v1.5"
     model_kwargs = {"device": "cpu"}
@@ -40,9 +92,9 @@ def load_bge_embeddings():
 
 def load_pinecone_vectorstore():
     emb_fn = load_bge_embeddings()
-    pc = PineconeClient(api_key=st.secrets["pinecone_api_key"])
+    pc = Pinecone(api_key=st.secrets["pinecone_api_key"])
     index = pc.Index(st.secrets["pinecone_index_name"])
-    vectorstore = Pinecone(
+    vectorstore = PineconeVectorStore(
         index=index,
         embedding=emb_fn,
         text_key="text",
@@ -51,8 +103,255 @@ def load_pinecone_vectorstore():
     return vectorstore
 
 
+def write_outreach_links():
+    nomic_base_url = "https://atlas.nomic.ai/data/gabrielhyperdemocracy"
+    nomic_map_name = "us-congressional-legislation-s1024o256nomic"
+    nomic_url = f"{nomic_base_url}/{nomic_map_name}/map"
+    hf_url = "https://huggingface.co/hyperdemocracy"
+    st.subheader(":brain: Learn about [hyperdemocracy](https://hyperdemocracy.us)")
+    st.subheader(f":world_map: Visualize with [nomic atlas]({nomic_url})")
+    st.subheader(f":hugging_face: Explore the [huggingface datasets](hf_url)")
+
+
+def group_docs(docs) -> list[tuple[str, list[Document]]]:
+    doc_grps = defaultdict(list)
+
+    # create legis_id groups
+    for doc in docs:
+        doc_grps[doc.metadata["legis_id"]].append(doc)
+
+    # sort docs in each group by start index
+    for legis_id in doc_grps.keys():
+        doc_grps[legis_id] = sorted(
+            doc_grps[legis_id],
+            key=lambda x: x.metadata["start_index"],
+        )
+
+    # sort groups by number of docs
+    doc_grps = sorted(
+        tuple(doc_grps.items()),
+        key=lambda x: -len(x[1]),
+    )
+
+    return doc_grps
+
+
+def format_docs_v1(docs):
+    """Simple double new line join"""
+    return "\n\n".join([doc.page_content for doc in docs])
+
+
+def format_docs_v2(docs):
+    """Format with snippet_num, legis_id, and title"""
+
+    def format_doc(idoc, doc):
+        return "snippet_num: {}\nlegis_id: {}\ntitle: {}\n... {} ...\n".format(
+            idoc,
+            doc.metadata["legis_id"],
+            doc.metadata["title"],
+            doc.page_content,
+        )
+
+    snips = []
+    for idoc, doc in enumerate(docs):
+        txt = format_doc(idoc, doc)
+        snips.append(txt)
+
+    return "\n===\n".join(snips)
+
+
+def format_docs_v3(docs):
+
+    def format_header(doc):
+        return "legis_id: {}\ntitle: {}".format(
+            doc.metadata["legis_id"],
+            doc.metadata["title"],
+        )
+
+    def format_content(doc):
+        return "... {} ...\n".format(
+            doc.page_content,
+        )
+
+    snips = []
+    doc_grps = group_docs(docs)
+    for legis_id, doc_grp in doc_grps:
+        first_doc = doc_grp[0]
+        head = format_header(first_doc)
+        contents = []
+        for idoc, doc in enumerate(doc_grp):
+            txt = format_content(doc)
+            contents.append(txt)
+        snips.append("{}\n\n{}".format(head, "\n".join(contents)))
+
+    return "\n===\n".join(snips)
+
+
+def format_docs_v4(docs):
+    """JSON grouped"""
+
+    doc_grps = group_docs(docs)
+    out = []
+    for legis_id, doc_grp in doc_grps:
+        dd = {
+            "legis_id": doc_grp[0].metadata["legis_id"],
+            "title": doc_grp[0].metadata["title"],
+            "snippets": [doc.page_content for doc in doc_grp],
+        }
+        out.append(dd)
+    return json.dumps(out, indent=4)
+
+
+DOC_FORMATTERS = {
+    "v1": format_docs_v1,
+    "v2": format_docs_v2,
+    "v3": format_docs_v3,
+    "v4": format_docs_v4,
+}
+
+
+def escape_markdown(text):
+    MD_SPECIAL_CHARS = r"\`*_{}[]()#+-.!$"
+    for char in MD_SPECIAL_CHARS:
+        text = text.replace(char, "\\" + char)
+    return text
+
+
+with st.sidebar:
+
+    with st.container(border=True):
+        write_outreach_links()
+
+    st.checkbox("escape markdown in answer", key="response_escape_markdown")
+
+    with st.expander("Generative Config"):
+        st.selectbox(label="model name", options=OPENAI_CHAT_MODELS, key="model_name")
+        st.slider(
+            "temperature", min_value=0.0, max_value=2.0, value=0.0, key="temperature"
+        )
+        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=40,
+            value=10,
+            key="n_ret_docs",
+        )
+        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.text_input("Congress (e.g. 118)", key="filter_congress_num")
+
+    with st.expander("Prompt Config"):
+        st.selectbox(
+            label="prompt version",
+            options=PROMPT_TEMPLATES.keys(),
+            index=3,
+            key="prompt_version",
+        )
+        st.text_area(
+            "prompt template",
+            PROMPT_TEMPLATES[SS["prompt_version"]],
+            height=300,
+            key="prompt_template",
+        )
+
+
+llm = ChatOpenAI(
+    model_name=SS["model_name"],
+    temperature=SS["temperature"],
+    openai_api_key=st.secrets["openai_api_key"],
+    model_kwargs={"top_p": SS["top_p"], "seed": SEED},
+)
 
 vectorstore = load_pinecone_vectorstore()
-query = st.text_area("Enter query")
-docs = vectorstore.similarity_search_with_score(query)
-st.write(docs)
+format_docs = DOC_FORMATTERS[SS["prompt_version"]]
+
+with st.form("my_form"):
+    st.text_area("Enter question:", key="query")
+    query_submitted = st.form_submit_button("Submit")
+
+
+def get_vectorstore_filter():
+    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"]
+    if SS["filter_congress_num"] != "":
+        vs_filter["congress_num"] = int(SS["filter_congress_num"])
+    return vs_filter
+
+
+if query_submitted:
+
+    vs_filter = get_vectorstore_filter()
+    retriever = vectorstore.as_retriever(
+        search_kwargs={"k": SS["n_ret_docs"], "filter": vs_filter},
+    )
+    prompt = PromptTemplate.from_template(SS["prompt_template"])
+    rag_chain_from_docs = (
+        RunnablePassthrough.assign(context=(lambda x: format_docs(x["context"])))
+        | prompt
+        | llm
+        | StrOutputParser()
+    )
+    rag_chain_with_source = RunnableParallel(
+        {"context": retriever, "question": RunnablePassthrough()}
+    ).assign(answer=rag_chain_from_docs)
+    out = rag_chain_with_source.invoke(SS["query"])
+    SS["out"] = out
+
+
+def write_doc_grp(legis_id: str, doc_grp: list[Document]):
+    first_doc = doc_grp[0]
+
+    congress_gov_url = get_congress_gov_url(
+        first_doc.metadata["congress_num"],
+        first_doc.metadata["legis_type"],
+        first_doc.metadata["legis_num"],
+    )
+    congress_gov_link = f"[congress.gov]({congress_gov_url})"
+
+    gov_track_url = get_govtrack_url(
+        first_doc.metadata["congress_num"],
+        first_doc.metadata["legis_type"],
+        first_doc.metadata["legis_num"],
+    )
+    gov_track_link = f"[govtrack.us]({gov_track_url})"
+
+    ref = "{} chunks from {}\n\n{}\n\n{} | {}\n\n[{} ({}) ]({})".format(
+        len(doc_grp),
+        first_doc.metadata["legis_id"],
+        first_doc.metadata["title"],
+        congress_gov_link,
+        gov_track_link,
+        first_doc.metadata["sponsor_full_name"],
+        first_doc.metadata["sponsor_bioguide_id"],
+        get_sponsor_url(first_doc.metadata["sponsor_bioguide_id"]),
+    )
+    doc_contents = [
+        "[start_index={}] ".format(int(doc.metadata["start_index"])) + doc.page_content
+        for doc in doc_grp
+    ]
+    with st.expander(ref):
+        st.write(escape_markdown("\n\n...\n\n".join(doc_contents)))
+
+
+out = SS.get("out")
+if out:
+
+    if SS["response_escape_markdown"]:
+        st.info(escape_markdown(out["answer"]))
+    else:
+        st.info(out["answer"])
+
+    doc_grps = group_docs(out["context"])
+    for legis_id, doc_grp in doc_grps:
+        write_doc_grp(legis_id, doc_grp)
+
+    with st.expander("Debug doc format"):
+
+        st.text_area("formatted docs", value=format_docs(out["context"]), height=600)
+        #    st.write(json.loads(format_docs(out["context"])))

requirements.txt

@@ -33,6 +33,7 @@ jsonschema-specifications==2023.12.1
 langchain-community==0.0.24
 langchain-core==0.1.26
 langchain-openai==0.0.7
+langchain-pinecone==0.0.3
 langsmith==0.1.7
 markdown-it-py==3.0.0
 MarkupSafe==2.1.5