Quick “hello world” style ChatBot based on Ollama running locally

Ollama-ChatDocs-cli.py

from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_community.document_loaders import SeleniumURLLoader
from langchain_community.embeddings.sentence_transformer import (SentenceTransformerEmbeddings)
from langchain.embeddings import OllamaEmbeddings
from langchain_community.vectorstores import Chroma
from langchain.llms import Ollama
from langchain.chains import RetrievalQA

print("Loading story and splitting it into documents…")
loader = SeleniumURLLoader(urls=["https://americanliterature.com/author/o-henry/short-story/the-gift-of-the-magi#google_vignette"])
data = loader.load()
text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=20)
docs = text_splitter.split_documents(data)

print("Loading", docs.__len__(), "documents into database…")
embedding_function = OllamaEmbeddings(base_url = "http://localhost:11434", model = "mistral")
# Quicker, but less precise: embedding_function = SentenceTransformerEmbeddings(model_name="all-MiniLM-L6-v2")
db = Chroma.from_documents(docs, embedding_function)

print("I'm ready for your questions. Enter an empty question to leave.")
ollama = Ollama(base_url = "http://localhost:11434", model = "mistral")
while query := input(">>> ").strip():
    qachain = RetrievalQA.from_chain_type(llm = ollama, chain_type = "stuff", retriever = db.as_retriever())
    result = qachain({"query": query})
    print(result["result"])