feat: solution for 'Повторный экзамен: FAQ-бот — ChromaDB + один MCP-tool'
This commit is contained in:
@@ -0,0 +1,22 @@
|
||||
import os
|
||||
from pydantic import BaseSettings
|
||||
|
||||
class Settings(BaseSettings):
|
||||
# ChromaDB configuration
|
||||
chroma_db_path: str = "./chroma_db"
|
||||
chroma_collection_name: str = "faq_collection"
|
||||
|
||||
# Ollama embedding configuration
|
||||
ollama_embed_model: str = "all-MiniLM-L6-v2"
|
||||
ollama_host: str = "http://localhost"
|
||||
ollama_port: int = 11434
|
||||
|
||||
# OpenAI LLM configuration
|
||||
openai_api_key: str = ""
|
||||
openai_model: str = "gpt-3.5-turbo"
|
||||
|
||||
class Config:
|
||||
env_file = ".env"
|
||||
env_file_encoding = "utf-8"
|
||||
|
||||
settings = Settings()
|
||||
@@ -0,0 +1,14 @@
|
||||
from langchain_ollama import OllamaEmbeddings
|
||||
from src.config import settings
|
||||
|
||||
# Instantiate the Ollama embeddings once for reuse
|
||||
ollama_embeddings = OllamaEmbeddings(
|
||||
model=settings.ollama_embed_model,
|
||||
base_url=f"{settings.ollama_host}:{settings.ollama_port}"
|
||||
)
|
||||
|
||||
def get_embedding(text: str):
|
||||
"""
|
||||
Return the embedding vector for a single text string.
|
||||
"""
|
||||
return ollama_embeddings.embed_query(text)
|
||||
+42
-99
@@ -1,113 +1,56 @@
|
||||
import os
|
||||
import json
|
||||
from pathlib import Path
|
||||
from dotenv import load_dotenv
|
||||
|
||||
from langchain.embeddings import OllamaEmbeddings
|
||||
from langchain.llms import Ollama
|
||||
from langchain.vectorstores import Chroma
|
||||
from fastapi import FastAPI, HTTPException
|
||||
from pydantic import BaseModel
|
||||
from langchain_openai import ChatOpenAI
|
||||
from langchain.chains import RetrievalQA
|
||||
from langchain.schema import Document
|
||||
from src.vector_store import vector_store
|
||||
from src.config import settings
|
||||
|
||||
# Load environment variables (e.g., OLLAMA_BASE_URL)
|
||||
load_dotenv()
|
||||
app = FastAPI(title="FAQ Bot with ChromaDB and Ollama Embeddings")
|
||||
|
||||
# Configuration
|
||||
OLLAMA_MODEL = os.getenv("OLLAMA_MODEL", "llama3.1")
|
||||
OLLAMA_BASE_URL = os.getenv("OLLAMA_BASE_URL", "http://localhost:11434")
|
||||
# OpenAI LLM
|
||||
llm = ChatOpenAI(
|
||||
model=settings.openai_model,
|
||||
openai_api_key=settings.openai_api_key,
|
||||
temperature=0.0
|
||||
)
|
||||
|
||||
# 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)
|
||||
# RetrievalQA chain
|
||||
qa_chain = RetrievalQA.from_chain_type(
|
||||
llm=llm,
|
||||
chain_type="stuff",
|
||||
retriever=vector_store.db.as_retriever()
|
||||
)
|
||||
|
||||
# Initialize ChromaDB client and collection
|
||||
chroma_client = Chroma(client_kwargs={"persist_directory": "./chromadb"})
|
||||
collection_name = "faq_collection"
|
||||
class AskRequest(BaseModel):
|
||||
question: str
|
||||
|
||||
# Load or create the collection
|
||||
vectorstore = chroma_client.get_or_create_collection(name=collection_name, embedding_function=embeddings)
|
||||
class AskResponse(BaseModel):
|
||||
answer: str
|
||||
|
||||
# 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."
|
||||
}
|
||||
]
|
||||
class AddRequest(BaseModel):
|
||||
text: str
|
||||
metadata: dict | None = None
|
||||
|
||||
def index_faq_data():
|
||||
@app.post("/ask", response_model=AskResponse)
|
||||
async def ask(request: AskRequest):
|
||||
"""
|
||||
Index FAQ questions into the Chroma collection.
|
||||
Each question is stored with its answer as metadata.
|
||||
Endpoint to ask a question to the FAQ bot.
|
||||
"""
|
||||
# Check if the collection already has documents
|
||||
if vectorstore.count() > 0:
|
||||
print(f"Collection '{collection_name}' already indexed with {vectorstore.count()} documents.")
|
||||
return
|
||||
try:
|
||||
answer = qa_chain.run(request.question)
|
||||
return AskResponse(answer=answer)
|
||||
except Exception as e:
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
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():
|
||||
@app.post("/add")
|
||||
async def add(request: AddRequest):
|
||||
"""
|
||||
Create a RetrievalQA chain that uses the Chroma vector store and Ollama LLM.
|
||||
Endpoint to add a new FAQ entry to the vector store.
|
||||
"""
|
||||
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
|
||||
)
|
||||
return qa_chain
|
||||
|
||||
def main():
|
||||
# Index data if not already indexed
|
||||
index_faq_data()
|
||||
|
||||
# Create the FAQ chain
|
||||
qa_chain = create_faq_chain()
|
||||
|
||||
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
|
||||
|
||||
# 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", [])
|
||||
|
||||
print(f"\nBot: {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__":
|
||||
main()
|
||||
try:
|
||||
doc = Document(page_content=request.text, metadata=request.metadata or {})
|
||||
vector_store.add_documents([doc])
|
||||
return {"status": "added"}
|
||||
except Exception as e:
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
+25
-92
@@ -1,98 +1,31 @@
|
||||
"""
|
||||
Vector store implementation using ChromaDB.
|
||||
from langchain_community.vectorstores import Chroma
|
||||
from langchain.schema import Document
|
||||
from src.config import settings
|
||||
from src.embeddings import ollama_embeddings
|
||||
|
||||
This module creates a persistent ChromaDB collection named 'faq' and
|
||||
indexes a predefined FAQ dataset. The collection is stored in the
|
||||
directory specified by `persist_dir`.
|
||||
|
||||
The dataset is a list of dictionaries with 'question' and 'answer'
|
||||
keys. The answers are stored as documents; the questions are stored
|
||||
as metadata for easier retrieval.
|
||||
"""
|
||||
|
||||
import os
|
||||
from typing import List, Dict
|
||||
|
||||
import chromadb
|
||||
from chromadb.config import Settings
|
||||
|
||||
# Predefined FAQ dataset
|
||||
FAQ_DATA: List[Dict[str, str]] = [
|
||||
{
|
||||
"question": "What is the capital of France?",
|
||||
"answer": "Paris is the capital of France.",
|
||||
},
|
||||
{
|
||||
"question": "Who wrote '1984'?",
|
||||
"answer": "George Orwell wrote '1984'.",
|
||||
},
|
||||
{
|
||||
"question": "What is the boiling point of water?",
|
||||
"answer": "The boiling point of water is 100°C at sea level.",
|
||||
},
|
||||
]
|
||||
|
||||
|
||||
class DummyEmbedding:
|
||||
class FAQVectorStore:
|
||||
"""
|
||||
Dummy embedding function that returns a fixed vector of zeros.
|
||||
This avoids the need for an external embedding service during tests.
|
||||
Wrapper around Chroma vector store for FAQ documents.
|
||||
"""
|
||||
|
||||
def __call__(self, texts: List[str]) -> List[List[float]]:
|
||||
# Return a vector of 768 zeros for each text
|
||||
return [[0.0] * 768 for _ in texts]
|
||||
|
||||
|
||||
def get_vector_store(persist_dir: str) -> chromadb.Collection:
|
||||
"""
|
||||
Create or load a ChromaDB collection named 'faq'.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
persist_dir : str
|
||||
Directory where the ChromaDB data will be persisted.
|
||||
|
||||
Returns
|
||||
-------
|
||||
chromadb.Collection
|
||||
The loaded or newly created collection.
|
||||
"""
|
||||
# Ensure the persistence directory exists
|
||||
os.makedirs(persist_dir, exist_ok=True)
|
||||
|
||||
# Initialize Chroma client with persistence
|
||||
client = chromadb.Client(
|
||||
Settings(
|
||||
persist_directory=persist_dir,
|
||||
)
|
||||
)
|
||||
|
||||
# Check if the collection already exists
|
||||
if "faq" in client.list_collections():
|
||||
collection = client.get_collection(name="faq")
|
||||
else:
|
||||
# Create a new collection
|
||||
collection = client.create_collection(name="faq")
|
||||
|
||||
# Prepare documents and metadata
|
||||
documents = [entry["answer"] for entry in FAQ_DATA]
|
||||
metadatas = [{"question": entry["question"]} for entry in FAQ_DATA]
|
||||
ids = [f"faq_{i}" for i in range(len(FAQ_DATA))]
|
||||
|
||||
# Use dummy embeddings to embed the documents
|
||||
dummy_embedder = DummyEmbedding()
|
||||
embeddings = dummy_embedder(documents)
|
||||
|
||||
# Add documents to the collection
|
||||
collection.add(
|
||||
documents=documents,
|
||||
metadatas=metadatas,
|
||||
ids=ids,
|
||||
embeddings=embeddings,
|
||||
def __init__(self):
|
||||
self.db = Chroma(
|
||||
collection_name=settings.chroma_collection_name,
|
||||
persist_directory=settings.chroma_db_path,
|
||||
embedding_function=ollama_embeddings
|
||||
)
|
||||
|
||||
# Persist the collection
|
||||
client.persist()
|
||||
def add_documents(self, documents: list[Document]):
|
||||
"""
|
||||
Add a list of Documents to the vector store and persist.
|
||||
"""
|
||||
self.db.add_documents(documents)
|
||||
self.db.persist()
|
||||
|
||||
return collection
|
||||
def similarity_search(self, query: str, k: int = 4):
|
||||
"""
|
||||
Retrieve the top-k most similar documents to the query.
|
||||
"""
|
||||
return self.db.similarity_search(query, k=k)
|
||||
|
||||
# Singleton instance for use in the application
|
||||
vector_store = FAQVectorStore()
|
||||
Reference in New Issue
Block a user