agent update

app.py

@@ -1,9 +1,13 @@
 from collections import defaultdict
 import json
-from operator import itemgetter
 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
@@ -14,6 +18,7 @@ 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
@@ -28,6 +33,8 @@ 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"]
+os.environ["TOKENIZERS_PARALLELISM"] = "false"
+
 
 SS = st.session_state
 SEED = 292764
@@ -54,34 +61,6 @@ ANTHROPIC_CHAT_MODELS = [
 ]
 CHAT_MODELS = OPENAI_CHAT_MODELS + ANTHROPIC_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: {query}""",
-    "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: {query}""",
-    "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: {query}""",
-    "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: {query}""",
-}
-
 
 def get_sponsor_url(bioguide_id: str) -> str:
     return f"https://bioguide.congress.gov/search/bio/{bioguide_id}"
@@ -92,10 +71,6 @@ def get_congress_gov_url(congress_num: int, legis_type: str, legis_num: int) ->
     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"}
@@ -156,58 +131,7 @@ def group_docs(docs) -> list[tuple[str, list[Document]]]:
     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):
+def format_docs(docs):
     """JSON grouped"""
 
     doc_grps = group_docs(docs)
@@ -216,20 +140,13 @@ def format_docs_v4(docs):
         dd = {
             "legis_id": doc_grp[0].metadata["legis_id"],
             "title": doc_grp[0].metadata["title"],
+            "sponsor": doc_grp[0].metadata["sponsor_full_name"],
             "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:
@@ -258,12 +175,6 @@ def render_doc_grp(legis_id: str, doc_grp: list[Document]):
     )
     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),
@@ -392,26 +303,31 @@ def render_sidebar():
             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",
-        )
-
 
-def render_query_tab():
+def render_query_rag_tab():
 
     render_example_queries()
 
-    with st.form("my_form"):
+    QUERY_TEMPLATE = """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", "sponsor", and "snippets" keys. If a snippet is useful in writing part of your response, then cite the "title", "legis_id", and "sponsor" in the response. If you don't know how to respond, just tell the user.
+
+---
+
+Congressional Legislation Excerpts:
+
+{context}
+
+---
+
+Query: {query}"""
+
+    prompt = ChatPromptTemplate.from_messages(
+        [
+            ("system", "You are an expert legislative analyst."),
+            ("human", QUERY_TEMPLATE),
+        ]
+    )
+
+    with st.form("query_form"):
         st.text_area("Enter query:", key="query")
         query_submitted = st.form_submit_button("Submit")
 
@@ -421,7 +337,7 @@ def render_query_tab():
         retriever = vectorstore.as_retriever(
             search_kwargs={"k": SS["n_ret_docs"], "filter": vs_filter},
         )
-        prompt = PromptTemplate.from_template(SS["prompt_template"])
+
         rag_chain = (
             RunnableParallel(
                 {
@@ -430,7 +346,7 @@ def render_query_tab():
                 }
             )
             .assign(context=(lambda x: format_docs(x["docs"])))
-            .assign(answer=prompt | llm | StrOutputParser())
+            .assign(output=prompt | llm | StrOutputParser())
         )
 
         if SS["model_name"] in OPENAI_CHAT_MODELS:
@@ -442,7 +358,7 @@ def render_query_tab():
 
     if "out" in SS:
 
-        out_display = SS["out"]["answer"]
+        out_display = SS["out"]["output"]
         if SS["response_escape_markdown"]:
             out_display = escape_markdown(out_display)
         if SS["response_add_legis_urls"]:
@@ -451,7 +367,7 @@ def render_query_tab():
             st.write("Response")
             st.info(out_display)
 
-        if SS["model_name"] in OPENAI_CHAT_MODELS:
+        if "cb" in SS:
             with st.container(border=True):
                 st.write("API Usage")
                 st.warning(SS["cb"])
@@ -468,24 +384,82 @@ def render_query_tab():
             st.write(SS["out"])
 
 
-def render_chat_tab():
+def render_query_agent_tab():
+
+    from custom_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.agents import create_react_agent
+    from langchain import hub
 
-    contextualize_q_system_prompt = """Given a chat history and the latest user question \
-which might reference context in the chat history, formulate a standalone question \
-which can be understood without the chat history. Do NOT answer the question, \
-just reformulate it if needed and otherwise return it as is."""
-    contextualize_q_prompt = ChatPromptTemplate.from_messages(
+    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},
+    )
+    legis_retrieval_tool = get_retriever_tool(
+        retriever,
+        "search_legislation",
+        "Searches and returns excerpts from congressional legislation. Always call this tool first.",
+        format_docs,
+    )
+
+    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", contextualize_q_system_prompt),
-            MessagesPlaceholder(variable_name="chat_history"),
-            ("human", "{query}"),
+            ("system", "You are a helpful assistant."),
+            ("human", "{input}"),
+            MessagesPlaceholder(variable_name="agent_scratchpad"),
         ]
     )
-    contextualize_q_chain = contextualize_q_prompt | llm | StrOutputParser()
+    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")
 
-    pass
+
+
 
 
 ##################
@@ -521,16 +495,22 @@ else:
 
 
 vectorstore = load_pinecone_vectorstore()
-format_docs = DOC_FORMATTERS[SS["prompt_version"]]
 
+query_rag_tab, query_agent_tab, chat_agent_tab, guide_tab = st.tabs([
+    "query_rag",
+    "query_agent",
+    "chat_agent",
+    "guide",
+])
 
-query_tab, chat_tab, guide_tab = st.tabs(["query", "chat", "guide"])
+with query_rag_tab:
+    render_query_rag_tab()
 
-with guide_tab:
-    render_guide()
+with query_agent_tab:
+    render_query_agent_tab()
 
-with query_tab:
-    render_query_tab()
+with chat_agent_tab:
+    render_chat_agent_tab()
 
-with chat_tab:
-    render_chat_tab()
+with guide_tab:
+    render_guide()

custom_tools.py

@@ -0,0 +1,98 @@
+"""
+TODO clean all this up
+modified from https://github.com/langchain-ai/langchain/blob/master/libs/langchain/langchain/tools/retriever.py
+"""
+
+from functools import partial
+from typing import Optional
+
+from langchain_core.callbacks.manager import Callbacks
+from langchain_core.prompts import BasePromptTemplate, PromptTemplate
+from langchain_core.pydantic_v1 import BaseModel, Field
+from langchain_core.retrievers import BaseRetriever
+from langchain.tools import Tool
+
+
+def get_retriever_tool(
+    retriever,
+    name,
+    description,
+    format_docs,
+    *,
+    document_prompt: Optional[BasePromptTemplate] = None,
+    document_separator: str = "\n\n",
+):
+
+    class RetrieverInput(BaseModel):
+        """Input to the retriever."""
+
+        query: str = Field(description="query to look up in retriever")
+
+
+    def _get_relevant_documents(
+        query: str,
+        retriever: BaseRetriever,
+        document_prompt: BasePromptTemplate,
+        document_separator: str,
+        callbacks: Callbacks = None,
+    ) -> str:
+        docs = retriever.get_relevant_documents(query, callbacks=callbacks)
+        return format_docs(docs)
+
+    async def _aget_relevant_documents(
+        query: str,
+        retriever: BaseRetriever,
+        document_prompt: BasePromptTemplate,
+        document_separator: str,
+        callbacks: Callbacks = None,
+    ) -> str:
+        docs = await retriever.aget_relevant_documents(query, callbacks=callbacks)
+        return format_docs(docs)
+
+    def create_retriever_tool(
+        retriever: BaseRetriever,
+        name: str,
+        description: str,
+        *,
+        document_prompt: Optional[BasePromptTemplate] = None,
+        document_separator: str = "\n\n",
+    ) -> Tool:
+        """Create a tool to do retrieval of documents.
+
+        Args:
+            retriever: The retriever to use for the retrieval
+            name: The name for the tool. This will be passed to the language model,
+                so should be unique and somewhat descriptive.
+            description: The description for the tool. This will be passed to the language
+                model, so should be descriptive.
+
+        Returns:
+            Tool class to pass to an agent
+        """
+        document_prompt = document_prompt or PromptTemplate.from_template("{page_content}")
+        func = partial(
+            _get_relevant_documents,
+            retriever=retriever,
+            document_prompt=document_prompt,
+            document_separator=document_separator,
+        )
+        afunc = partial(
+            _aget_relevant_documents,
+            retriever=retriever,
+            document_prompt=document_prompt,
+            document_separator=document_separator,
+        )
+        return Tool(
+            name=name,
+            description=description,
+            func=func,
+            coroutine=afunc,
+            args_schema=RetrieverInput,
+        )
+
+
+    return create_retriever_tool(
+        retriever,
+        name,
+        description,
+    )

requirements.txt

@@ -37,12 +37,14 @@ jsonpatch==1.33
 jsonpointer==2.4
 jsonschema==4.21.1
 jsonschema-specifications==2023.12.1
+langchain==0.1.13
 langchain-anthropic==0.1.1
-langchain-community==0.0.24
-langchain-core==0.1.26
+langchain-community==0.0.29
+langchain-core==0.1.36
 langchain-openai==0.0.7
 langchain-pinecone==0.0.3
-langsmith==0.1.7
+langchain-text-splitters==0.0.1
+langsmith==0.1.38
 markdown-it-py==3.0.0
 MarkupSafe==2.1.5
 marshmallow==3.20.2