from collections import defaultdict import json import os import re 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.callbacks import get_openai_callback from langchain_community.embeddings import HuggingFaceBgeEmbeddings from langchain_community.vectorstores.utils import DistanceStrategy from langchain_openai import ChatOpenAI from langchain_pinecone import PineconeVectorStore from pinecone import Pinecone import streamlit as st 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"] SS = st.session_state SEED = 292764 CONGRESS_NUMBERS = [113, 114, 115, 116, 117, 118] SPONSOR_PARTIES = ["D", "R", "L", "I"] CONGRESS_GOV_TYPE_MAP = { "hconres": "house-concurrent-resolution", "hjres": "house-joint-resolution", "hr": "house-bill", "hres": "house-resolution", "s": "senate-bill", "sconres": "senate-concurrent-resolution", "sjres": "senate-joint-resolution", "sres": "senate-resolution", } OPENAI_CHAT_MODELS = [ "gpt-3.5-turbo-0125", "gpt-4-0125-preview", ] 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 cite 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 cite 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 cite the "title" and "legis_id" in the response. 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"} encode_kwargs = {"normalize_embeddings": True} emb_fn = HuggingFaceBgeEmbeddings( model_name=model_name, model_kwargs=model_kwargs, encode_kwargs=encode_kwargs, query_instruction="Represent this question for searching relevant passages: ", ) return emb_fn def load_pinecone_vectorstore(): emb_fn = load_bge_embeddings() vectorstore = PineconeVectorStore( embedding=emb_fn, text_key="text", distance_strategy=DistanceStrategy.COSINE, pinecone_api_key=st.secrets["pinecone_api_key"], index_name=st.secrets["pinecone_index_name"], ) 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" pc_url = "https://www.pinecone.io/blog/serverless" st.subheader(":brain: About [hyperdemocracy](https://hyperdemocracy.us)") st.subheader(f":world_map: Visualize [nomic atlas]({nomic_url})") st.subheader(f":hugging_face: Raw [huggingface datasets]({hf_url})") st.subheader(f":evergreen_tree: Index [pinecone serverless]({pc_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 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"] vs_filter = {**vs_filter, "congress_num": {"$in": SS["filter_congress_nums"]}} vs_filter = {**vs_filter, "sponsor_party": {"$in": SS["filter_sponsor_parties"]}} return vs_filter 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, 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))) def legis_id_to_link(legis_id: str) -> str: congress_num, legis_type, legis_num = legis_id.split("-") return get_congress_gov_url(congress_num, legis_type, legis_num) def legis_id_match_to_link(matchobj): mstring = matchobj.string[matchobj.start() : matchobj.end()] url = legis_id_to_link(mstring) link = f"[{mstring}]({url})" return link def replace_legis_ids_with_urls(text): pattern = "11[345678]-[a-z]+-\d{1,5}" rtext = re.sub(pattern, legis_id_match_to_link, text) return rtext def write_guide(): st.write( """ When you send a query to LegisQA, it will attempt to retrieve relevant content from the past six congresses ([113th-118th](https://en.wikipedia.org/wiki/List_of_United_States_Congresses)) covering 2013 to the present, pass it to a [large language model (LLM)](https://en.wikipedia.org/wiki/Large_language_model), and generate a response. This technique is known as Retrieval Augmented Generation (RAG). You can read [an academic paper](https://proceedings.neurips.cc/paper/2020/hash/6b493230205f780e1bc26945df7481e5-Abstract.html) or [a high level summary](https://research.ibm.com/blog/retrieval-augmented-generation-RAG) to get more details. Once the response is generated, the retrieved content will be available for inspection with links to the bills and sponsors. ## Disclaimer This is a research project. The RAG technique helps to ground the LLM response by providing context from a trusted source, but it does not guarantee a high quality response. We encourage you to play around, find questions that work and find questions that fail. There is a small monthly budget dedicated to the OpenAI endpoints. Once that is used up each month, queries will no longer work. ## Sidebar Config Use the `Generative Config` to change LLM parameters. Use the `Retrieval Config` to change the number of chunks retrieved from our congress corpus and to apply various filters to the content before it is retrieved (e.g. filter to a specific set of congresses). Use the `Prompt Config` to try out different document formatting and prompting strategies. """ ) def write_example_queries(): with st.expander("Example Queries"): st.write( """ ``` What are the themes around artificial intelligence? ``` ``` Write a well cited 3 paragraph essay on food insecurity. ``` ``` Create a table summarizing the major climate change ideas with columns legis_id, title, idea. ``` """ ) ################## st.title(":classical_building: LegisQA :classical_building:") with st.sidebar: with st.container(border=True): write_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") 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=32, value=8, 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.multiselect( "Congress Numbers", CONGRESS_NUMBERS, default=CONGRESS_NUMBERS, key="filter_congress_nums", ) st.multiselect( "Sponsor Party", SPONSOR_PARTIES, default=SPONSOR_PARTIES, key="filter_sponsor_parties", ) 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() format_docs = DOC_FORMATTERS[SS["prompt_version"]] vs_filter = get_vectorstore_filter() query_tab, guide_tab = st.tabs(["query", "guide"]) with guide_tab: write_guide() with query_tab: write_example_queries() with st.form("my_form"): st.text_area("Enter query:", key="query") query_submitted = st.form_submit_button("Submit") if query_submitted: retriever = vectorstore.as_retriever( search_kwargs={"k": SS["n_ret_docs"], "filter": vs_filter}, ) prompt = PromptTemplate.from_template(SS["prompt_template"]) # # takes in a dict. adds context key with formatted docs # rag_chain_from_docs = ( # RunnablePassthrough.assign(context=(lambda x: format_docs(x["context"]))) # | prompt # | llm # | StrOutputParser() # ) # # takes in a query string. # # passes to retriever and passthru # # assign answer # rag_chain_with_source = RunnableParallel( # {"context": retriever, "question": RunnablePassthrough()} # ).assign(answer=rag_chain_from_docs) # takes in a dict. adds context key with formatted docs rag_chain_from_docs = ( RunnablePassthrough.assign(context=(lambda x: format_docs(x["context"]))) | prompt | llm | StrOutputParser() ) # takes in a query string. # passes to retriever and passthru # assign answer rag_chain_with_source = RunnableParallel( {"context": retriever, "question": RunnablePassthrough()} ).assign(answer=rag_chain_from_docs) print(rag_chain_with_source) with get_openai_callback() as cb: SS["out"] = rag_chain_with_source.invoke(SS["query"]) SS["cb"] = cb if "out" in SS: out_display = SS["out"]["answer"] if SS["response_escape_markdown"]: out_display = escape_markdown(out_display) if SS["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") st.warning(SS["cb"]) with st.container(border=True): doc_grps = group_docs(SS["out"]["context"]) 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: write_doc_grp(legis_id, doc_grp) # with st.expander("Debug doc format"): # st.text_area("formatted docs", value=format_docs(SS["out"]["context"]), height=600)