added sub-agent for document retrieval

langchain-streamlit-demo/app.py

@@ -23,6 +23,7 @@ from streamlit_feedback import streamlit_feedback
 from defaults import default_values
 from llm_resources import (
     get_agent,
+    get_doc_agent,
     get_llm,
     get_runnable,
     get_texts_and_multiretriever,
@@ -494,17 +495,29 @@ if st.session_state.llm:
                         chat_prompt,
                         prompt,
                     )
+
                     doc_chain_tool = Tool.from_function(
                         func=lambda s: st.session_state.doc_chain.invoke(
                             s,
                             # config=get_config(callbacks),
                         ),
                         name="user-document-chat",
+                        description="this assistant returns a response based on the user's custom context. ",
+                    )
+                    doc_chain_agent = get_doc_agent(
+                        [doc_chain_tool],
+                    )
+                    doc_agent_tool = Tool.from_function(
+                        func=lambda s: doc_chain_agent.invoke(
+                            s,
+                        ),
+                        name="document-agent",
                         description="this assistant returns a response based on the user's custom context. "
                         "if the user's meaning is unclear, perhaps the answer is here. "
-                        "generally speaking, try this tool before conducting web research.",
+                        "generally speaking, try this tool before conducting web research."
+                        "it is best to send this tool a question, as it will attempt to break complex questions down into several, simpler questions.",
                     )
-                    TOOLS = [doc_chain_tool, research_assistant_tool] + default_tools
+                    TOOLS = [doc_agent_tool, research_assistant_tool] + default_tools
 
                 st.session_state.chain = get_agent(
                     TOOLS,

langchain-streamlit-demo/llm_resources.py

@@ -2,7 +2,7 @@ import uuid
 from tempfile import NamedTemporaryFile
 from typing import Tuple, List, Optional, Dict
 
-from langchain.agents import AgentExecutor
+from langchain.agents import AgentExecutor, AgentType, initialize_agent
 from langchain.agents.openai_functions_agent.agent_token_buffer_memory import (
     AgentTokenBufferMemory,
 )
@@ -19,7 +19,7 @@ from langchain.chat_models import (
 from langchain.document_loaders import PyPDFLoader
 from langchain.embeddings import AzureOpenAIEmbeddings, OpenAIEmbeddings
 from langchain.llms.base import BaseLLM
-from langchain.prompts import MessagesPlaceholder
+from langchain.prompts import MessagesPlaceholder, ChatPromptTemplate
 from langchain.retrievers import EnsembleRetriever
 from langchain.retrievers.multi_query import MultiQueryRetriever
 from langchain.retrievers.multi_vector import MultiVectorRetriever
@@ -78,6 +78,48 @@ def get_agent(
     )
 
 
+def get_doc_agent(
+    tools: list[BaseTool],
+    llm: Optional[BaseLLM] = None,
+    agent_type: AgentType = AgentType.STRUCTURED_CHAT_ZERO_SHOT_REACT_DESCRIPTION,
+):
+    if llm is None:
+        llm = ChatOpenAI(
+            model_name="gpt-4-1106-preview",
+            temperature=0.0,
+            streaming=True,
+        )
+    prompt = ChatPromptTemplate.from_messages(
+        [
+            (
+                "system",
+                """
+                You assist a chatbot with answering questions about a document.
+                If necessary, break up incoming questions into multiple parts,
+                and use the tools provided to answer smaller questions before
+                answering the larger question.
+                """,
+            ),
+            ("user", "{input}"),
+            MessagesPlaceholder(variable_name="agent_scratchpad"),
+        ],
+    )
+    agent_executor = initialize_agent(
+        tools,
+        llm,
+        agent=agent_type,
+        verbose=True,
+        memory=None,
+        handle_parsing_errors=True,
+        prompt=prompt,
+    )
+    return (
+        {"input": RunnablePassthrough()}
+        | agent_executor
+        | (lambda output: output["output"])
+    )
+
+
 def get_runnable(
     use_document_chat: bool,
     document_chat_chain_type: str,