feat: solution for 'Повторный экзамен: FAQ-бот — ChromaDB + один MCP-tool'
This commit is contained in:
@@ -0,0 +1,115 @@
|
||||
"""
|
||||
Chromadb client wrapper for storing and querying FAQ documents.
|
||||
"""
|
||||
|
||||
import os
|
||||
from typing import List, Dict, Any
|
||||
|
||||
import chromadb
|
||||
from chromadb.config import Settings
|
||||
from chromadb.utils import embedding_functions
|
||||
|
||||
import openai
|
||||
from dotenv import load_dotenv
|
||||
|
||||
load_dotenv()
|
||||
|
||||
# Ensure OpenAI API key is set
|
||||
OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
|
||||
if not OPENAI_API_KEY:
|
||||
raise RuntimeError("OPENAI_API_KEY not set in environment")
|
||||
|
||||
openai.api_key = OPENAI_API_KEY
|
||||
|
||||
|
||||
class ChromadbClient:
|
||||
"""
|
||||
A simple wrapper around ChromaDB for storing FAQ documents and performing similarity searches.
|
||||
"""
|
||||
|
||||
def __init__(self, collection_name: str = "faq_collection", persist_directory: str = "chromadb"):
|
||||
"""
|
||||
Initialize the ChromaDB client and collection.
|
||||
|
||||
:param collection_name: Name of the collection to use.
|
||||
:param persist_directory: Directory to persist the database.
|
||||
"""
|
||||
self.client = chromadb.Client(Settings(
|
||||
chroma_db_impl="duckdb+parquet",
|
||||
persist_directory=persist_directory,
|
||||
))
|
||||
self.collection_name = collection_name
|
||||
self.collection = self.client.get_or_create_collection(name=collection_name)
|
||||
|
||||
def _embed_text(self, text: str) -> List[float]:
|
||||
"""
|
||||
Generate embeddings for a given text using OpenAI embeddings.
|
||||
|
||||
:param text: Text to embed.
|
||||
:return: List of floats representing the embedding.
|
||||
"""
|
||||
response = openai.Embedding.create(
|
||||
input=text,
|
||||
model="text-embedding-ada-002",
|
||||
)
|
||||
return response["data"][0]["embedding"]
|
||||
|
||||
def add_documents(self, documents: List[Dict[str, Any]]) -> None:
|
||||
"""
|
||||
Add a list of documents to the collection.
|
||||
|
||||
Each document should be a dict with keys:
|
||||
- id: unique identifier
|
||||
- text: the content of the document
|
||||
- metadata: optional dict of metadata
|
||||
|
||||
:param documents: List of document dicts.
|
||||
"""
|
||||
ids = []
|
||||
embeddings = []
|
||||
metadatas = []
|
||||
texts = []
|
||||
|
||||
for doc in documents:
|
||||
doc_id = str(doc["id"])
|
||||
text = doc["text"]
|
||||
metadata = doc.get("metadata", {})
|
||||
|
||||
ids.append(doc_id)
|
||||
embeddings.append(self._embed_text(text))
|
||||
metadatas.append(metadata)
|
||||
texts.append(text)
|
||||
|
||||
self.collection.add(
|
||||
ids=ids,
|
||||
embeddings=embeddings,
|
||||
metadatas=metadatas,
|
||||
documents=texts,
|
||||
)
|
||||
|
||||
def query(self, query_text: str, top_k: int = 3) -> List[Dict[str, Any]]:
|
||||
"""
|
||||
Query the collection for the most similar documents to the query_text.
|
||||
|
||||
:param query_text: The query string.
|
||||
:param top_k: Number of top results to return.
|
||||
:return: List of dicts containing id, score, metadata, and document text.
|
||||
"""
|
||||
query_embedding = self._embed_text(query_text)
|
||||
results = self.collection.query(
|
||||
query_embeddings=[query_embedding],
|
||||
n_results=top_k,
|
||||
include=["documents", "metadatas", "distances"],
|
||||
)
|
||||
|
||||
# ChromaDB returns lists; we flatten them
|
||||
hits = []
|
||||
for i in range(len(results["ids"][0])):
|
||||
hit = {
|
||||
"id": results["ids"][0][i],
|
||||
"score": 1 - results["distances"][0][i], # convert distance to similarity
|
||||
"metadata": results["metadatas"][0][i],
|
||||
"document": results["documents"][0][i],
|
||||
}
|
||||
hits.append(hit)
|
||||
return hits
|
||||
+142
-45
@@ -1,56 +1,153 @@
|
||||
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
|
||||
"""
|
||||
Main entry point for the FAQ bot using ChromaDB and a single MCP-tool.
|
||||
"""
|
||||
|
||||
app = FastAPI(title="FAQ Bot with ChromaDB and Ollama Embeddings")
|
||||
import os
|
||||
import json
|
||||
import sys
|
||||
from typing import List, Dict, Any
|
||||
|
||||
# OpenAI LLM
|
||||
llm = ChatOpenAI(
|
||||
model=settings.openai_model,
|
||||
openai_api_key=settings.openai_api_key,
|
||||
temperature=0.0
|
||||
)
|
||||
import openai
|
||||
from dotenv import load_dotenv
|
||||
|
||||
# RetrievalQA chain
|
||||
qa_chain = RetrievalQA.from_chain_type(
|
||||
llm=llm,
|
||||
chain_type="stuff",
|
||||
retriever=vector_store.db.as_retriever()
|
||||
)
|
||||
from chromadb_client import ChromadbClient
|
||||
from mcp_tool import MCPTool
|
||||
|
||||
class AskRequest(BaseModel):
|
||||
question: str
|
||||
load_dotenv()
|
||||
|
||||
class AskResponse(BaseModel):
|
||||
answer: str
|
||||
# Ensure OpenAI API key is set
|
||||
OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
|
||||
if not OPENAI_API_KEY:
|
||||
print("Error: OPENAI_API_KEY not set in environment.", file=sys.stderr)
|
||||
sys.exit(1)
|
||||
|
||||
class AddRequest(BaseModel):
|
||||
text: str
|
||||
metadata: dict | None = None
|
||||
openai.api_key = OPENAI_API_KEY
|
||||
|
||||
@app.post("/ask", response_model=AskResponse)
|
||||
async def ask(request: AskRequest):
|
||||
# Initialize the ChromaDB client
|
||||
db_client = ChromadbClient()
|
||||
|
||||
# Load FAQ documents from a local file (JSON lines format)
|
||||
FAQ_FILE = os.getenv("FAQ_FILE", "data/faq.jsonl")
|
||||
|
||||
def load_faq_documents(file_path: str) -> List[Dict[str, Any]]:
|
||||
"""
|
||||
Endpoint to ask a question to the FAQ bot.
|
||||
"""
|
||||
try:
|
||||
answer = qa_chain.run(request.question)
|
||||
return AskResponse(answer=answer)
|
||||
except Exception as e:
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
Load FAQ documents from a JSON lines file.
|
||||
|
||||
@app.post("/add")
|
||||
async def add(request: AddRequest):
|
||||
Each line should be a JSON object with keys:
|
||||
- id: unique identifier
|
||||
- text: the content of the FAQ
|
||||
- metadata: optional dict
|
||||
"""
|
||||
Endpoint to add a new FAQ entry to the vector store.
|
||||
docs = []
|
||||
if not os.path.exists(file_path):
|
||||
print(f"FAQ file {file_path} not found. Skipping load.", file=sys.stderr)
|
||||
return docs
|
||||
|
||||
with open(file_path, "r", encoding="utf-8") as f:
|
||||
for line in f:
|
||||
try:
|
||||
doc = json.loads(line.strip())
|
||||
docs.append(doc)
|
||||
except json.JSONDecodeError:
|
||||
continue
|
||||
return docs
|
||||
|
||||
# Load and add documents to the collection if not already present
|
||||
if not db_client.collection.count():
|
||||
print("Loading FAQ documents into ChromaDB...")
|
||||
faq_docs = load_faq_documents(FAQ_FILE)
|
||||
if faq_docs:
|
||||
db_client.add_documents(faq_docs)
|
||||
print(f"Added {len(faq_docs)} documents.")
|
||||
else:
|
||||
print("No FAQ documents loaded.", file=sys.stderr)
|
||||
|
||||
# Instantiate the MCP-tool
|
||||
mcp_tool = MCPTool()
|
||||
|
||||
# Define the function schema for OpenAI function calling
|
||||
function_schema = {
|
||||
"name": mcp_tool.name,
|
||||
"description": mcp_tool.description,
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {},
|
||||
"required": [],
|
||||
},
|
||||
}
|
||||
|
||||
def ask_question(question: str) -> str:
|
||||
"""
|
||||
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))
|
||||
Ask a question to the bot. The bot will:
|
||||
1. Retrieve relevant FAQ documents from ChromaDB.
|
||||
2. Use OpenAI LLM to generate an answer, possibly invoking the MCP-tool.
|
||||
"""
|
||||
# Retrieve top 3 relevant documents
|
||||
hits = db_client.query(question, top_k=3)
|
||||
|
||||
# Build context from hits
|
||||
context = "\n\n".join([f"Document {hit['id']}:\n{hit['document']}" for hit in hits])
|
||||
|
||||
# Construct the prompt for the LLM
|
||||
messages = [
|
||||
{"role": "system", "content": "You are an FAQ assistant. Use the provided documents to answer questions."},
|
||||
{"role": "user", "content": f"Question: {question}\n\nContext:\n{context}"},
|
||||
]
|
||||
|
||||
# Call OpenAI with function calling enabled
|
||||
response = openai.ChatCompletion.create(
|
||||
model="gpt-4o-mini",
|
||||
messages=messages,
|
||||
functions=[function_schema],
|
||||
function_call="auto",
|
||||
)
|
||||
|
||||
# Parse the response
|
||||
reply = response["choices"][0]["message"]
|
||||
if reply.get("function_call"):
|
||||
# The model wants to call the MCP-tool
|
||||
func_name = reply["function_call"]["name"]
|
||||
if func_name == mcp_tool.name:
|
||||
# Execute the tool
|
||||
tool_response = mcp_tool({})
|
||||
# Send the tool response back to the model
|
||||
tool_message = {
|
||||
"role": "tool",
|
||||
"name": func_name,
|
||||
"content": json.dumps(tool_response),
|
||||
}
|
||||
# Re-send the conversation with the tool response
|
||||
messages.append(reply)
|
||||
messages.append(tool_message)
|
||||
# Get the final answer
|
||||
final_response = openai.ChatCompletion.create(
|
||||
model="gpt-4o-mini",
|
||||
messages=messages,
|
||||
)
|
||||
return final_response["choices"][0]["message"]["content"]
|
||||
else:
|
||||
return f"Unknown function call: {func_name}"
|
||||
else:
|
||||
return reply["content"]
|
||||
|
||||
def main():
|
||||
print("FAQ Bot (ChromaDB + MCP-tool). Type 'exit' to quit.")
|
||||
while True:
|
||||
try:
|
||||
user_input = input("\nYou: ").strip()
|
||||
except (EOFError, KeyboardInterrupt):
|
||||
print("\nGoodbye!")
|
||||
break
|
||||
|
||||
if user_input.lower() in {"exit", "quit"}:
|
||||
print("Goodbye!")
|
||||
break
|
||||
|
||||
if not user_input:
|
||||
continue
|
||||
|
||||
answer = ask_question(user_input)
|
||||
print(f"\nBot: {answer}")
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
+15
-60
@@ -1,69 +1,24 @@
|
||||
"""
|
||||
MCP-tool: Simple embedding generator.
|
||||
|
||||
This module provides a single function `get_embedding` that returns a vector
|
||||
representation of a given text. The implementation first tries to use the
|
||||
OpenAI embeddings API. If no API key is available or the request fails,
|
||||
a deterministic dummy embedding is returned so that the rest of the
|
||||
application can continue to work without external dependencies.
|
||||
A single MCP-tool implementation for the FAQ bot.
|
||||
"""
|
||||
|
||||
import os
|
||||
import hashlib
|
||||
from typing import List
|
||||
import datetime
|
||||
from typing import Dict, Any
|
||||
|
||||
try:
|
||||
import openai
|
||||
except ImportError:
|
||||
openai = None
|
||||
|
||||
|
||||
def _hash_embedding(text: str, dim: int = 1536) -> List[float]:
|
||||
class MCPTool:
|
||||
"""
|
||||
Create a deterministic dummy embedding from a hash of the text.
|
||||
The values are in the range [0, 1).
|
||||
Example MCP-tool that returns the current UTC datetime.
|
||||
"""
|
||||
h = hashlib.sha256(text.encode("utf-8")).digest()
|
||||
# Expand the hash to the required dimension
|
||||
values = []
|
||||
idx = 0
|
||||
while len(values) < dim:
|
||||
# Take 4 bytes at a time
|
||||
chunk = h[idx : idx + 4]
|
||||
if len(chunk) < 4:
|
||||
chunk = chunk.ljust(4, b"\0")
|
||||
val = int.from_bytes(chunk, "big") / 2**32
|
||||
values.append(val)
|
||||
idx += 4
|
||||
return values
|
||||
|
||||
name = "get_current_utc_time"
|
||||
description = "Returns the current UTC datetime in ISO 8601 format."
|
||||
|
||||
def get_embedding(text: str) -> List[float]:
|
||||
"""
|
||||
Return an embedding vector for the given text.
|
||||
def __call__(self, arguments: Dict[str, Any]) -> Dict[str, Any]:
|
||||
"""
|
||||
Execute the tool.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
text : str
|
||||
The input text to embed.
|
||||
|
||||
Returns
|
||||
-------
|
||||
List[float]
|
||||
The embedding vector.
|
||||
"""
|
||||
api_key = os.getenv("OPENAI_API_KEY")
|
||||
if api_key and openai:
|
||||
openai.api_key = api_key
|
||||
try:
|
||||
response = openai.Embedding.create(
|
||||
input=text,
|
||||
model="text-embedding-ada-002",
|
||||
)
|
||||
return response["data"][0]["embedding"]
|
||||
except Exception:
|
||||
# Fall back to dummy embedding on any error
|
||||
pass
|
||||
|
||||
# Dummy deterministic embedding
|
||||
return _hash_embedding(text)
|
||||
:param arguments: Dictionary of arguments (unused in this simple tool).
|
||||
:return: Dictionary with the result.
|
||||
"""
|
||||
now = datetime.datetime.utcnow().isoformat() + "Z"
|
||||
return {"current_time": now}
|
||||
Reference in New Issue
Block a user