main.py•3.69 kB
import os
import json
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from typing import List, Dict, Any, Optional
from dotenv import load_dotenv
from mcp_server.models.vector_store import FAISSVectorStore
from mcp_server.utils.document_processor import DocumentProcessor
import mcp_server.mcp_api as mcp_api
# Load environment variables
load_dotenv()
# Get configuration from environment variables
PORT = int(os.getenv("PORT", "8000"))
HOST = os.getenv("HOST", "0.0.0.0")
INDEX_FILE = os.getenv("INDEX_FILE", "./data/faiss_index.bin")
EMBEDDING_MODEL = os.getenv("EMBEDDING_MODEL", "all-MiniLM-L6-v2")
app = FastAPI(title="MCP RAG Server")
# Initialize vector store
vector_store = None
doc_processor = None
class QueryRequest(BaseModel):
query: str
top_k: int = 5
class QueryResponse(BaseModel):
documents: List[Dict[str, Any]]
metadata: Optional[Dict[str, Any]] = None
@app.on_event("startup")
async def startup_event():
global vector_store, doc_processor
# Initialize document processor
doc_processor = DocumentProcessor(model_name=EMBEDDING_MODEL)
# Load or create vector store
vector_store = FAISSVectorStore()
# Check if index exists
if os.path.exists(INDEX_FILE):
print(f"Loading existing FAISS index from {INDEX_FILE}...")
vector_store.load(INDEX_FILE)
else:
print("No existing FAISS index found. Please index documents first.")
# Configure MCP API with our components
mcp_api.configure(vector_store, doc_processor)
# Include the MCP API router
app.include_router(mcp_api.router)
@app.post("/query", response_model=QueryResponse)
async def query_documents(request: QueryRequest):
if vector_store is None:
raise HTTPException(status_code=500, detail="Vector store not initialized")
# Get query embeddings
query_embedding = doc_processor.get_embedding(request.query)
# Query vector store
results = vector_store.search(query_embedding, request.top_k)
return {"documents": results}
@app.post("/index")
async def index_documents(directory_path: str = "docs"):
global vector_store, doc_processor
if vector_store is None or doc_processor is None:
raise HTTPException(status_code=500, detail="Server components not initialized")
# Process documents
documents = doc_processor.process_documents(directory_path)
# Create embeddings and index
vector_store.index_documents(documents)
# Save index
vector_store.save(INDEX_FILE)
return {"status": "success", "message": f"Indexed {len(documents)} documents"}
def run_server():
"""Entry point for pipx to run the server"""
import uvicorn
import argparse
parser = argparse.ArgumentParser(description="Run the MCP RAG Server")
parser.add_argument("--host", default=HOST, help=f"Host address (default: {HOST})")
parser.add_argument("--port", type=int, default=PORT, help=f"Port number (default: {PORT})")
parser.add_argument("--index-file", default=INDEX_FILE, help=f"Path to FAISS index file (default: {INDEX_FILE})")
parser.add_argument("--no-reload", action="store_true", help="Disable auto-reload")
args = parser.parse_args()
# Update environment variables based on args
global HOST, PORT, INDEX_FILE
HOST = args.host
PORT = args.port
INDEX_FILE = args.index_file
print(f"Starting MCP Server on {HOST}:{PORT}")
print(f"Using index file: {INDEX_FILE}")
# Run the server
uvicorn.run("mcp_server.main:app", host=HOST, port=PORT, reload=not args.no_reload)
if __name__ == "__main__":
run_server()