refactored api

api.py

@@ -5,35 +5,22 @@ import json
 import logging
 import time
 import shutil
+from typing import List, Optional
 
-from models import OptimizeRequest, QARequest, AutotuneRequest
 from fastapi import FastAPI, HTTPException, UploadFile, File, Form
 from fastapi.middleware.cors import CORSMiddleware
+from dotenv import load_dotenv
 
-try:
-    from ragmint.autotuner import AutoRAGTuner
-    from ragmint.qa_generator import generate_validation_qa
-    from ragmint.explainer import explain_results
-    from ragmint.leaderboard import Leaderboard
-    from ragmint.tuner import RAGMint
-except Exception as e:
-    AutoRAGTuner = None
-    generate_validation_qa = None
-    explain_results = None
-    Leaderboard = None
-    RAGMint = None
-    _import_error = e
-else:
-    _import_error = None
+from models import OptimizeRequest, QARequest, AutotuneRequest
 
-from dotenv import load_dotenv
+# Load environment
 load_dotenv()
 
 # Logging
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger("ragmint_mcp_server")
 
-# FastAPI app (exported for mounting)
+# FastAPI app
 app = FastAPI(title="Ragmint MCP Server", version="0.1.0")
 app.add_middleware(
     CORSMiddleware,
@@ -43,14 +30,30 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
-# Use repo-local data folder (not parent dirs)
+# Directories
 DEFAULT_DATA_DIR = "data/docs"
 LEADERBOARD_STORAGE = "experiments/leaderboard.jsonl"
-
-# ensure folders exist
 os.makedirs(DEFAULT_DATA_DIR, exist_ok=True)
 os.makedirs("experiments", exist_ok=True)
 
+# Try importing ragmint modules
+try:
+    from ragmint.autotuner import AutoRAGTuner
+    from ragmint.qa_generator import generate_validation_qa
+    from ragmint.explainer import explain_results
+    from ragmint.leaderboard import Leaderboard
+    from ragmint.tuner import RAGMint
+except Exception as e:
+    AutoRAGTuner = None
+    generate_validation_qa = None
+    explain_results = None
+    Leaderboard = None
+    RAGMint = None
+    _import_error = e
+else:
+    _import_error = None
+
+
 @app.get("/health")
 def health():
     return {
@@ -59,10 +62,11 @@ def health():
         "import_error": str(_import_error) if _import_error else None,
     }
 
+
 @app.post("/upload_docs")
 async def upload_docs(
-    docs_path: str = Form(...),
-    files: list[UploadFile] = File(...)
+        docs_path: str = Form(...),
+        files: List[UploadFile] = File(...)
 ):
     os.makedirs(docs_path, exist_ok=True)
     saved_files = []
@@ -74,6 +78,34 @@ async def upload_docs(
     return {"status": "ok", "uploaded_files": saved_files, "docs_path": docs_path}
 
 
+def handle_validation_choice(docs_path: str, validation_choice: Optional[str], llm_model: str) -> Optional[str]:
+    """Determine which validation QA set to use or generate one."""
+    validation_choice = (validation_choice or "").strip()
+    default_path = os.path.join(docs_path, "validation_qa.json")
+
+    if not validation_choice:
+        if os.path.exists(default_path):
+            logger.info("Using default validation QA: %s", default_path)
+            return default_path
+        return None
+
+    if validation_choice.lower() == "generate":
+        generate_validation_qa(
+            docs_path=docs_path,
+            output_path=default_path,
+            llm_model=llm_model
+        )
+        logger.info("Generated validation QA at: %s", default_path)
+        return default_path
+
+    if os.path.exists(validation_choice) or "/" in validation_choice:
+        logger.info("Using specified validation dataset: %s", validation_choice)
+        return validation_choice
+
+    logger.warning("Validation choice provided but not found: %s", validation_choice)
+    return None
+
+
 @app.post("/optimize_rag")
 def optimize_rag(req: OptimizeRequest):
     logger.info("Received optimize_rag request: %s", req.json())
@@ -89,40 +121,14 @@ def optimize_rag(req: OptimizeRequest):
             docs_path=docs_path,
             retrievers=req.retriever,
             embeddings=req.embedding_model,
-            rerankers=(req.rerankers or ["mmr"]),
+            rerankers=req.rerankers or ["mmr"],
             chunk_sizes=req.chunk_sizes,
             overlaps=req.overlaps,
             strategies=req.strategy,
         )
 
-        # validation set handling
-        validation_set = None
-        validation_choice = (req.validation_choice or "").strip()
-        default_val_path = os.path.join(docs_path, "validation_qa.json")
-
-        if not validation_choice:
-            if os.path.exists(default_val_path):
-                validation_set = default_val_path
-                logger.info("Using default validation set: %s", validation_set)
-            else:
-                logger.warning("No validation_choice provided and no default found.")
-                validation_set = None
-        elif "/" in validation_choice and not os.path.exists(validation_choice):
-            validation_set = validation_choice
-            logger.info("Using HF dataset as validation: %s", validation_set)
-        elif os.path.exists(validation_choice):
-            validation_set = validation_choice
-            logger.info("Using local validation dataset: %s", validation_set)
-        elif validation_choice.lower() == "generate":
-            gen_path = os.path.join(docs_path, "validation_qa.json")
-            generate_validation_qa(
-                docs_path=docs_path,
-                output_path=gen_path,
-                llm_model=req.llm_model if hasattr(req, "llm_model") else "gemini-2.5-flash-lite"
-            )
-            validation_set = gen_path
-            logger.info("Generated validation QA at: %s", validation_set)
-
+        validation_set = handle_validation_choice(docs_path, req.validation_choice,
+                                                  getattr(req, "llm_model", "gemini-2.5-flash-lite"))
         start_time = time.time()
         best, results = rag.optimize(
             validation_set=validation_set,
@@ -133,29 +139,23 @@ def optimize_rag(req: OptimizeRequest):
         elapsed = time.time() - start_time
         run_id = f"opt_{int(time.time())}"
 
-        try:
-            corpus_stats = {
-                "num_docs": len(rag.documents),
-                "avg_len": sum(len(d.split()) for d in rag.documents) / max(1, len(rag.documents)),
-                "corpus_size": sum(len(d) for d in rag.documents),
-            }
-        except Exception:
-            corpus_stats = None
-
-        try:
-            if Leaderboard:
-                lb = Leaderboard()
-                lb.upload(
-                    run_id=run_id,
-                    best_config=best,
-                    best_score=best.get("faithfulness", best.get("score", 0.0)),
-                    all_results=results,
-                    documents=os.listdir(docs_path),
-                    model=best.get("embedding_model", req.embedding_model),
-                    corpus_stats=corpus_stats,
-                )
-        except Exception:
-            logger.exception("Leaderboard persistence failed for optimize_rag")
+        corpus_stats = {
+            "num_docs": len(rag.documents),
+            "avg_len": sum(len(d.split()) for d in rag.documents) / max(1, len(rag.documents)),
+            "corpus_size": sum(len(d) for d in rag.documents),
+        }
+
+        if Leaderboard:
+            lb = Leaderboard()
+            lb.upload(
+                run_id=run_id,
+                best_config=best,
+                best_score=best.get("faithfulness", best.get("score", 0.0)),
+                all_results=results,
+                documents=os.listdir(docs_path),
+                model=best.get("embedding_model", req.embedding_model),
+                corpus_stats=corpus_stats,
+            )
 
         return {
             "status": "finished",
@@ -165,7 +165,6 @@ def optimize_rag(req: OptimizeRequest):
             "results": results,
             "corpus_stats": corpus_stats,
         }
-
     except Exception as exc:
         logger.exception("optimize_rag failed")
         raise HTTPException(status_code=500, detail=str(exc))
@@ -191,7 +190,6 @@ def autotune_rag(req: AutotuneRequest):
             num_pairs=int(req.num_chunk_pairs),
             step=20
         )
-
         chunk_sizes = sorted({c for c, _ in chunk_candidates})
         overlaps = sorted({o for _, o in chunk_candidates})
 
@@ -205,27 +203,8 @@ def autotune_rag(req: AutotuneRequest):
             strategies=[rec["strategy"]],
         )
 
-        validation_set = None
-        validation_choice = (req.validation_choice or "").strip()
-        default_val_path = os.path.join(docs_path, "validation_qa.jsonl")
-        if not validation_choice:
-            if os.path.exists(default_val_path):
-                validation_set = default_val_path
-            else:
-                validation_set = None
-        elif "/" in validation_choice and not os.path.exists(validation_choice):
-            validation_set = validation_choice
-        elif os.path.exists(validation_choice):
-            validation_set = validation_choice
-        elif validation_choice.lower() == "generate":
-            gen_path = os.path.join(docs_path, "validation_qa.json")
-            generate_validation_qa(
-                docs_path=docs_path,
-                output_path=gen_path,
-                llm_model=req.llm_model if hasattr(req, "llm_model") else "gemini-2.5-flash-lite",
-            )
-            validation_set = gen_path
-
+        validation_set = handle_validation_choice(docs_path, req.validation_choice,
+                                                  getattr(req, "llm_model", "gemini-2.5-flash-lite"))
         best, results = rag.optimize(
             validation_set=validation_set,
             metric=req.metric,
@@ -235,29 +214,23 @@ def autotune_rag(req: AutotuneRequest):
         elapsed = time.time() - start_time
         run_id = f"autotune_{int(time.time())}"
 
-        try:
-            corpus_stats = {
-                "num_docs": len(rag.documents),
-                "avg_len": sum(len(d.split()) for d in rag.documents) / max(1, len(rag.documents)),
-                "corpus_size": sum(len(d) for d in rag.documents),
-            }
-        except Exception:
-            corpus_stats = None
-
-        try:
-            if Leaderboard:
-                lb = Leaderboard()
-                lb.upload(
-                    run_id=run_id,
-                    best_config=best,
-                    best_score=best.get("faithfulness", best.get("score", 0.0)),
-                    all_results=results,
-                    documents=os.listdir(docs_path),
-                    model=best.get("embedding_model", rec.get("embedding_model")),
-                    corpus_stats=corpus_stats,
-                )
-        except Exception:
-            logger.exception("Leaderboard persistence failed for autotune_rag")
+        corpus_stats = {
+            "num_docs": len(rag.documents),
+            "avg_len": sum(len(d.split()) for d in rag.documents) / max(1, len(rag.documents)),
+            "corpus_size": sum(len(d) for d in rag.documents),
+        }
+
+        if Leaderboard:
+            lb = Leaderboard()
+            lb.upload(
+                run_id=run_id,
+                best_config=best,
+                best_score=best.get("faithfulness", best.get("score", 0.0)),
+                all_results=results,
+                documents=os.listdir(docs_path),
+                model=best.get("embedding_model", rec.get("embedding_model")),
+                corpus_stats=corpus_stats,
+            )
 
         return {
             "status": "finished",
@@ -276,13 +249,13 @@ def autotune_rag(req: AutotuneRequest):
 
 
 @app.post("/generate_validation_qa")
-def generate_qa(req: QARequest):
+def generate_validation_qa_endpoint(req: QARequest):
     logger.info("Received generate_validation_qa request: %s", req.json())
     if generate_validation_qa is None:
         raise HTTPException(status_code=500, detail=f"Ragmint imports failed: {_import_error}")
 
     try:
-        out_path = os.path.join("data", "docs", "validation_qa.json")
+        out_path = os.path.join(req.docs_path or DEFAULT_DATA_DIR, "validation_qa.json")
         os.makedirs(os.path.dirname(out_path), exist_ok=True)
 
         generate_validation_qa(
@@ -297,7 +270,12 @@ def generate_qa(req: QARequest):
         with open(out_path, "r", encoding="utf-8") as f:
             data = json.load(f)
 
-        return {"status": "finished", "output_path": out_path, "preview_count": len(data), "sample": data[:5]}
+        return {
+            "status": "finished",
+            "output_path": out_path,
+            "preview_count": len(data),
+            "sample": data[:5]
+        }
 
     except Exception as exc:
         logger.exception("generate_validation_qa failed")