feat: solution for 'Повторный экзамен: FAQ-бот — ChromaDB + один MCP-tool'
This commit is contained in:
+89
-117
@@ -1,141 +1,113 @@
|
||||
import os
|
||||
import re
|
||||
import click
|
||||
import pandas as pd
|
||||
import json
|
||||
from pathlib import Path
|
||||
from datetime import datetime
|
||||
from dotenv import load_dotenv
|
||||
|
||||
from langchain_ollama import Ollama, OllamaEmbeddings
|
||||
from langchain.embeddings import OllamaEmbeddings
|
||||
from langchain.llms import Ollama
|
||||
from langchain.vectorstores import Chroma
|
||||
from langchain.chains import RetrievalQA
|
||||
from langchain.prompts import PromptTemplate
|
||||
from langchain.tools import Tool
|
||||
from langchain.schema import Document
|
||||
|
||||
# Load environment variables
|
||||
# Load environment variables (e.g., OLLAMA_BASE_URL)
|
||||
load_dotenv()
|
||||
OLLAMA_MODEL = os.getenv("OLLAMA_MODEL", "llama3")
|
||||
CHROMA_DB_PATH = os.getenv("CHROMA_DB_PATH", "./chromadb")
|
||||
FAQ_DATA_PATH = Path("data/faq.csv")
|
||||
|
||||
# Initialize embeddings and LLM
|
||||
embeddings = OllamaEmbeddings(model=OLLAMA_MODEL)
|
||||
llm = Ollama(model=OLLAMA_MODEL)
|
||||
# Configuration
|
||||
OLLAMA_MODEL = os.getenv("OLLAMA_MODEL", "llama3.1")
|
||||
OLLAMA_BASE_URL = os.getenv("OLLAMA_BASE_URL", "http://localhost:11434")
|
||||
|
||||
# Initialize Chroma client and collection
|
||||
from chromadb import Client
|
||||
client = Client(path=CHROMA_DB_PATH)
|
||||
collection = client.get_or_create_collection(name="faq")
|
||||
# Initialize embeddings and LLM using Ollama
|
||||
embeddings = OllamaEmbeddings(model=OLLAMA_MODEL, base_url=OLLAMA_BASE_URL)
|
||||
llm = Ollama(model=OLLAMA_MODEL, base_url=OLLAMA_BASE_URL)
|
||||
|
||||
# Create vector store
|
||||
vectorstore = Chroma(collection=collection, embedding=embeddings)
|
||||
# Initialize ChromaDB client and collection
|
||||
chroma_client = Chroma(client_kwargs={"persist_directory": "./chromadb"})
|
||||
collection_name = "faq_collection"
|
||||
|
||||
# Prompt template for RetrievalQA
|
||||
prompt = PromptTemplate(
|
||||
input_variables=["context", "question"],
|
||||
template=(
|
||||
"You are a helpful FAQ bot. Use the following context to answer the question.\n"
|
||||
"Context: {context}\n"
|
||||
"Question: {question}\n"
|
||||
"Answer:"
|
||||
),
|
||||
)
|
||||
# Load or create the collection
|
||||
vectorstore = chroma_client.get_or_create_collection(name=collection_name, embedding_function=embeddings)
|
||||
|
||||
# RetrievalQA chain
|
||||
retrieval_chain = RetrievalQA.from_chain_type(
|
||||
llm=llm,
|
||||
chain_type="stuff",
|
||||
retriever=vectorstore.as_retriever(),
|
||||
chain_type_kwargs={"prompt": prompt},
|
||||
)
|
||||
# Sample FAQ data (could be loaded from a file or database)
|
||||
FAQ_DATA = [
|
||||
{
|
||||
"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 offer international shipping?",
|
||||
"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 reset my password?",
|
||||
"answer": "Click on 'Forgot password' at the login page and follow the instructions."
|
||||
}
|
||||
]
|
||||
|
||||
# MCP-tool: Current Time Tool
|
||||
def get_current_time(_input: str) -> str:
|
||||
"""Return the current system time."""
|
||||
return datetime.now().strftime("%Y-%m-%d %H:%M:%S")
|
||||
|
||||
time_tool = Tool(
|
||||
name="CurrentTime",
|
||||
description="Returns the current system time. Useful when the user asks about the time or date.",
|
||||
func=get_current_time,
|
||||
)
|
||||
|
||||
def ingest_faq():
|
||||
"""Read FAQ data from CSV and upsert into Chroma collection."""
|
||||
if not FAQ_DATA_PATH.exists():
|
||||
click.echo(f"FAQ data file not found at {FAQ_DATA_PATH}")
|
||||
def index_faq_data():
|
||||
"""
|
||||
Index FAQ questions into the Chroma collection.
|
||||
Each question is stored with its answer as metadata.
|
||||
"""
|
||||
# Check if the collection already has documents
|
||||
if vectorstore.count() > 0:
|
||||
print(f"Collection '{collection_name}' already indexed with {vectorstore.count()} documents.")
|
||||
return
|
||||
df = pd.read_csv(FAQ_DATA_PATH)
|
||||
if "question" not in df.columns or "answer" not in df.columns:
|
||||
click.echo("FAQ CSV must contain 'question' and 'answer' columns.")
|
||||
return
|
||||
# Prepare documents
|
||||
docs = df["answer"].tolist()
|
||||
metadatas = [{"question": q} for q in df["question"]]
|
||||
ids = [str(i) for i in range(len(docs))]
|
||||
# Upsert into collection
|
||||
collection.upsert(
|
||||
documents=docs,
|
||||
metadatas=metadatas,
|
||||
ids=ids,
|
||||
|
||||
texts = [item["question"] for item in FAQ_DATA]
|
||||
metadatas = [{"answer": item["answer"]} for item in FAQ_DATA]
|
||||
|
||||
# Add documents to the collection
|
||||
vectorstore.add_texts(texts=texts, metadatas=metadatas)
|
||||
print(f"Indexed {len(texts)} FAQ entries into '{collection_name}'.")
|
||||
|
||||
def create_faq_chain():
|
||||
"""
|
||||
Create a RetrievalQA chain that uses the Chroma vector store and Ollama LLM.
|
||||
"""
|
||||
retriever = vectorstore.as_retriever(search_kwargs={"k": 3})
|
||||
qa_chain = RetrievalQA.from_chain_type(
|
||||
llm=llm,
|
||||
chain_type="stuff",
|
||||
retriever=retriever,
|
||||
return_source_documents=True
|
||||
)
|
||||
click.echo(f"Ingested {len(docs)} FAQ entries into Chroma collection.")
|
||||
return qa_chain
|
||||
|
||||
def is_collection_empty() -> bool:
|
||||
"""Check if the Chroma collection has any documents."""
|
||||
return len(collection.get(ids=None)["ids"]) == 0
|
||||
def main():
|
||||
# Index data if not already indexed
|
||||
index_faq_data()
|
||||
|
||||
def answer_query(question: str) -> str:
|
||||
"""Determine whether to use the time tool or the retrieval chain."""
|
||||
if re.search(r"\b(time|date)\b", question, re.I):
|
||||
return time_tool.run(question)
|
||||
else:
|
||||
return retrieval_chain.run(question)
|
||||
# Create the FAQ chain
|
||||
qa_chain = create_faq_chain()
|
||||
|
||||
# CLI implementation
|
||||
@click.group()
|
||||
def cli():
|
||||
"""FAQ Bot CLI."""
|
||||
pass
|
||||
print("\nFAQ Bot is ready! Type your question (or 'exit' to quit).")
|
||||
while True:
|
||||
user_input = input("\nYou: ").strip()
|
||||
if user_input.lower() in {"exit", "quit"}:
|
||||
print("Goodbye!")
|
||||
break
|
||||
|
||||
@cli.command()
|
||||
@click.argument("question", nargs=-1, required=True)
|
||||
@click.option("--init", is_flag=True, help="Ingest FAQ data before answering.")
|
||||
def ask(question, init):
|
||||
"""Ask a question to the FAQ bot."""
|
||||
if init or is_collection_empty():
|
||||
ingest_faq()
|
||||
query = " ".join(question)
|
||||
answer = answer_query(query)
|
||||
click.echo(f"Answer: {answer}")
|
||||
# Get answer from the chain
|
||||
result = qa_chain({"query": user_input})
|
||||
answer = result.get("result", "Sorry, I couldn't find an answer.")
|
||||
sources = result.get("source_documents", [])
|
||||
|
||||
@cli.command()
|
||||
@click.option("--init", is_flag=True, help="Ingest FAQ data before starting the server.")
|
||||
def serve(init):
|
||||
"""Start the FastAPI web server."""
|
||||
if init or is_collection_empty():
|
||||
ingest_faq()
|
||||
import uvicorn
|
||||
uvicorn.run("src.main:app", host="0.0.0.0", port=8000, reload=True)
|
||||
print(f"\nBot: {answer}")
|
||||
|
||||
# FastAPI web interface
|
||||
from fastapi import FastAPI, HTTPException
|
||||
from pydantic import BaseModel
|
||||
|
||||
app = FastAPI(title="FAQ Bot API")
|
||||
|
||||
class QuestionRequest(BaseModel):
|
||||
question: str
|
||||
|
||||
class AnswerResponse(BaseModel):
|
||||
answer: str
|
||||
|
||||
@app.post("/ask", response_model=AnswerResponse)
|
||||
async def ask_endpoint(req: QuestionRequest):
|
||||
if not req.question:
|
||||
raise HTTPException(status_code=400, detail="Question cannot be empty.")
|
||||
answer = answer_query(req.question)
|
||||
return AnswerResponse(answer=answer)
|
||||
if sources:
|
||||
print("\nSources:")
|
||||
for doc in sources:
|
||||
# Each doc is a Document with metadata containing the answer
|
||||
source_answer = doc.metadata.get("answer", "No answer metadata.")
|
||||
print(f"- {source_answer}")
|
||||
|
||||
if __name__ == "__main__":
|
||||
cli()
|
||||
main()
|
||||
Reference in New Issue
Block a user