update

app.py

@@ -1,3 +1,8 @@
+"""
+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
 import json
 import os
@@ -140,6 +145,7 @@ def format_docs(docs):
         dd = {
             "legis_id": doc_grp[0].metadata["legis_id"],
             "title": doc_grp[0].metadata["title"],
+            "introduced_date": doc_grp[0].metadata["introduced_date"],
             "sponsor": doc_grp[0].metadata["sponsor_full_name"],
             "snippets": [doc.page_content for doc in doc_grp],
         }
@@ -308,7 +314,7 @@ def render_query_rag_tab():
 
     render_example_queries()
 
-    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.
+    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", "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.
 
 ---
 
@@ -328,7 +334,7 @@ Query: {query}"""
     )
 
     with st.form("query_form"):
-        st.text_area("Enter query:", key="query")
+        st.text_area("Enter a query that can be answered with congressional legislation:", key="query")
         query_submitted = st.form_submit_button("Submit")
 
     if query_submitted:
@@ -354,6 +360,7 @@ Query: {query}"""
                 SS["out"] = rag_chain.invoke(SS["query"])
                 SS["cb"] = cb
         else:
+            SS.pop("cb", None)
             SS["out"] = rag_chain.invoke(SS["query"])
 
     if "out" in SS:
@@ -386,7 +393,7 @@ Query: {query}"""
 
 def render_query_agent_tab():
 
-    from custom_tools import get_retriever_tool
+    from retriever_tools import get_retriever_tool
 
     from langchain_community.tools import WikipediaQueryRun
     from langchain_community.utilities import WikipediaAPIWrapper
@@ -465,9 +472,8 @@ def render_chat_agent_tab():
 ##################
 
 
-st.title(
-    ":classical_building: LegisQA - Chat With Congressional Bills :classical_building:"
-)
+st.title(":classical_building: LegisQA :classical_building:")
+st.header("Chat With Congressional Bills")
 
 
 with st.sidebar:

custom_tools.py

@@ -1,98 +0,0 @@
-"""
-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,
-    )

retriever_tools.py

@@ -0,0 +1,79 @@
+"""
+modified from https://github.com/langchain-ai/langchain/blob/master/libs/langchain/langchain/tools/retriever.py
+"""
+
+from functools import partial
+from typing import Callable
+from typing import Iterable
+from typing import Optional
+
+from langchain.schema import Document
+from langchain.tools import Tool
+from langchain_core.callbacks.manager import Callbacks
+from langchain_core.pydantic_v1 import BaseModel
+from langchain_core.pydantic_v1 import Field
+from langchain_core.retrievers import BaseRetriever
+
+
+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,
+    format_docs: Callable[[Iterable[Document]], 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,
+    format_docs: Callable[[Iterable[Document]], str],
+    callbacks: Callbacks = None,
+) -> str:
+    docs = await retriever.aget_relevant_documents(query, callbacks=callbacks)
+    return format_docs(docs)
+
+
+def get_retriever_tool(
+    retriever: BaseRetriever,
+    name: str,
+    description: str,
+    format_docs: Callable[[Iterable[Document]], str],
+) -> 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.
+        format_docs: A function to turn an iterable of docs into a string.
+
+        Returns:
+            Tool class to pass to an agent
+    """
+    func = partial(
+        _get_relevant_documents,
+        retriever=retriever,
+        format_docs=format_docs,
+    )
+    afunc = partial(
+        _aget_relevant_documents,
+        retriever=retriever,
+        format_docs=format_docs,
+    )
+    return Tool(
+        name=name,
+        description=description,
+        func=func,
+        coroutine=afunc,
+        args_schema=RetrieverInput,
+    )