153 lines
4.4 KiB
Python
153 lines
4.4 KiB
Python
"""
|
|
Main entry point for the FAQ bot using ChromaDB and a single MCP-tool.
|
|
"""
|
|
|
|
import os
|
|
import json
|
|
import sys
|
|
from typing import List, Dict, Any
|
|
|
|
import openai
|
|
from dotenv import load_dotenv
|
|
|
|
from chromadb_client import ChromadbClient
|
|
from mcp_tool import MCPTool
|
|
|
|
load_dotenv()
|
|
|
|
# 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)
|
|
|
|
openai.api_key = OPENAI_API_KEY
|
|
|
|
# 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]]:
|
|
"""
|
|
Load FAQ documents from a JSON lines file.
|
|
|
|
Each line should be a JSON object with keys:
|
|
- id: unique identifier
|
|
- text: the content of the FAQ
|
|
- metadata: optional dict
|
|
"""
|
|
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:
|
|
"""
|
|
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() |