server.py

"""
Server — FastAPI OpenAI-compatible proxy that translates between
the OpenAI ChatCompletion protocol and the Gemini web API.

This is the core translation middleware for BrowserOS integration.
"""

import json
import logging
import re
import time
import uuid
from typing import Any, AsyncGenerator

from fastapi import FastAPI, Request
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse, StreamingResponse

from gemini_client import GeminiClientWrapper
from schema_sanitizer import sanitize_tools
from json_fixer import fix_tool_arguments

logger = logging.getLogger(__name__)

# ─── App Setup ───────────────────────────────────────────────────────────────

app = FastAPI(
    title="Gemini-BrowserOS Gateway",
    description="OpenAI-compatible proxy for Gemini web API → BrowserOS agentic workflows",
    version="1.0.0",
)

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# Will be set by run.py at startup
gemini: GeminiClientWrapper | None = None
config: dict[str, Any] = {}


# ─── Models Endpoint ─────────────────────────────────────────────────────────

@app.get("/v1/models")
async def list_models():
    """Return available models in OpenAI format."""
    models = GeminiClientWrapper.get_models()
    return {"object": "list", "data": models}


# ─── Chat Completions Endpoint ───────────────────────────────────────────────

@app.post("/v1/chat/completions")
async def chat_completions(request: Request):
    """
    Main proxy endpoint — accepts OpenAI ChatCompletion requests,
    translates to Gemini, and returns OpenAI-formatted responses.
    """
    body = await request.json()

    model = body.get("model", config.get("gemini", {}).get("default_model", "gemini-3-flash"))
    messages = body.get("messages", [])
    tools = body.get("tools", [])
    tool_choice = body.get("tool_choice", "auto")
    stream = body.get("stream", False)

    # Log incoming request summary
    logger.info("━" * 60)
    logger.info("INCOMING REQUEST — model=%s, stream=%s, tools=%d, tool_choice=%s",
                model, stream, len(tools), tool_choice)
    for i, msg in enumerate(messages[-3:]):  # Log last 3 messages
        role = msg.get("role", "?")
        content = msg.get("content", "")
        tc = msg.get("tool_calls")
        if tc:
            logger.info("  msg[%d] role=%s tool_calls=%s", len(messages) - 3 + i, role, json.dumps(tc)[:200])
        else:
            content_preview = (content[:150] + "...") if isinstance(content, str) and len(content) > 150 else content
            if isinstance(content_preview, list):
                content_preview = "[multimodal]"
            logger.info("  msg[%d] role=%s content=%s", len(messages) - 3 + i, role, content_preview)
            
    if tools:
        logger.debug("FULL TOOLS PAYLOAD: %s", json.dumps(tools, indent=2))

    # Extract system prompt
    system_parts = []
    conversation_parts = []
    for msg in messages:
        role = msg.get("role", "")
        content = msg.get("content", "")
        if role == "system":
            system_parts.append(content if isinstance(content, str) else json.dumps(content))
        else:
            conversation_parts.append(msg)

    # Sanitize tool schemas for Gemini compatibility
    tool_schemas = {}
    if tools:
        sanitized = sanitize_tools(tools)
        for tool in sanitized:
            if tool.get("type") == "function" and "function" in tool:
                func = tool["function"]
                tool_schemas[func["name"]] = func.get("parameters", {})

    # Build prompt
    prompt = _build_prompt(conversation_parts, tools, tool_choice, system_parts)
    logger.debug("PROMPT (last 500 chars): ...%s", prompt[-500:])

    try:
        assert gemini is not None
        logger.debug("Calling gemini.generate with prompt length: %d", len(prompt))
        result = await gemini.generate(prompt=prompt, model=model)
        logger.debug("gemini.generate returned result keys: %s", list(result.keys()))
    except Exception as e:
        logger.error("Gemini generation failed: %s", e)
        return JSONResponse(
            status_code=502,
            content={
                "error": {
                    "message": f"Gemini backend error: {str(e)}",
                    "type": "upstream_error",
                    "code": 502,
                }
            },
        )

    text = result.get("text", "")
    thoughts = result.get("thoughts")

    # Log raw model output
    logger.info("RAW MODEL OUTPUT (%d chars): %s", len(text), text[:500])

    # Detect and parse function calls in the response
    tool_calls = _extract_tool_calls(text, tool_schemas)

    logger.info("EXTRACTED TOOL CALLS: %d — %s",
                len(tool_calls),
                json.dumps([tc["function"]["name"] for tc in tool_calls]) if tool_calls else "none")

    if stream:
        return StreamingResponse(
            _stream_response(text, model, tool_calls, thoughts),
            media_type="text/event-stream",
        )

    resp = _build_completion_response(text, model, tool_calls, thoughts)
    logger.info("RESPONSE finish_reason=%s",
                "tool_calls" if tool_calls else "stop")
    logger.info("━" * 60)
    return resp


# ─── Prompt Construction ─────────────────────────────────────────────────────

_TOOL_SYSTEM_PROMPT = """You are a strict proxy agent meant to translate a single user request into exactly one tool call if an action is needed, or a plain text response if the action is already complete.

RESPONSE FORMAT:
If you need to take an action, you MUST use this EXACT format:
<tool_call>{"name": "TOOL_NAME", "arguments": {ARGS}}</tool_call>

CRITICAL RULES:
1. EXECUTE ONLY THE SINGLE IMMEDIATELY REQUESTED ACTION. Do NOT try to complete a full workflow.
2. If the user's requested action is already complete or the page is already in the correct state, DO NOT output any <tool_call> tags. Instead, reply directly to the user in plain text.
3. Replying in plain text without a <tool_call> will stop the automation loop. Do this immediately after the single requested action is observed.
4. Do NOT explore the page, take extra actions, or try to be "helpful" by completing unrequested steps.
5. ALL argument values MUST match their specified types EXACTLY.
6. Do NOT wrap tool calls in markdown code blocks. Use ONLY raw <tool_call> tags."""


def _build_prompt(
    messages: list[dict],
    tools: list[dict],
    tool_choice: str | dict,
    system_parts: list[str],
) -> str:
    """
    Compose a single prompt string from OpenAI messages, tool defs, and system prompt.
    The Gemini web API doesn't natively support structured tool schemas,
    so we embed them as structured instructions in the prompt.
    """
    parts = []

    # System instructions — inject tool-calling enforcement first
    if tools:
        parts.append(_TOOL_SYSTEM_PROMPT)

    if system_parts:
        parts.append("## Additional System Instructions\n" + "\n".join(system_parts))

    # Tool definitions
    if tools:
        sanitized = sanitize_tools(tools)
        tool_block = _format_tools_as_prompt(sanitized, tool_choice)
        parts.append(tool_block)

    # Conversation history
    for msg in messages:
        role = msg.get("role", "user")
        content = msg.get("content", "")

        if role == "tool":
            # Tool response — format as structured feedback
            tool_call_id = msg.get("tool_call_id", "unknown")
            parts.append(f"[Tool Result for call {tool_call_id}]\n{content}")
        elif role == "assistant":
            # Check for tool_calls in assistant messages
            if "tool_calls" in msg and msg["tool_calls"]:
                tc_strs = []
                for tc in msg["tool_calls"]:
                    func = tc.get("function", {})
                    tc_strs.append(
                        f'<tool_call>{{"name": "{func.get("name", "")}", '
                        f'"arguments": {func.get("arguments", "{}")}}}</tool_call>'
                    )
                parts.append(f"Assistant:\n{''.join(tc_strs)}")
            elif content:
                parts.append(f"Assistant: {content}")
        elif role == "user":
            if isinstance(content, list):
                # Multimodal content (text + images)
                text_parts = []
                for item in content:
                    if isinstance(item, dict) and item.get("type") == "text":
                        text_parts.append(item.get("text", ""))
                    elif isinstance(item, dict) and item.get("type") == "image_url":
                        text_parts.append("[Image provided]")
                parts.append(f"User: {' '.join(text_parts)}")
            else:
                parts.append(f"User: {content}")

    # Final reinforcement for tool-calling
    if tools:
        parts.append(
            "CRITICAL INSTRUCTION: You are a DUMB API proxy. Execute ONLY the specific action requested by the user. "
            "If the user asked to navigate somewhere, navigate and then STOP immediately by replying in plain text. "
            "DO NOT click buttons, DO NOT fill forms, and DO NOT log in unless explicitly instructed in the User prompt. "
            "Remember: respond with <tool_call> tags IF an action is needed, otherwise reply in plain text to stop."
        )

    return "\n\n".join(parts)


def _format_tools_as_prompt(tools: list[dict], tool_choice: str | dict) -> str:
    """Format sanitized tool definitions as a structured prompt block."""
    lines = [
        "## Available Tools",
        "",
    ]

    for tool in tools:
        if tool.get("type") != "function" or "function" not in tool:
            continue
        func = tool["function"]
        name = func.get("name", "unknown")
        desc = func.get("description", "")
        params = func.get("parameters", {})
        lines.append(f"### Tool: {name}")
        if desc:
            lines.append(f"Description: {desc}")
        if params.get("properties"):
            lines.append(f"Parameters: {json.dumps(params, indent=2)}")
        lines.append("")

    # Tool choice instructions
    if isinstance(tool_choice, str):
        if tool_choice in ("required", "any"):
            lines.append("⚠️ You are strongly encouraged to call one of the tools above if an action is needed.")
        elif tool_choice == "none":
            lines.append("Do NOT use any tools. Respond with plain text only.")
    elif isinstance(tool_choice, dict):
        forced_name = tool_choice.get("function", {}).get("name", "")
        if forced_name:
            lines.append(f"⚠️ MANDATORY: You MUST call the '{forced_name}' tool.")

    return "\n".join(lines)


# ─── Tool Call Extraction ─────────────────────────────────────────────────────

# Multiple patterns the model might use to express tool calls
_TOOL_CALL_PATTERNS = [
    # Standard: <tool_call>{"name": ..., "arguments": ...}</tool_call>
    re.compile(r'<tool_call>\s*(.*?)\s*</tool_call>', re.DOTALL),
    # Variants: ```tool_call or ```json with name/arguments structure
    re.compile(r'```(?:tool_call|json)?\s*\n?\s*(\{[^`]*?"name"\s*:\s*"[^"]+?"[^`]*?\})\s*```', re.DOTALL),
    # Bare JSON with name+arguments keys (last resort)
    re.compile(r'(\{\s*"name"\s*:\s*"[^"]+?"\s*,\s*"arguments"\s*:\s*\{[^}]*?\}\s*\})', re.DOTALL),
]


def _extract_tool_calls(
    text: str, tool_schemas: dict[str, dict]
) -> list[dict[str, Any]]:
    """
    Parse tool calls from model output using multiple patterns.
    Apply JSON repair + schema coercion on arguments.
    """
    tool_calls = []
    matches = []

    # Pre-clean: strip markdown formatting from the raw text
    # Gemini web app embeds markdown links [text](url) in tool arguments
    cleaned_text = re.sub(r'\[([^\]]*?)\]\(([^)]+?)\)', r'\2', text)
    # Strip bold/italic markdown
    cleaned_text = re.sub(r'\*{1,2}([^*]+?)\*{1,2}', r'\1', cleaned_text)

    # Try patterns in order of specificity
    for pattern in _TOOL_CALL_PATTERNS:
        found = pattern.findall(cleaned_text)
        if found:
            matches = found
            logger.debug("Matched tool calls with pattern: %s", pattern.pattern[:60])
            break

    # If still no matches, try to detect if the model output looks like
    # it's trying to call a known tool but without proper formatting
    if not matches and tool_schemas:
        for tool_name in tool_schemas:
            # Check if the model mentioned a tool name with arguments nearby
            loose_pattern = re.compile(
                rf'["\']?{re.escape(tool_name)}["\']?\s*[:=]\s*\{{(.*?)\}}',
                re.DOTALL
            )
            loose_match = loose_pattern.search(cleaned_text)
            if loose_match:
                # Construct a synthetic tool call
                args_str = loose_match.group(1)
                synthetic = json.dumps({"name": tool_name, "arguments": "{" + args_str + "}"})
                matches.append(synthetic)
                logger.info("Recovered loose tool call for '%s' from unstructured text", tool_name)

    repair_cfg = config.get("json_repair", {})
    schema_mode = repair_cfg.get("schema_mode", "salvage")
    repair_enabled = repair_cfg.get("enabled", True)

    for match in matches:
        match = match.strip()
        logger.debug("Processing tool call match: %s", match[:200])
        try:
            # Try direct parse first
            try:
                call_data = json.loads(match)
            except json.JSONDecodeError:
                from json_repair import repair_json
                call_data = repair_json(match, return_objects=True)

            if not isinstance(call_data, dict):
                logger.warning("Tool call parsed to non-dict: %s", type(call_data))
                continue

            name = call_data.get("name", "")
            if not name:
                logger.warning("Tool call missing 'name' field: %s", match[:100])
                continue

            raw_args = call_data.get("arguments", {})

            # If arguments is a string, repair it
            if isinstance(raw_args, str):
                schema = tool_schemas.get(name, {}) if repair_enabled else None
                raw_args = fix_tool_arguments(raw_args, schema, schema_mode)
            elif isinstance(raw_args, dict) and repair_enabled:
                # Even if it's already a dict, run type coercion
                schema = tool_schemas.get(name)
                if schema:
                    raw_args = fix_tool_arguments(
                        json.dumps(raw_args), schema, schema_mode
                    )

            tool_calls.append({
                "id": f"call_{uuid.uuid4().hex[:24]}",
                "type": "function",
                "function": {
                    "name": name,
                    "arguments": json.dumps(raw_args) if isinstance(raw_args, dict) else str(raw_args),
                },
            })
            logger.info("✅ Parsed tool call: %s(%s)", name,
                        json.dumps(raw_args)[:100] if isinstance(raw_args, dict) else str(raw_args)[:100])

        except Exception as e:
            logger.warning("Failed to parse tool call: %s — %s", match[:100], e)
            continue

    return tool_calls


# ─── Response Formatting ─────────────────────────────────────────────────────

def _build_completion_response(
    text: str,
    model: str,
    tool_calls: list[dict],
    thoughts: str | None = None,
) -> JSONResponse:
    """Build a standard OpenAI ChatCompletion JSON response."""
    completion_id = f"chatcmpl-{uuid.uuid4().hex[:29]}"
    created = int(time.time())

    # Build the message
    message: dict[str, Any] = {"role": "assistant"}

    if tool_calls:
        # Strip tool_call tags and JSON blocks from content
        clean_text = re.sub(r'<tool_call>.*?</tool_call>', '', text, flags=re.DOTALL).strip()
        clean_text = re.sub(r'```(?:json|tool_call)?\s*\n?\{.*?\}\s*```', '', clean_text, flags=re.DOTALL).strip()
        # Also strip bare JSON tool-call-like blocks
        clean_text = re.sub(r'\{\s*"name"\s*:.*?"arguments"\s*:.*?\}', '', clean_text, flags=re.DOTALL).strip()
        message["content"] = clean_text if clean_text else None
        message["tool_calls"] = tool_calls
        finish_reason = "tool_calls"
    else:
        message["content"] = text
        finish_reason = "stop"

    response = {
        "id": completion_id,
        "object": "chat.completion",
        "created": created,
        "model": model,
        "choices": [
            {
                "index": 0,
                "message": message,
                "finish_reason": finish_reason,
            }
        ],
        "usage": {
            "prompt_tokens": 0,
            "completion_tokens": 0,
            "total_tokens": 0,
        },
    }

    # Include thoughts as a custom field
    if thoughts:
        response["_thoughts"] = thoughts

    return JSONResponse(content=response)


async def _stream_response(
    text: str,
    model: str,
    tool_calls: list[dict],
    thoughts: str | None = None,
) -> AsyncGenerator[str, None]:
    """Yield SSE chunks in OpenAI ChatCompletionChunk format."""
    completion_id = f"chatcmpl-{uuid.uuid4().hex[:29]}"
    created = int(time.time())

    if tool_calls:
        # Emit tool calls as a single chunk
        chunk = {
            "id": completion_id,
            "object": "chat.completion.chunk",
            "created": created,
            "model": model,
            "choices": [
                {
                    "index": 0,
                    "delta": {
                        "role": "assistant",
                        "content": None,
                        "tool_calls": [
                            {
                                "index": i,
                                "id": tc["id"],
                                "type": "function",
                                "function": tc["function"],
                            }
                            for i, tc in enumerate(tool_calls)
                        ],
                    },
                    "finish_reason": None,
                }
            ],
        }
        yield f"data: {json.dumps(chunk)}\n\n"

        # Final chunk with finish reason
        final_chunk = {
            "id": completion_id,
            "object": "chat.completion.chunk",
            "created": created,
            "model": model,
            "choices": [
                {
                    "index": 0,
                    "delta": {},
                    "finish_reason": "tool_calls",
                }
            ],
        }
        yield f"data: {json.dumps(final_chunk)}\n\n"
    else:
        # Stream text in chunks
        clean_text = text
        chunk_size = 20  # characters per chunk

        # First chunk with role
        first_chunk = {
            "id": completion_id,
            "object": "chat.completion.chunk",
            "created": created,
            "model": model,
            "choices": [
                {
                    "index": 0,
                    "delta": {"role": "assistant", "content": ""},
                    "finish_reason": None,
                }
            ],
        }
        yield f"data: {json.dumps(first_chunk)}\n\n"

        # Content chunks
        for i in range(0, len(clean_text), chunk_size):
            chunk_text = clean_text[i : i + chunk_size]
            chunk = {
                "id": completion_id,
                "object": "chat.completion.chunk",
                "created": created,
                "model": model,
                "choices": [
                    {
                        "index": 0,
                        "delta": {"content": chunk_text},
                        "finish_reason": None,
                    }
                ],
            }
            yield f"data: {json.dumps(chunk)}\n\n"

        # Final chunk
        final_chunk = {
            "id": completion_id,
            "object": "chat.completion.chunk",
            "created": created,
            "model": model,
            "choices": [
                {
                    "index": 0,
                    "delta": {},
                    "finish_reason": "stop",
                }
            ],
        }
        yield f"data: {json.dumps(final_chunk)}\n\n"

    yield "data: [DONE]\n\n"


# ─── Health Check ─────────────────────────────────────────────────────────────

@app.get("/health")
async def health():
    return {"status": "ok", "timestamp": int(time.time())}