Upload app.py with huggingface_hub

app.py

@@ -1,13 +1,14 @@
 """
 Anthropic-Compatible API Endpoint
 Lightweight CPU-based implementation for Hugging Face Spaces
-Full Anthropic API parameter compatibility
+Full Anthropic API parameter compatibility with Extended Thinking support
 """
 
 import os
 import time
 import uuid
 import logging
+import re
 from datetime import datetime
 from logging.handlers import RotatingFileHandler
 from typing import List, Optional, Union, Dict, Any, Literal
@@ -83,7 +84,7 @@ async def lifespan(app: FastAPI):
 
 app = FastAPI(
     title="Anthropic-Compatible API",
-    description="Lightweight CPU-based API with full Anthropic Messages API compatibility",
+    description="Lightweight CPU-based API with full Anthropic Messages API compatibility including Extended Thinking",
     version="1.0.0",
     lifespan=lifespan
 )
@@ -184,6 +185,16 @@ class SystemContent(BaseModel):
     text: str
     cache_control: Optional[Dict[str, str]] = None
 
+# ============== Extended Thinking (ThinkingConfig) ==============
+class ThinkingConfig(BaseModel):
+    """
+    Extended thinking configuration (matching Anthropic's ThinkingConfig)
+    Enables Claude to think through complex problems before responding
+    """
+    type: Literal["enabled", "disabled"] = "enabled"
+    # Budget tokens for thinking (Anthropic uses budget_tokens)
+    budget_tokens: Optional[int] = Field(default=1024, ge=1, le=128000)
+
 # Main request model (matching Anthropic exactly)
 class MessageRequest(BaseModel):
     # Required parameters
@@ -202,25 +213,33 @@ class MessageRequest(BaseModel):
     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)
+    # Extended Thinking (ThinkingConfig)
+    thinking: Optional[ThinkingConfig] = None
+
+# Usage model (matching Anthropic exactly with thinking tokens)
 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
+# Response content blocks
 class ResponseTextBlock(BaseModel):
     type: Literal["text"] = "text"
     text: str
 
+class ResponseThinkingBlock(BaseModel):
+    """Thinking block in response (matching Anthropic's thinking content block)"""
+    type: Literal["thinking"] = "thinking"
+    thinking: str
+
 class ResponseToolUseBlock(BaseModel):
     type: Literal["tool_use"] = "tool_use"
     id: str
     name: str
     input: Dict[str, Any]
 
-ResponseContentBlock = Union[ResponseTextBlock, ResponseToolUseBlock]
+ResponseContentBlock = Union[ResponseTextBlock, ResponseThinkingBlock, ResponseToolUseBlock]
 
 # Main response model (matching Anthropic exactly)
 class MessageResponse(BaseModel):
@@ -248,6 +267,7 @@ class TokenCountRequest(BaseModel):
     messages: List[Message]
     system: Optional[Union[str, List[SystemContent]]] = None
     tools: Optional[List[Tool]] = None
+    thinking: Optional[ThinkingConfig] = None
 
 class TokenCountResponse(BaseModel):
     input_tokens: int
@@ -281,11 +301,36 @@ def extract_system_content(system: Optional[Union[str, List[SystemContent]]]) ->
             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"""
+def format_messages_with_thinking(
+    messages: List[Message],
+    system: Optional[Union[str, List[SystemContent]]] = None,
+    thinking_enabled: bool = False,
+    budget_tokens: int = 1024
+) -> str:
+    """Format messages with optional thinking prompt"""
     formatted_messages = []
 
     system_text = extract_system_content(system)
+
+    # Add thinking instructions to system prompt if enabled
+    if thinking_enabled:
+        thinking_instruction = f"""You are a helpful AI assistant with extended thinking capabilities.
+
+When responding to complex problems:
+1. First, think through the problem step by step inside <thinking>...</thinking> tags
+2. Consider multiple approaches and evaluate them
+3. Show your reasoning process clearly
+4. After thinking, provide your final answer outside the thinking tags
+
+Budget for thinking: up to {budget_tokens} tokens for reasoning.
+
+Think deeply and thoroughly before responding."""
+
+        if system_text:
+            system_text = f"{thinking_instruction}\n\n{system_text}"
+        else:
+            system_text = thinking_instruction
+
     if system_text:
         formatted_messages.append({"role": "system", "content": system_text})
 
@@ -305,6 +350,22 @@ def format_messages(messages: List[Message], system: Optional[Union[str, List[Sy
     prompt += "Assistant: "
     return prompt
 
+def parse_thinking_response(text: str) -> tuple:
+    """
+    Parse response to extract thinking and final answer
+    Returns: (thinking_text, answer_text)
+    """
+    thinking_pattern = r'<thinking>(.*?)</thinking>'
+    thinking_matches = re.findall(thinking_pattern, text, re.DOTALL)
+
+    if thinking_matches:
+        thinking_text = "\n".join(thinking_matches).strip()
+        # Remove thinking blocks from response
+        answer_text = re.sub(thinking_pattern, '', text, flags=re.DOTALL).strip()
+        return thinking_text, answer_text
+    else:
+        return None, text.strip()
+
 def generate_id() -> str:
     return f"msg_{uuid.uuid4().hex[:24]}"
 
@@ -318,6 +379,7 @@ async def root():
         "model": MODEL_ID,
         "api_version": "2023-06-01",
         "compatibility": "anthropic-messages-api",
+        "features": ["extended-thinking", "streaming", "tool-use"],
         "log_file": LOG_FILE
     }
 
@@ -331,7 +393,8 @@ async def list_models():
             "object": "model",
             "created": int(time.time()),
             "owned_by": "huggingface",
-            "display_name": "SmolLM2 135M Instruct"
+            "display_name": "SmolLM2 135M Instruct",
+            "supports_thinking": True
         }]
     }
 
@@ -357,13 +420,27 @@ async def create_message(
     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 with Extended Thinking)"""
     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}")
+
+    # Check if thinking is enabled
+    thinking_enabled = False
+    budget_tokens = 1024
+    if request.thinking:
+        thinking_enabled = request.thinking.type == "enabled"
+        budget_tokens = request.thinking.budget_tokens or 1024
+
+    logger.info(f"[{message_id}] Creating message - model: {request.model}, max_tokens: {request.max_tokens}, stream: {request.stream}, thinking: {thinking_enabled}")
+    logger.debug(f"[{message_id}] Request params - temp: {request.temperature}, top_p: {request.top_p}, top_k: {request.top_k}, thinking_budget: {budget_tokens}")
 
     try:
-        prompt = format_messages(request.messages, request.system)
+        # Format prompt with thinking if enabled
+        prompt = format_messages_with_thinking(
+            request.messages,
+            request.system,
+            thinking_enabled=thinking_enabled,
+            budget_tokens=budget_tokens
+        )
         logger.debug(f"[{message_id}] Prompt length: {len(prompt)} chars")
 
         inputs = tokenizer(prompt, return_tensors="pt").to(DEVICE)
@@ -372,29 +449,28 @@ async def create_message(
 
         if request.stream:
             logger.info(f"[{message_id}] Starting streaming response")
-            return await stream_response(request, inputs, input_token_count, message_id)
+            return await stream_response(request, inputs, input_token_count, message_id, thinking_enabled, budget_tokens)
 
-        # Build generation kwargs matching Anthropic params
+        # Calculate max tokens (include thinking budget if enabled)
+        total_max_tokens = request.max_tokens
+        if thinking_enabled:
+            total_max_tokens += budget_tokens
+
+        # Build generation kwargs
         gen_kwargs = {
-            "max_new_tokens": request.max_tokens,
+            "max_new_tokens": total_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:
@@ -413,17 +489,27 @@ async def create_message(
         generated_text = tokenizer.decode(generated_tokens, skip_special_tokens=True)
         output_token_count = len(generated_tokens)
 
+        # Parse thinking from response if enabled
+        content_blocks = []
+        if thinking_enabled:
+            thinking_text, answer_text = parse_thinking_response(generated_text)
+            if thinking_text:
+                logger.info(f"[{message_id}] Thinking extracted: {len(thinking_text)} chars")
+                content_blocks.append(ResponseThinkingBlock(type="thinking", thinking=thinking_text))
+            content_blocks.append(ResponseTextBlock(type="text", text=answer_text))
+        else:
+            content_blocks.append(ResponseTextBlock(type="text", text=generated_text.strip()))
+
         # Determine stop reason
         stop_reason = "end_turn"
         stop_sequence = None
-        if output_token_count >= request.max_tokens:
+        if output_token_count >= total_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
@@ -431,7 +517,7 @@ async def create_message(
 
         response = MessageResponse(
             id=message_id,
-            content=[ResponseTextBlock(type="text", text=generated_text.strip())],
+            content=content_blocks,
             model=request.model,
             stop_reason=stop_reason,
             stop_sequence=stop_sequence,
@@ -446,8 +532,15 @@ async def create_message(
         logger.error(f"[{message_id}] Error creating message: {e}", exc_info=True)
         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) - Anthropic format"""
+async def stream_response(
+    request: MessageRequest,
+    inputs,
+    input_token_count: int,
+    message_id: str,
+    thinking_enabled: bool = False,
+    budget_tokens: int = 1024
+):
+    """Stream response using SSE (Server-Sent Events) - Anthropic format with thinking support"""
 
     async def generate():
         # message_start event
@@ -466,22 +559,24 @@ async def stream_response(request: MessageRequest, inputs, input_token_count: in
         }
         yield f"event: message_start\ndata: {json.dumps(start_event)}\n\n"
 
-        # content_block_start event
-        block_start = {
-            "type": "content_block_start",
-            "index": 0,
-            "content_block": {"type": "text", "text": ""}
-        }
-        yield f"event: content_block_start\ndata: {json.dumps(block_start)}\n\n"
+        # If thinking is enabled, we'll track thinking vs text blocks
+        block_index = 0
+        in_thinking = False
+        thinking_started = False
+        text_block_started = False
 
-        # ping event (Anthropic sends these)
+        # ping event
         yield f"event: ping\ndata: {json.dumps({'type': 'ping'})}\n\n"
 
         streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
 
+        total_max_tokens = request.max_tokens
+        if thinking_enabled:
+            total_max_tokens += budget_tokens
+
         gen_kwargs = {
             **inputs,
-            "max_new_tokens": request.max_tokens,
+            "max_new_tokens": total_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,
@@ -500,26 +595,89 @@ async def stream_response(request: MessageRequest, inputs, input_token_count: in
         thread.start()
 
         output_tokens = 0
+        accumulated_text = ""
+
         for text in streamer:
             if text:
                 output_tokens += len(tokenizer.encode(text, add_special_tokens=False))
-                delta_event = {
-                    "type": "content_block_delta",
-                    "index": 0,
-                    "delta": {"type": "text_delta", "text": text}
-                }
-                yield f"event: content_block_delta\ndata: {json.dumps(delta_event)}\n\n"
+                accumulated_text += text
+
+                if thinking_enabled:
+                    # Check for thinking tags
+                    if "<thinking>" in accumulated_text and not thinking_started:
+                        # Start thinking block
+                        thinking_started = True
+                        in_thinking = True
+                        block_start = {
+                            "type": "content_block_start",
+                            "index": block_index,
+                            "content_block": {"type": "thinking", "thinking": ""}
+                        }
+                        yield f"event: content_block_start\ndata: {json.dumps(block_start)}\n\n"
+
+                    if in_thinking:
+                        # Stream thinking content
+                        clean_text = text.replace("<thinking>", "").replace("</thinking>", "")
+                        if clean_text:
+                            delta_event = {
+                                "type": "content_block_delta",
+                                "index": block_index,
+                                "delta": {"type": "thinking_delta", "thinking": clean_text}
+                            }
+                            yield f"event: content_block_delta\ndata: {json.dumps(delta_event)}\n\n"
+
+                        if "</thinking>" in accumulated_text:
+                            # End thinking block
+                            in_thinking = False
+                            yield f"event: content_block_stop\ndata: {json.dumps({'type': 'content_block_stop', 'index': block_index})}\n\n"
+                            block_index += 1
+
+                            # Start text block
+                            text_block_started = True
+                            block_start = {
+                                "type": "content_block_start",
+                                "index": block_index,
+                                "content_block": {"type": "text", "text": ""}
+                            }
+                            yield f"event: content_block_start\ndata: {json.dumps(block_start)}\n\n"
+
+                    elif text_block_started:
+                        # Stream text content
+                        delta_event = {
+                            "type": "content_block_delta",
+                            "index": block_index,
+                            "delta": {"type": "text_delta", "text": text}
+                        }
+                        yield f"event: content_block_delta\ndata: {json.dumps(delta_event)}\n\n"
+
+                else:
+                    # No thinking - just stream text
+                    if not text_block_started:
+                        text_block_started = True
+                        block_start = {
+                            "type": "content_block_start",
+                            "index": 0,
+                            "content_block": {"type": "text", "text": ""}
+                        }
+                        yield f"event: content_block_start\ndata: {json.dumps(block_start)}\n\n"
+
+                    delta_event = {
+                        "type": "content_block_delta",
+                        "index": 0,
+                        "delta": {"type": "text_delta", "text": text}
+                    }
+                    yield f"event: content_block_delta\ndata: {json.dumps(delta_event)}\n\n"
 
         thread.join()
         gen_time = time.time() - gen_start
         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)")
 
-        # content_block_stop event
-        yield f"event: content_block_stop\ndata: {json.dumps({'type': 'content_block_stop', 'index': 0})}\n\n"
+        # content_block_stop for final block
+        yield f"event: content_block_stop\ndata: {json.dumps({'type': 'content_block_stop', 'index': block_index})}\n\n"
 
         # message_delta event
-        stop_reason = "max_tokens" if output_tokens >= request.max_tokens else "end_turn"
+        stop_reason = "max_tokens" if output_tokens >= total_max_tokens else "end_turn"
         delta = {
             "type": "message_delta",
             "delta": {"stop_reason": stop_reason, "stop_sequence": None},
@@ -543,14 +701,27 @@ async def stream_response(request: MessageRequest, inputs, input_token_count: in
 @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)
+    thinking_enabled = request.thinking and request.thinking.type == "enabled"
+    budget_tokens = request.thinking.budget_tokens if request.thinking else 1024
+
+    prompt = format_messages_with_thinking(
+        request.messages,
+        request.system,
+        thinking_enabled=thinking_enabled,
+        budget_tokens=budget_tokens
+    )
     tokens = tokenizer.encode(prompt)
-    logger.debug(f"Token count request: {len(tokens)} tokens")
+    logger.debug(f"Token count request: {len(tokens)} tokens (thinking: {thinking_enabled})")
     return TokenCountResponse(input_tokens=len(tokens))
 
 @app.get("/health")
 async def health():
-    return {"status": "ok", "model_loaded": model is not None, "log_file": LOG_FILE}
+    return {
+        "status": "ok",
+        "model_loaded": model is not None,
+        "log_file": LOG_FILE,
+        "features": ["extended-thinking", "streaming"]
+    }
 
 if __name__ == "__main__":
     import uvicorn