import ipaddress import uuid from pathlib import Path from typing import Any, AsyncGenerator, Literal from fastapi import APIRouter, Depends, Header, HTTPException, Request from fastapi.responses import StreamingResponse from pydantic import BaseModel, Field from mlflow.assistant import clear_project_path_cache, get_project_path from mlflow.assistant.config import AssistantConfig, PermissionsConfig, ProjectConfig from mlflow.assistant.providers import list_providers from mlflow.assistant.providers.base import ( CLINotInstalledError, NotAuthenticatedError, ProviderNotConfiguredError, clear_config_cache, ) from mlflow.assistant.skill_installer import install_skills, list_installed_skills from mlflow.assistant.types import EventType from mlflow.server.assistant.session import SessionManager, terminate_session_process def _get_provider(name: str): for p in list_providers(): if p.name == name: return p return None def _get_selected_provider(): config = AssistantConfig.load() for provider_name, provider_config in config.providers.items(): if provider_config.selected: return _get_provider(provider_name) return None _BLOCK_REMOTE_ACCESS_ERROR_MSG = ( "Assistant API is only accessible from the same host where the MLflow server is running." ) async def _require_localhost(request: Request) -> None: """ Dependency that restricts access to localhost only. Uses ipaddress library for robust loopback detection. Raises: HTTPException: If request is not from localhost """ client_host = request.client.host if request.client else None if not client_host: raise HTTPException(status_code=403, detail=_BLOCK_REMOTE_ACCESS_ERROR_MSG) try: ip = ipaddress.ip_address(client_host) except ValueError: raise HTTPException(status_code=403, detail=_BLOCK_REMOTE_ACCESS_ERROR_MSG) if not ip.is_loopback: raise HTTPException(status_code=403, detail=_BLOCK_REMOTE_ACCESS_ERROR_MSG) assistant_router = APIRouter( prefix="/ajax-api/3.0/mlflow/assistant", tags=["assistant"], dependencies=[Depends(_require_localhost)], ) class MessageRequest(BaseModel): message: str session_id: str | None = None # empty for the first message experiment_id: str | None = None context: dict[str, Any] = Field(default_factory=dict) class MessageResponse(BaseModel): session_id: str stream_url: str # Config-related models class ConfigResponse(BaseModel): providers: dict[str, Any] = Field(default_factory=dict) projects: dict[str, Any] = Field(default_factory=dict) class ConfigUpdateRequest(BaseModel): providers: dict[str, Any] | None = None projects: dict[str, Any] | None = None class SessionPatchRequest(BaseModel): status: Literal["cancelled"] class SessionPatchResponse(BaseModel): message: str class PermissionDecision(BaseModel): request_id: str # the paused tool_call's id decision: Literal["allow", "deny"] # Skills-related models class SkillsInstallRequest(BaseModel): type: Literal["global", "project", "custom"] = "global" custom_path: str | None = None # Required if type="custom" experiment_id: str | None = None # Used to get project_path for type="project" class SkillsInstallResponse(BaseModel): installed_skills: list[str] skills_directory: str @assistant_router.post("/message") async def send_message(request: MessageRequest) -> MessageResponse: """ Send a message to the assistant and get a session for streaming the response. Args: request: MessageRequest with message, context, and optional session_id Returns: MessageResponse with session_id and stream_url """ # Generate or use existing session ID session_id = request.session_id or str(uuid.uuid4()) project_path = get_project_path(request.experiment_id) if request.experiment_id else None # Create or update session session = SessionManager.load(session_id) if session is None: session = SessionManager.create( context=request.context, working_dir=Path(project_path) if project_path else None ) elif request.context: session.update_context(request.context) # Store the pending message with role session.set_pending_message(role="user", content=request.message) session.add_message(role="user", content=request.message) SessionManager.save(session_id, session) return MessageResponse( session_id=session_id, stream_url=f"/ajax-api/3.0/mlflow/assistant/stream/{session_id}", ) @assistant_router.get("/sessions/{session_id}/stream") async def stream_response(request: Request, session_id: str) -> StreamingResponse: """ Stream the assistant's response via Server-Sent Events. Args: request: The FastAPI request object session_id: The session ID returned from /message Returns: StreamingResponse with SSE events """ session = SessionManager.load(session_id) if session is None: raise HTTPException(status_code=404, detail="Session not found") # A turn is driven by either a pending user message (a new turn) or pending # tool-call decisions (resuming a turn paused at a permission prompt). Both # are consumed here so the stream is replay-safe. pending_message = session.clear_pending_message() tool_decisions = session.pending_tool_decisions session.pending_tool_decisions = {} if not pending_message and not tool_decisions: raise HTTPException(status_code=400, detail="No pending message to process") SessionManager.save(session_id, session) prompt = pending_message.content if pending_message else "" # On resume the decision rides in the context; the provider detects the # pending tool_calls in history and applies it instead of starting a turn. # A new message supersedes a pending decision: if both are present (e.g. a # resume stream never consumed the decision and the user typed again), # forwarding the stale tool_decisions would make the provider resume the # abandoned turn and silently drop the new message. Prefer the message. context = dict(session.context) if tool_decisions and not pending_message: context["tool_decisions"] = tool_decisions # Extract the MLflow server URL from the request for the assistant to use. # This assumes the assistant is accessing the same MLflow server that serves this API, # which works because the assistant endpoint is localhost-only. # TODO: Extend this to support remote/proxy scenarios where the tracking URI may differ. tracking_uri = str(request.base_url).rstrip("/") async def event_generator() -> AsyncGenerator[str, None]: nonlocal session provider = _get_selected_provider() if provider is None: from mlflow.assistant.types import Event yield Event.from_error( "No assistant provider is configured or available." ).to_sse_event() return async for event in provider.astream( prompt=prompt, tracking_uri=tracking_uri, session_id=session.provider_session_id, mlflow_session_id=session_id, cwd=session.working_dir, context=context, ): # Store provider session ID if returned (for conversation continuity). # On a paused turn this persists the history with the unanswered # tool_call so a later resume can continue from it. if event.type == EventType.DONE: session.provider_session_id = event.data.get("session_id") SessionManager.save(session_id, session) yield event.to_sse_event() return StreamingResponse( event_generator(), media_type="text/event-stream", headers={ "Cache-Control": "no-cache", "Connection": "keep-alive", "X-Accel-Buffering": "no", }, ) @assistant_router.patch("/sessions/{session_id}") async def patch_session(session_id: str, request: SessionPatchRequest) -> SessionPatchResponse: """ Update session status. Currently supports cancelling an active session, which terminates the running assistant process. Args: session_id: The session ID request: SessionPatchRequest with status to set Returns: SessionPatchResponse indicating success """ session = SessionManager.load(session_id) if session is None: raise HTTPException(status_code=404, detail="Session not found") if request.status == "cancelled": # Terminate any associated subprocess. The OpenAI-compatible provider # holds no in-process state to release (the turn ends at each prompt). # Drop any tool permissions so later stream doesn't see stale decisions. session.pending_tool_decisions = {} SessionManager.save(session_id, session) terminated = terminate_session_process(session_id) msg = "Session cancelled and process terminated" if terminated else "Session cancelled" return SessionPatchResponse(message=msg) # This branch is unreachable due to Literal type, but satisfies type checker raise HTTPException(status_code=400, detail=f"Unknown status: {request.status}") @assistant_router.post("/sessions/{session_id}/permission") async def resolve_permission(session_id: str, request: PermissionDecision) -> MessageResponse: """Deliver a tool-call permission decision and resume the paused turn on a new stream. The decision is stored on the session and consumed by the next stream, which re-enters the provider with the choice in context. Stateless across requests: any worker can serve the decision because the pending state lives in the session, not process memory. """ try: SessionManager.validate_session_id(session_id) except ValueError as e: raise HTTPException(status_code=400, detail=str(e)) session = SessionManager.load(session_id) if session is None: raise HTTPException(status_code=404, detail="Session not found") session.pending_tool_decisions = {request.request_id: request.decision} SessionManager.save(session_id, session) return MessageResponse( session_id=session_id, stream_url=f"/ajax-api/3.0/mlflow/assistant/sessions/{session_id}/stream", ) @assistant_router.get("/providers/{provider}/health") async def provider_health_check(provider: str) -> dict[str, str]: p = _get_provider(provider) if p is None: raise HTTPException(status_code=404, detail=f"Provider '{provider}' not found") try: p.check_connection() except NotImplementedError as e: # Presets that delegate verification to the frontend (e.g. the # in-server MLflow AI Gateway). Returning a clear 501 prevents the # wizard from claiming a successful probe that never ran. raise HTTPException(status_code=501, detail=str(e)) except CLINotInstalledError as e: raise HTTPException(status_code=412, detail=str(e)) except NotAuthenticatedError as e: raise HTTPException(status_code=401, detail=str(e)) return {"status": "ok"} @assistant_router.get("/config") async def get_config() -> ConfigResponse: """ Get the current assistant configuration. Returns: Current configuration including providers and projects. """ config = AssistantConfig.load() return ConfigResponse( providers={name: p.model_dump() for name, p in config.providers.items()}, projects={exp_id: p.model_dump() for exp_id, p in config.projects.items()}, ) @assistant_router.put("/config") async def update_config(request: ConfigUpdateRequest) -> ConfigResponse: """ Update the assistant configuration. Args: request: Partial configuration update. Returns: Updated configuration. """ config = AssistantConfig.load() # Update providers if request.providers: for name, provider_data in request.providers.items(): existing = config.providers.get(name) model = provider_data.get("model") or (existing.model if existing else "default") base_url = provider_data.get("base_url") api_key = provider_data.get("api_key") permissions = None if "permissions" in provider_data: perm_data = provider_data["permissions"] permissions = PermissionsConfig( allow_edit_files=perm_data.get("allow_edit_files", True), allow_read_docs=perm_data.get("allow_read_docs", True), full_access=perm_data.get("full_access", False), ) selected = provider_data.get("selected", False) if selected: config.set_provider(name, model, permissions, base_url=base_url, api_key=api_key) else: config.update_provider( name, model=model, permissions=permissions, base_url=base_url, api_key=api_key, ) # Update projects if request.projects: for exp_id, project_data in request.projects.items(): if project_data is None: # Remove project mapping config.projects.pop(exp_id, None) else: location = project_data.get("location", "") project_path = Path(location).expanduser() if not project_path or not project_path.exists(): raise HTTPException( status_code=400, detail=f"Project path does not exist: {location}", ) config.projects[exp_id] = ProjectConfig( type=project_data.get("type", "local"), location=str(project_path), ) config.save() # Clear caches so provider and project path lookups pick up new settings clear_config_cache() clear_project_path_cache() return ConfigResponse( providers={name: p.model_dump() for name, p in config.providers.items()}, projects={exp_id: p.model_dump() for exp_id, p in config.projects.items()}, ) @assistant_router.post("/skills/install") async def install_skills_endpoint(request: SkillsInstallRequest) -> SkillsInstallResponse: """ Install skills bundled with MLflow. This endpoint only handles installation. Config updates should be done via PUT /config. Args: request: SkillsInstallRequest with type, custom_path, and experiment_id. Returns: SkillsInstallResponse with installed skill names and directory. Raises: HTTPException 400: If custom type without custom_path or project type without experiment_id. """ config = AssistantConfig.load() project_path: Path | None = None if request.type == "project": if not request.experiment_id: raise HTTPException(status_code=400, detail="experiment_id required for 'project' type") project_location = config.get_project_path(request.experiment_id) if not project_location: raise HTTPException( status_code=400, detail=f"No project path configured for experiment {request.experiment_id}", ) project_path = Path(project_location) provider = _get_selected_provider() if provider is None: raise HTTPException( status_code=412, detail="No assistant provider is configured or available.", ) match request.type: case "global": destination = provider.resolve_skills_path(Path.home()) case "project": destination = provider.resolve_skills_path(project_path) case "custom": if not request.custom_path: raise HTTPException( status_code=400, detail="custom_path is required when type='custom'.", ) destination = Path(request.custom_path).expanduser() # Check if skills already exist - skip re-installation if destination.exists(): if current_skills := list_installed_skills(destination): return SkillsInstallResponse( installed_skills=current_skills, skills_directory=str(destination) ) installed = install_skills(destination) return SkillsInstallResponse(installed_skills=installed, skills_directory=str(destination)) @assistant_router.get("/providers/{provider}/models") async def list_provider_models( provider: str, base_url: str | None = None, x_api_key: str | None = Header(default=None, alias="X-API-Key"), ) -> dict[str, Any]: # api_key is read from the X-API-Key header (not a query param) so the # bearer token doesn't land in access logs, browser history, or referer # headers. Localhost-only gating mitigates remote exposure but not # local logging. api_key = x_api_key p = _get_provider(provider) if p is None: raise HTTPException( status_code=404, detail=f"Provider '{provider}' not found", ) try: models = p.list_models(base_url, api_key) return {"models": models} except NotImplementedError: raise HTTPException( status_code=404, detail=f"Model listing is not supported for provider '{provider}'", ) except CLINotInstalledError as e: raise HTTPException(status_code=412, detail=str(e)) except ProviderNotConfiguredError as e: raise HTTPException( status_code=503, detail=str(e), )