server/api.py

"""FastAPI app wiring Copilot onto the OpenAI Chat Completions API."""

import os
import threading
import time

from typing import Optional

from fastapi import FastAPI
from fastapi.responses import HTMLResponse, JSONResponse, StreamingResponse

from copilot import CopilotClient
from copilot.auth import SESSION_ENV_VAR, _load_session_from_env

from .config import MODEL_NAME, RATE_LIMIT_BURST, RATE_LIMIT_RPM
from .openai_format import (
    completion_response,
    new_id,
    sse_event,
    stream_chunk,
)
from .prompt import messages_to_prompt
from .ratelimit import TokenBucket
from .schemas import ChatCompletionRequest

app = FastAPI(title="Copilot OpenAI-compatible API", version="1.0.0")

# The client is created lazily on first use so that server/__init__ can inject an
# anonymous one when ANONYMOUS=true is set.  If no injection happens, the
# default signed-in client is used (requires `python -m copilot login`).
_client: Optional[CopilotClient] = None


def _get_client() -> CopilotClient:
    global _client
    if _client is not None:
        return _client
    # Check if an anonymous client was injected via app.state (set by server/__init__).
    injected = getattr(app.state, "copilot_client", None)
    if injected is not None:
        _client = injected
        return _client
    _client = CopilotClient()
    return _client

# Self-imposed rate limit on top of the concurrency lock below: this caps
# requests-per-minute, the lock caps requests-in-flight. See server/ratelimit.py.
_rate_limiter = TokenBucket(RATE_LIMIT_RPM, RATE_LIMIT_BURST)


def _rate_limited_response():
    """Spend a token; return an OpenAI-shaped 429 if none left, else ``None``."""
    allowed, wait = _rate_limiter.try_acquire()
    if allowed:
        return None
    secs = max(1, round(wait))
    return JSONResponse(
        status_code=429,
        headers={"Retry-After": str(secs)},
        content={"error": {
            "message": (
                f"Rate limit exceeded (>{RATE_LIMIT_RPM:g} req/min). "
                f"Retry in {secs}s."
            ),
            "type": "rate_limit_error",
            "code": "rate_limit_exceeded",
        }},
    )

# Copilot's per-account chat socket doesn't tolerate concurrent conversations
# from one process (parallel requests error out or hang). This server bridges a
# single signed-in account, so we serialize upstream calls: concurrent HTTP
# requests queue here and run one at a time. Predictable, at the cost of
# parallelism — fine for a personal bridge.
_upstream_lock = threading.Lock()


def _stream(prompt: str, model: str, conversation_id=None):
    """Yield OpenAI ``chat.completion.chunk`` SSE events for ``prompt``.

    ``conversation_id`` continues an existing Copilot thread; ``None`` starts a
    fresh one (its id is emitted on the final chunk).
    """
    cid = new_id()
    created = int(time.time())
    try:
        with _upstream_lock:  # one upstream chat at a time (released on disconnect)
            yield sse_event(stream_chunk(cid, created, model, {"role": "assistant"}))
            stream = _get_client().stream(prompt, conversation_id=conversation_id)
            for piece in stream:
                if isinstance(piece, str) and piece:
                    yield sse_event(stream_chunk(cid, created, model, {"content": piece}))
            # Copilot's conversation id is known once the stream has run; emit it
            # on the final chunk so callers can track the upstream thread.
            yield sse_event(
                stream_chunk(
                    cid, created, model, {}, finish="stop",
                    conversation_id=stream.conversation_id,
                )
            )
    except Exception as exc:  # surface errors to the client instead of hanging
        yield sse_event(
            stream_chunk(cid, created, model, {"content": f"\n[error: {exc}]"}, finish="error")
        )
    yield "data: [DONE]\n\n"


@app.get("/v1/models")
def list_models():
    return {
        "object": "list",
        "data": [
            {"id": MODEL_NAME, "object": "model", "created": 0, "owned_by": "microsoft"}
        ],
    }


@app.post("/v1/chat/completions")
def chat_completions(req: ChatCompletionRequest):
    prompt = messages_to_prompt(req.messages)
    if not prompt.strip():
        return JSONResponse(
            status_code=400,
            content={"error": {"message": "no text content in messages", "type": "invalid_request_error"}},
        )
    model = req.model or MODEL_NAME

    # Enforce the per-minute ceiling before touching the upstream lock, so excess
    # callers get a fast 429 instead of piling up behind the serialized queue.
    limited = _rate_limited_response()
    if limited is not None:
        return limited

    if req.stream:
        return StreamingResponse(
            _stream(prompt, model, req.conversation_id), media_type="text/event-stream"
        )

    try:
        with _upstream_lock:  # serialize: one upstream chat at a time
            reply = _get_client().chat(prompt, conversation_id=req.conversation_id)
    except Exception as exc:
        return JSONResponse(
            status_code=502,
            content={"error": {"message": str(exc), "type": "upstream_error"}},
        )
    return completion_response(reply.text, model, reply.conversation_id)


@app.get("/", response_class=HTMLResponse)
def root():
    """Status landing page — shows auth mode and quick-start instructions."""
    env_session = _load_session_from_env()
    anonymous = os.environ.get("ANONYMOUS", "").lower() in ("1", "true", "yes")

    if env_session is not None:
        auth_mode = "Signed-in (via COPILOT_SESSION)"
        auth_color = "#22c55e"
        auth_detail = "Session loaded from environment variable / HF Secret."
    elif anonymous:
        auth_mode = "Anonymous"
        auth_color = "#eab308"
        auth_detail = (
            "Running without sign-in. May be geo-blocked in some regions. "
            f"Set <code>{SESSION_ENV_VAR}</code> to use your signed-in account."
        )
    else:
        auth_mode = "Signed-in (local session)"
        auth_color = "#22c55e"
        auth_detail = "Session from local browser login."

    return f"""<!DOCTYPE html>
<html lang="en">
<head>
<meta charset="utf-8"/>
<meta name="viewport" content="width=device-width, initial-scale=1"/>
<title>Copilot OpenAI-compatible API</title>
<style>
  body {{ font-family: -apple-system, BlinkMacSystemFont, "Segoe UI", Helvetica, Arial, sans-serif;
         max-width: 720px; margin: 40px auto; padding: 0 20px; color: #1f2937; line-height: 1.6; }}
  h1 {{ border-bottom: 2px solid #3b82f6; padding-bottom: 8px; }}
  code {{ background: #f3f4f6; padding: 2px 6px; border-radius: 4px; font-size: 90%; }}
  pre {{ background: #1e293b; color: #e2e8f0; padding: 16px; border-radius: 8px; overflow-x: auto; }}
  .badge {{ display: inline-block; padding: 4px 12px; border-radius: 9999px;
            font-weight: 600; font-size: 14px; color: #fff; background: {auth_color}; }}
  .endpoint {{ margin: 8px 0; padding: 8px 12px; background: #f9fafb; border-radius: 6px; }}
  .method {{ font-weight: 700; color: #3b82f6; margin-right: 8px; }}
</style>
</head>
<body>
  <h1>Copilot OpenAI-compatible API</h1>
  <p>Auth mode: <span class="badge">{auth_mode}</span></p>
  <p style="color:#6b7280;font-size:14px;">{auth_detail}</p>

  <h2>Endpoints</h2>
  <div class="endpoint"><span class="method">POST</span><code>/v1/chat/completions</code> &mdash; Chat completion (supports streaming)</div>
  <div class="endpoint"><span class="method">GET</span><code>/v1/models</code> &mdash; List available models</div>

  <h2>Quick start</h2>
  <pre>from openai import OpenAI

client = OpenAI(
    base_url="YOUR_SPACE_URL/v1",
    api_key="unused",
)

resp = client.chat.completions.create(
    model="copilot",
    messages=[{{"role": "user", "content": "Hello!"}}],
)
print(resp.choices[0].message.content)</pre>

  <h2>Sign in with your account</h2>
  <p>To use your signed-in Microsoft Copilot account on this Space:</p>
  <ol>
    <li>Run this project locally and sign in:<br><code>python -m copilot login</code></li>
    <li>Copy the contents of <code>session/token.json</code></li>
    <li>Go to your Space <b>Settings &rarr; Variables and secrets</b></li>
    <li>Add a new Secret named <code>{SESSION_ENV_VAR}</code> and paste the JSON</li>
    <li>If <code>ANONYMOUS=true</code> was set, remove it (or set <code>ANONYMOUS=false</code>)</li>
    <li>Restart the Space</li>
  </ol>
  <p style="color:#6b7280;font-size:13px;">Note: Microsoft access tokens expire after ~60-90 minutes. When the token expires you will need to re-login locally and update the Secret. For long-running use, consider running the project on a host where you can keep a persistent browser profile.</p>
</body>
</html>"""