Create openai_server.py

openai_server.py

@@ -0,0 +1,178 @@
+import json
+import re
+import time
+import uuid
+from typing import List, Optional, Dict, Any
+
+from fastapi import FastAPI, Request, HTTPException
+from fastapi.responses import JSONResponse
+from sse_starlette.sse import EventSourceResponse
+from pydantic import BaseModel, Field
+
+# Import the core from the cloned Grok-Api repository
+from core import Grok
+
+app = FastAPI(title="Grok OpenAI Wrapper (Agent & MCP Compatible)")
+
+# --- OpenAI Pydantic Models ---
+class ChatMessage(BaseModel):
+    role: str
+    content: Optional[str] = None
+    name: Optional[str] = None
+    tool_calls: Optional[List[Dict[Any, Any]]] = None
+    tool_call_id: Optional[str] = None
+
+class ChatCompletionRequest(BaseModel):
+    model: str = "grok-3-fast"
+    messages: List[ChatMessage]
+    stream: Optional[bool] = False
+    tools: Optional[List[Dict[Any, Any]]] = None
+    tool_choice: Optional[Any] = None
+    temperature: Optional[float] = 0.7
+
+# --- 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"""
+    prompt = ""
+    
+    # 1. Inject Tools via System Prompt Strategy if tools exist
+    if tools:
+        prompt += (
+            "SYSTEM INSTRUCTION: You are an intelligent AI acting as an API. You have access to tools. "
+            "If you need to call a tool, you MUST reply ONLY with a JSON block in the exact format below, and no other text.\n"
+            '```json\n{"tool_calls":[{"name": "function_name", "arguments": {"arg_name": "arg_value"}}]}\n```\n'
+            "Available tools:\n" + json.dumps(tools, indent=2) + "\n\n"
+        )
+
+    # 2. Append Message History
+    for msg in messages:
+        if msg.role == "system":
+            prompt += f"System: {msg.content}\n\n"
+        elif msg.role == "user":
+            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"
+            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
+            prompt += f"TOOL RESULT (for {msg.tool_call_id or msg.name}): {msg.content}\n\n"
+            
+    prompt += "Assistant: "
+    return prompt
+
+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
+    match = re.search(r'```json\s*(.*?)\s*```', text, re.DOTALL)
+    json_str = match.group(1) if match else text.strip()
+    
+    try:
+        parsed = json.loads(json_str)
+        if "tool_calls" in parsed and isinstance(parsed["tool_calls"], list):
+            formatted_calls = []
+            for tc in parsed["tool_calls"]:
+                formatted_calls.append({
+                    "id": f"call_{uuid.uuid4().hex[:8]}",
+                    "type": "function",
+                    "function": {
+                        "name": tc.get("name"),
+                        "arguments": json.dumps(tc.get("arguments", {})) # OpenAI expects a stringified JSON here
+                    }
+                })
+            return formatted_calls
+    except json.JSONDecodeError:
+        pass
+    return None
+
+# --- API Endpoints ---
+@app.post("/v1/chat/completions")
+async def chat_completions(request: ChatCompletionRequest):
+    # 1. Prepare Prompt
+    mega_prompt = format_messages_and_tools(request.messages, request.tools)
+    
+    try:
+        # 2. Call the underlying Grok Wrapper (Stateless, passing entire context in prompt)
+        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 Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
+    # 3. Check if response is a tool call
+    tool_calls = extract_tool_calls(response_text) if request.tools else None
+
+    # 4. 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
+            if tool_calls:
+                chunk = {
+                    "id": f"chatcmpl-{uuid.uuid4()}",
+                    "object": "chat.completion.chunk",
+                    "model": request.model,
+                    "choices":[{"index": 0, "delta": {"tool_calls": tool_calls}, "finish_reason": "tool_calls"}]
+                }
+                yield {"data": json.dumps(chunk)}
+            else:
+                # Fake the stream using the token array returned by the API
+                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}]
+                    }
+                    yield {"data": json.dumps(chunk)}
+                    time.sleep(0.01) # slight delay to emulate natural streaming
+                
+                # Final finish reason chunk
+                yield {
+                    "data": json.dumps({
+                        "id": f"chatcmpl-{uuid.uuid4()}",
+                        "object": "chat.completion.chunk",
+                        "model": request.model,
+                        "choices":[{"index": 0, "delta": {}, "finish_reason": "stop"}]
+                    })
+                }
+            yield {"data": "[DONE]"}
+            
+        return EventSourceResponse(event_generator())
+
+    # 5. 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
+
+    return JSONResponse(content={
+        "id": f"chatcmpl-{uuid.uuid4()}",
+        "object": "chat.completion",
+        "created": int(time.time()),
+        "model": request.model,
+        "choices":[{
+            "index": 0,
+            "message": response_msg,
+            "finish_reason": "tool_calls" if tool_calls else "stop"
+        }],
+        "usage": {
+            "prompt_tokens": len(mega_prompt) // 4,
+            "completion_tokens": len(response_text) // 4,
+            "total_tokens": (len(mega_prompt) + len(response_text)) // 4
+        }
+    })
+
+@app.get("/v1/models")
+async def list_models():
+    return {
+        "object": "list",
+        "data":[
+            {"id": "grok-3-auto", "object": "model", "created": int(time.time()), "owned_by": "xai"},
+            {"id": "grok-3-fast", "object": "model", "created": int(time.time()), "owned_by": "xai"},
+            {"id": "grok-4", "object": "model", "created": int(time.time()), "owned_by": "xai"},
+            {"id": "grok-4-mini-thinking-tahoe", "object": "model", "created": int(time.time()), "owned_by": "xai"}
+        ]
+    }