feat: solution for 'Повторный экзамен: FAQ-бот — ChromaDB + один MCP-tool'
This commit is contained in:
+103
-61
@@ -1,78 +1,120 @@
|
||||
import os
|
||||
import sys
|
||||
from pathlib import Path
|
||||
from typing import List
|
||||
|
||||
from dotenv import load_dotenv
|
||||
from langchain_community.document_loaders import TextLoader
|
||||
from langchain_community.embeddings import OllamaEmbeddings
|
||||
from langchain_community.vectorstores.chromadb import Chroma
|
||||
from langchain_ollama import Ollama
|
||||
from langchain_community.tools.mcp_tool import MCPTool
|
||||
from langchain.schema import Document
|
||||
from langchain.vectorstores import Chroma
|
||||
from langchain.chains import RetrievalQA
|
||||
from langchain_ollama import OllamaEmbeddings, Ollama
|
||||
import chromadb
|
||||
|
||||
# Load environment variables (if any)
|
||||
load_dotenv()
|
||||
|
||||
# Directory containing FAQ documents (plain text files)
|
||||
DATA_DIR = Path("data")
|
||||
# Directory where ChromaDB will persist its data
|
||||
CHROMA_DIR = Path("chroma_db")
|
||||
|
||||
# Ensure directories exist
|
||||
DATA_DIR.mkdir(parents=True, exist_ok=True)
|
||||
CHROMA_DIR.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
# Load all text files from the data directory
|
||||
documents = []
|
||||
for txt_file in DATA_DIR.glob("*.txt"):
|
||||
loader = TextLoader(str(txt_file))
|
||||
documents.extend(loader.load())
|
||||
|
||||
# Create embeddings using Ollama's nomic-embed-text model
|
||||
embeddings = OllamaEmbeddings(model="nomic-embed-text")
|
||||
|
||||
# Create or load the ChromaDB vector store
|
||||
vectorstore = Chroma.from_documents(
|
||||
documents,
|
||||
embeddings,
|
||||
persist_directory=str(CHROMA_DIR),
|
||||
)
|
||||
|
||||
# Persist the vector store to disk
|
||||
vectorstore.persist()
|
||||
|
||||
# Initialize the Ollama LLM for generation (e.g., llama3)
|
||||
llm = Ollama(model="llama3")
|
||||
|
||||
# Instantiate the MCPTool with the LLM and vector store
|
||||
mcp_tool = MCPTool(llm=llm, vectorstore=vectorstore)
|
||||
|
||||
def answer_question(question: str) -> str:
|
||||
def load_faq_data() -> List[Document]:
|
||||
"""
|
||||
Answer a question using the MCPTool, which internally retrieves relevant
|
||||
documents from the ChromaDB vector store and generates a response with
|
||||
the Ollama LLM.
|
||||
Load FAQ data. In a real application this could read from a file or database.
|
||||
Here we use a hard-coded list for demonstration purposes.
|
||||
"""
|
||||
return mcp_tool.run(question)
|
||||
faq_pairs = [
|
||||
{
|
||||
"question": "What is the return policy?",
|
||||
"answer": "You can return any item within 30 days of purchase with a receipt."
|
||||
},
|
||||
{
|
||||
"question": "How do I track my order?",
|
||||
"answer": "After placing an order, you will receive a tracking number via email."
|
||||
},
|
||||
{
|
||||
"question": "Do you ship internationally?",
|
||||
"answer": "Yes, we ship to most countries worldwide. Shipping fees apply."
|
||||
},
|
||||
{
|
||||
"question": "What payment methods are accepted?",
|
||||
"answer": "We accept credit cards, debit cards, and PayPal."
|
||||
},
|
||||
{
|
||||
"question": "How can I contact customer support?",
|
||||
"answer": "You can reach us at support@example.com or call 1-800-123-4567."
|
||||
},
|
||||
]
|
||||
|
||||
def main() -> None:
|
||||
documents = []
|
||||
for pair in faq_pairs:
|
||||
# Store the answer as the document content and the question as metadata
|
||||
doc = Document(
|
||||
page_content=pair["answer"],
|
||||
metadata={"source": pair["question"]}
|
||||
)
|
||||
documents.append(doc)
|
||||
return documents
|
||||
|
||||
def create_vectorstore(embeddings, persist_directory: str = "chroma_db") -> Chroma:
|
||||
"""
|
||||
Simple command-line interface for the FAQ bot.
|
||||
Create or load a Chroma vector store with the given embeddings function.
|
||||
"""
|
||||
print("FAQ Bot powered by ChromaDB and Ollama.")
|
||||
print("Type 'exit' to quit.")
|
||||
# Ensure the persistence directory exists
|
||||
os.makedirs(persist_directory, exist_ok=True)
|
||||
|
||||
# Create a persistent Chroma client
|
||||
client = chromadb.PersistentClient(path=persist_directory)
|
||||
|
||||
# Create or get the collection named "faq"
|
||||
collection = client.get_or_create_collection(name="faq")
|
||||
|
||||
# Wrap the collection in LangChain's Chroma wrapper
|
||||
vectorstore = Chroma(
|
||||
client=client,
|
||||
collection_name="faq",
|
||||
embedding_function=embeddings
|
||||
)
|
||||
return vectorstore
|
||||
|
||||
def main():
|
||||
# 1. Set up embeddings using Ollama's "nomic-embed-text" model
|
||||
embeddings = OllamaEmbeddings(model="nomic-embed-text")
|
||||
|
||||
# 2. Load FAQ data
|
||||
documents = load_faq_data()
|
||||
|
||||
# 3. Create or load the vector store
|
||||
vectorstore = create_vectorstore(embeddings)
|
||||
|
||||
# 4. Add documents to the vector store if not already present
|
||||
# We check if the collection is empty by attempting a simple query
|
||||
try:
|
||||
# Try retrieving a dummy query; if it returns nothing, we add documents
|
||||
dummy_query = "dummy"
|
||||
results = vectorstore.similarity_search(dummy_query, k=1)
|
||||
if not results:
|
||||
vectorstore.add_documents(documents)
|
||||
except Exception:
|
||||
# If any error occurs (e.g., collection not found), add documents
|
||||
vectorstore.add_documents(documents)
|
||||
|
||||
# 5. Set up the LLM for generation (any Ollama model suitable for text generation)
|
||||
llm = Ollama(model="llama3") # You can replace "llama3" with another model if desired
|
||||
|
||||
# 6. Build the RetrievalQA chain
|
||||
qa_chain = RetrievalQA.from_chain_type(
|
||||
llm=llm,
|
||||
chain_type="stuff",
|
||||
retriever=vectorstore.as_retriever()
|
||||
)
|
||||
|
||||
# 7. Interactive loop
|
||||
print("FAQ Bot is ready. Type your question (or 'exit' to quit).")
|
||||
while True:
|
||||
try:
|
||||
user_input = input("\nYour question: ").strip()
|
||||
except (KeyboardInterrupt, EOFError):
|
||||
print("\nExiting.")
|
||||
break
|
||||
user_input = input("\nYou: ").strip()
|
||||
if user_input.lower() in {"exit", "quit"}:
|
||||
print("Goodbye!")
|
||||
break
|
||||
if not user_input:
|
||||
print("Please enter a question.")
|
||||
continue
|
||||
answer = answer_question(user_input)
|
||||
print(f"\nAnswer: {answer}")
|
||||
|
||||
# Retrieve answer
|
||||
try:
|
||||
result = qa_chain.run(user_input)
|
||||
print(f"Bot: {result}")
|
||||
except Exception as e:
|
||||
print(f"Error: {e}")
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user