DEV Community

Apify for Apify

Posted on • Originally published at blog.apify.com on

How to create a custom AI chatbot with Python

In this tutorial, were going to build a custom AI chatbot. Our chatbot is going to work on top of data that will be fed to a large language model (LLM). In other words, well be developing a retrieval-augmented chatbot. The main tools well use are Streamlit and LangChain.

  • Streamlit is a tool for the quick creation of web apps. Well use it to implement the chat interface.

  • LangChain is a framework that simplifies the building of LLM apps. It mostly acts as the glue between vector databases, LLMs, and your custom code.

Well split this tutorial into 3 steps:

  1. First, well get some data that can be used as context for the LLM.

  2. Second, well use Streamlit to create the chat interface.

  3. Lastly, well connect everything together using LangChain.

The code is available at https://github.com/apify/chat-with-a-website.

Related: What is retrieval-augmented generation, and why use it for chatbots?

Obtaining the data and saving it in a vector database

First, we want to collect some data. We'll later use this as the context provided to the LLM when chatting. Our example code will use Apifys Website Content Crawler to scrape the selected website and store it in a local vector database.

First, lets create an .env file that will contain the website we want to chat with and API tokens for Apify and OpenAI:

OPENAI_API_KEY=your_api_key APIFY_API_TOKEN=your_api_key WEBSITE_URL="<https://docs.apify.com/platform>" 
Enter fullscreen mode Exit fullscreen mode

Next, lets install all the required packages:

pip install apify-client chromadb langchain openai python-dotenv streamlit tiktoken 
Enter fullscreen mode Exit fullscreen mode

Our environments all set, so lets write some Python code!

Lets create a new file called scrape.py. First, we want to import the necessary packages and load our .env file:

import os from apify_client import ApifyClient from dotenv import load_dotenv from langchain.document_loaders import ApifyDatasetLoader from langchain.document_loaders.base import Document from langchain.embeddings.openai import OpenAIEmbeddings from langchain.text_splitter import RecursiveCharacterTextSplitter from langchain.vectorstores import Chroma # Load environment variables from a .env file load_dotenv() 
Enter fullscreen mode Exit fullscreen mode

Next, well write the main function:

if __name__ == ' __main__': apify_client = ApifyClient(os.environ.get('APIFY_API_TOKEN')) website_url = os.environ.get('WEBSITE_URL') print(f'Extracting data from "{website_url}". Please wait...') actor_run_info = apify_client.actor('apify/website-content-crawler').call( run_input={'startUrls': [{'url': website_url}]} ) print('Saving data into the vector database. Please wait...') loader = ApifyDatasetLoader( dataset_id=actor_run_info['defaultDatasetId'], dataset_mapping_function=lambda item: Document( page_content=item['text'] or '', metadata={'source': item['url']} ), ) documents = loader.load() text_splitter = RecursiveCharacterTextSplitter(chunk_size=1500, chunk_overlap=100) docs = text_splitter.split_documents(documents) embedding = OpenAIEmbeddings() vectordb = Chroma.from_documents( documents=docs, embedding=embedding, persist_directory='db2', ) vectordb.persist() print('All done!') 
Enter fullscreen mode Exit fullscreen mode

We'll run the Website Content Crawler Actor on Apify to scrape the target website, then use the ApifyDatasetLoader that is integrated into LangChain to load the scraped documents.

Then, we use the RecursiveCharacterTextSplitter to chunk the documents, and finally, we use OpenAIs embeddings to convert our documents into vectors that get stored in the db directory.

Creating the chat interface

We're gonna use Streamlit to create the interface. Well base it on examples provided at https://github.com/langchain-ai/streamlit-agent.

Lets start with the imports and some settings:

import os import streamlit as st from dotenv import load_dotenv from langchain.callbacks.base import BaseCallbackHandler from langchain.chains import ConversationalRetrievalChain from langchain.chat_models import ChatOpenAI from langchain.embeddings import OpenAIEmbeddings from langchain.memory import ConversationBufferMemory from langchain.memory.chat_message_histories import StreamlitChatMessageHistory from langchain.vectorstores import Chroma load_dotenv() website_url = os.environ.get('WEBSITE_URL', 'a website') st.set_page_config(page_title=f'Chat with {website_url}') st.title('Chat with a website') 
Enter fullscreen mode Exit fullscreen mode

Next, we'll implement some helpers. The get_retriever function will create a retriever based on data we extracted in the previous step using scrape.py. The StreamHandler class will be used for streaming the responses from ChatGPT to our application.

@st.cache_resource(ttl='1h') def get_retriever(): embeddings = OpenAIEmbeddings() vectordb = Chroma(persist_directory='db', embedding_function=embeddings) retriever = vectordb.as_retriever(search_type='mmr') return retriever class StreamHandler(BaseCallbackHandler): def __init__ (self, container: st.delta_generator.DeltaGenerator, initial_text: str = ''): self.container = container self.text = initial_text def on_llm_new_token(self, token: str, **kwargs) -> None: self.text += token self.container.markdown(self.text) 
Enter fullscreen mode Exit fullscreen mode

Finally, lets add the main code. We use the ConversationalRetrievalChain utility provided by LangChain along with OpenAIs gpt-3.5-turbo. The rest of the code sets up the Streamlit chat interface.

retriever = get_retriever() msgs = StreamlitChatMessageHistory() memory = ConversationBufferMemory(memory_key='chat_history', chat_memory=msgs, return_messages=True) llm = ChatOpenAI(model_name='gpt-3.5-turbo', temperature=0, streaming=True) qa_chain = ConversationalRetrievalChain.from_llm( llm, retriever=retriever, memory=memory, verbose=False ) if st.sidebar.button('Clear message history') or len(msgs.messages) == 0: msgs.clear() msgs.add_ai_message(f'Ask me anything about {website_url}!') avatars = {'human': 'user', 'ai': 'assistant'} for msg in msgs.messages: st.chat_message(avatars[msg.type]).write(msg.content) if user_query := st.chat_input(placeholder='Ask me anything!'): st.chat_message('user').write(user_query) with st.chat_message('assistant'): stream_handler = StreamHandler(st.empty()) response = qa_chain.run(user_query, callbacks=[stream_handler]) 
Enter fullscreen mode Exit fullscreen mode

Connecting everything together

If youve followed along with this tutorial, then by now, you should have three files: .env, [scrape.py](<http://scrape.py>)and chat.py. Lets take what weve created and use it to chat with a website!

First, run python scrape.py to extract the relevant data from the target website. Note that this step may take a while since the website might be pretty big. You can check the progress at https://console.apify.com/actors/runs.

After the data extraction is done, you can start chatting with the website by running streamlit run chat.py!

How to create a custom AI chatbot with Python and Streamlit

Top comments (0)