mirror of
https://github.com/hwchase17/langchain
synced 2024-10-31 15:20:26 +00:00
32 lines
1.0 KiB
Python
32 lines
1.0 KiB
Python
|
"""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],
|
||
|
|
)
|