#!/usr/bin/env python3 """OpenAI-compatible proxy backed by the Mistral key manager. The proxy accepts OpenAI-style `/v1/*` requests, chooses a usable key from the key manager, forwards the request to Mistral, and reports bad keys so the pool can rotate away from them. """ from __future__ import annotations import argparse import json import os import random import re import sys import threading import time import traceback import urllib.error import urllib.parse import urllib.request import uuid from dataclasses import dataclass from datetime import datetime, timezone from http import HTTPStatus from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer from typing import Any, Dict, List, Optional, Tuple DEFAULT_KEY_MANAGER_URL = "http://127.0.0.1:8097" DEFAULT_MISTRAL_API_BASE = "https://api.mistral.ai/v1" DEFAULT_CONVERSATION_MODELS = "glm-5-2" HOP_BY_HOP_HEADERS = { "connection", "keep-alive", "proxy-authenticate", "proxy-authorization", "te", "trailer", "transfer-encoding", "upgrade", } OPENAI_ROOT_PATHS = { "/models", "/chat/completions", # Mistral-native endpoint with OpenAI-style request normalization. "/conversations", "/responses", "/embeddings", "/moderations", "/rerank", } BAD_KEY_STATUS_CODES = {401, 403} RETRYABLE_STATUS_CODES = {408, 409, 429, 500, 502, 503, 504} # Mistral can use 401/403 for a model entitlement failure. That says nothing # about the API key itself and must not remove the key from the shared pool. MODEL_ACCESS_ERROR_PATTERNS = ( re.compile(r"\bthis model is not available in your subscription\b", re.I), re.compile(r"\bmodel\b.*\bnot available\b.*\bsubscription\b", re.I | re.S), re.compile(r"\bsubscription\b.*\bmodel\b.*\bnot available\b", re.I | re.S), re.compile(r"\bupgrade\b.*\b(subscription|plan|tier)\b", re.I | re.S), re.compile(r"\bdoes not have access to (the )?model\b", re.I), ) CONVERSATION_COMPLETION_ARG_FIELDS = { "stop", "presence_penalty", "frequency_penalty", "temperature", "top_p", "max_tokens", "random_seed", "prediction", "response_format", "tool_choice", "reasoning_effort", } CONVERSATION_COMPLETION_ARG_ALIASES = { # OpenAI's newer name maps to Mistral's supported completion argument. "max_completion_tokens": "max_tokens", "seed": "random_seed", } def utc_now() -> str: return datetime.now(timezone.utc).isoformat(timespec="seconds") def json_bytes(data: Any) -> bytes: return json.dumps(data, ensure_ascii=False, separators=(",", ":")).encode("utf-8") def is_model_access_error(status_code: Optional[int], error: str) -> bool: """Return whether an auth-looking response is about model entitlement.""" if status_code not in BAD_KEY_STATUS_CODES: return False text = str(error or "") return any(pattern.search(text) for pattern in MODEL_ACCESS_ERROR_PATTERNS) def read_env_int(name: str, default: int) -> int: raw = os.getenv(name, "").strip() if not raw: return default try: return int(raw) except ValueError: return default def read_env_float(name: str, default: float) -> float: raw = os.getenv(name, "").strip() if not raw: return default try: return float(raw) except ValueError: return default def read_env_bool(name: str, default: bool = False) -> bool: raw = os.getenv(name) if raw is None or raw.strip() == "": return default return raw.strip().lower() in {"1", "true", "yes", "y", "on"} def parse_model_list(raw: str) -> frozenset[str]: return frozenset(model.strip() for model in raw.split(",") if model.strip()) def default_max_concurrency() -> int: cpu_count = os.cpu_count() or 8 return max(4, min(16, cpu_count)) def mask_key(api_key: str) -> str: clean = api_key.strip() if len(clean) <= 10: return "*" * len(clean) return f"{clean[:6]}...{clean[-4:]}" class HTTPClientError(RuntimeError): def __init__(self, message: str, *, status_code: Optional[int] = None, body: bytes = b"") -> None: super().__init__(message) self.status_code = status_code self.body = body def parse_retry_after(value: str) -> Optional[float]: clean = (value or "").strip() if not clean: return None try: return max(0.0, min(30.0, float(clean))) except ValueError: return None def is_stream_request(headers: Dict[str, str], body: bytes) -> bool: accept = headers.get("accept", "") if "text/event-stream" in accept.lower(): return True if not body: return False try: payload = json.loads(body.decode("utf-8")) except Exception: return False return bool(payload.get("stream")) def normalize_proxy_path(path: str) -> str: parsed = urllib.parse.urlsplit(path) route = parsed.path if route.startswith("/v1/"): route = route[3:] elif route == "/v1": route = "/" if not route.startswith("/"): route = "/" + route return urllib.parse.urlunsplit(("", "", route, parsed.query, "")) def is_proxy_path(path: str) -> bool: if path == "/v1" or path.startswith("/v1/"): return True return path in OPENAI_ROOT_PATHS or any(path.startswith(prefix + "/") for prefix in OPENAI_ROOT_PATHS) def is_conversation_start_path(method: str, path: str) -> bool: return method.upper() == "POST" and urllib.parse.urlsplit(normalize_proxy_path(path)).path == "/conversations" def openai_instruction_text(content: Any) -> str: if isinstance(content, str): return content if isinstance(content, list): parts: List[str] = [] for part in content: if isinstance(part, str): parts.append(part) continue if not isinstance(part, dict): raise HTTPClientError( "system and developer message content must contain text", status_code=HTTPStatus.BAD_REQUEST, ) if part.get("type") not in {"text", "input_text", "output_text"} or not isinstance(part.get("text"), str): raise HTTPClientError( "system and developer message content must contain text", status_code=HTTPStatus.BAD_REQUEST, ) parts.append(part["text"]) return "".join(parts) raise HTTPClientError( "system and developer message content must be text", status_code=HTTPStatus.BAD_REQUEST, ) def openai_tool_result_content(content: Any) -> str: if isinstance(content, str): return content if content is None: return "" return json.dumps(content, ensure_ascii=False, separators=(",", ":")) def translate_openai_messages(messages: List[Any]) -> Tuple[List[Dict[str, Any]], List[str]]: inputs: List[Dict[str, Any]] = [] instruction_parts: List[str] = [] for message in messages: if not isinstance(message, dict): raise HTTPClientError( "OpenAI-compatible messages must contain objects", status_code=HTTPStatus.BAD_REQUEST, ) role = message.get("role") if role in {"system", "developer"}: instruction_parts.append(openai_instruction_text(message.get("content"))) continue if role == "tool": tool_call_id = message.get("tool_call_id") if not isinstance(tool_call_id, str) or not tool_call_id: raise HTTPClientError( "tool messages require tool_call_id", status_code=HTTPStatus.BAD_REQUEST, ) inputs.append( { "tool_call_id": tool_call_id, "result": openai_tool_result_content(message.get("content")), } ) continue if role == "assistant" and isinstance(message.get("tool_calls"), list): content = message.get("content") if content is not None and content != "": inputs.append({"role": "assistant", "content": content}) for tool_call in message["tool_calls"]: if not isinstance(tool_call, dict): raise HTTPClientError( "assistant tool_calls must contain objects", status_code=HTTPStatus.BAD_REQUEST, ) function = tool_call.get("function") tool_call_id = tool_call.get("id") if not isinstance(function, dict) or not isinstance(tool_call_id, str) or not tool_call_id: raise HTTPClientError( "assistant tool_calls require id and function", status_code=HTTPStatus.BAD_REQUEST, ) name = function.get("name") if not isinstance(name, str) or not name: raise HTTPClientError( "assistant tool_calls require function.name", status_code=HTTPStatus.BAD_REQUEST, ) arguments = function.get("arguments", "") if not isinstance(arguments, str): arguments = json.dumps(arguments, ensure_ascii=False, separators=(",", ":")) inputs.append( { "tool_call_id": tool_call_id, "name": name, "arguments": arguments, } ) continue if role not in {"user", "assistant"}: raise HTTPClientError( f"unsupported OpenAI message role for Conversations: {role!r}", status_code=HTTPStatus.BAD_REQUEST, ) inputs.append({"role": role, "content": message.get("content")}) return inputs, instruction_parts def translate_openai_tool_choice(payload: Dict[str, Any]) -> None: tool_choice = payload.get("tool_choice") if not isinstance(tool_choice, dict): return function = tool_choice.get("function") name = function.get("name") if isinstance(function, dict) else None if tool_choice.get("type") != "function" or not isinstance(name, str) or not name: raise HTTPClientError( "only OpenAI function tool_choice is supported", status_code=HTTPStatus.BAD_REQUEST, ) tools = payload.get("tools") if not isinstance(tools, list): raise HTTPClientError( "named tool_choice requires a tools array", status_code=HTTPStatus.BAD_REQUEST, ) selected_tools = [ tool for tool in tools if isinstance(tool, dict) and isinstance(tool.get("function"), dict) and tool["function"].get("name") == name ] if not selected_tools: raise HTTPClientError( f"tool_choice references unknown function: {name}", status_code=HTTPStatus.BAD_REQUEST, ) # Mistral only supports auto/none/required. Restricting the tool list # preserves OpenAI's named-function behavior. payload["tools"] = selected_tools payload["tool_choice"] = "required" def translate_openai_conversation_body(body: bytes) -> bytes: """Convert OpenAI chat-style fields into Mistral Conversations request fields. A request using native `inputs` is deliberately left unchanged. A request using `messages` opts into compatibility mode, where completion parameters are moved into Mistral's `completion_args` object. """ if not body: return body try: payload = json.loads(body.decode("utf-8")) except (UnicodeDecodeError, json.JSONDecodeError) as exc: raise HTTPClientError( f"invalid JSON for /v1/conversations: {exc}", status_code=HTTPStatus.BAD_REQUEST, ) from exc if not isinstance(payload, dict): raise HTTPClientError( "/v1/conversations JSON body must be an object", status_code=HTTPStatus.BAD_REQUEST, ) if "messages" not in payload: return body if "inputs" in payload: raise HTTPClientError( "use either OpenAI-compatible messages or native Mistral inputs, not both", status_code=HTTPStatus.BAD_REQUEST, ) messages = payload.pop("messages") # This is an OpenAI streaming control field. Mistral Conversations rejects # it, and the proxy handles include_usage when translating the response. payload.pop("stream_options", None) payload.pop("parallel_tool_calls", None) if not isinstance(messages, list): raise HTTPClientError( "OpenAI-compatible messages must be an array", status_code=HTTPStatus.BAD_REQUEST, ) inputs, instruction_parts = translate_openai_messages(messages) if instruction_parts: existing_instructions = payload.get("instructions") if existing_instructions is not None and not isinstance(existing_instructions, str): raise HTTPClientError( "instructions must be a string", status_code=HTTPStatus.BAD_REQUEST, ) if isinstance(existing_instructions, str) and existing_instructions: instruction_parts.append(existing_instructions) payload["instructions"] = "\n\n".join(part for part in instruction_parts if part) translate_openai_tool_choice(payload) completion_args_raw = payload.pop("completion_args", None) if completion_args_raw is None: completion_args: Dict[str, Any] = {} elif isinstance(completion_args_raw, dict): completion_args = dict(completion_args_raw) else: raise HTTPClientError( "completion_args must be an object", status_code=HTTPStatus.BAD_REQUEST, ) completion_fields = set(CONVERSATION_COMPLETION_ARG_FIELDS) completion_fields.update(CONVERSATION_COMPLETION_ARG_ALIASES) for source_field in completion_fields: if source_field not in payload: continue target_field = CONVERSATION_COMPLETION_ARG_ALIASES.get(source_field, source_field) if target_field in completion_args: raise HTTPClientError( f"{target_field} was supplied both at the top level and in completion_args", status_code=HTTPStatus.BAD_REQUEST, ) completion_args[target_field] = payload.pop(source_field) payload["inputs"] = inputs if completion_args: payload["completion_args"] = completion_args return json_bytes(payload) def request_model_id(body: bytes) -> str: try: payload = json.loads(body.decode("utf-8")) except (UnicodeDecodeError, json.JSONDecodeError): return "" if not isinstance(payload, dict): return "" model = payload.get("model") return model.strip() if isinstance(model, str) else "" def openai_stream_include_usage(body: bytes) -> bool: """Return OpenAI's requested final usage chunk, if present.""" try: payload = json.loads(body.decode("utf-8")) except (UnicodeDecodeError, json.JSONDecodeError): return False if not isinstance(payload, dict): return False stream_options = payload.get("stream_options") return isinstance(stream_options, dict) and stream_options.get("include_usage") is True def is_openai_chat_completions_path(method: str, path: str) -> bool: return method.upper() == "POST" and urllib.parse.urlsplit(normalize_proxy_path(path)).path == "/chat/completions" def merge_conversation_output_content(contents: List[Any]) -> Any: if not contents: return None if all(isinstance(content, str) for content in contents): return "".join(contents) if len(contents) == 1: return contents[0] merged: List[Any] = [] for content in contents: if isinstance(content, list): merged.extend(content) elif isinstance(content, str): merged.append({"type": "text", "text": content}) else: merged.append(content) return merged def conversation_text_content(content: Any) -> Optional[str]: if isinstance(content, str): return content if isinstance(content, dict) and content.get("type") == "text" and isinstance(content.get("text"), str): return content["text"] if isinstance(content, list): parts = [conversation_text_content(part) for part in content] if all(part is not None for part in parts): return "".join(part for part in parts if part is not None) return None def conversation_thinking_content_to_openai(content: Any) -> Optional[str]: if not isinstance(content, dict) or content.get("type") != "thinking": return None thinking = conversation_text_content(content.get("thinking")) return thinking if thinking is not None else "" def split_conversation_content(content: Any) -> Tuple[str, str]: """Separate Mistral visible text from its thinking blocks.""" thinking = conversation_thinking_content_to_openai(content) if thinking is not None: return "", thinking text = conversation_text_content(content) if text is not None: return text, "" if isinstance(content, list): visible_parts: List[str] = [] reasoning_parts: List[str] = [] for part in content: visible, reasoning = split_conversation_content(part) if visible: visible_parts.append(visible) if reasoning: reasoning_parts.append(reasoning) return "".join(visible_parts), "".join(reasoning_parts) # The local OpenAI clients use string content. Unknown structured content # is not a valid text completion, so do not leak it into message.content. return "", "" def conversation_usage_to_openai(usage: Any) -> Optional[Dict[str, int]]: if not isinstance(usage, dict): return None return { "prompt_tokens": int(usage.get("prompt_tokens") or 0), "completion_tokens": int(usage.get("completion_tokens") or 0), "total_tokens": int(usage.get("total_tokens") or 0), } def conversation_response_to_openai_chat_response(response: "ProxyResponse", model: str) -> "ProxyResponse": if response.status >= 400: return response try: payload = json.loads(response.body.decode("utf-8")) except (UnicodeDecodeError, json.JSONDecodeError): return response if not isinstance(payload, dict): return response conversation_id = payload.get("conversation_id") outputs = payload.get("outputs") if not isinstance(conversation_id, str) or not isinstance(outputs, list): return response contents: List[str] = [] reasoning_contents: List[str] = [] tool_calls: List[Dict[str, Any]] = [] for output in outputs: if not isinstance(output, dict): continue output_type = output.get("type") if output_type == "message.output" or ( output_type is None and output.get("role") == "assistant" and "content" in output ): if output.get("content") is not None: visible_content, reasoning_content = split_conversation_content(output["content"]) if visible_content: contents.append(visible_content) if reasoning_content: reasoning_contents.append(reasoning_content) elif output_type == "function.call": arguments = output.get("arguments", "") if not isinstance(arguments, str): arguments = json.dumps(arguments, ensure_ascii=False, separators=(",", ":")) tool_calls.append( { "id": str(output.get("tool_call_id") or output.get("id") or uuid.uuid4().hex), "type": "function", "function": { "name": str(output.get("name") or ""), "arguments": arguments, }, } ) message: Dict[str, Any] = { "role": "assistant", "content": "".join(contents), } if tool_calls: message["tool_calls"] = tool_calls if reasoning_contents: message["reasoning_content"] = "".join(reasoning_contents) usage = conversation_usage_to_openai(payload.get("usage")) or {} body = json_bytes( { "id": f"chatcmpl-{conversation_id}", "object": "chat.completion", "created": int(time.time()), "model": model, "choices": [ { "index": 0, "message": message, "finish_reason": "tool_calls" if tool_calls else "stop", } ], "usage": usage, } ) return ProxyResponse( status=response.status, reason=response.reason, headers=response.headers, body=body, ) def conversation_stream_content_to_openai(content: Any) -> Optional[str]: visible_content, _ = split_conversation_content(content) return visible_content or None @dataclass class KeyCandidate: api_key: str item: Dict[str, Any] @dataclass class ProxyResponse: status: int reason: str headers: List[Tuple[str, str]] body: bytes def retry_after(self) -> Optional[float]: for key, value in self.headers: if key.lower() == "retry-after": return parse_retry_after(value) return None @dataclass(frozen=True) class RoutedRequest: path: str body: bytes openai_chat_model: str = "" openai_stream_include_usage: bool = False @dataclass class UpstreamStream: status: int reason: str headers: List[Tuple[str, str]] response: Any class KeyManagerClient: def __init__( self, base_url: str, token: str, timeout_sec: float, cache_ttl_sec: float, stale_cache_ttl_sec: float, refresh_backoff_sec: float, ) -> None: self.base_url = base_url.rstrip("/") self.token = token.strip() self.timeout_sec = max(1.0, timeout_sec) self.cache_ttl_sec = max(0.0, cache_ttl_sec) self.stale_cache_ttl_sec = max(self.cache_ttl_sec, stale_cache_ttl_sec) self.refresh_backoff_sec = max(0.5, refresh_backoff_sec) self._active_cache: List[KeyCandidate] = [] self._active_cache_loaded = False self._active_cache_expires_at = 0.0 self._active_cache_stale_until = 0.0 self._active_cache_retry_at = 0.0 self._active_cache_lock = threading.Lock() self._active_cache_refresh_lock = threading.Lock() self._last_stale_cache_warning_at = 0.0 def headers(self) -> Dict[str, str]: result = {"Accept": "application/json", "User-Agent": "yunwo-mistral-openai-proxy/1.0"} if self.token: result["X-Admin-Token"] = self.token return result def request_json(self, method: str, path: str, payload: Optional[Dict[str, Any]] = None) -> Dict[str, Any]: data = json_bytes(payload) if payload is not None else None headers = self.headers() if data is not None: headers["Content-Type"] = "application/json" request = urllib.request.Request(f"{self.base_url}{path}", data=data, headers=headers, method=method) try: with urllib.request.urlopen(request, timeout=self.timeout_sec) as response: raw = response.read().decode("utf-8", errors="replace") except urllib.error.HTTPError as exc: raw = exc.read(2000).decode("utf-8", errors="replace") raise HTTPClientError(f"key manager HTTP {exc.code}: {raw}", status_code=exc.code) from exc except TimeoutError as exc: raise HTTPClientError(f"key manager timed out after {self.timeout_sec:g}s") from exc except urllib.error.URLError as exc: raise HTTPClientError(f"key manager unreachable: {exc.reason}") from exc return json.loads(raw or "{}") def fetch_active_keys(self) -> List[KeyCandidate]: envelope = self.request_json("GET", "/api/keys/active") if not envelope.get("success"): raise HTTPClientError(str(envelope.get("error") or "key manager returned failure")) keys = [str(item or "").strip() for item in envelope.get("keys") or [] if str(item or "").strip()] items = list(envelope.get("items") or []) candidates: List[KeyCandidate] = [] for index, api_key in enumerate(keys): item = items[index] if index < len(items) and isinstance(items[index], dict) else {} candidates.append(KeyCandidate(api_key=api_key, item=item)) return candidates def active_keys(self, *, force_refresh: bool = False) -> List[KeyCandidate]: if self.cache_ttl_sec <= 0: return self.fetch_active_keys() now = time.monotonic() with self._active_cache_lock: cached = list(self._active_cache) cache_loaded = self._active_cache_loaded cache_fresh = cache_loaded and now < self._active_cache_expires_at refresh_deferred = cache_loaded and now < self._active_cache_retry_at and now < self._active_cache_stale_until if not force_refresh and (cache_fresh or refresh_deferred): return cached with self._active_cache_refresh_lock: now = time.monotonic() with self._active_cache_lock: cached = list(self._active_cache) cache_loaded = self._active_cache_loaded cache_fresh = cache_loaded and now < self._active_cache_expires_at refresh_deferred = cache_loaded and now < self._active_cache_retry_at and now < self._active_cache_stale_until if not force_refresh and (cache_fresh or refresh_deferred): return cached try: candidates = self.fetch_active_keys() except Exception as exc: with self._active_cache_lock: cached = list(self._active_cache) cache_loaded = self._active_cache_loaded stale_cache_usable = cache_loaded and now < self._active_cache_stale_until if stale_cache_usable: self._active_cache_retry_at = time.monotonic() + self.refresh_backoff_sec if stale_cache_usable: self.warn_stale_cache(exc) return cached raise refreshed_at = time.monotonic() with self._active_cache_lock: self._active_cache = list(candidates) self._active_cache_loaded = True self._active_cache_expires_at = refreshed_at + self.cache_ttl_sec self._active_cache_stale_until = refreshed_at + self.stale_cache_ttl_sec self._active_cache_retry_at = 0.0 return candidates def warn_stale_cache(self, exc: Exception) -> None: now = time.monotonic() with self._active_cache_lock: if now < self._last_stale_cache_warning_at + 30.0: return self._last_stale_cache_warning_at = now sys.stderr.write(f"[{utc_now()}] key manager unavailable; using stale active-key cache: {exc}\n") def drop_cached_key(self, api_key: str) -> None: with self._active_cache_lock: self._active_cache = [candidate for candidate in self._active_cache if candidate.api_key != api_key] def report(self, api_key: str, status_code: Optional[int], error: str, *, status: str = "") -> None: payload = { "api_key": api_key, "status_code": status_code, "error": error[:1000], } if status: payload["status"] = status try: self.request_json("POST", "/api/keys/report", payload) except Exception as exc: sys.stderr.write(f"[{utc_now()}] key report failed for {mask_key(api_key)}: {exc}\n") class MistralForwarder: def __init__(self, base_url: str, proxy: str, timeout_sec: float) -> None: self.base_url = base_url.rstrip("/") self.proxy = proxy.strip() self.timeout_sec = max(1.0, timeout_sec) self._opener = self.build_opener() def build_opener(self) -> urllib.request.OpenerDirector: if not self.proxy: return urllib.request.build_opener() return urllib.request.build_opener(urllib.request.ProxyHandler({"http": self.proxy, "https": self.proxy})) def opener(self) -> urllib.request.OpenerDirector: return self._opener def forward( self, *, method: str, path: str, headers: Dict[str, str], body: bytes, api_key: str, stream: bool, ) -> ProxyResponse: url = f"{self.base_url}{normalize_proxy_path(path)}" outgoing = self.forward_headers(headers, api_key, stream) data = body if method.upper() not in {"GET", "HEAD"} else None request = urllib.request.Request(url, data=data, headers=outgoing, method=method) try: with self.opener().open(request, timeout=self.timeout_sec) as response: return ProxyResponse( status=response.status, reason=response.reason, headers=list(response.headers.items()), body=response.read(), ) except urllib.error.HTTPError as exc: return ProxyResponse( status=exc.code, reason=exc.reason, headers=list(exc.headers.items()), body=exc.read(), ) except TimeoutError as exc: raise HTTPClientError(f"upstream timed out after {self.timeout_sec:g}s") from exc except urllib.error.URLError as exc: raise HTTPClientError(f"upstream unreachable: {exc.reason}") from exc def open_stream( self, *, method: str, path: str, headers: Dict[str, str], body: bytes, api_key: str, ) -> UpstreamStream | ProxyResponse: url = f"{self.base_url}{normalize_proxy_path(path)}" outgoing = self.forward_headers(headers, api_key, True) data = body if method.upper() not in {"GET", "HEAD"} else None request = urllib.request.Request(url, data=data, headers=outgoing, method=method) try: response = self.opener().open(request, timeout=self.timeout_sec) return UpstreamStream( status=response.status, reason=response.reason, headers=list(response.headers.items()), response=response, ) except urllib.error.HTTPError as exc: return ProxyResponse( status=exc.code, reason=exc.reason, headers=list(exc.headers.items()), body=exc.read(), ) except TimeoutError as exc: raise HTTPClientError(f"upstream timed out after {self.timeout_sec:g}s") from exc except urllib.error.URLError as exc: raise HTTPClientError(f"upstream unreachable: {exc.reason}") from exc @staticmethod def forward_headers(headers: Dict[str, str], api_key: str, stream: bool) -> Dict[str, str]: result: Dict[str, str] = {} for key, value in headers.items(): lower = key.lower() if lower in HOP_BY_HOP_HEADERS or lower in {"host", "authorization", "content-length", "accept-encoding"}: continue result[key] = value result["Authorization"] = f"Bearer {api_key}" result.setdefault("User-Agent", "yunwo-mistral-openai-proxy/1.0") if stream: result.setdefault("Accept", "text/event-stream") return result class ProxyApp: def __init__( self, *, key_manager: KeyManagerClient, forwarder: MistralForwarder, service_token: str, max_attempts: int, hard_max_attempts: int, shuffle_keys: bool, retry_backoff_sec: float, access_log: bool, conversation_models: frozenset[str], ) -> None: self.key_manager = key_manager self.forwarder = forwarder self.service_token = service_token.strip() self.max_attempts = max_attempts self.hard_max_attempts = hard_max_attempts self.shuffle_keys = shuffle_keys self.retry_backoff_sec = max(0.0, retry_backoff_sec) self.access_log = access_log self.conversation_models = conversation_models def authorized(self, handler: BaseHTTPRequestHandler) -> bool: if not self.service_token: return True auth = handler.headers.get("Authorization", "") if not auth.startswith("Bearer "): return False return auth[7:].strip() == self.service_token def attempt_count(self, pool_size: int) -> int: if pool_size <= 0: return 0 attempts = pool_size if self.max_attempts <= 0 else min(self.max_attempts, pool_size) if self.hard_max_attempts > 0: attempts = min(attempts, self.hard_max_attempts) return max(1, attempts) def choose_keys(self) -> List[KeyCandidate]: candidates = self.key_manager.active_keys() attempts = self.attempt_count(len(candidates)) if attempts <= 0: return [] if attempts >= len(candidates): selected = list(candidates) if self.shuffle_keys: random.shuffle(selected) return selected if self.shuffle_keys: return random.sample(candidates, attempts) return candidates[:attempts] def report_bad_response(self, candidate: KeyCandidate, response: ProxyResponse) -> bool: error = response.body[:1000].decode("utf-8", errors="replace") if is_model_access_error(response.status, error): # The subscription does not include this model. Retrying with a # different account only repeats the same request and can poison # every otherwise valid key in the pool. return False self.key_manager.drop_cached_key(candidate.api_key) status = "quota_exhausted" if response.status == 401 else "error" self.key_manager.report(candidate.api_key, response.status, error, status=status) return True def report_network_error(self, candidate: KeyCandidate, exc: Exception) -> None: sys.stderr.write(f"[{utc_now()}] upstream error with {mask_key(candidate.api_key)}: {exc}\n") def route_request(self, *, method: str, path: str, body: bytes) -> RoutedRequest: if not self.conversation_models or not is_openai_chat_completions_path(method, path): return RoutedRequest(path=path, body=body) model = request_model_id(body) if model not in self.conversation_models: return RoutedRequest(path=path, body=body) return RoutedRequest( path="/v1/conversations", body=translate_openai_conversation_body(body), openai_chat_model=model, openai_stream_include_usage=openai_stream_include_usage(body), ) def proxy_request( self, *, method: str, path: str, headers: Dict[str, str], body: bytes, stream: bool, ) -> ProxyResponse: candidates = self.choose_keys() if not candidates: raise HTTPClientError("no active Mistral API keys are available", status_code=HTTPStatus.SERVICE_UNAVAILABLE) attempts = len(candidates) last_response: Optional[ProxyResponse] = None last_error: Optional[Exception] = None for index, candidate in enumerate(candidates[:attempts], start=1): try: response = self.forwarder.forward( method=method, path=path, headers=headers, body=body, api_key=candidate.api_key, stream=stream, ) except Exception as exc: last_error = exc if index < attempts: time.sleep(self.retry_backoff_sec) continue break if response.status < 400: return response last_response = response if response.status in BAD_KEY_STATUS_CODES: if not self.report_bad_response(candidate, response): # This key is valid but cannot use the requested model. # Try another pooled key without changing this key's # manager state. if index < attempts: continue return response if index < attempts: continue if response.status == 429 or response.status in RETRYABLE_STATUS_CODES: if index < attempts: delay = response.retry_after() if delay is None: delay = self.retry_backoff_sec if delay > 0: time.sleep(delay) continue return response if last_response is not None: return last_response raise HTTPClientError(str(last_error or "proxy request failed"), status_code=HTTPStatus.BAD_GATEWAY) def proxy_stream( self, handler: BaseHTTPRequestHandler, *, method: str, path: str, headers: Dict[str, str], body: bytes, openai_chat_model: str = "", openai_stream_include_usage: bool = False, ) -> None: candidates = self.choose_keys() if not candidates: raise HTTPClientError("no active Mistral API keys are available", status_code=HTTPStatus.SERVICE_UNAVAILABLE) attempts = len(candidates) last_response: Optional[ProxyResponse] = None last_error: Optional[Exception] = None for index, candidate in enumerate(candidates[:attempts], start=1): try: upstream = self.forwarder.open_stream( method=method, path=path, headers=headers, body=body, api_key=candidate.api_key, ) except Exception as exc: last_error = exc if index < attempts: time.sleep(self.retry_backoff_sec) continue break if isinstance(upstream, ProxyResponse): last_response = upstream if upstream.status in BAD_KEY_STATUS_CODES: if not self.report_bad_response(candidate, upstream): # This key is valid but cannot use the requested # model. Try another key without reporting it as bad. if index < attempts: continue handler.send_upstream_response(upstream) return if index < attempts: continue if upstream.status == 429 or upstream.status in RETRYABLE_STATUS_CODES: if index < attempts: delay = upstream.retry_after() if delay is None: delay = self.retry_backoff_sec if delay > 0: time.sleep(delay) continue handler.send_upstream_response(upstream) return if openai_chat_model: handler.send_openai_chat_stream_response( upstream, openai_chat_model, include_usage=openai_stream_include_usage, ) else: handler.send_stream_response(upstream) return if last_response is not None: handler.send_upstream_response(last_response) return raise HTTPClientError(str(last_error or "stream proxy request failed"), status_code=HTTPStatus.BAD_GATEWAY) APP: Optional[ProxyApp] = None class Handler(BaseHTTPRequestHandler): server_version = "YunwoMistralOpenAIProxy/1.0" def log_message(self, fmt: str, *args: Any) -> None: app = APP if app is not None and app.access_log: sys.stderr.write("[%s] %s\n" % (utc_now(), fmt % args)) def require_app(self) -> ProxyApp: if APP is None: raise RuntimeError("app is not initialized") return APP def send_json(self, data: Any, status: int = 200) -> None: body = json_bytes(data) self.send_response(status) self.send_header("Content-Type", "application/json; charset=utf-8") self.send_header("Content-Length", str(len(body))) self.send_header("Access-Control-Allow-Origin", "*") self.send_header("Access-Control-Allow-Headers", "Authorization, Content-Type") self.send_header("Access-Control-Allow-Methods", "GET, POST, OPTIONS") self.end_headers() self.wfile.write(body) def send_error_json(self, status: int, message: str) -> None: self.send_json( { "error": { "message": message, "type": "yunwo_mistral_proxy_error", "code": status, } }, status=status, ) def do_OPTIONS(self) -> None: self.send_response(204) self.send_header("Access-Control-Allow-Origin", "*") self.send_header("Access-Control-Allow-Headers", "Authorization, Content-Type") self.send_header("Access-Control-Allow-Methods", "GET, POST, OPTIONS") self.end_headers() def do_GET(self) -> None: self.handle_proxy("GET") def do_POST(self) -> None: self.handle_proxy("POST") def do_HEAD(self) -> None: self.handle_proxy("HEAD") def read_body(self) -> bytes: length = int(self.headers.get("Content-Length") or 0) if length <= 0: return b"" return self.rfile.read(length) def handle_proxy(self, method: str) -> None: try: app = self.require_app() parsed = urllib.parse.urlsplit(self.path) if parsed.path == "/health": self.send_json( { "success": True, "time": utc_now(), "key_manager_url": app.key_manager.base_url, "mistral_api_base": app.forwarder.base_url, "upstream_proxy_configured": bool(app.forwarder.proxy), "max_attempts": app.max_attempts, "hard_max_attempts": app.hard_max_attempts, "max_concurrency": getattr(self.server, "max_workers", None), "key_cache_ttl_sec": app.key_manager.cache_ttl_sec, "requires_token": bool(app.service_token), "conversation_models": sorted(app.conversation_models), } ) return if is_proxy_path(parsed.path): if not app.authorized(self): self.send_error_json(HTTPStatus.UNAUTHORIZED, "unauthorized") return body = self.read_body() route = app.route_request(method=method, path=self.path, body=body) if is_conversation_start_path(method, route.path): body = route.body body = translate_openai_conversation_body(body) else: body = route.body headers = {key: value for key, value in self.headers.items()} stream = is_stream_request({key.lower(): value for key, value in headers.items()}, body) if stream: app.proxy_stream( self, method=method, path=route.path, headers=headers, body=body, openai_chat_model=route.openai_chat_model, openai_stream_include_usage=route.openai_stream_include_usage, ) else: response = app.proxy_request( method=method, path=route.path, headers=headers, body=body, stream=False, ) if route.openai_chat_model: response = conversation_response_to_openai_chat_response(response, route.openai_chat_model) self.send_upstream_response(response, head_only=method == "HEAD") return self.send_error_json(HTTPStatus.NOT_FOUND, "not found") except HTTPClientError as exc: self.send_error_json(exc.status_code or HTTPStatus.BAD_GATEWAY, str(exc)) except Exception as exc: traceback.print_exc() self.send_error_json(HTTPStatus.INTERNAL_SERVER_ERROR, str(exc)) def send_upstream_response(self, response: ProxyResponse, *, head_only: bool = False) -> None: self.send_response(response.status, response.reason) for key, value in response.headers: lower = key.lower() if lower in HOP_BY_HOP_HEADERS or lower in {"content-length", "content-encoding"}: continue self.send_header(key, value) self.send_header("Content-Length", str(len(response.body))) self.end_headers() if not head_only: self.wfile.write(response.body) def send_stream_response(self, upstream: UpstreamStream) -> None: self.send_response(upstream.status, upstream.reason) for key, value in upstream.headers: lower = key.lower() if lower in HOP_BY_HOP_HEADERS or lower in {"content-length", "content-encoding"}: continue self.send_header(key, value) self.send_header("Cache-Control", "no-cache") self.send_header("X-Accel-Buffering", "no") self.end_headers() try: while True: chunk = upstream.response.readline() if not chunk: break self.wfile.write(chunk) self.wfile.flush() finally: upstream.response.close() def send_openai_chat_stream_response( self, upstream: UpstreamStream, model: str, *, include_usage: bool = False, ) -> None: self.send_response(upstream.status, upstream.reason) for key, value in upstream.headers: lower = key.lower() if lower in HOP_BY_HOP_HEADERS or lower in {"content-length", "content-encoding", "content-type"}: continue self.send_header(key, value) self.send_header("Content-Type", "text/event-stream; charset=utf-8") self.send_header("Cache-Control", "no-cache") self.send_header("X-Accel-Buffering", "no") self.end_headers() response_id = f"chatcmpl-{uuid.uuid4().hex}" created = int(time.time()) role_sent = False finished = False used_tools = False tool_indexes: Dict[str, int] = {} pending_data: List[str] = [] final_usage: Optional[Dict[str, int]] = None def emit(data: Any) -> None: self.wfile.write(b"data: " + json_bytes(data) + b"\n\n") self.wfile.flush() def emit_done() -> None: nonlocal finished if not finished: self.wfile.write(b"data: [DONE]\n\n") self.wfile.flush() finished = True def emit_delta(delta: Dict[str, Any], finish_reason: Optional[str] = None) -> None: nonlocal role_sent if not role_sent: delta = {"role": "assistant", **delta} role_sent = True emit( { "id": response_id, "object": "chat.completion.chunk", "created": created, "model": model, "choices": [{"index": 0, "delta": delta, "finish_reason": finish_reason}], } ) def emit_usage() -> None: if final_usage is None: return emit( { "id": response_id, "object": "chat.completion.chunk", "created": created, "model": model, "choices": [], "usage": final_usage, } ) def process_event(raw_data: str) -> None: nonlocal response_id, used_tools, final_usage if not raw_data or raw_data == "[DONE]": return try: event = json.loads(raw_data) except json.JSONDecodeError: return if not isinstance(event, dict): return event_type = event.get("type") if event_type == "conversation.response.started": conversation_id = event.get("conversation_id") if isinstance(conversation_id, str) and conversation_id: response_id = f"chatcmpl-{conversation_id}" return if event_type == "message.output.delta": raw_content = event.get("content") _, thinking = split_conversation_content(raw_content) if thinking: emit_delta({"reasoning_content": thinking}) content = conversation_stream_content_to_openai(raw_content) if content is not None: emit_delta({"content": content}) return if event_type == "function.call.delta": used_tools = True tool_call_id = str(event.get("tool_call_id") or event.get("id") or uuid.uuid4().hex) tool_index = tool_indexes.setdefault(tool_call_id, len(tool_indexes)) emit_delta( { "tool_calls": [ { "index": tool_index, "id": tool_call_id, "type": "function", "function": { "name": str(event.get("name") or ""), "arguments": str(event.get("arguments") or ""), }, } ] } ) return if event_type == "conversation.response.done": final_usage = conversation_usage_to_openai(event.get("usage")) emit_delta({}, "tool_calls" if used_tools else "stop") if include_usage: emit_usage() emit_done() return if event_type == "conversation.response.error": emit({"error": {"message": str(event.get("message") or "conversation stream failed"), "code": event.get("code")}}) emit_done() try: while True: chunk = upstream.response.readline() if not chunk: break line = chunk.decode("utf-8", errors="replace").rstrip("\r\n") if line.startswith("data:"): pending_data.append(line[5:].lstrip()) elif not line: if pending_data: process_event("\n".join(pending_data)) pending_data = [] if pending_data: process_event("\n".join(pending_data)) if not finished: emit_delta({}, "tool_calls" if used_tools else "stop") if include_usage: emit_usage() emit_done() finally: upstream.response.close() class BoundedThreadingHTTPServer(ThreadingHTTPServer): daemon_threads = True request_queue_size = 64 def __init__(self, server_address: tuple[str, int], RequestHandlerClass: type[BaseHTTPRequestHandler], *, max_workers: int) -> None: self.max_workers = max(1, max_workers) self._worker_semaphore = threading.BoundedSemaphore(self.max_workers) super().__init__(server_address, RequestHandlerClass) def process_request(self, request: Any, client_address: tuple[str, int]) -> None: if not self._worker_semaphore.acquire(blocking=False): self._send_overloaded_response(request) self.shutdown_request(request) return try: super().process_request(request, client_address) except Exception: self._worker_semaphore.release() raise def process_request_thread(self, request: Any, client_address: tuple[str, int]) -> None: try: super().process_request_thread(request, client_address) finally: self._worker_semaphore.release() @staticmethod def _send_overloaded_response(request: Any) -> None: body = json_bytes( { "error": { "message": "proxy overloaded", "type": "yunwo_mistral_proxy_error", "code": HTTPStatus.SERVICE_UNAVAILABLE, } } ) response = ( f"HTTP/1.1 {HTTPStatus.SERVICE_UNAVAILABLE.value} {HTTPStatus.SERVICE_UNAVAILABLE.phrase}\r\n" "Content-Type: application/json; charset=utf-8\r\n" "Connection: close\r\n" f"Content-Length: {len(body)}\r\n\r\n" ).encode("ascii") + body try: request.sendall(response) except Exception: pass def build_app(args: argparse.Namespace) -> ProxyApp: key_manager_url = args.key_manager_url or os.getenv("MISTRAL_KEY_MANAGER_URL") or DEFAULT_KEY_MANAGER_URL key_manager_token = args.key_manager_token if args.key_manager_token is not None else os.getenv("MISTRAL_KEY_MANAGER_TOKEN", "") mistral_api_base = args.mistral_api_base or os.getenv("MISTRAL_OPENAI_PROXY_MISTRAL_API_BASE") or os.getenv("MISTRAL_API_BASE") or DEFAULT_MISTRAL_API_BASE upstream_proxy = args.upstream_proxy or os.getenv("MISTRAL_OPENAI_PROXY_UPSTREAM_PROXY") or os.getenv("MISTRAL_API_PROXY") or os.getenv("MISTRAL_OCR_PROXY", "") service_token = args.token if args.token is not None else os.getenv("MISTRAL_OPENAI_PROXY_TOKEN", "") timeout_sec = read_env_float("MISTRAL_OPENAI_PROXY_TIMEOUT_SEC", 120.0) key_manager_timeout_sec = read_env_float("MISTRAL_OPENAI_PROXY_KEY_MANAGER_TIMEOUT_SEC", min(10.0, timeout_sec)) cache_ttl_sec = read_env_float("MISTRAL_OPENAI_PROXY_KEY_CACHE_TTL_SEC", 30.0) stale_cache_ttl_sec = read_env_float("MISTRAL_OPENAI_PROXY_KEY_STALE_CACHE_TTL_SEC", 300.0) refresh_backoff_sec = read_env_float("MISTRAL_OPENAI_PROXY_KEY_REFRESH_BACKOFF_SEC", 5.0) access_log = read_env_bool("MISTRAL_OPENAI_PROXY_ACCESS_LOG", False) conversation_models = parse_model_list(args.conversation_models) return ProxyApp( key_manager=KeyManagerClient( key_manager_url, key_manager_token, key_manager_timeout_sec, cache_ttl_sec, stale_cache_ttl_sec, refresh_backoff_sec, ), forwarder=MistralForwarder(mistral_api_base, upstream_proxy, timeout_sec), service_token=service_token or "", max_attempts=args.max_attempts, hard_max_attempts=args.hard_max_attempts, shuffle_keys=not args.no_shuffle, retry_backoff_sec=args.retry_backoff_sec, access_log=access_log, conversation_models=conversation_models, ) def parse_args(argv: Optional[List[str]] = None) -> argparse.Namespace: parser = argparse.ArgumentParser(description="OpenAI-compatible Mistral proxy backed by key manager") parser.add_argument("--host", default=os.getenv("MISTRAL_OPENAI_PROXY_HOST", "127.0.0.1")) parser.add_argument("--port", type=int, default=int(os.getenv("MISTRAL_OPENAI_PROXY_PORT", "8099"))) parser.add_argument("--token", default=None, help="Optional proxy bearer token") parser.add_argument("--key-manager-url", default=os.getenv("MISTRAL_KEY_MANAGER_URL", DEFAULT_KEY_MANAGER_URL)) parser.add_argument("--key-manager-token", default=os.getenv("MISTRAL_KEY_MANAGER_TOKEN", "")) parser.add_argument("--mistral-api-base", default=os.getenv("MISTRAL_OPENAI_PROXY_MISTRAL_API_BASE", os.getenv("MISTRAL_API_BASE", DEFAULT_MISTRAL_API_BASE))) parser.add_argument("--upstream-proxy", default=os.getenv("MISTRAL_OPENAI_PROXY_UPSTREAM_PROXY", os.getenv("MISTRAL_API_PROXY", os.getenv("MISTRAL_OCR_PROXY", "")))) parser.add_argument("--max-attempts", type=int, default=read_env_int("MISTRAL_OPENAI_PROXY_MAX_ATTEMPTS", 8)) parser.add_argument("--hard-max-attempts", type=int, default=read_env_int("MISTRAL_OPENAI_PROXY_HARD_MAX_ATTEMPTS", 16)) parser.add_argument("--max-concurrency", type=int, default=read_env_int("MISTRAL_OPENAI_PROXY_MAX_CONCURRENCY", default_max_concurrency())) parser.add_argument("--retry-backoff-sec", type=float, default=read_env_float("MISTRAL_OPENAI_PROXY_RETRY_BACKOFF_SEC", 0.5)) parser.add_argument("--no-shuffle", action="store_true", default=read_env_bool("MISTRAL_OPENAI_PROXY_NO_SHUFFLE", False)) parser.add_argument( "--conversation-models", default=os.getenv("MISTRAL_OPENAI_PROXY_CONVERSATION_MODELS", DEFAULT_CONVERSATION_MODELS), help="Comma-separated models routed from /v1/chat/completions to /v1/conversations", ) return parser.parse_args(argv) def main(argv: Optional[List[str]] = None) -> int: global APP args = parse_args(argv) APP = build_app(args) max_workers = max(1, args.max_concurrency) server = BoundedThreadingHTTPServer((args.host, args.port), Handler, max_workers=max_workers) print( f"Mistral OpenAI proxy listening on http://{args.host}:{args.port} " f"key_manager={APP.key_manager.base_url} mistral={APP.forwarder.base_url} " f"upstream_proxy={'yes' if APP.forwarder.proxy else 'no'} max_workers={max_workers} " f"conversation_models={','.join(sorted(APP.conversation_models)) or 'none'}", flush=True, ) try: server.serve_forever() except KeyboardInterrupt: print("\nshutting down", flush=True) finally: server.server_close() return 0 if __name__ == "__main__": raise SystemExit(main())