StreamlitCallbackHandler (#6315)

A new implementation of `StreamlitCallbackHandler`. It formats Agent
thoughts into Streamlit expanders.

You can see the handler in action here:
https://langchain-mrkl.streamlit.app/

Per a discussion with Harrison, we'll be adding a
`StreamlitCallbackHandler` implementation to an upcoming
[Streamlit](https://github.com/streamlit/streamlit) release as well, and
will be updating it as we add new LLM- and LangChain-specific features
to Streamlit.

The idea with this PR is that the LangChain `StreamlitCallbackHandler`
will "auto-update" in a way that keeps it forward- (and backward-)
compatible with Streamlit. If the user has an older Streamlit version
installed, the LangChain `StreamlitCallbackHandler` will be used; if
they have a newer Streamlit version that has an updated
`StreamlitCallbackHandler`, that implementation will be used instead.

(I'm opening this as a draft to get the conversation going and make sure
we're on the same page. We're really excited to land this into
LangChain!)

#### Who can review?

@agola11, @hwchase17
This commit is contained in:
Tim Conkling
2023-06-22 16:14:28 -04:00
committed by GitHub
parent 74ac6fb6b9
commit c28990d871
9 changed files with 939 additions and 638 deletions

View File

@@ -0,0 +1,31 @@
"""Integration tests for the StreamlitCallbackHandler module."""
import pytest
from langchain.agents import AgentType, initialize_agent, load_tools
# Import the internal StreamlitCallbackHandler from its module - and not from
# the `langchain.callbacks.streamlit` package - so that we don't end up using
# Streamlit's externally-provided callback handler.
from langchain.callbacks.streamlit.streamlit_callback_handler import (
StreamlitCallbackHandler,
)
from langchain.llms import OpenAI
@pytest.mark.requires("streamlit")
def test_streamlit_callback_agent() -> None:
import streamlit as st
streamlit_callback = StreamlitCallbackHandler(st.container())
llm = OpenAI(temperature=0)
tools = load_tools(["serpapi", "llm-math"], llm=llm)
agent = initialize_agent(
tools, llm, agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION, verbose=True
)
agent.run(
"Who is Olivia Wilde's boyfriend? "
"What is his current age raised to the 0.23 power?",
callbacks=[streamlit_callback],
)

View File

@@ -0,0 +1,86 @@
import builtins
import unittest
from typing import Any
from unittest import mock
from unittest.mock import MagicMock
from langchain.callbacks.streamlit import StreamlitCallbackHandler
class TestImport(unittest.TestCase):
"""Test the StreamlitCallbackHandler 'auto-updating' API"""
def setUp(self) -> None:
self.builtins_import = builtins.__import__
def tearDown(self) -> None:
builtins.__import__ = self.builtins_import
@mock.patch("langchain.callbacks.streamlit._InternalStreamlitCallbackHandler")
def test_create_internal_handler(self, mock_internal_handler: Any) -> None:
"""If we're using a Streamlit that does not expose its own
StreamlitCallbackHandler, use our own implementation.
"""
def external_import_error(
name: str, globals: Any, locals: Any, fromlist: Any, level: int
) -> Any:
if name == "streamlit.external.langchain":
raise ImportError
return self.builtins_import(name, globals, locals, fromlist, level)
builtins.__import__ = external_import_error # type: ignore[assignment]
parent_container = MagicMock()
thought_labeler = MagicMock()
StreamlitCallbackHandler(
parent_container,
max_thought_containers=1,
expand_new_thoughts=True,
collapse_completed_thoughts=False,
thought_labeler=thought_labeler,
)
# Our internal handler should be created
mock_internal_handler.assert_called_once_with(
parent_container,
max_thought_containers=1,
expand_new_thoughts=True,
collapse_completed_thoughts=False,
thought_labeler=thought_labeler,
)
def test_create_external_handler(self) -> None:
"""If we're using a Streamlit that *does* expose its own callback handler,
delegate to that implementation.
"""
mock_streamlit_module = MagicMock()
def external_import_success(
name: str, globals: Any, locals: Any, fromlist: Any, level: int
) -> Any:
if name == "streamlit.external.langchain":
return mock_streamlit_module
return self.builtins_import(name, globals, locals, fromlist, level)
builtins.__import__ = external_import_success # type: ignore[assignment]
parent_container = MagicMock()
thought_labeler = MagicMock()
StreamlitCallbackHandler(
parent_container,
max_thought_containers=1,
expand_new_thoughts=True,
collapse_completed_thoughts=False,
thought_labeler=thought_labeler,
)
# Streamlit's handler should be created
mock_streamlit_module.StreamlitCallbackHandler.assert_called_once_with(
parent_container,
max_thought_containers=1,
expand_new_thoughts=True,
collapse_completed_thoughts=False,
thought_labeler=thought_labeler,
)