callbacks (still not working 100%

.pre-commit-config.yaml

@@ -40,24 +40,20 @@ repos:
     -   id: trailing-whitespace
     -   id: mixed-line-ending
     -   id: requirements-txt-fixer
--   repo: https://github.com/pre-commit/mirrors-mypy
-    rev: v1.5.1
+-   repo: https://github.com/psf/black
+    rev: 23.9.1
     hooks:
-    -   id: mypy
-        additional_dependencies:
-            - types-requests
+    -   id: black
 -   repo: https://github.com/asottile/add-trailing-comma
     rev: v3.1.0
     hooks:
     -   id: add-trailing-comma
-#-   repo: https://github.com/dannysepler/rm_unneeded_f_str
-#    rev: v0.2.0
-#    hooks:
-#    -   id: rm-unneeded-f-str
--   repo: https://github.com/psf/black
-    rev: 23.9.1
+-   repo: https://github.com/pre-commit/mirrors-mypy
+    rev: v1.5.1
     hooks:
-    -   id: black
+    -   id: mypy
+        additional_dependencies:
+            - types-requests
 -   repo: https://github.com/PyCQA/bandit
     rev: 1.7.5
     hooks:

langchain-streamlit-demo/app.py

@@ -6,6 +6,7 @@ import langsmith.utils
 import openai
 import streamlit as st
 from langchain.callbacks import StreamlitCallbackHandler
+from langchain.callbacks.base import BaseCallbackHandler
 from langchain.callbacks.tracers.langchain import LangChainTracer, wait_for_all_tracers
 from langchain.callbacks.tracers.run_collector import RunCollectorCallbackHandler
 from langchain.memory import ConversationBufferMemory, StreamlitChatMessageHistory
@@ -20,6 +21,7 @@ from defaults import default_values
 from llm_resources import (
     get_agent,
     get_llm,
+    get_runnable,
     get_texts_and_multiretriever,
 )
 from research_assistant.chain import chain as research_assistant_chain
@@ -379,15 +381,6 @@ st.session_state.llm = get_llm(
     },
 )
 
-research_assistant_tool = Tool.from_function(
-    func=lambda s: research_assistant_chain.invoke({"question": s}),
-    name="web-research-assistant",
-    description="this assistant returns a report based on web research",
-)
-
-TOOLS = [research_assistant_tool]
-st.session_state.agent = get_agent(TOOLS, STMEMORY, st.session_state.llm)
-
 # --- Chat History ---
 for msg in STMEMORY.messages:
     st.chat_message(
@@ -424,12 +417,16 @@ if st.session_state.llm:
             if st.session_state.ls_tracer:
                 callbacks.append(st.session_state.ls_tracer)
 
-            config: Dict[str, Any] = dict(
-                callbacks=callbacks,
-                tags=["Streamlit Chat"],
-            )
-            if st.session_state.provider == "Anthropic":
-                config["max_concurrency"] = 5
+            def get_config(callbacks: list[BaseCallbackHandler]) -> dict[str, Any]:
+                config: Dict[str, Any] = dict(
+                    callbacks=callbacks,
+                    tags=["Streamlit Chat"],
+                    verbose=True,
+                    return_intermediate_steps=True,
+                )
+                if st.session_state.provider == "Anthropic":
+                    config["max_concurrency"] = 5
+                return config
 
             use_document_chat = all(
                 [
@@ -439,32 +436,66 @@ if st.session_state.llm:
             )
 
             full_response: Union[str, None] = None
-
             # stream_handler = StreamHandler(message_placeholder)
             # callbacks.append(stream_handler)
+            message_placeholder = st.empty()
 
-            st_callback = StreamlitCallbackHandler(st.container())
-            callbacks.append(st_callback)
+            if st.session_state.provider in ("Azure OpenAI", "OpenAI"):
+                st_callback = StreamlitCallbackHandler(st.container())
+                callbacks.append(st_callback)
+                research_assistant_tool = Tool.from_function(
+                    func=lambda s: research_assistant_chain.invoke(
+                        {"question": s},
+                        config=get_config(callbacks),
+                    ),
+                    name="web-research-assistant",
+                    description="this assistant returns a report based on web research",
+                )
 
-            message_placeholder = st.empty()
-            # TODO use agent if openai or azure openai
-            # otherwise use runnable
-            # for agent + runnable, add to tools
-
-            # st.session_state.chain = get_runnable(
-            #     use_document_chat,
-            #     document_chat_chain_type,
-            #     st.session_state.llm,
-            #     st.session_state.retriever,
-            #     MEMORY,
-            #     chat_prompt,
-            #     prompt,
-            #     STMEMORY,
-            # )
+                TOOLS = [research_assistant_tool]
+                if use_document_chat:
+                    st.session_state.doc_chain = get_runnable(
+                        use_document_chat,
+                        document_chat_chain_type,
+                        st.session_state.llm,
+                        st.session_state.retriever,
+                        MEMORY,
+                        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. if the user's meaning is unclear, perhaps the answer is here. generally speaking, try this tool before conducting web research.",
+                    )
+                    TOOLS = [doc_chain_tool, research_assistant_tool]
+
+                st.session_state.chain = get_agent(
+                    TOOLS,
+                    STMEMORY,
+                    st.session_state.llm,
+                    callbacks,
+                )
+            else:
+                st.session_state.chain = get_runnable(
+                    use_document_chat,
+                    document_chat_chain_type,
+                    st.session_state.llm,
+                    st.session_state.retriever,
+                    MEMORY,
+                    chat_prompt,
+                    prompt,
+                )
 
             # --- LLM call ---
             try:
-                full_response = st.session_state.agent.invoke(prompt, config)
+                full_response = st.session_state.chain.invoke(
+                    prompt,
+                    config=get_config(callbacks),
+                )
 
             except (openai.AuthenticationError, anthropic.AuthenticationError):
                 st.error(

langchain-streamlit-demo/llm_resources.py

@@ -3,13 +3,13 @@ from tempfile import NamedTemporaryFile
 from typing import Tuple, List, Optional, Dict
 
 from langchain.agents import AgentExecutor
-from langchain.agents.agent_toolkits import create_retriever_tool
 from langchain.agents.openai_functions_agent.agent_token_buffer_memory import (
     AgentTokenBufferMemory,
 )
 from langchain.agents.openai_functions_agent.base import OpenAIFunctionsAgent
 from langchain.callbacks.base import BaseCallbackHandler
 from langchain.chains import LLMChain
+from langchain.chains import RetrievalQA
 from langchain.chat_models import (
     AzureChatOpenAI,
     ChatOpenAI,
@@ -18,29 +18,30 @@ 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.retrievers import EnsembleRetriever
 from langchain.retrievers.multi_query import MultiQueryRetriever
 from langchain.retrievers.multi_vector import MultiVectorRetriever
 from langchain.schema import Document, BaseRetriever
+from langchain.schema.chat_history import BaseChatMessageHistory
 from langchain.schema.runnable import RunnablePassthrough
 from langchain.storage import InMemoryStore
 from langchain.text_splitter import RecursiveCharacterTextSplitter
+from langchain.tools.base import BaseTool
 from langchain.vectorstores import FAISS
 from langchain_core.messages import SystemMessage
 
 from defaults import DEFAULT_CHUNK_SIZE, DEFAULT_CHUNK_OVERLAP, DEFAULT_RETRIEVER_K
 from qagen import get_rag_qa_gen_chain
 from summarize import get_rag_summarization_chain
-from langchain.tools.base import BaseTool
-from langchain.schema.chat_history import BaseChatMessageHistory
-from langchain.llms.base import BaseLLM
 
 
 def get_agent(
     tools: list[BaseTool],
     chat_history: BaseChatMessageHistory,
     llm: BaseLLM,
+    callbacks,
 ):
     memory_key = "agent_history"
     system_message = SystemMessage(
@@ -68,6 +69,7 @@ def get_agent(
         memory=agent_memory,
         verbose=True,
         return_intermediate_steps=True,
+        callbacks=callbacks,
     )
     return (
         {"input": RunnablePassthrough()}
@@ -84,7 +86,6 @@ def get_runnable(
     memory,
     chat_prompt,
     summarization_prompt,
-    chat_history,
 ):
     if not use_document_chat:
         return LLMChain(
@@ -105,14 +106,12 @@ def get_runnable(
             llm,
         )
     else:
-        tool = create_retriever_tool(
-            retriever,
-            "search_user_document",
-            "Retrieves custom context provided by the user for this conversation. Use this if you cannot answer immediately and confidently.",
-        )
-        tools = [tool]
-
-        return get_agent(tools, chat_history, llm)
+        return RetrievalQA.from_chain_type(
+            llm=llm,
+            chain_type=document_chat_chain_type,
+            retriever=retriever,
+            output_key="output_text",
+        ) | (lambda output: output["output_text"])
 
 
 def get_llm(