synth-ai 0.2.9.dev11__py3-none-any.whl → 0.4.1__py3-none-any.whl

examples/warming_up_to_rl/task_app/synth_envs_hosted/inference/openai_client.py

@@ -1,618 +0,0 @@
-from __future__ import annotations
-
-import asyncio
-import contextlib
-import logging
-from typing import Any
-
-import httpx
-
-logger = logging.getLogger(__name__)
-
-
-class OpenAIClient:
-    """Async HTTP client for OpenAI-compatible inference servers (vLLM)."""
-
-    def __init__(
-        self,
-        base_url: str,
-        api_key: str | None = None,
-        timeout_s: float = 120.0,
-    ) -> None:
-        self.base_url = base_url.rstrip("/")
-        self.api_key = api_key
-        self.timeout_s = timeout_s
-        self.headers = {}
-
-        if api_key:
-            self.headers["Authorization"] = f"Bearer {api_key}"
-
-    def _fix_model_parameters(
-        self, request: dict[str, Any], target_url: str | None = None
-    ) -> dict[str, Any]:
-        """
-        Fix parameter compatibility for newer OpenAI models.
-
-        Newer models like gpt-5-nano use 'max_completion_tokens' instead of 'max_tokens'.
-        """
-        if not request:
-            return request
-
-        # Make a copy to avoid modifying the original
-        fixed_request = request.copy()
-
-        # Determine if target is OpenAI-compatible (OpenAI, Azure OpenAI, Groq);
-        # strip fields those endpoints don't accept
-        is_openai = False
-        try:
-            if isinstance(target_url, str):
-                low = target_url.lower()
-                is_openai = (
-                    ("openai.com" in low)
-                    or ("azure" in low and ".openai." in low)
-                    or ("groq.com" in low)
-                    or ("/openai" in low)
-                )
-        except Exception:
-            is_openai = False
-
-        model = fixed_request.get("model", "")
-
-        if is_openai:
-            # Remove fields OpenAI/Groq don't accept
-            for k in (
-                "stop_after_tool_calls",
-                "thinking_mode",
-                "thinking_budget",
-                "reasoning",
-                "extra_body",
-                "parallel_tool_calls",
-                "function_call",
-            ):
-                if k in fixed_request:
-                    fixed_request.pop(k, None)
-
-            # GPT-5 family specifics
-            if "gpt-5" in model or "gpt-4.1" in model:
-                # Convert max_tokens to max_completion_tokens for newer models
-                if "max_tokens" in fixed_request:
-                    if "max_completion_tokens" not in fixed_request:
-                        fixed_request["max_completion_tokens"] = fixed_request.pop("max_tokens")
-                        logger.info(
-                            f"Converted max_tokens to max_completion_tokens for model {model}"
-                        )
-                    else:
-                        fixed_request.pop("max_tokens")
-                        logger.info(f"Removed conflicting max_tokens parameter for model {model}")
-                # Some OpenAI endpoints ignore/deny sampling fields for reasoning models
-                for k in ("temperature", "top_p"):
-                    if k in fixed_request:
-                        fixed_request.pop(k, None)
-                # If tools are present, force single tool choice to our function
-                try:
-                    tools = fixed_request.get("tools")
-                    if isinstance(tools, list) and tools:
-                        # Choose the first provided function name from tools schema (e.g., run_command)
-                        func_name = None
-                        for t in tools:
-                            try:
-                                cand = None
-                                if isinstance(t, dict):
-                                    f = t.get("function")
-                                    if isinstance(f, dict):
-                                        cand = f.get("name")
-                                if isinstance(cand, str) and cand:
-                                    func_name = cand
-                                    break
-                            except Exception:
-                                continue
-                        if not func_name:
-                            func_name = "run_command"
-                        fixed_request["tool_choice"] = {
-                            "type": "function",
-                            "function": {"name": func_name},
-                        }
-                        fixed_request["parallel_tool_calls"] = False
-                except Exception:
-                    pass
-
-        return fixed_request
-
-    async def generate(
-        self,
-        request: dict[str, Any],
-        base_url: str | None = None,
-        timeout_s: float | None = None,
-        extra_headers: dict[str, str] | None = None,
-    ) -> dict[str, Any]:
-        """
-        Send a chat completion request to the inference server.
-
-        Args:
-            request: OpenAI-compatible chat completion request
-            base_url: Override base URL for this request
-            timeout_s: Override timeout for this request
-            extra_headers: Additional headers to include (e.g., X-Policy-Name)
-
-        Returns:
-            OpenAI-compatible chat completion response
-        """
-        url = (base_url or self.base_url).rstrip("/") + "/v1/chat/completions"
-        timeout = timeout_s or self.timeout_s
-
-        # Merge headers
-        headers = self.headers.copy()
-        if extra_headers:
-            headers.update(extra_headers)
-
-        # Fix parameter compatibility for newer models
-        processed_request = self._fix_model_parameters(request, target_url=url)
-
-        # Log request (redact messages in production)
-        logger.info(f"Inference POST target: {url}")
-        if extra_headers:
-            logger.info(f"Extra headers: {extra_headers}")
-        with contextlib.suppress(Exception):
-            keys_preview = sorted(processed_request.keys())
-            logger.info(f"Request keys: {keys_preview}")
-
-        # Final hard-guard for OpenAI: ensure unsupported field is not present
-        try:
-            if "openai" in url.lower() and "stop_after_tool_calls" in processed_request:
-                processed_request.pop("stop_after_tool_calls", None)
-                logger.info("Removed stop_after_tool_calls for OpenAI request")
-            # Groq-specific requirement: when using JSON mode, one of the messages must contain the word 'json'
-            low_url = url.lower()
-            if ("groq.com" in low_url or "/openai" in low_url) and isinstance(
-                processed_request, dict
-            ):
-                rf = processed_request.get("response_format")
-                rf_type = None
-                if isinstance(rf, dict):
-                    rf_type = str(rf.get("type") or "").lower()
-                if rf_type in {"json_object", "json_schema"}:
-                    msgs = processed_request.get("messages")
-                    has_json_word = False
-                    if isinstance(msgs, list):
-                        for m in msgs:
-                            try:
-                                content = m.get("content") if isinstance(m, dict) else None
-                                text = None
-                                if isinstance(content, str):
-                                    text = content
-                                elif isinstance(content, list):
-                                    # Join any text segments
-                                    parts = []
-                                    for seg in content:
-                                        if isinstance(seg, dict) and isinstance(
-                                            seg.get("text"), str
-                                        ):
-                                            parts.append(seg["text"])
-                                    text = "\n".join(parts)
-                                if isinstance(text, str) and ("json" in text.lower()):
-                                    has_json_word = True
-                                    break
-                            except Exception:
-                                continue
-                    if not has_json_word:
-                        try:
-                            instruction = (
-                                "Respond in strict JSON only. Output a single valid JSON object."
-                            )
-                            if not isinstance(msgs, list):
-                                msgs = []
-                            # Prepend a system message to satisfy Groq requirement without changing user intent
-                            prepend = {"role": "system", "content": instruction}
-                            processed_request["messages"] = [prepend] + list(msgs)
-                            logger.info(
-                                "Injected JSON-mode system instruction for Groq response_format compliance"
-                            )
-                        except Exception:
-                            pass
-        except Exception:
-            pass
-
-        async with httpx.AsyncClient(timeout=timeout) as client:
-            try:
-                response = await client.post(
-                    url,
-                    json=processed_request,
-                    headers=headers,
-                )
-                response.raise_for_status()
-
-                # Rich response diagnostics
-                content_type = response.headers.get("content-type")
-                body_text = response.text
-                logger.info(
-                    f"Inference response status=200, content-type={content_type}, bytes={len(body_text)}"
-                )
-                if body_text:
-                    preview_len = min(800, len(body_text))
-                    logger.info(
-                        f"Inference response preview ({preview_len} bytes): {body_text[:preview_len]}"
-                    )
-
-                result = response.json()
-                logger.info(f"Inference response parsed_type={type(result).__name__}")
-                return result
-
-            except httpx.TimeoutException:
-                logger.error(f"Request to {url} timed out after {timeout}s")
-                raise
-            except httpx.HTTPStatusError as e:
-                status = e.response.status_code if e.response is not None else None
-                text = e.response.text if e.response is not None else str(e)
-                # Log full body for debugging remote failures
-                try:
-                    logger.error(
-                        {
-                            "openai_http_error": True,
-                            "status": status,
-                            "url": url,
-                            "body": text,
-                        }
-                    )
-                except Exception:
-                    logger.error(f"HTTP error from {url}: {status} - {text}")
-                # For 4xx/5xx, print full sanitized request to aid debugging (especially Groq 400s)
-                try:
-                    redacted_headers = dict(headers)
-                    if "Authorization" in redacted_headers:
-                        redacted_headers["Authorization"] = "***REDACTED***"
-                    logger.error(
-                        {
-                            "request_debug": True,
-                            "status": status,
-                            "target": url,
-                            "headers": redacted_headers,
-                            "payload": processed_request,
-                        }
-                    )
-                except Exception:
-                    pass
-                # Special case: token budget exceeded (OpenAI-compatible error schema)
-                try:
-                    if status == 400 and e.response is not None:
-                        data = e.response.json()
-                        detail = data.get("detail") if isinstance(data, dict) else None
-                        err_code = (detail or {}).get("error") if isinstance(detail, dict) else None
-                        if err_code == "token_budget_exceeded":
-                            info = (detail or {}).get("details") or {}
-                            messages_tokens = int(info.get("messages_tokens") or 0)
-                            model_limit = int(info.get("model_limit") or 0)
-                            safety = 64
-                            # Compute a conservative new max_tokens
-                            new_max = max(16, model_limit - messages_tokens - safety)
-                            try:
-                                # Update request and retry once immediately with smaller budget
-                                if isinstance(processed_request, dict):
-                                    processed_request = dict(processed_request)
-                                    if "max_completion_tokens" in processed_request:
-                                        processed_request["max_completion_tokens"] = new_max
-                                        processed_request.pop("max_tokens", None)
-                                    else:
-                                        processed_request["max_tokens"] = new_max
-                                    # Remove optional fields that some servers reject
-                                    for k in ("thinking_mode", "thinking_budget", "reasoning"):
-                                        processed_request.pop(k, None)
-                                    # Force structured tool choice
-                                    if processed_request.get("tool_choice") == "required":
-                                        func_name = "run_command"
-                                        try:
-                                            tools_arr = processed_request.get("tools") or []
-                                            if isinstance(tools_arr, list) and tools_arr:
-                                                f = (
-                                                    tools_arr[0].get("function")
-                                                    if isinstance(tools_arr[0], dict)
-                                                    else None
-                                                )
-                                                cand = (
-                                                    (f or {}).get("name")
-                                                    if isinstance(f, dict)
-                                                    else None
-                                                )
-                                                if isinstance(cand, str) and cand:
-                                                    func_name = cand
-                                        except Exception:
-                                            pass
-                                        processed_request["tool_choice"] = {
-                                            "type": "function",
-                                            "function": {"name": func_name},
-                                        }
-                                        processed_request["parallel_tool_calls"] = False
-                                    logger.warning(
-                                        {
-                                            "token_budget_recovery": True,
-                                            "messages_tokens": messages_tokens,
-                                            "model_limit": model_limit,
-                                            "retry_max_tokens": new_max,
-                                        }
-                                    )
-                                    # Retry once with reduced budget
-                                    async with httpx.AsyncClient(timeout=timeout) as client2:
-                                        r2 = await client2.post(
-                                            url, json=processed_request, headers=headers
-                                        )
-                                        r2.raise_for_status()
-                                        return r2.json()
-                            except Exception:
-                                pass
-                except Exception:
-                    pass
-                # Gracefully degrade on 422 so rollouts can still produce a trajectory
-                if status == 422:
-                    try:
-                        # Best-effort parse of error for diagnostics
-                        err = None
-                        try:
-                            err = e.response.json()
-                        except Exception:
-                            err = {"error": "unprocessable", "detail": (text or "")[:200]}
-                        logger.warning(
-                            {
-                                "inference_422_recovered": True,
-                                "detail": err,
-                            }
-                        )
-                    except Exception:
-                        pass
-                    # Return a minimal OpenAI-compatible response with no tool_calls/content
-                    import time as _t
-
-                    return {
-                        "id": f"cmpl-{int(_t.time())}",
-                        "object": "chat.completion",
-                        "created": int(_t.time()),
-                        "model": processed_request.get("model") or "unknown",
-                        "choices": [
-                            {
-                                "index": 0,
-                                "message": {"role": "assistant", "content": "", "tool_calls": []},
-                                "finish_reason": "stop",
-                            }
-                        ],
-                        "usage": {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0},
-                    }
-                raise
-            except Exception as e:
-                logger.error(f"Unexpected error calling {url}: {e}")
-                raise
-
-    async def check_health(
-        self,
-        base_url: str | None = None,
-        timeout_s: float | None = None,
-    ) -> dict[str, Any]:
-        """
-        Check if the inference service is healthy.
-
-        Args:
-            base_url: Override base URL for this request
-            timeout_s: Override timeout for this request
-
-        Returns:
-            Health status dict with 'status' field
-        """
-        url = (base_url or self.base_url).rstrip("/") + "/health"
-        timeout = timeout_s or 10.0
-
-        try:
-            async with httpx.AsyncClient(timeout=timeout) as client:
-                response = await client.get(url, headers=self.headers)
-                response.raise_for_status()
-                return response.json()
-        except httpx.HTTPStatusError as e:
-            if e.response.status_code == 400:
-                # Service is overloaded but still responding
-                try:
-                    data = e.response.json()
-                    if data.get("status") == "overloaded":
-                        return {"status": "overloaded", "retry_after": data.get("retry_after", 1)}
-                except Exception:
-                    pass
-            return {"status": "unhealthy", "error": str(e)}
-        except Exception as e:
-            return {"status": "unhealthy", "error": str(e)}
-
-    async def generate_with_retries(
-        self,
-        request: dict[str, Any],
-        base_url: str | None = None,
-        timeout_s: float | None = None,
-        max_retries: int = 4,
-        backoff_factor: float = 2.0,
-        extra_headers: dict[str, str] | None = None,
-    ) -> dict[str, Any]:
-        """
-        Generate with exponential backoff retries for transient errors.
-
-        Args:
-            request: OpenAI-compatible chat completion request
-            base_url: Override base URL
-            timeout_s: Override timeout
-            max_retries: Maximum number of retry attempts
-            backoff_factor: Exponential backoff multiplier
-            extra_headers: Additional headers to include (e.g., X-Policy-Name)
-
-        Returns:
-            OpenAI-compatible chat completion response
-        """
-        last_error = None
-        wait_time = 1.0
-
-        for attempt in range(max_retries + 1):
-            try:
-                # Apply parameter fixes to the request
-                processed_request = self._fix_model_parameters(
-                    request,
-                    target_url=(base_url or self.base_url).rstrip("/") + "/v1/chat/completions",
-                )
-                return await self.generate(
-                    request=processed_request,
-                    base_url=base_url,
-                    timeout_s=timeout_s,
-                    extra_headers=extra_headers,
-                )
-            except httpx.HTTPStatusError as e:
-                # Retry on 400 (overloaded), 429 (rate limit), 500 (internal error), 503 (service unavailable)
-                if e.response.status_code not in [400, 429, 500, 503]:
-                    raise
-                last_error = e
-                if e.response.status_code == 400:
-                    # Check if this is an overload error by looking at response content
-                    try:
-                        response_data = e.response.json()
-                        if response_data.get("status") == "overloaded":
-                            retry_after = response_data.get("retry_after", 1)
-                            # Use the suggested retry_after time instead of exponential backoff for overload
-                            wait_time = max(wait_time, float(retry_after))
-                            logger.warning(
-                                f"Inference service overloaded (400). {response_data} Retrying after {wait_time}s..."
-                            )
-                        else:
-                            # This is a different type of 400 error, don't retry
-                            try:
-                                redacted_headers = {}
-                                try:
-                                    redacted_headers = dict(self.headers)
-                                    if "Authorization" in redacted_headers:
-                                        redacted_headers["Authorization"] = "***REDACTED***"
-                                except Exception:
-                                    redacted_headers = {}
-                                logger.error(
-                                    {
-                                        "non_overload_400": True,
-                                        "target": (base_url or self.base_url),
-                                        "payload": processed_request,
-                                        "headers": redacted_headers,
-                                        "body": e.response.text if e.response is not None else None,
-                                    }
-                                )
-                            except Exception:
-                                pass
-                            raise RuntimeError(
-                                f"Inference 400 response: {e.response.text if e.response is not None else 'Bad Request'}"
-                            ) from e
-                    except Exception:
-                        # If we can't parse the response, don't retry 400 errors
-                        with contextlib.suppress(Exception):
-                            logger.error(
-                                {
-                                    "non_overload_400_unparsed": True,
-                                    "target": (base_url or self.base_url),
-                                    "payload": processed_request,
-                                }
-                            )
-                        raise RuntimeError(
-                            f"Inference 400 response (unparsed): {e.response.text if e.response is not None else 'Bad Request'}"
-                        ) from e
-                elif e.response.status_code == 503:
-                    # Avoid referencing undefined response_data
-                    try:
-                        preview = (e.response.text or "")[:200]
-                    except Exception:
-                        preview = ""
-                    logger.warning(
-                        f"Flash returned 503; container may be cold starting. Retrying... body={preview}"
-                    )
-                elif e.response.status_code == 500:
-                    try:
-                        preview = (e.response.text or "")[:200]
-                    except Exception:
-                        preview = ""
-                    logger.warning(
-                        f"Flash returned 500; inference service error. Retrying... body={preview}"
-                    )
-            except httpx.TimeoutException as e:
-                last_error = e
-
-            if attempt < max_retries:
-                logger.warning(
-                    f"Inference request failed (attempt {attempt + 1}/{max_retries + 1}), "
-                    f"retrying in {wait_time}s..."
-                )
-                await asyncio.sleep(wait_time)
-                wait_time *= backoff_factor
-
-        raise last_error
-
-
-def create_inference_client(
-    task_app: Any,
-    api_key: str | None = None,
-) -> OpenAIClient:
-    """
-    Create an inference client using TaskApp configuration.
-
-    Args:
-        task_app: TaskApp instance with vllm_base_url
-        api_key: Optional API key for authentication
-
-    Returns:
-        Configured OpenAIClient instance
-    """
-    # Fallback to environment if caller didn't provide an API key
-    if api_key is None:
-        try:
-            import os as _os  # local import to avoid module-level side effects
-
-            api_key = _os.getenv("OPENAI_API_KEY") or getattr(task_app, "openai_api_key", None)
-        except Exception:
-            api_key = None
-
-    import json as _json
-    import os as _os
-    import time as _time
-
-    if _os.getenv("SYNTH_FAKE_INFERENCE", "").strip():
-
-        class _DummyClient:
-            async def generate_with_retries(
-                self,
-                request: dict[str, Any],
-                base_url: str | None = None,
-                max_retries: int = 0,
-                backoff_factor: float = 1.0,
-                extra_headers: dict[str, str] | None = None,
-            ) -> dict[str, Any]:
-                tool_call = {
-                    "id": "call_dummy",
-                    "type": "function",
-                    "function": {
-                        "name": "interact_many",
-                        "arguments": _json.dumps({"actions": ["move_right"]}),
-                    },
-                }
-                return {
-                    "id": f"cmpl-{int(_time.time())}",
-                    "object": "chat.completion",
-                    "created": int(_time.time()),
-                    "model": request.get("model") or "dummy-model",
-                    "choices": [
-                        {
-                            "index": 0,
-                            "message": {
-                                "role": "assistant",
-                                "content": "",
-                                "tool_calls": [tool_call],
-                            },
-                            "finish_reason": "tool_calls",
-                        }
-                    ],
-                    "usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15},
-                }
-
-            async def check_health(
-                self,
-                base_url: str | None = None,
-                timeout_s: float | None = None,
-            ) -> dict[str, Any]:
-                return {"status": "ok", "dummy": True}
-
-        return _DummyClient()
-
-    return OpenAIClient(
-        base_url=task_app.vllm_base_url,
-        api_key=api_key,
-    )