Update openai_server.py

openai_server.py

@@ -2,9 +2,11 @@ import json
 import re
 import time
 import uuid
-from typing import List, Optional, Dict, Any
+import os
+import traceback
+from typing import List, Optional, Dict, Any, Union
 
-from fastapi import FastAPI, Request, HTTPException
+from fastapi import FastAPI, HTTPException
 from fastapi.responses import JSONResponse
 from sse_starlette.sse import EventSourceResponse
 from pydantic import BaseModel, Field
@@ -15,11 +17,20 @@ from core import Grok
 app = FastAPI(title="Grok OpenAI Wrapper (Agent & MCP Compatible)")
 
 # --- OpenAI Pydantic Models ---
+class FunctionCall(BaseModel):
+    name: str
+    arguments: str
+
+class ToolCall(BaseModel):
+    id: str
+    type: str = "function"
+    function: FunctionCall
+
 class ChatMessage(BaseModel):
     role: str
     content: Optional[str] = None
     name: Optional[str] = None
-    tool_calls: Optional[List[Dict[Any, Any]]] = None
+    tool_calls: Optional[List[ToolCall]] = None
     tool_call_id: Optional[str] = None
 
 class ChatCompletionRequest(BaseModel):
@@ -32,7 +43,7 @@ class ChatCompletionRequest(BaseModel):
 
 # --- Helper Functions ---
 def format_messages_and_tools(messages: List[ChatMessage], tools: Optional[List[Dict]]) -> str:
-    """Translates the standard OpenAI message history into a single string for the web scraper"""
+    """Translates the standard OpenAI message history into a single string for Grok"""
     prompt = ""
     
     # 1. Inject Tools via System Prompt Strategy if tools exist
@@ -52,11 +63,12 @@ def format_messages_and_tools(messages: List[ChatMessage], tools: Optional[List[
             prompt += f"User: {msg.content}\n\n"
         elif msg.role == "assistant":
             if msg.tool_calls:
-                prompt += f"Assistant called tools: {json.dumps(msg.tool_calls)}\n\n"
+                # Convert tool calls to dicts to cleanly dump them
+                tc_dicts =[{"name": tc.function.name, "arguments": json.loads(tc.function.arguments)} for tc in msg.tool_calls]
+                prompt += f"Assistant called tools: {json.dumps(tc_dicts)}\n\n"
             if msg.content:
                 prompt += f"Assistant: {msg.content}\n\n"
-        elif msg.role == "tool" or msg.role == "function":
-            # Pass tool results back to the model
+        elif msg.role in ["tool", "function"]:
             prompt += f"TOOL RESULT (for {msg.tool_call_id or msg.name}): {msg.content}\n\n"
             
     prompt += "Assistant: "
@@ -64,7 +76,7 @@ def format_messages_and_tools(messages: List[ChatMessage], tools: Optional[List[
 
 def extract_tool_calls(text: str):
     """Parses Grok's response to check if it emitted our forced JSON tool call"""
-    # Look for a markdown JSON block
+    # Look for a markdown JSON block, or fall back to raw text
     match = re.search(r'```json\s*(.*?)\s*```', text, re.DOTALL)
     json_str = match.group(1) if match else text.strip()
     
@@ -78,7 +90,7 @@ def extract_tool_calls(text: str):
                     "type": "function",
                     "function": {
                         "name": tc.get("name"),
-                        "arguments": json.dumps(tc.get("arguments", {})) # OpenAI expects a stringified JSON here
+                        "arguments": json.dumps(tc.get("arguments", {})) # OpenAI expects a stringified JSON
                     }
                 })
             return formatted_calls
@@ -92,24 +104,42 @@ async def chat_completions(request: ChatCompletionRequest):
     # 1. Prepare Prompt
     mega_prompt = format_messages_and_tools(request.messages, request.tools)
     
+    # 2. Check for Proxy in Environment Variables
+    # If Hugging Face IPs are blocked by Cloudflare, setting this in HF Secrets fixes it.
+    proxy_url = os.environ.get("GROK_PROXY", None)
+    
     try:
-        # 2. Call the underlying Grok Wrapper (Stateless, passing entire context in prompt)
-        grok_client = Grok(request.model)
+        # 3. Call Grok
+        if proxy_url:
+            grok_client = Grok(request.model, proxy_url)
+        else:
+            grok_client = Grok(request.model)
+            
         raw_response = grok_client.start_convo(mega_prompt)
         
         response_text = raw_response.get("response", "")
         stream_array = raw_response.get("stream_response",[])
         
+    except UnboundLocalError as e:
+        # This catches the specific "local variable 'script_content1' referenced before assignment" error
+        error_msg = (
+            "Grok API Scraper failed to find session tokens. "
+            "This usually means Hugging Face's IP is blocked by Grok's Cloudflare, or Grok updated their website DOM. "
+            "Fix: Add a 'GROK_PROXY' secret in your HF Space settings (e.g., http://user:pass@ip:port)."
+        )
+        print(f"Scraper Error: {traceback.format_exc()}")
+        raise HTTPException(status_code=500, detail=error_msg)
     except Exception as e:
-        raise HTTPException(status_code=500, detail=str(e))
+        print(f"Unknown Error: {traceback.format_exc()}")
+        raise HTTPException(status_code=500, detail=f"Upstream Grok Error: {str(e)}")
 
-    # 3. Check if response is a tool call
+    # 4. Parse Tool Calls
     tool_calls = extract_tool_calls(response_text) if request.tools else None
 
-    # 4. Handle Streaming Response
+    # 5. Handle Streaming Response
     if request.stream:
         async def event_generator():
-            # If it's a tool call, we typically don't stream it, but send it as one chunk
+            # Tool calls are emitted as one chunk to prevent breaking JSON parsers in agents
             if tool_calls:
                 chunk = {
                     "id": f"chatcmpl-{uuid.uuid4()}",
@@ -119,18 +149,18 @@ async def chat_completions(request: ChatCompletionRequest):
                 }
                 yield {"data": json.dumps(chunk)}
             else:
-                # Fake the stream using the token array returned by the API
+                # Simulate streaming using the token array
                 for token in stream_array:
                     chunk = {
                         "id": f"chatcmpl-{uuid.uuid4()}",
                         "object": "chat.completion.chunk",
                         "model": request.model,
-                        "choices":[{"index": 0, "delta": {"content": token}, "finish_reason": None}]
+                        "choices": [{"index": 0, "delta": {"content": token}, "finish_reason": None}]
                     }
                     yield {"data": json.dumps(chunk)}
-                    time.sleep(0.01) # slight delay to emulate natural streaming
+                    time.sleep(0.01) # Small delay for smooth streaming
                 
-                # Final finish reason chunk
+                # Final STOP chunk
                 yield {
                     "data": json.dumps({
                         "id": f"chatcmpl-{uuid.uuid4()}",
@@ -143,7 +173,7 @@ async def chat_completions(request: ChatCompletionRequest):
             
         return EventSourceResponse(event_generator())
 
-    # 5. Handle Standard Sync Response
+    # 6. Handle Standard Sync Response
     response_msg = {"role": "assistant", "content": None if tool_calls else response_text}
     if tool_calls:
         response_msg["tool_calls"] = tool_calls