Feat/mcp interface (#260)

* add fastapi app

* add test for http

* add mcp app

* add deps

* add tests for mcp

* `api` -> `server`

* properly load dataset

* remove unnecessary data models

* run code formatter

* configure mcp server running

* configure http server running

* fix typing

* fix typing

Author: voorhs

autointent/_dump_tools/unit_dumpers.py

@@ -12,7 +12,7 @@
 from peft import PeftModel
 from pydantic import BaseModel
 from sklearn.base import BaseEstimator
-from transformers import (  # type: ignore[attr-defined]
+from transformers import (
     AutoModelForSequenceClassification,
     AutoTokenizer,
     PreTrainedModel,

autointent/context/data_handler/_stratification.py

@@ -13,7 +13,7 @@
 from numpy import typing as npt
 from sklearn.model_selection import train_test_split
 from skmultilearn.model_selection import IterativeStratification
-from transformers import set_seed  # type: ignore[attr-defined]
+from transformers import set_seed
 
 from autointent import Dataset
 from autointent.custom_types import LabelType

autointent/modules/scoring/_bert.py

@@ -9,7 +9,7 @@
 import torch
 from datasets import Dataset, DatasetDict
 from sklearn.model_selection import train_test_split
-from transformers import (  # type: ignore[attr-defined]
+from transformers import (
     AutoModelForSequenceClassification,
     AutoTokenizer,
     DataCollatorWithPadding,

autointent/server/__init__.py

@@ -0,0 +1 @@
+

autointent/server/http.py

@@ -0,0 +1,114 @@
+"""FastAPI application for AutoIntent pipeline inference."""
+
+import logging
+from collections.abc import AsyncGenerator
+from contextlib import asynccontextmanager
+from functools import lru_cache
+from pathlib import Path
+
+from fastapi import FastAPI, HTTPException
+from pydantic import BaseModel, Field
+from pydantic_settings import BaseSettings, SettingsConfigDict
+
+from autointent import Pipeline
+from autointent.custom_types import ListOfLabelsWithOOS
+
+
+class Settings(BaseSettings):
+    """Application settings loaded from environment variables."""
+
+    model_config = SettingsConfigDict(env_file=".env", env_prefix="AUTOINTENT_")
+    path: str = Field(..., description="Path to the optimized pipeline assets")
+    host: str = "127.0.0.1"
+    port: int = 8013
+
+
+class PredictRequest(BaseModel):
+    """Request model for the predict endpoint."""
+
+    utterances: list[str] = Field(..., description="List of text utterances to classify")
+
+
+class PredictResponse(BaseModel):
+    """Response model for the predict endpoint."""
+
+    predictions: ListOfLabelsWithOOS = Field(..., description="List of predicted class labels")
+
+
+settings = Settings()
+logger = logging.getLogger(__name__)
+
+
+@lru_cache(maxsize=1)
+def load_pipeline() -> Pipeline:
+    """Load the optimized pipeline from disk."""
+    pipeline_path = Path(settings.path)
+    if not pipeline_path.exists():
+        msg = f"Pipeline path does not exist: {pipeline_path}"
+        logger.error(msg)
+        raise HTTPException(status_code=404, detail=msg)
+
+    try:
+        msg = f"Loading pipeline from: {pipeline_path}"
+        logger.info(msg)
+        pipeline = Pipeline.load(pipeline_path)
+        logger.info("Pipeline loaded successfully")
+
+    except Exception:
+        logger.exception("Failed to load pipeline")
+        raise
+    else:
+        return pipeline
+
+
+@asynccontextmanager
+async def lifespan(_: FastAPI) -> AsyncGenerator[None, None]:
+    """Load pipe."""
+    load_pipeline()
+    yield
+
+
+app = FastAPI(
+    title="AutoIntent Pipeline API",
+    description="API for serving AutoIntent predictions",
+    version="0.0.1",
+    lifespan=lifespan,
+)
+
+
+@app.get("/health")
+async def health_check() -> dict[str, str]:
+    """Health check endpoint."""
+    return {"status": "healthy"}
+
+
+@app.post("/predict")
+async def predict(request: PredictRequest) -> PredictResponse:
+    """Predict class labels for the given utterances.
+
+    Args:
+        request: Request containing list of utterances to classify
+
+    Returns:
+        Response containing predicted class labels
+    """
+    current_pipeline = load_pipeline()
+
+    if not request.utterances:
+        return PredictResponse(predictions=[])
+
+    predictions = current_pipeline.predict(request.utterances)
+
+    return PredictResponse(predictions=predictions)
+
+
+def main() -> None:
+    """Main entry point for the HTTP server."""
+    import uvicorn
+
+    uvicorn.run(
+        "autointent.server.http:app",
+        host=settings.host,
+        port=settings.port,
+        reload=False,
+    )