feat: solution for 'Повторный экзамен: FAQ-бот — ChromaDB + один MCP-tool'
This commit is contained in:
+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()
|
||||
Reference in New Issue
Block a user