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(