Skip to content

Streamlit chatbot

StreamlitChatPack #

Bases: BaseLlamaPack

Streamlit 聊天机器人包。

Source code in llama_index/packs/streamlit_chatbot/base.py
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
class StreamlitChatPack(BaseLlamaPack):
    """Streamlit 聊天机器人包。"""

    def __init__(
        self,
        wikipedia_page: str = "Snowflake Inc.",
        run_from_main: bool = False,
        **kwargs: Any,
    ) -> None:
        """初始化参数。"""
        if not run_from_main:
            raise ValueError(
                "Please run this llama-pack directly with "
                "`streamlit run [download_dir]/streamlit_chatbot/base.py`"
            )

        self.wikipedia_page = wikipedia_page

    def get_modules(self) -> Dict[str, Any]:
        """获取模块。"""
        return {}

    def run(self, *args: Any, **kwargs: Any) -> Any:
        """运行流水线。"""
        import streamlit as st
        from streamlit_pills import pills

        st.set_page_config(
            page_title=f"Chat with {self.wikipedia_page}'s Wikipedia page, powered by LlamaIndex",
            page_icon="🦙",
            layout="centered",
            initial_sidebar_state="auto",
            menu_items=None,
        )

        if "messages" not in st.session_state:  # Initialize the chat messages history
            st.session_state["messages"] = [
                {"role": "assistant", "content": "Ask me a question about Snowflake!"}
            ]

        st.title(
            f"Chat with {self.wikipedia_page}'s Wikipedia page, powered by LlamaIndex 💬🦙"
        )
        st.info(
            "This example is powered by the **[Llama Hub Wikipedia Loader](https://llamahub.ai/l/wikipedia)**. Use any of [Llama Hub's many loaders](https://llamahub.ai/) to retrieve and chat with your data via a Streamlit app.",
            icon="ℹ️",
        )

        def add_to_message_history(role, content):
            message = {"role": role, "content": str(content)}
            st.session_state["messages"].append(
                message
            )  # Add response to message history

        @st.cache_resource
        def load_index_data():
            loader = WikipediaReader()
            docs = loader.load_data(pages=[self.wikipedia_page])
            service_context = ServiceContext.from_defaults(
                llm=OpenAI(model="gpt-3.5-turbo", temperature=0.5)
            )
            return VectorStoreIndex.from_documents(
                docs, service_context=service_context
            )

        index = load_index_data()

        selected = pills(
            "Choose a question to get started or write your own below.",
            [
                "What is Snowflake?",
                "What company did Snowflake announce they would acquire in October 2023?",
                "What company did Snowflake acquire in March 2022?",
                "When did Snowflake IPO?",
            ],
            clearable=True,
            index=None,
        )

        if "chat_engine" not in st.session_state:  # Initialize the query engine
            st.session_state["chat_engine"] = index.as_chat_engine(
                chat_mode="context", verbose=True
            )

        for message in st.session_state["messages"]:  # Display the prior chat messages
            with st.chat_message(message["role"]):
                st.write(message["content"])

        # To avoid duplicated display of answered pill questions each rerun
        if selected and selected not in st.session_state.get(
            "displayed_pill_questions", set()
        ):
            st.session_state.setdefault("displayed_pill_questions", set()).add(selected)
            with st.chat_message("user"):
                st.write(selected)
            with st.chat_message("assistant"):
                response = st.session_state["chat_engine"].stream_chat(selected)
                response_str = ""
                response_container = st.empty()
                for token in response.response_gen:
                    response_str += token
                    response_container.write(response_str)
                add_to_message_history("user", selected)
                add_to_message_history("assistant", response)

        if prompt := st.chat_input(
            "Your question"
        ):  # Prompt for user input and save to chat history
            add_to_message_history("user", prompt)

            # Display the new question immediately after it is entered
            with st.chat_message("user"):
                st.write(prompt)

            # If last message is not from assistant, generate a new response
            # if st.session_state["messages"][-1]["role"] != "assistant":
            with st.chat_message("assistant"):
                response = st.session_state["chat_engine"].stream_chat(prompt)
                response_str = ""
                response_container = st.empty()
                for token in response.response_gen:
                    response_str += token
                    response_container.write(response_str)
                # st.write(response.response)
                add_to_message_history("assistant", response.response)

            # Save the state of the generator
            st.session_state["response_gen"] = response.response_gen

get_modules #

get_modules() -> Dict[str, Any]

获取模块。

Source code in llama_index/packs/streamlit_chatbot/base.py
37
38
39
def get_modules(self) -> Dict[str, Any]:
    """获取模块。"""
    return {}

run #

run(*args: Any, **kwargs: Any) -> Any

运行流水线。

Source code in llama_index/packs/streamlit_chatbot/base.py
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
def run(self, *args: Any, **kwargs: Any) -> Any:
    """运行流水线。"""
    import streamlit as st
    from streamlit_pills import pills

    st.set_page_config(
        page_title=f"Chat with {self.wikipedia_page}'s Wikipedia page, powered by LlamaIndex",
        page_icon="🦙",
        layout="centered",
        initial_sidebar_state="auto",
        menu_items=None,
    )

    if "messages" not in st.session_state:  # Initialize the chat messages history
        st.session_state["messages"] = [
            {"role": "assistant", "content": "Ask me a question about Snowflake!"}
        ]

    st.title(
        f"Chat with {self.wikipedia_page}'s Wikipedia page, powered by LlamaIndex 💬🦙"
    )
    st.info(
        "This example is powered by the **[Llama Hub Wikipedia Loader](https://llamahub.ai/l/wikipedia)**. Use any of [Llama Hub's many loaders](https://llamahub.ai/) to retrieve and chat with your data via a Streamlit app.",
        icon="ℹ️",
    )

    def add_to_message_history(role, content):
        message = {"role": role, "content": str(content)}
        st.session_state["messages"].append(
            message
        )  # Add response to message history

    @st.cache_resource
    def load_index_data():
        loader = WikipediaReader()
        docs = loader.load_data(pages=[self.wikipedia_page])
        service_context = ServiceContext.from_defaults(
            llm=OpenAI(model="gpt-3.5-turbo", temperature=0.5)
        )
        return VectorStoreIndex.from_documents(
            docs, service_context=service_context
        )

    index = load_index_data()

    selected = pills(
        "Choose a question to get started or write your own below.",
        [
            "What is Snowflake?",
            "What company did Snowflake announce they would acquire in October 2023?",
            "What company did Snowflake acquire in March 2022?",
            "When did Snowflake IPO?",
        ],
        clearable=True,
        index=None,
    )

    if "chat_engine" not in st.session_state:  # Initialize the query engine
        st.session_state["chat_engine"] = index.as_chat_engine(
            chat_mode="context", verbose=True
        )

    for message in st.session_state["messages"]:  # Display the prior chat messages
        with st.chat_message(message["role"]):
            st.write(message["content"])

    # To avoid duplicated display of answered pill questions each rerun
    if selected and selected not in st.session_state.get(
        "displayed_pill_questions", set()
    ):
        st.session_state.setdefault("displayed_pill_questions", set()).add(selected)
        with st.chat_message("user"):
            st.write(selected)
        with st.chat_message("assistant"):
            response = st.session_state["chat_engine"].stream_chat(selected)
            response_str = ""
            response_container = st.empty()
            for token in response.response_gen:
                response_str += token
                response_container.write(response_str)
            add_to_message_history("user", selected)
            add_to_message_history("assistant", response)

    if prompt := st.chat_input(
        "Your question"
    ):  # Prompt for user input and save to chat history
        add_to_message_history("user", prompt)

        # Display the new question immediately after it is entered
        with st.chat_message("user"):
            st.write(prompt)

        # If last message is not from assistant, generate a new response
        # if st.session_state["messages"][-1]["role"] != "assistant":
        with st.chat_message("assistant"):
            response = st.session_state["chat_engine"].stream_chat(prompt)
            response_str = ""
            response_container = st.empty()
            for token in response.response_gen:
                response_str += token
                response_container.write(response_str)
            # st.write(response.response)
            add_to_message_history("assistant", response.response)

        # Save the state of the generator
        st.session_state["response_gen"] = response.response_gen