feat: solution for 'Повторный экзамен: FAQ-бот — ChromaDB + один MCP-tool'
This commit is contained in:
+123
-104
@@ -1,120 +1,139 @@
|
||||
import os
|
||||
from typing import List
|
||||
from typing import List, Dict
|
||||
|
||||
from langchain.schema import Document
|
||||
from langchain.vectorstores import Chroma
|
||||
from langchain.chains import RetrievalQA
|
||||
from langchain_ollama import OllamaEmbeddings, Ollama
|
||||
import chromadb
|
||||
import openai
|
||||
from qdrant_client import QdrantClient
|
||||
from qdrant_client.http import models as qdrant_models
|
||||
from dotenv import load_dotenv
|
||||
|
||||
def load_faq_data() -> List[Document]:
|
||||
load_dotenv()
|
||||
|
||||
# Configuration
|
||||
QDRANT_HOST = os.getenv("QDRANT_HOST", "localhost")
|
||||
QDRANT_PORT = int(os.getenv("QDRANT_PORT", "6333"))
|
||||
QDRANT_API_KEY = os.getenv("QDRANT_API_KEY", None)
|
||||
OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
|
||||
|
||||
if not OPENAI_API_KEY:
|
||||
raise RuntimeError("OPENAI_API_KEY environment variable not set")
|
||||
|
||||
openai.api_key = OPENAI_API_KEY
|
||||
|
||||
# Collection name
|
||||
COLLECTION_NAME = "faq_collection"
|
||||
|
||||
# Embedding dimension for text-embedding-ada-002
|
||||
EMBEDDING_DIM = 1536
|
||||
|
||||
# Sample FAQ data
|
||||
FAQ_DATA = [
|
||||
{
|
||||
"question": "What is the return policy?",
|
||||
"answer": "You can return any item within 30 days of purchase."
|
||||
},
|
||||
{
|
||||
"question": "How do I track my order?",
|
||||
"answer": "Use the tracking link sent to your email after shipping."
|
||||
},
|
||||
{
|
||||
"question": "Do you offer international shipping?",
|
||||
"answer": "Yes, we ship to most countries worldwide."
|
||||
},
|
||||
{
|
||||
"question": "What payment methods are accepted?",
|
||||
"answer": "We accept credit cards, PayPal, and bank transfers."
|
||||
},
|
||||
{
|
||||
"question": "How can I contact customer support?",
|
||||
"answer": "Email us at support@example.com or call 1-800-123-4567."
|
||||
}
|
||||
]
|
||||
|
||||
def get_embedding(text: str) -> List[float]:
|
||||
"""
|
||||
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.
|
||||
Generate an embedding for the given text using OpenAI's embedding model.
|
||||
"""
|
||||
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."
|
||||
},
|
||||
]
|
||||
response = openai.Embedding.create(
|
||||
input=text,
|
||||
model="text-embedding-ada-002"
|
||||
)
|
||||
return response["data"][0]["embedding"]
|
||||
|
||||
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"]}
|
||||
def create_or_recreate_collection(client: QdrantClient) -> None:
|
||||
"""
|
||||
Create a new collection or recreate it if it already exists.
|
||||
"""
|
||||
client.recreate_collection(
|
||||
collection_name=COLLECTION_NAME,
|
||||
vectors_config=qdrant_models.VectorParams(
|
||||
size=EMBEDDING_DIM,
|
||||
distance=qdrant_models.Distance.COSINE
|
||||
)
|
||||
documents.append(doc)
|
||||
return documents
|
||||
|
||||
def create_vectorstore(embeddings, persist_directory: str = "chroma_db") -> Chroma:
|
||||
"""
|
||||
Create or load a Chroma vector store with the given embeddings function.
|
||||
"""
|
||||
# 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).")
|
||||
def ingest_faqs(client: QdrantClient, faqs: List[Dict[str, str]]) -> None:
|
||||
"""
|
||||
Ingest FAQ data into Qdrant.
|
||||
"""
|
||||
points = []
|
||||
for idx, faq in enumerate(faqs):
|
||||
vector = get_embedding(faq["question"])
|
||||
point = qdrant_models.PointStruct(
|
||||
id=idx,
|
||||
vector=vector,
|
||||
payload={
|
||||
"question": faq["question"],
|
||||
"answer": faq["answer"]
|
||||
}
|
||||
)
|
||||
points.append(point)
|
||||
|
||||
# Upsert points in batches
|
||||
batch_size = 100
|
||||
for i in range(0, len(points), batch_size):
|
||||
batch = points[i:i+batch_size]
|
||||
client.upsert(
|
||||
collection_name=COLLECTION_NAME,
|
||||
points=batch
|
||||
)
|
||||
|
||||
def query_faq(client: QdrantClient, question: str, top_k: int = 1) -> str:
|
||||
"""
|
||||
Query the FAQ collection for the most relevant answer.
|
||||
"""
|
||||
query_vector = get_embedding(question)
|
||||
search_result = client.search(
|
||||
collection_name=COLLECTION_NAME,
|
||||
query_vector=query_vector,
|
||||
limit=top_k,
|
||||
with_payload=True
|
||||
)
|
||||
if not search_result:
|
||||
return "Sorry, I couldn't find an answer to your question."
|
||||
# Return the answer from the top result
|
||||
return search_result[0].payload.get("answer", "Answer not found.")
|
||||
|
||||
def main() -> None:
|
||||
client = QdrantClient(
|
||||
host=QDRANT_HOST,
|
||||
port=QDRANT_PORT,
|
||||
api_key=QDRANT_API_KEY
|
||||
)
|
||||
|
||||
# Ingest FAQs (only if collection is empty or you want to refresh)
|
||||
print("Ingesting FAQ data into Qdrant...")
|
||||
create_or_recreate_collection(client)
|
||||
ingest_faqs(client, FAQ_DATA)
|
||||
print("Ingestion complete.")
|
||||
|
||||
print("\nFAQ Bot is ready. Type your question (or 'exit' to quit).")
|
||||
while True:
|
||||
user_input = input("\nYou: ").strip()
|
||||
user_input = input("\nYour question: ").strip()
|
||||
if user_input.lower() in {"exit", "quit"}:
|
||||
print("Goodbye!")
|
||||
break
|
||||
if not user_input:
|
||||
print("Please enter a question.")
|
||||
continue
|
||||
|
||||
# Retrieve answer
|
||||
try:
|
||||
result = qa_chain.run(user_input)
|
||||
print(f"Bot: {result}")
|
||||
except Exception as e:
|
||||
print(f"Error: {e}")
|
||||
answer = query_faq(client, user_input)
|
||||
print(f"Answer: {answer}")
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user