@rulemetric/proxy 0.2.2

package/addon/providers.py

@@ -0,0 +1,483 @@
+"""Provider detection and request body parsing for mitmproxy addon.
+
+Mirrors the TypeScript provider adapters. Uses only Python stdlib.
+"""
+
+from __future__ import annotations
+
+import hashlib
+import json
+import re
+from typing import Any
+
+from context_classifier import classify_context
+
+# ---------------------------------------------------------------------------
+# Provider detection
+# ---------------------------------------------------------------------------
+
+_KNOWN_OPENAI_COMPAT_HOSTS: dict[str, str] = {
+    "api.openai.com": "openai",
+    "api.groq.com": "groq",
+    "api.together.xyz": "together",
+    "api.fireworks.ai": "fireworks",
+    "api.deepseek.com": "deepseek",
+    "api.perplexity.ai": "perplexity",
+    "api.cerebras.ai": "cerebras",
+    "api.sambanova.ai": "sambanova",
+    "api.x.ai": "xai",
+    "api.mistral.ai": "mistral",
+    "integrate.api.nvidia.com": "nvidia_nim",
+    "api.ai21.com": "ai21",
+    "api.cohere.com": "cohere",
+    "openrouter.ai": "openrouter",
+    "openai-proxy.replicate.com": "replicate",
+    "api.githubcopilot.com": "github_copilot_chat",
+    "api.individual.githubcopilot.com": "github_copilot_chat",
+    "api.business.githubcopilot.com": "github_copilot_chat",
+    "api.enterprise.githubcopilot.com": "github_copilot_chat",
+}
+
+
+def detect_provider(host: str, path: str) -> str | None:
+    """Return a canonical provider name based on the request host and path.
+
+    Specific providers with non-OpenAI formats are matched first.
+    Then known OpenAI-compatible hosts are checked.
+    Finally, any request to a ``chat/completions`` path is treated as
+    a generic OpenAI-compatible provider, with the name derived from the host.
+    """
+    # --- Non-OpenAI formats (must match first) ---
+    if host == "api.anthropic.com" and "/v1/messages" in path:
+        return "anthropic"
+    if host == "generativelanguage.googleapis.com" and ":generateContent" in path:
+        return "gemini"
+    if "bedrock-runtime" in host and "/converse" in path:
+        return "bedrock"
+    if host.endswith(".snowflakecomputing.com") and "/api/v2/cortex/inference:complete" in path:
+        return "cortex"
+
+    # --- GitHub Copilot FIM completions (non-chat format) ---
+    if host == "copilot-proxy.githubusercontent.com":
+        if "/v1/engines/" in path and "/completions" in path:
+            return "github_copilot_fim"
+        if "/chat/completions" in path:
+            return "github_copilot_chat"
+        return None  # telemetry or other traffic — skip
+
+    # --- Azure OpenAI (host-based, before generic catch-all) ---
+    if host.endswith(".openai.azure.com") and "/chat/completions" in path:
+        return "azure_openai"
+
+    # --- Azure AI Model Catalog ---
+    if host.endswith(".models.ai.azure.com") and "/chat/completions" in path:
+        return "azure_ai"
+
+    # --- Cloudflare Workers AI (OpenAI-compat mode) ---
+    if host == "api.cloudflare.com" and "/ai/v1/chat/completions" in path:
+        return "cloudflare"
+
+    # --- Databricks Foundation Model API ---
+    if host.endswith(".cloud.databricks.com") and "/invocations" in path:
+        return "databricks"
+
+    # NOTE: VS Code Copilot also POSTs to `/models/session` and
+    # `/agents/session` on api.githubcopilot.com — these are session-init /
+    # model-list polling endpoints, NOT chat completions. Routing them
+    # through parse_openai produces empty `messages: []` snapshots and 0-len
+    # assistant_response events that pollute the dashboard. Detection for
+    # those endpoints is intentionally removed; only `/chat/completions`
+    # gets captured (handled by the _KNOWN_OPENAI_COMPAT_HOSTS block below).
+
+    # --- Known OpenAI-compatible hosts ---
+    for known_host, provider_name in _KNOWN_OPENAI_COMPAT_HOSTS.items():
+        if host == known_host or host.endswith("." + known_host):
+            if "chat/completions" in path or "/completions" in path:
+                return provider_name
+
+    # --- Ollama native API ---
+    if "localhost" in host and "/api/chat" in path:
+        return "ollama"
+
+    # --- Ollama OpenAI-compat endpoint ---
+    if "localhost" in host and "/v1/chat/completions" in path:
+        return "ollama"
+
+    # --- Generic OpenAI-compatible catch-all ---
+    if "chat/completions" in path:
+        return _provider_from_host(host)
+
+    return None
+
+
+def _provider_from_host(host: str) -> str:
+    """Derive a short provider name from a hostname.
+
+    Examples:
+        api.foo.com       → foo
+        foo.example.com   → foo
+        localhost:1234    → localhost
+    """
+    # Strip port
+    h = host.split(":")[0]
+    parts = h.split(".")
+    if len(parts) >= 2:
+        # Use the second-to-last part (e.g., "foo" from "api.foo.com")
+        # Unless it's a known TLD part like "co" in "api.foo.co.uk"
+        return parts[-2] if parts[-2] not in ("co", "com", "net", "org", "io", "ai") else parts[-3] if len(parts) >= 3 else parts[0]
+    return parts[0]
+
+
+# ---------------------------------------------------------------------------
+# System-reminder extraction & hashing
+# ---------------------------------------------------------------------------
+
+_SYSTEM_REMINDER_RE = re.compile(
+    r"<system-reminder>(.*?)</system-reminder>", re.DOTALL
+)
+
+
+def extract_system_reminders(text: str) -> list[str]:
+    """Extract all fully-closed <system-reminder> blocks from *text*."""
+    return _SYSTEM_REMINDER_RE.findall(text)
+
+
+def compute_system_hash(text: str) -> str:
+    """Return the SHA-256 hex digest of *text*."""
+    return hashlib.sha256(text.encode("utf-8")).hexdigest()
+
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+def _estimate_tokens(text: str) -> int:
+    """Rough token estimate (~4 chars per token)."""
+    return max(1, len(text) // 4)
+
+
+def _flatten_text(parts: Any) -> str:
+    """Collapse a string-or-list-of-blocks into a single string."""
+    if isinstance(parts, str):
+        return parts
+    if isinstance(parts, list):
+        return "\n".join(
+            block.get("text", "") if isinstance(block, dict) else str(block)
+            for block in parts
+        )
+    return str(parts) if parts else ""
+
+
+def _base_snapshot() -> dict[str, Any]:
+    return {
+        "provider": "unknown",
+        "model": "unknown",
+        "system_prompt": "",
+        "system_blocks": [],
+        "messages": [],
+        "messages_raw": [],  # Full raw message objects from the API request
+        "tools": [],
+        "max_tokens": None,
+        "temperature": None,
+        "top_p": None,
+        "top_k": None,
+        "stop_sequences": [],
+        "stream": False,
+        "tool_choice": None,
+        "thinking": None,
+        "injected_instructions": [],
+        "system_hash": "",
+        "request_hash": "",  # Hash of the entire request (system + messages + tools)
+        "message_count": 0,
+        "tool_count": 0,
+        "estimated_tokens": 0,
+    }
+
+
+def _finalize(snapshot: dict[str, Any]) -> dict[str, Any]:
+    """Compute derived fields on a snapshot dict."""
+    snapshot["message_count"] = len(snapshot["messages"])
+    snapshot["tool_count"] = len(snapshot["tools"])
+
+    blob = json.dumps(
+        {
+            "system_prompt": snapshot["system_prompt"],
+            "messages": snapshot["messages"],
+            "tools": snapshot["tools"],
+        },
+        separators=(",", ":"),
+    )
+    snapshot["estimated_tokens"] = _estimate_tokens(blob)
+
+    if snapshot["system_prompt"]:
+        snapshot["system_hash"] = compute_system_hash(snapshot["system_prompt"])
+
+    # Hash the entire request for per-turn dedup
+    full_blob = json.dumps(
+        {
+            "system_prompt": snapshot["system_prompt"],
+            "messages_raw": snapshot["messages_raw"],
+            "tools": snapshot["tools"],
+        },
+        separators=(",", ":"),
+        sort_keys=True,
+    )
+    snapshot["request_hash"] = compute_system_hash(full_blob)
+
+    # Extract system-reminders from ALL messages (not just system prompt)
+    all_injected = []
+    if snapshot["system_prompt"]:
+        all_injected.extend(extract_system_reminders(snapshot["system_prompt"]))
+    for msg in snapshot["messages"]:
+        content = msg.get("content", "")
+        if content:
+            all_injected.extend(extract_system_reminders(content))
+    if all_injected:
+        snapshot["injected_instructions"] = all_injected
+
+    # Classify all context items
+    snapshot["context_items"] = classify_context(snapshot)
+
+    return snapshot
+
+
+# ---------------------------------------------------------------------------
+# Anthropic
+# ---------------------------------------------------------------------------
+
+def parse_anthropic(body: dict[str, Any]) -> dict[str, Any]:
+    """Parse an Anthropic ``/v1/messages`` request body."""
+    snap = _base_snapshot()
+    snap["provider"] = "anthropic"
+    snap["model"] = body.get("model", "unknown")
+    snap["max_tokens"] = body.get("max_tokens")
+    snap["temperature"] = body.get("temperature")
+    snap["top_p"] = body.get("top_p")
+    snap["top_k"] = body.get("top_k")
+    snap["stop_sequences"] = body.get("stop_sequences", [])
+    snap["stream"] = body.get("stream", False)
+    snap["tool_choice"] = body.get("tool_choice")
+    snap["thinking"] = body.get("thinking")
+    # Metadata (user_id, etc.)
+    snap["api_metadata"] = body.get("metadata")
+
+    # System — string or array of content blocks
+    raw_system = body.get("system", "")
+    if isinstance(raw_system, list):
+        snap["system_blocks"] = raw_system
+        snap["system_prompt"] = _flatten_text(raw_system)
+    else:
+        snap["system_prompt"] = str(raw_system)
+        if raw_system:
+            snap["system_blocks"] = [{"type": "text", "text": raw_system}]
+
+    # Messages — keep both normalized (for display) and raw (for fidelity)
+    snap["messages_raw"] = body.get("messages", [])
+    for msg in body.get("messages", []):
+        snap["messages"].append({
+            "role": msg.get("role", "user"),
+            "content": _flatten_text(msg.get("content", "")),
+        })
+
+    # Tools — rename input_schema → parameters
+    for tool in body.get("tools", []):
+        normalized: dict[str, Any] = {"name": tool.get("name", "")}
+        if "description" in tool:
+            normalized["description"] = tool["description"]
+        if "input_schema" in tool:
+            normalized["parameters"] = tool["input_schema"]
+        snap["tools"].append(normalized)
+
+    return _finalize(snap)
+
+
+# ---------------------------------------------------------------------------
+# OpenAI (and Azure OpenAI)
+# ---------------------------------------------------------------------------
+
+def parse_openai(body: dict[str, Any]) -> dict[str, Any]:
+    """Parse an OpenAI ``/v1/chat/completions`` request body."""
+    snap = _base_snapshot()
+    snap["provider"] = "openai"
+    snap["model"] = body.get("model", "unknown")
+    snap["max_tokens"] = body.get("max_completion_tokens") or body.get("max_tokens")
+    snap["temperature"] = body.get("temperature")
+    snap["top_p"] = body.get("top_p")
+    snap["stop_sequences"] = body.get("stop") if isinstance(body.get("stop"), list) else ([body["stop"]] if body.get("stop") else [])
+    snap["stream"] = body.get("stream", False)
+    snap["tool_choice"] = body.get("tool_choice")
+    snap["api_metadata"] = {"user": body.get("user")} if body.get("user") else None
+
+    system_parts: list[str] = []
+    developer_parts: list[str] = []
+    messages: list[dict[str, str]] = []
+
+    snap["messages_raw"] = body.get("messages", [])
+    for msg in body.get("messages", []):
+        role = msg.get("role", "user")
+        content = _flatten_text(msg.get("content", ""))
+        if role == "system":
+            system_parts.append(content)
+        elif role == "developer":
+            developer_parts.append(content)
+        else:
+            messages.append({"role": role, "content": content})
+
+    snap["system_prompt"] = "\n".join(system_parts)
+    snap["injected_instructions"] = developer_parts
+    snap["messages"] = messages
+
+    # Tools — unwrap function wrapper
+    for tool in body.get("tools", []):
+        func = tool.get("function", tool)
+        normalized: dict[str, Any] = {"name": func.get("name", "")}
+        if "description" in func:
+            normalized["description"] = func["description"]
+        if "parameters" in func:
+            normalized["parameters"] = func["parameters"]
+        snap["tools"].append(normalized)
+
+    return _finalize(snap)
+
+
+# ---------------------------------------------------------------------------
+# Gemini
+# ---------------------------------------------------------------------------
+
+def parse_gemini(body: dict[str, Any]) -> dict[str, Any]:
+    """Parse a Gemini ``generateContent`` request body."""
+    snap = _base_snapshot()
+    snap["provider"] = "gemini"
+    snap["model"] = body.get("model", "unknown")
+
+    gen_config = body.get("generationConfig", {})
+    snap["max_tokens"] = gen_config.get("maxOutputTokens")
+    snap["temperature"] = gen_config.get("temperature")
+
+    # System instruction
+    sys_instr = body.get("systemInstruction", {})
+    parts = sys_instr.get("parts", [])
+    sys_texts = [p.get("text", "") for p in parts if isinstance(p, dict)]
+    snap["system_prompt"] = "\n".join(sys_texts)
+
+    # Contents
+    snap["messages_raw"] = body.get("contents", [])
+    for content in body.get("contents", []):
+        role = content.get("role", "user")
+        if role == "model":
+            role = "assistant"
+        text_parts = content.get("parts", [])
+        text = "\n".join(
+            p.get("text", "") for p in text_parts if isinstance(p, dict) and "text" in p
+        )
+        snap["messages"].append({"role": role, "content": text})
+
+    # Tools — functionDeclarations
+    for tool_group in body.get("tools", []):
+        for decl in tool_group.get("functionDeclarations", []):
+            normalized: dict[str, Any] = {"name": decl.get("name", "")}
+            if "description" in decl:
+                normalized["description"] = decl["description"]
+            if "parameters" in decl:
+                normalized["parameters"] = decl["parameters"]
+            snap["tools"].append(normalized)
+
+    return _finalize(snap)
+
+
+# ---------------------------------------------------------------------------
+# Bedrock
+# ---------------------------------------------------------------------------
+
+def parse_bedrock(body: dict[str, Any]) -> dict[str, Any]:
+    """Parse an AWS Bedrock ``/converse`` request body."""
+    snap = _base_snapshot()
+    snap["provider"] = "bedrock"
+    snap["model"] = body.get("modelId", "unknown")
+
+    inf_config = body.get("inferenceConfig", {})
+    snap["max_tokens"] = inf_config.get("maxTokens")
+    snap["temperature"] = inf_config.get("temperature")
+
+    # System — list of text blocks
+    sys_blocks = body.get("system", [])
+    sys_texts = [
+        b.get("text", "") for b in sys_blocks if isinstance(b, dict) and "text" in b
+    ]
+    snap["system_prompt"] = "\n".join(sys_texts)
+    snap["system_blocks"] = sys_blocks
+
+    # Messages
+    snap["messages_raw"] = body.get("messages", [])
+    for msg in body.get("messages", []):
+        role = msg.get("role", "user")
+        content_blocks = msg.get("content", [])
+        text = "\n".join(
+            b.get("text", "")
+            for b in content_blocks
+            if isinstance(b, dict) and "text" in b
+        )
+        snap["messages"].append({"role": role, "content": text})
+
+    # Tools — toolConfig.tools[].toolSpec
+    tool_config = body.get("toolConfig", {})
+    for tool in tool_config.get("tools", []):
+        spec = tool.get("toolSpec", {})
+        normalized: dict[str, Any] = {"name": spec.get("name", "")}
+        if "description" in spec:
+            normalized["description"] = spec["description"]
+        input_schema = spec.get("inputSchema", {})
+        if "json" in input_schema:
+            normalized["parameters"] = input_schema["json"]
+        snap["tools"].append(normalized)
+
+    return _finalize(snap)
+
+
+# ---------------------------------------------------------------------------
+# GitHub Copilot FIM (Fill-in-the-Middle)
+# ---------------------------------------------------------------------------
+
+def parse_copilot_fim(body: dict[str, Any]) -> dict[str, Any]:
+    """Parse a Copilot FIM ``/v1/engines/.../completions`` request body.
+
+    FIM requests use ``prompt`` + ``suffix`` instead of ``messages``.
+    """
+    snap = _base_snapshot()
+    snap["provider"] = "github_copilot_fim"
+    snap["model"] = body.get("model", "unknown")
+    snap["max_tokens"] = body.get("max_tokens")
+    snap["temperature"] = body.get("temperature")
+    snap["top_p"] = body.get("top_p")
+    snap["stop_sequences"] = body.get("stop", [])
+    snap["stream"] = body.get("stream", False)
+
+    # Prompt is the code context before the cursor
+    prompt = body.get("prompt", "")
+    snap["system_prompt"] = prompt
+
+    # Suffix is the code after the cursor
+    suffix = body.get("suffix", "")
+    if suffix:
+        snap["messages"].append({"role": "context", "content": suffix})
+
+    # Raw body for hashing
+    snap["messages_raw"] = [{"prompt": prompt, "suffix": suffix}]
+
+    # Extra metadata (language, repo, token counts)
+    extra = body.get("extra", {})
+    metadata: dict[str, Any] = {}
+    if extra.get("language"):
+        metadata["language"] = extra["language"]
+    if body.get("nwo"):
+        metadata["repo"] = body["nwo"]
+    if extra.get("prompt_tokens"):
+        metadata["prompt_tokens"] = extra["prompt_tokens"]
+    if extra.get("suffix_tokens"):
+        metadata["suffix_tokens"] = extra["suffix_tokens"]
+    if body.get("n"):
+        metadata["n"] = body["n"]
+    if metadata:
+        snap["api_metadata"] = metadata
+
+    return _finalize(snap)