""" Docker Model Runner - Anthropic API Compatible Full compatibility with Anthropic Messages API format Optimized for: 2 vCPU, 16GB RAM """ from fastapi import FastAPI, HTTPException, Header, Request from fastapi.responses import StreamingResponse from pydantic import BaseModel, Field from typing import Optional, List, Union, Literal, Any, Dict import torch from transformers import AutoTokenizer, AutoModelForCausalLM import os from datetime import datetime from contextlib import asynccontextmanager import uuid import time import json import asyncio # CPU-optimized lightweight models GENERATOR_MODEL = os.getenv("GENERATOR_MODEL", "distilgpt2") MODEL_DISPLAY_NAME = os.getenv("MODEL_NAME", "MiniMax-M2") # Set CPU threading torch.set_num_threads(2) # Global model cache models = {} def load_models(): """Pre-load models for faster inference""" global models print("Loading models for CPU inference...") models["tokenizer"] = AutoTokenizer.from_pretrained(GENERATOR_MODEL) models["model"] = AutoModelForCausalLM.from_pretrained(GENERATOR_MODEL) models["model"].eval() if models["tokenizer"].pad_token is None: models["tokenizer"].pad_token = models["tokenizer"].eos_token print("✅ All models loaded successfully!") @asynccontextmanager async def lifespan(app: FastAPI): load_models() yield models.clear() app = FastAPI( title="Docker Model Runner", description="Anthropic API Compatible Endpoint", version="1.0.0", lifespan=lifespan ) # ============== Anthropic API Models ============== class TextBlock(BaseModel): type: Literal["text"] = "text" text: str class ThinkingBlock(BaseModel): type: Literal["thinking"] = "thinking" thinking: str class ToolUseBlock(BaseModel): type: Literal["tool_use"] = "tool_use" id: str name: str input: Dict[str, Any] class ToolResultContent(BaseModel): type: Literal["tool_result"] = "tool_result" tool_use_id: str content: Union[str, List[TextBlock]] is_error: Optional[bool] = False class ImageSource(BaseModel): type: Literal["base64", "url"] media_type: Optional[str] = None data: Optional[str] = None url: Optional[str] = None class ImageBlock(BaseModel): type: Literal["image"] = "image" source: ImageSource ContentBlock = Union[TextBlock, ThinkingBlock, ToolUseBlock, ToolResultContent, ImageBlock, str] class MessageParam(BaseModel): role: Literal["user", "assistant"] content: Union[str, List[ContentBlock]] class ToolInputSchema(BaseModel): type: str = "object" properties: Optional[Dict[str, Any]] = None required: Optional[List[str]] = None class Tool(BaseModel): name: str description: str input_schema: ToolInputSchema class ToolChoice(BaseModel): type: Literal["auto", "any", "tool"] = "auto" name: Optional[str] = None class ThinkingConfig(BaseModel): type: Literal["enabled", "disabled"] = "disabled" budget_tokens: Optional[int] = None class Metadata(BaseModel): user_id: Optional[str] = None class AnthropicRequest(BaseModel): model: str = "MiniMax-M2" messages: List[MessageParam] max_tokens: int = 1024 temperature: Optional[float] = Field(default=1.0, gt=0.0, le=1.0) top_p: Optional[float] = Field(default=1.0, gt=0.0, le=1.0) top_k: Optional[int] = None # Ignored stop_sequences: Optional[List[str]] = None # Ignored stream: Optional[bool] = False system: Optional[Union[str, List[TextBlock]]] = None tools: Optional[List[Tool]] = None tool_choice: Optional[ToolChoice] = None metadata: Optional[Metadata] = None thinking: Optional[ThinkingConfig] = None service_tier: Optional[str] = None # Ignored class Usage(BaseModel): input_tokens: int output_tokens: int cache_creation_input_tokens: Optional[int] = 0 cache_read_input_tokens: Optional[int] = 0 class AnthropicResponse(BaseModel): id: str type: Literal["message"] = "message" role: Literal["assistant"] = "assistant" content: List[Union[TextBlock, ThinkingBlock, ToolUseBlock]] model: str stop_reason: Optional[Literal["end_turn", "max_tokens", "stop_sequence", "tool_use"]] = "end_turn" stop_sequence: Optional[str] = None usage: Usage # Streaming Event Models class StreamEvent(BaseModel): type: str index: Optional[int] = None content_block: Optional[Dict[str, Any]] = None delta: Optional[Dict[str, Any]] = None message: Optional[Dict[str, Any]] = None usage: Optional[Dict[str, Any]] = None # ============== Helper Functions ============== def extract_text_from_content(content: Union[str, List[ContentBlock]]) -> str: """Extract text from content which may be string or list of blocks""" if isinstance(content, str): return content texts = [] for block in content: if isinstance(block, str): texts.append(block) elif hasattr(block, 'text'): texts.append(block.text) elif hasattr(block, 'thinking'): texts.append(block.thinking) elif isinstance(block, dict): if block.get('type') == 'text': texts.append(block.get('text', '')) elif block.get('type') == 'thinking': texts.append(block.get('thinking', '')) return " ".join(texts) def format_system_prompt(system: Optional[Union[str, List[TextBlock]]]) -> str: """Format system prompt from string or list of blocks""" if system is None: return "" if isinstance(system, str): return system return " ".join([block.text for block in system if hasattr(block, 'text')]) def format_messages_to_prompt(messages: List[MessageParam], system: Optional[Union[str, List[TextBlock]]] = None) -> str: """Convert chat messages to a single prompt string""" prompt_parts = [] system_text = format_system_prompt(system) if system_text: prompt_parts.append(f"System: {system_text}\n\n") for msg in messages: role = msg.role content_text = extract_text_from_content(msg.content) if role == "user": prompt_parts.append(f"Human: {content_text}\n\n") elif role == "assistant": prompt_parts.append(f"Assistant: {content_text}\n\n") prompt_parts.append("Assistant:") return "".join(prompt_parts) def generate_text(prompt: str, max_tokens: int, temperature: float, top_p: float) -> tuple: """Generate text and return (text, input_tokens, output_tokens)""" tokenizer = models["tokenizer"] model = models["model"] inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=512) input_tokens = inputs["input_ids"].shape[1] with torch.no_grad(): outputs = model.generate( **inputs, max_new_tokens=min(max_tokens, 256), # Limit for CPU temperature=temperature if temperature > 0 else 1.0, top_p=top_p, do_sample=temperature > 0, pad_token_id=tokenizer.pad_token_id, eos_token_id=tokenizer.eos_token_id ) generated_tokens = outputs[0][input_tokens:] output_tokens = len(generated_tokens) generated_text = tokenizer.decode(generated_tokens, skip_special_tokens=True) return generated_text.strip(), input_tokens, output_tokens async def generate_stream(prompt: str, max_tokens: int, temperature: float, top_p: float, message_id: str, model_name: str): """Generate streaming response in Anthropic SSE format""" tokenizer = models["tokenizer"] model = models["model"] inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=512) input_tokens = inputs["input_ids"].shape[1] # Send message_start event message_start = { "type": "message_start", "message": { "id": message_id, "type": "message", "role": "assistant", "content": [], "model": model_name, "stop_reason": None, "stop_sequence": None, "usage": {"input_tokens": input_tokens, "output_tokens": 0} } } yield f"event: message_start\ndata: {json.dumps(message_start)}\n\n" # Send content_block_start event content_block_start = { "type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""} } yield f"event: content_block_start\ndata: {json.dumps(content_block_start)}\n\n" # Generate tokens with torch.no_grad(): outputs = model.generate( **inputs, max_new_tokens=min(max_tokens, 256), temperature=temperature if temperature > 0 else 1.0, top_p=top_p, do_sample=temperature > 0, pad_token_id=tokenizer.pad_token_id, eos_token_id=tokenizer.eos_token_id ) generated_tokens = outputs[0][input_tokens:] generated_text = tokenizer.decode(generated_tokens, skip_special_tokens=True).strip() output_tokens = len(generated_tokens) # Stream text in chunks chunk_size = 5 for i in range(0, len(generated_text), chunk_size): chunk = generated_text[i:i+chunk_size] content_block_delta = { "type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": chunk} } yield f"event: content_block_delta\ndata: {json.dumps(content_block_delta)}\n\n" await asyncio.sleep(0.01) # Small delay for realistic streaming # Send content_block_stop event content_block_stop = {"type": "content_block_stop", "index": 0} yield f"event: content_block_stop\ndata: {json.dumps(content_block_stop)}\n\n" # Send message_delta event message_delta = { "type": "message_delta", "delta": {"stop_reason": "end_turn", "stop_sequence": None}, "usage": {"output_tokens": output_tokens} } yield f"event: message_delta\ndata: {json.dumps(message_delta)}\n\n" # Send message_stop event message_stop = {"type": "message_stop"} yield f"event: message_stop\ndata: {json.dumps(message_stop)}\n\n" def handle_tool_call(tools: List[Tool], messages: List[MessageParam], generated_text: str) -> Optional[ToolUseBlock]: """Check if the response should trigger a tool call""" if not tools: return None # Simple heuristic: check if response mentions tool names for tool in tools: if tool.name.lower() in generated_text.lower(): return ToolUseBlock( type="tool_use", id=f"toolu_{uuid.uuid4().hex[:24]}", name=tool.name, input={} ) return None # ============== Anthropic API Endpoints ============== @app.post("/v1/messages") async def create_message(request: AnthropicRequest): """ Anthropic Messages API compatible endpoint POST /v1/messages Supports: - Text messages - System prompts - Streaming responses - Tool/function calling - Thinking/reasoning blocks """ try: message_id = f"msg_{uuid.uuid4().hex[:24]}" # Format messages to prompt prompt = format_messages_to_prompt(request.messages, request.system) # Handle streaming if request.stream: return StreamingResponse( generate_stream( prompt=prompt, max_tokens=request.max_tokens, temperature=request.temperature or 1.0, top_p=request.top_p or 1.0, message_id=message_id, model_name=request.model ), media_type="text/event-stream", headers={ "Cache-Control": "no-cache", "Connection": "keep-alive", "X-Accel-Buffering": "no" } ) # Non-streaming response generated_text, input_tokens, output_tokens = generate_text( prompt=prompt, max_tokens=request.max_tokens, temperature=request.temperature or 1.0, top_p=request.top_p or 1.0 ) # Build content blocks content_blocks = [] # Add thinking block if enabled if request.thinking and request.thinking.type == "enabled": thinking_text = f"Analyzing the user's request and formulating a response..." content_blocks.append(ThinkingBlock(type="thinking", thinking=thinking_text)) # Check for tool calls tool_use = handle_tool_call(request.tools, request.messages, generated_text) if request.tools else None if tool_use: content_blocks.append(TextBlock(type="text", text=generated_text)) content_blocks.append(tool_use) stop_reason = "tool_use" else: content_blocks.append(TextBlock(type="text", text=generated_text)) stop_reason = "end_turn" return AnthropicResponse( id=message_id, content=content_blocks, model=request.model, stop_reason=stop_reason, usage=Usage(input_tokens=input_tokens, output_tokens=output_tokens) ) except Exception as e: raise HTTPException(status_code=500, detail=str(e)) # ============== OpenAI Compatible Endpoints ============== class ChatMessage(BaseModel): role: str content: str class ChatCompletionRequest(BaseModel): model: str = "distilgpt2" messages: List[ChatMessage] max_tokens: Optional[int] = 1024 temperature: Optional[float] = 0.7 top_p: Optional[float] = 1.0 stream: Optional[bool] = False @app.post("/v1/chat/completions") async def chat_completions(request: ChatCompletionRequest): """OpenAI Chat Completions API compatible endpoint""" try: # Convert to Anthropic format anthropic_messages = [ MessageParam(role=msg.role if msg.role in ["user", "assistant"] else "user", content=msg.content) for msg in request.messages if msg.role in ["user", "assistant"] ] prompt = format_messages_to_prompt(anthropic_messages) generated_text, input_tokens, output_tokens = generate_text( prompt=prompt, max_tokens=request.max_tokens or 1024, temperature=request.temperature or 0.7, top_p=request.top_p or 1.0 ) return { "id": f"chatcmpl-{uuid.uuid4().hex[:24]}", "object": "chat.completion", "created": int(time.time()), "model": request.model, "choices": [{ "index": 0, "message": {"role": "assistant", "content": generated_text}, "finish_reason": "stop" }], "usage": { "prompt_tokens": input_tokens, "completion_tokens": output_tokens, "total_tokens": input_tokens + output_tokens } } except Exception as e: raise HTTPException(status_code=500, detail=str(e)) @app.get("/v1/models") async def list_models(): """List available models""" return { "object": "list", "data": [ {"id": "MiniMax-M2", "object": "model", "created": int(time.time()), "owned_by": "local"}, {"id": "MiniMax-M2-Stable", "object": "model", "created": int(time.time()), "owned_by": "local"}, {"id": GENERATOR_MODEL, "object": "model", "created": int(time.time()), "owned_by": "local"} ] } # ============== Utility Endpoints ============== @app.get("/") async def root(): """Welcome endpoint""" return { "message": "Docker Model Runner API (Anthropic Compatible)", "hardware": "CPU Basic: 2 vCPU · 16 GB RAM", "docs": "/docs", "api_endpoints": { "anthropic_messages": "POST /v1/messages", "openai_chat": "POST /v1/chat/completions", "models": "GET /v1/models" }, "supported_features": [ "text messages", "system prompts", "streaming responses", "tool/function calling", "thinking blocks", "metadata" ] } @app.get("/health") async def health(): """Health check endpoint""" return { "status": "healthy", "timestamp": datetime.utcnow().isoformat(), "hardware": "CPU Basic: 2 vCPU · 16 GB RAM", "models_loaded": len(models) > 0 } @app.get("/info") async def info(): """API information""" return { "name": "Docker Model Runner", "version": "1.0.0", "api_compatibility": ["anthropic", "openai"], "supported_models": ["MiniMax-M2", "MiniMax-M2-Stable"], "supported_parameters": { "fully_supported": ["model", "messages", "max_tokens", "stream", "system", "temperature", "top_p", "tools", "tool_choice", "metadata", "thinking"], "ignored": ["top_k", "stop_sequences", "service_tier"] }, "message_types": { "supported": ["text", "tool_use", "tool_result", "thinking"], "not_supported": ["image", "document"] } } if __name__ == "__main__": import uvicorn uvicorn.run(app, host="0.0.0.0", port=7860)