Upload app.py with huggingface_hub

app.py

@@ -1,6 +1,7 @@
 """
 Anthropic-Compatible API Endpoint
 Lightweight CPU-based implementation for Hugging Face Spaces
+Full Anthropic API parameter compatibility
 """
 
 import os
@@ -9,7 +10,7 @@ import uuid
 import logging
 from datetime import datetime
 from logging.handlers import RotatingFileHandler
-from typing import List, Optional, Union
+from typing import List, Optional, Union, Dict, Any, Literal
 from contextlib import asynccontextmanager
 
 from fastapi import FastAPI, HTTPException, Header, Request
@@ -26,32 +27,24 @@ LOG_DIR = "/tmp/logs"
 os.makedirs(LOG_DIR, exist_ok=True)
 LOG_FILE = os.path.join(LOG_DIR, "api.log")
 
-# Create formatters
 log_format = logging.Formatter(
     '%(asctime)s | %(levelname)-8s | %(name)s | %(message)s',
     datefmt='%Y-%m-%d %H:%M:%S'
 )
 
-# File handler with rotation (10MB max, keep 5 backups)
 file_handler = RotatingFileHandler(
-    LOG_FILE,
-    maxBytes=10*1024*1024,
-    backupCount=5,
-    encoding='utf-8'
+    LOG_FILE, maxBytes=10*1024*1024, backupCount=5, encoding='utf-8'
 )
 file_handler.setFormatter(log_format)
 file_handler.setLevel(logging.DEBUG)
 
-# Console handler
 console_handler = logging.StreamHandler()
 console_handler.setFormatter(log_format)
 console_handler.setLevel(logging.INFO)
 
-# Root logger
 logging.basicConfig(level=logging.DEBUG, handlers=[file_handler, console_handler])
 logger = logging.getLogger("anthropic-api")
 
-# Also capture uvicorn logs
 for uvicorn_logger in ["uvicorn", "uvicorn.error", "uvicorn.access"]:
     uv_log = logging.getLogger(uvicorn_logger)
     uv_log.handlers = [file_handler, console_handler]
@@ -62,29 +55,21 @@ logger.info(f"Log file: {LOG_FILE}")
 logger.info("=" * 60)
 
 # ============== Configuration ==============
-MODEL_ID = "HuggingFaceTB/SmolLM2-135M-Instruct"  # Ultra-lightweight 135M model
-MAX_TOKENS_DEFAULT = 1024
+MODEL_ID = "HuggingFaceTB/SmolLM2-135M-Instruct"
 DEVICE = "cpu"
 
-# Global model and tokenizer
 model = None
 tokenizer = None
 
 @asynccontextmanager
 async def lifespan(app: FastAPI):
-    """Load model on startup"""
     global model, tokenizer
     logger.info(f"Loading model: {MODEL_ID}")
-
     try:
         tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
         logger.info("Tokenizer loaded successfully")
-
         model = AutoModelForCausalLM.from_pretrained(
-            MODEL_ID,
-            torch_dtype=torch.float32,
-            device_map=DEVICE,
-            low_cpu_mem_usage=True
+            MODEL_ID, torch_dtype=torch.float32, device_map=DEVICE, low_cpu_mem_usage=True
         )
         model.eval()
         logger.info("Model loaded successfully!")
@@ -92,21 +77,17 @@ async def lifespan(app: FastAPI):
     except Exception as e:
         logger.error(f"Failed to load model: {e}", exc_info=True)
         raise
-
     yield
-
-    # Cleanup
     logger.info("Shutting down, cleaning up model...")
     del model, tokenizer
 
 app = FastAPI(
     title="Anthropic-Compatible API",
-    description="Lightweight CPU-based API with Anthropic Messages API compatibility",
+    description="Lightweight CPU-based API with full Anthropic Messages API compatibility",
     version="1.0.0",
     lifespan=lifespan
 )
 
-# CORS middleware
 app.add_middleware(
     CORSMiddleware,
     allow_origins=["*"],
@@ -115,14 +96,11 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
-# Request logging middleware
 @app.middleware("http")
 async def log_requests(request: Request, call_next):
     request_id = str(uuid.uuid4())[:8]
     start_time = time.time()
-
     logger.info(f"[{request_id}] {request.method} {request.url.path} - Started")
-
     try:
         response = await call_next(request)
         duration = (time.time() - start_time) * 1000
@@ -133,69 +111,193 @@ async def log_requests(request: Request, call_next):
         logger.error(f"[{request_id}] {request.method} {request.url.path} - Error: {e} ({duration:.2f}ms)")
         raise
 
-# ============== Pydantic Models (Anthropic-Compatible) ==============
+# ============== Anthropic-Compatible Pydantic Models ==============
 
-class ContentBlock(BaseModel):
-    type: str = "text"
+# Content block types (matching Anthropic exactly)
+class TextBlock(BaseModel):
+    type: Literal["text"] = "text"
     text: str
 
+class ImageSource(BaseModel):
+    type: Literal["base64", "url"] = "base64"
+    media_type: Optional[str] = None
+    data: Optional[str] = None
+    url: Optional[str] = None
+
+class ImageBlock(BaseModel):
+    type: Literal["image"] = "image"
+    source: ImageSource
+
+class ToolUseBlock(BaseModel):
+    type: Literal["tool_use"] = "tool_use"
+    id: str
+    name: str
+    input: Dict[str, Any]
+
+class ToolResultBlock(BaseModel):
+    type: Literal["tool_result"] = "tool_result"
+    tool_use_id: str
+    content: Optional[Union[str, List[TextBlock]]] = None
+    is_error: Optional[bool] = False
+
+ContentBlock = Union[TextBlock, ImageBlock, ToolUseBlock, ToolResultBlock]
+
+# Message structure (matching Anthropic exactly)
 class Message(BaseModel):
-    role: str
+    role: Literal["user", "assistant"]
     content: Union[str, List[ContentBlock]]
 
+# Tool definition (matching Anthropic exactly)
+class ToolInputSchema(BaseModel):
+    type: Literal["object"] = "object"
+    properties: Optional[Dict[str, Any]] = None
+    required: Optional[List[str]] = None
+
+class Tool(BaseModel):
+    name: str
+    description: Optional[str] = None
+    input_schema: ToolInputSchema
+
+# Tool choice (matching Anthropic exactly)
+class ToolChoiceAuto(BaseModel):
+    type: Literal["auto"] = "auto"
+    disable_parallel_tool_use: Optional[bool] = None
+
+class ToolChoiceAny(BaseModel):
+    type: Literal["any"] = "any"
+    disable_parallel_tool_use: Optional[bool] = None
+
+class ToolChoiceTool(BaseModel):
+    type: Literal["tool"] = "tool"
+    name: str
+    disable_parallel_tool_use: Optional[bool] = None
+
+ToolChoice = Union[ToolChoiceAuto, ToolChoiceAny, ToolChoiceTool]
+
+# Metadata (matching Anthropic exactly)
+class Metadata(BaseModel):
+    user_id: Optional[str] = None
+
+# System content (matching Anthropic exactly)
+class SystemContent(BaseModel):
+    type: Literal["text"] = "text"
+    text: str
+    cache_control: Optional[Dict[str, str]] = None
+
+# Main request model (matching Anthropic exactly)
 class MessageRequest(BaseModel):
+    # Required parameters
     model: str
+    max_tokens: int
     messages: List[Message]
-    max_tokens: int = MAX_TOKENS_DEFAULT
-    temperature: Optional[float] = 0.7
-    top_p: Optional[float] = 0.9
-    top_k: Optional[int] = 50
-    stream: Optional[bool] = False
-    system: Optional[str] = None
-    stop_sequences: Optional[List[str]] = None
 
+    # Optional parameters (matching Anthropic exactly)
+    metadata: Optional[Metadata] = None
+    stop_sequences: Optional[List[str]] = None
+    stream: Optional[bool] = False
+    system: Optional[Union[str, List[SystemContent]]] = None
+    temperature: Optional[float] = Field(default=1.0, ge=0.0, le=1.0)
+    tool_choice: Optional[ToolChoice] = None
+    tools: Optional[List[Tool]] = None
+    top_k: Optional[int] = Field(default=None, ge=0)
+    top_p: Optional[float] = Field(default=None, ge=0.0, le=1.0)
+
+# Usage model (matching Anthropic exactly)
 class Usage(BaseModel):
     input_tokens: int
     output_tokens: int
+    cache_creation_input_tokens: Optional[int] = None
+    cache_read_input_tokens: Optional[int] = None
 
+# Response content block
+class ResponseTextBlock(BaseModel):
+    type: Literal["text"] = "text"
+    text: str
+
+class ResponseToolUseBlock(BaseModel):
+    type: Literal["tool_use"] = "tool_use"
+    id: str
+    name: str
+    input: Dict[str, Any]
+
+ResponseContentBlock = Union[ResponseTextBlock, ResponseToolUseBlock]
+
+# Main response model (matching Anthropic exactly)
 class MessageResponse(BaseModel):
     id: str
-    type: str = "message"
-    role: str = "assistant"
-    content: List[ContentBlock]
+    type: Literal["message"] = "message"
+    role: Literal["assistant"] = "assistant"
+    content: List[ResponseContentBlock]
     model: str
-    stop_reason: str = "end_turn"
+    stop_reason: Optional[Literal["end_turn", "max_tokens", "stop_sequence", "tool_use"]] = None
     stop_sequence: Optional[str] = None
     usage: Usage
 
+# Error response (matching Anthropic exactly)
+class ErrorDetail(BaseModel):
+    type: str
+    message: str
+
 class ErrorResponse(BaseModel):
-    type: str = "error"
-    error: dict
+    type: Literal["error"] = "error"
+    error: ErrorDetail
+
+# Token count request/response (matching Anthropic exactly)
+class TokenCountRequest(BaseModel):
+    model: str
+    messages: List[Message]
+    system: Optional[Union[str, List[SystemContent]]] = None
+    tools: Optional[List[Tool]] = None
+
+class TokenCountResponse(BaseModel):
+    input_tokens: int
 
 # ============== Helper Functions ==============
 
-def format_messages(messages: List[Message], system: Optional[str] = None) -> str:
+def extract_text_content(content: Union[str, List[ContentBlock]]) -> str:
+    """Extract text from content (string or list of blocks)"""
+    if isinstance(content, str):
+        return content
+    texts = []
+    for block in content:
+        if isinstance(block, dict):
+            if block.get("type") == "text":
+                texts.append(block.get("text", ""))
+        elif hasattr(block, "type") and block.type == "text":
+            texts.append(block.text)
+    return " ".join(texts)
+
+def extract_system_content(system: Optional[Union[str, List[SystemContent]]]) -> Optional[str]:
+    """Extract system prompt from string or list of system content blocks"""
+    if system is None:
+        return None
+    if isinstance(system, str):
+        return system
+    texts = []
+    for block in system:
+        if isinstance(block, dict):
+            texts.append(block.get("text", ""))
+        elif hasattr(block, "text"):
+            texts.append(block.text)
+    return " ".join(texts)
+
+def format_messages(messages: List[Message], system: Optional[Union[str, List[SystemContent]]] = None) -> str:
     """Format messages into a prompt string"""
     formatted_messages = []
 
-    if system:
-        formatted_messages.append({"role": "system", "content": system})
+    system_text = extract_system_content(system)
+    if system_text:
+        formatted_messages.append({"role": "system", "content": system_text})
 
     for msg in messages:
-        content = msg.content
-        if isinstance(content, list):
-            content = " ".join([block.text for block in content if block.type == "text"])
+        content = extract_text_content(msg.content)
         formatted_messages.append({"role": msg.role, "content": content})
 
-    # Use chat template if available
     if tokenizer.chat_template:
         return tokenizer.apply_chat_template(
-            formatted_messages,
-            tokenize=False,
-            add_generation_prompt=True
+            formatted_messages, tokenize=False, add_generation_prompt=True
         )
 
-    # Fallback simple format
     prompt = ""
     for msg in formatted_messages:
         role = msg["role"].capitalize()
@@ -204,14 +306,12 @@ def format_messages(messages: List[Message], system: Optional[str] = None) -> st
     return prompt
 
 def generate_id() -> str:
-    """Generate a unique message ID"""
     return f"msg_{uuid.uuid4().hex[:24]}"
 
 # ============== API Endpoints ==============
 
 @app.get("/")
 async def root():
-    """Health check endpoint"""
     logger.debug("Root endpoint accessed")
     return {
         "status": "healthy",
@@ -223,24 +323,20 @@ async def root():
 
 @app.get("/v1/models")
 async def list_models():
-    """List available models (Anthropic-compatible)"""
     logger.debug("Models list requested")
     return {
         "object": "list",
-        "data": [
-            {
-                "id": "smollm2-135m",
-                "object": "model",
-                "created": int(time.time()),
-                "owned_by": "huggingface",
-                "display_name": "SmolLM2 135M Instruct"
-            }
-        ]
+        "data": [{
+            "id": "smollm2-135m",
+            "object": "model",
+            "created": int(time.time()),
+            "owned_by": "huggingface",
+            "display_name": "SmolLM2 135M Instruct"
+        }]
     }
 
 @app.get("/logs")
 async def get_logs(lines: int = 100):
-    """Get recent log entries"""
     try:
         with open(LOG_FILE, 'r') as f:
             all_lines = f.readlines()
@@ -254,24 +350,22 @@ async def get_logs(lines: int = 100):
     except FileNotFoundError:
         return {"error": "Log file not found", "log_file": LOG_FILE}
 
-@app.post("/v1/messages")
+@app.post("/v1/messages", response_model=MessageResponse)
 async def create_message(
     request: MessageRequest,
     x_api_key: Optional[str] = Header(None, alias="x-api-key"),
-    anthropic_version: Optional[str] = Header(None, alias="anthropic-version")
+    anthropic_version: Optional[str] = Header(None, alias="anthropic-version"),
+    anthropic_beta: Optional[str] = Header(None, alias="anthropic-beta")
 ):
-    """
-    Create a message (Anthropic Messages API compatible)
-    """
+    """Create a message (Anthropic Messages API compatible)"""
     message_id = generate_id()
     logger.info(f"[{message_id}] Creating message - model: {request.model}, max_tokens: {request.max_tokens}, stream: {request.stream}")
+    logger.debug(f"[{message_id}] Request params - temp: {request.temperature}, top_p: {request.top_p}, top_k: {request.top_k}")
 
     try:
-        # Format the prompt
         prompt = format_messages(request.messages, request.system)
         logger.debug(f"[{message_id}] Prompt length: {len(prompt)} chars")
 
-        # Tokenize
         inputs = tokenizer(prompt, return_tensors="pt").to(DEVICE)
         input_token_count = inputs.input_ids.shape[1]
         logger.info(f"[{message_id}] Input tokens: {input_token_count}")
@@ -280,41 +374,72 @@ async def create_message(
             logger.info(f"[{message_id}] Starting streaming response")
             return await stream_response(request, inputs, input_token_count, message_id)
 
-        # Generate
+        # Build generation kwargs matching Anthropic params
+        gen_kwargs = {
+            "max_new_tokens": request.max_tokens,
+            "do_sample": request.temperature > 0 if request.temperature else False,
+            "pad_token_id": tokenizer.eos_token_id,
+            "eos_token_id": tokenizer.eos_token_id,
+        }
+
+        # Temperature (Anthropic default: 1.0)
+        if request.temperature is not None and request.temperature > 0:
+            gen_kwargs["temperature"] = request.temperature
+
+        # Top-p (nucleus sampling)
+        if request.top_p is not None:
+            gen_kwargs["top_p"] = request.top_p
+
+        # Top-k sampling
+        if request.top_k is not None:
+            gen_kwargs["top_k"] = request.top_k
+
+        # Stop sequences
+        if request.stop_sequences:
+            stop_token_ids = []
+            for seq in request.stop_sequences:
+                tokens = tokenizer.encode(seq, add_special_tokens=False)
+                if tokens:
+                    stop_token_ids.extend(tokens)
+            if stop_token_ids:
+                gen_kwargs["eos_token_id"] = list(set([tokenizer.eos_token_id] + stop_token_ids))
+
         gen_start = time.time()
         with torch.no_grad():
-            outputs = model.generate(
-                **inputs,
-                max_new_tokens=request.max_tokens,
-                temperature=request.temperature if request.temperature > 0 else 1.0,
-                top_p=request.top_p,
-                top_k=request.top_k,
-                do_sample=request.temperature > 0,
-                pad_token_id=tokenizer.eos_token_id,
-                eos_token_id=tokenizer.eos_token_id,
-            )
+            outputs = model.generate(**inputs, **gen_kwargs)
         gen_time = time.time() - gen_start
 
-        # Decode only new tokens
         generated_tokens = outputs[0][input_token_count:]
         generated_text = tokenizer.decode(generated_tokens, skip_special_tokens=True)
         output_token_count = len(generated_tokens)
 
+        # Determine stop reason
+        stop_reason = "end_turn"
+        stop_sequence = None
+        if output_token_count >= request.max_tokens:
+            stop_reason = "max_tokens"
+        elif request.stop_sequences:
+            for seq in request.stop_sequences:
+                if seq in generated_text:
+                    stop_reason = "stop_sequence"
+                    stop_sequence = seq
+                    generated_text = generated_text.split(seq)[0]
+                    break
+
         tokens_per_sec = output_token_count / gen_time if gen_time > 0 else 0
         logger.info(f"[{message_id}] Generated {output_token_count} tokens in {gen_time:.2f}s ({tokens_per_sec:.1f} tok/s)")
 
-        # Build response
         response = MessageResponse(
             id=message_id,
-            content=[ContentBlock(type="text", text=generated_text.strip())],
+            content=[ResponseTextBlock(type="text", text=generated_text.strip())],
             model=request.model,
-            stop_reason="end_turn",
+            stop_reason=stop_reason,
+            stop_sequence=stop_sequence,
             usage=Usage(
                 input_tokens=input_token_count,
                 output_tokens=output_token_count
             )
         )
-
         return response
 
     except Exception as e:
@@ -322,10 +447,10 @@ async def create_message(
         raise HTTPException(status_code=500, detail=str(e))
 
 async def stream_response(request: MessageRequest, inputs, input_token_count: int, message_id: str):
-    """Stream response using SSE (Server-Sent Events)"""
+    """Stream response using SSE (Server-Sent Events) - Anthropic format"""
 
     async def generate():
-        # Send message_start event
+        # message_start event
         start_event = {
             "type": "message_start",
             "message": {
@@ -341,7 +466,7 @@ async def stream_response(request: MessageRequest, inputs, input_token_count: in
         }
         yield f"event: message_start\ndata: {json.dumps(start_event)}\n\n"
 
-        # Send content_block_start
+        # content_block_start event
         block_start = {
             "type": "content_block_start",
             "index": 0,
@@ -349,24 +474,29 @@ async def stream_response(request: MessageRequest, inputs, input_token_count: in
         }
         yield f"event: content_block_start\ndata: {json.dumps(block_start)}\n\n"
 
-        # Setup streamer
+        # ping event (Anthropic sends these)
+        yield f"event: ping\ndata: {json.dumps({'type': 'ping'})}\n\n"
+
         streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
 
-        generation_kwargs = {
+        gen_kwargs = {
             **inputs,
             "max_new_tokens": request.max_tokens,
-            "temperature": request.temperature if request.temperature > 0 else 1.0,
-            "top_p": request.top_p,
-            "top_k": request.top_k,
-            "do_sample": request.temperature > 0,
+            "do_sample": request.temperature > 0 if request.temperature else False,
             "pad_token_id": tokenizer.eos_token_id,
             "eos_token_id": tokenizer.eos_token_id,
             "streamer": streamer,
         }
 
-        # Run generation in a thread
+        if request.temperature is not None and request.temperature > 0:
+            gen_kwargs["temperature"] = request.temperature
+        if request.top_p is not None:
+            gen_kwargs["top_p"] = request.top_p
+        if request.top_k is not None:
+            gen_kwargs["top_k"] = request.top_k
+
         gen_start = time.time()
-        thread = Thread(target=model.generate, kwargs=generation_kwargs)
+        thread = Thread(target=model.generate, kwargs=gen_kwargs)
         thread.start()
 
         output_tokens = 0
@@ -385,19 +515,19 @@ async def stream_response(request: MessageRequest, inputs, input_token_count: in
         tokens_per_sec = output_tokens / gen_time if gen_time > 0 else 0
         logger.info(f"[{message_id}] Stream completed: {output_tokens} tokens in {gen_time:.2f}s ({tokens_per_sec:.1f} tok/s)")
 
-        # Send content_block_stop
-        block_stop = {"type": "content_block_stop", "index": 0}
-        yield f"event: content_block_stop\ndata: {json.dumps(block_stop)}\n\n"
+        # content_block_stop event
+        yield f"event: content_block_stop\ndata: {json.dumps({'type': 'content_block_stop', 'index': 0})}\n\n"
 
-        # Send message_delta
+        # message_delta event
+        stop_reason = "max_tokens" if output_tokens >= request.max_tokens else "end_turn"
         delta = {
             "type": "message_delta",
-            "delta": {"stop_reason": "end_turn", "stop_sequence": None},
+            "delta": {"stop_reason": stop_reason, "stop_sequence": None},
             "usage": {"output_tokens": output_tokens}
         }
         yield f"event: message_delta\ndata: {json.dumps(delta)}\n\n"
 
-        # Send message_stop
+        # message_stop event
         yield f"event: message_stop\ndata: {json.dumps({'type': 'message_stop'})}\n\n"
 
     return StreamingResponse(
@@ -410,16 +540,14 @@ async def stream_response(request: MessageRequest, inputs, input_token_count: in
         }
     )
 
-# Token counting endpoint
-@app.post("/v1/messages/count_tokens")
-async def count_tokens(request: MessageRequest):
-    """Count tokens for a message request"""
+@app.post("/v1/messages/count_tokens", response_model=TokenCountResponse)
+async def count_tokens(request: TokenCountRequest):
+    """Count tokens for a message request (Anthropic compatible)"""
     prompt = format_messages(request.messages, request.system)
     tokens = tokenizer.encode(prompt)
     logger.debug(f"Token count request: {len(tokens)} tokens")
-    return {"input_tokens": len(tokens)}
+    return TokenCountResponse(input_tokens=len(tokens))
 
-# Health check
 @app.get("/health")
 async def health():
     return {"status": "ok", "model_loaded": model is not None, "log_file": LOG_FILE}