coco-framework 0.1.0__py3-none-any.whl

coco_core/utils/_gateway.py

@@ -0,0 +1,274 @@
+"""Private helpers for resolving LangSmith gateway configuration.
+
+The [LangSmith LLM gateway](https://docs.langchain.com/langsmith/llm-gateway)
+lets a chat model reach a provider through a proxy configured via environment
+variables. These helpers centralize the (non-trivial) precedence rules so that
+each provider integration resolves its base URL and API key identically.
+
+This module is private: the API is not stable and may change without notice.
+"""
+
+from __future__ import annotations
+
+import json
+import os
+from collections.abc import Mapping
+from typing import TYPE_CHECKING, Any, NamedTuple
+
+from pydantic import SecretStr
+
+if TYPE_CHECKING:
+    from collections.abc import Sequence
+
+    from pydantic import BaseModel
+
+# Key under which integrations place parsed gateway metadata on a generation's
+# `generation_info`. This is the contract consumed by the LangSmith tracer, which
+# promotes it to run metadata. Defined here so integrations and the tracer agree
+# on a single name.
+GATEWAY_METADATA_RESPONSE_KEY = "lc_gateway_metadata"
+
+# HTTP response header the LangSmith gateway uses to return per-request metadata
+# (e.g. the resolved provider and model). The value is JSON-encoded.
+GATEWAY_METADATA_HEADER = "X-LangSmith-Gateway-Metadata"
+
+_LANGSMITH_GATEWAY_ENV = "LANGSMITH_GATEWAY"
+_LANGSMITH_GATEWAY_API_KEY_ENV = "LANGSMITH_GATEWAY_API_KEY"
+_LANGSMITH_API_KEY_ENV = "LANGSMITH_API_KEY"
+_LANGSMITH_GATEWAY_DEFAULT_BASE = "https://gateway.smith.langchain.com"
+
+_TRUE_VALUES = ("true", "1", "yes")
+_FALSE_VALUES = ("false", "0", "no")
+
+
+class GatewayConfig(NamedTuple):
+    """Resolved gateway configuration.
+
+    Attributes:
+        base_url: The base URL the client should use, or None to defer to the
+            provider SDK's own default.
+        api_key: The API key to use. A ``SecretStr`` when derived from the
+            environment; the caller's value returned unchanged when one was
+            passed explicitly; None when no key could be resolved.
+        base_url_from_gateway: Whether ``base_url`` was populated from the
+            LangSmith gateway (as opposed to an explicit value, a provider env
+            var, or a default).
+    """
+
+    base_url: str | None
+    api_key: Any
+    base_url_from_gateway: bool
+
+
+def _first_env(names: str | Sequence[str]) -> str | None:
+    """Return the first non-empty value among ``names``, or None."""
+    if isinstance(names, str):
+        names = (names,)
+    for name in names:
+        value = os.getenv(name)
+        if value:
+            return value
+    return None
+
+
+def _resolve_gateway_base_url(provider_path: str) -> str | None:
+    """Resolve the LangSmith gateway base URL for a provider.
+
+    ``LANGSMITH_GATEWAY`` accepts either a boolean-ish string or an explicit
+    base URL:
+
+    - ``true`` / ``1`` / ``yes`` -> the default gateway host.
+    - ``false`` / ``0`` / ``no`` / unset / empty -> the gateway is disabled
+      (None).
+    - anything else -> treated as a custom gateway base URL.
+
+    The provider-specific path is appended in all enabled cases.
+
+    Args:
+        provider_path: Path segment for the provider, e.g. ``"openai/v1"`` or
+            ``"anthropic"``.
+
+    Returns:
+        The provider base URL on the gateway, or None if the gateway is
+        disabled.
+    """
+    raw = os.getenv(_LANGSMITH_GATEWAY_ENV)
+    if not raw or raw.lower() in _FALSE_VALUES:
+        return None
+    base = (
+        _LANGSMITH_GATEWAY_DEFAULT_BASE
+        if raw.lower() in _TRUE_VALUES
+        else raw.rstrip("/")
+    )
+    return f"{base}/{provider_path}"
+
+
+def _resolve_gateway_config(
+    *,
+    base_url: str | None,
+    api_key: Any,
+    provider_path: str,
+    base_url_env: str | Sequence[str] = (),
+    api_key_env: str | Sequence[str] = (),
+    default_base_url: str | None = None,
+) -> GatewayConfig:
+    """Resolve a provider's base URL and API key, applying gateway settings.
+
+    Precedence:
+
+    - **base_url:** explicit ``base_url`` > ``base_url_env`` > LangSmith gateway
+      > ``default_base_url``.
+    - **api_key:** explicit ``api_key`` (returned unchanged) > the gateway key if
+      the base URL came from the gateway, otherwise the provider key > the other.
+      The gateway key uses ``LANGSMITH_GATEWAY_API_KEY`` when the gateway is
+      enabled, with ``LANGSMITH_API_KEY`` as a fallback only when the base URL came
+      from the gateway.
+
+    The provenance flip on the key means an ``OPENAI_API_KEY``-style provider key
+    is preferred whenever the caller pointed the base URL at a non-gateway
+    endpoint (so a stray gateway key is not sent to the provider), while the
+    gateway key wins for the common "just enable the gateway" setup even if a
+    provider key happens to be present in the environment.
+
+    Args:
+        base_url: Explicitly-provided base URL (e.g. a ``base_url`` kwarg), or
+            None if not set by the caller.
+        api_key: Explicitly-provided API key, or None if not set by the caller.
+            Returned unchanged when not None, so a caller-supplied value (secret
+            or callable) always wins.
+        provider_path: Path segment appended to the gateway host.
+        base_url_env: Env var name(s) for the provider base URL, in priority
+            order.
+        api_key_env: Env var name(s) for the provider API key, in priority order.
+        default_base_url: Base URL used when nothing else is set.
+
+    Returns:
+        The resolved `GatewayConfig`.
+    """
+    gateway_base_url = _resolve_gateway_base_url(provider_path)
+
+    resolved_base_url = base_url
+    base_url_from_gateway = False
+    if resolved_base_url is None:
+        resolved_base_url = _first_env(base_url_env)
+        if resolved_base_url is None:
+            if gateway_base_url is not None:
+                resolved_base_url = gateway_base_url
+                base_url_from_gateway = True
+            else:
+                resolved_base_url = default_base_url
+
+    if api_key is not None:
+        resolved_api_key: Any = api_key
+    else:
+        gateway_api_key = (
+            os.getenv(_LANGSMITH_GATEWAY_API_KEY_ENV)
+            if gateway_base_url is not None
+            else None
+        )
+        if not gateway_api_key and base_url_from_gateway:
+            gateway_api_key = os.getenv(_LANGSMITH_API_KEY_ENV)
+        provider_api_key = _first_env(api_key_env)
+        chosen = (
+            (gateway_api_key or provider_api_key)
+            if base_url_from_gateway
+            else (provider_api_key or gateway_api_key)
+        )
+        resolved_api_key = SecretStr(chosen) if chosen else None
+
+    return GatewayConfig(resolved_base_url, resolved_api_key, base_url_from_gateway)
+
+
+def _pop_provided(values: dict[str, Any], cls: type[BaseModel], field: str) -> Any:
+    """Pop a caller-provided field value by name or alias; None if absent.
+
+    Handles models with ``populate_by_name=True``, where a field may be supplied
+    under either its name or its alias. Both keys are removed so the resolved
+    value can be written back canonically under the field name. The alias is read
+    from the model rather than hard-coded, so callers pass only field names.
+    """
+    value = values.pop(field, None)
+    alias = cls.model_fields[field].alias
+    if alias is not None:
+        alias_value = values.pop(alias, None)
+        if value is None:
+            value = alias_value
+    return value
+
+
+def _apply_gateway_config(
+    values: dict[str, Any],
+    cls: type[BaseModel],
+    *,
+    base_url_field: str,
+    api_key_field: str,
+    provider_path: str,
+    base_url_env: str | Sequence[str] = (),
+    api_key_env: str | Sequence[str] = (),
+    default_base_url: str | None = None,
+) -> GatewayConfig:
+    """Resolve gateway settings from a model's raw input in a "before" validator.
+
+    Reads the caller-provided base URL and API key (by field name or alias),
+    resolves them against the gateway and provider env vars via
+    `_resolve_gateway_config`, then writes the results back into ``values`` under
+    the canonical field names. This lets a provider integration keep a
+    non-optional key field: the resolved value is injected before field
+    validation runs, so the field always receives a concrete value.
+
+    The resolved API key is written only when non-None, so the field's own
+    default applies when no key is found. Returns the `GatewayConfig` so the
+    caller can, for example, raise when a required key is missing.
+
+    Args:
+        values: The raw input mapping passed to the model, mutated in place.
+        cls: The model class, used to look up field aliases.
+        base_url_field: Name of the base URL field (e.g. ``"anthropic_api_url"``).
+        api_key_field: Name of the API key field (e.g. ``"anthropic_api_key"``).
+        provider_path: Path segment appended to the gateway host.
+        base_url_env: Env var name(s) for the provider base URL.
+        api_key_env: Env var name(s) for the provider API key.
+        default_base_url: Base URL used when nothing else is set.
+
+    Returns:
+        The resolved `GatewayConfig`.
+    """
+    config = _resolve_gateway_config(
+        base_url=_pop_provided(values, cls, base_url_field),
+        api_key=_pop_provided(values, cls, api_key_field),
+        provider_path=provider_path,
+        base_url_env=base_url_env,
+        api_key_env=api_key_env,
+        default_base_url=default_base_url,
+    )
+    values[base_url_field] = config.base_url
+    if config.api_key is not None:
+        values[api_key_field] = config.api_key
+    return config
+
+
+def _parse_gateway_metadata(headers: Any) -> dict[str, Any] | None:
+    """Parse LangSmith gateway metadata from HTTP response headers.
+
+    Provider integrations call this on the raw response headers so the parsed
+    result can be surfaced under `GATEWAY_METADATA_RESPONSE_KEY` and later
+    promoted to run metadata by the LangSmith tracer.
+
+    Args:
+        headers: The response headers.
+
+    Returns:
+        The deserialized gateway metadata, or None when the header is absent,
+        not a mapping, or not a JSON object.
+    """
+    if not isinstance(headers, Mapping):
+        return None
+
+    raw = headers.get(GATEWAY_METADATA_HEADER)
+    if not raw:
+        return None
+    try:
+        parsed = json.loads(raw)
+    except (json.JSONDecodeError, TypeError):
+        return None
+    return parsed if isinstance(parsed, dict) else None

coco_core/utils/_merge.py

@@ -0,0 +1,218 @@
+from __future__ import annotations
+
+from typing import Any
+
+
+def merge_dicts(left: dict[str, Any], *others: dict[str, Any]) -> dict[str, Any]:
+    r"""Merge dictionaries.
+
+    Merge many dicts, handling specific scenarios where a key exists in both
+    dictionaries but has a value of `None` in `'left'`. In such cases, the method uses
+    the value from `'right'` for that key in the merged dictionary.
+
+    Args:
+        left: The first dictionary to merge.
+        others: The other dictionaries to merge.
+
+    Returns:
+        The merged dictionary.
+
+    Raises:
+        TypeError: If the key exists in both dictionaries but has a different type.
+        TypeError: If the value has an unsupported type.
+
+    Example:
+        If `left = {"function_call": {"arguments": None}}` and
+        `right = {"function_call": {"arguments": "{\n"}}`, then, after merging, for the
+        key `'function_call'`, the value from `'right'` is used, resulting in
+        `merged = {"function_call": {"arguments": "{\n"}}`.
+    """
+    merged = left.copy()
+    for right in others:
+        for right_k, right_v in right.items():
+            if right_k not in merged or (
+                right_v is not None and merged[right_k] is None
+            ):
+                merged[right_k] = right_v
+            elif right_v is None:
+                continue
+            elif type(merged[right_k]) is not type(right_v):
+                msg = (
+                    f'additional_kwargs["{right_k}"] already exists in this message,'
+                    " but with a different type."
+                )
+                raise TypeError(msg)
+            elif isinstance(merged[right_k], str):
+                # TODO: Add below special handling for 'type' key in 0.3 and remove
+                # merge_lists 'type' logic.
+                #
+                # if right_k == "type":
+                #     if merged[right_k] == right_v:
+                #         continue
+                #     else:
+                #         raise ValueError(
+                #             "Unable to merge. Two different values seen for special "
+                #             f"key 'type': {merged[right_k]} and {right_v}. 'type' "
+                #             "should either occur once or have the same value across "
+                #             "all dicts."
+                #         )
+                if (right_k == "index" and merged[right_k].startswith("lc_")) or (
+                    right_k in {"id", "output_version", "model_provider"}
+                    and merged[right_k] == right_v
+                ):
+                    continue
+                merged[right_k] += right_v
+            elif isinstance(merged[right_k], dict):
+                merged[right_k] = merge_dicts(merged[right_k], right_v)
+            elif isinstance(merged[right_k], list):
+                merged[right_k] = merge_lists(merged[right_k], right_v)
+            elif merged[right_k] == right_v:
+                continue
+            elif isinstance(merged[right_k], bool):
+                # `bool` is a subclass of `int`, so without this check differing
+                # booleans would fall into the int branch below and get summed,
+                # silently turning e.g. `True + False` into the int `1`.
+                msg = (
+                    f"Additional kwargs key {right_k} already exists in left dict and "
+                    f"value has unsupported type {type(merged[right_k])}."
+                )
+                raise TypeError(msg)
+            elif isinstance(merged[right_k], int):
+                # Preserve identification and temporal fields using last-wins strategy
+                # instead of summing:
+                # - index: identifies which tool call a chunk belongs to
+                # - created/timestamp: temporal values that shouldn't be accumulated
+                if right_k in {"index", "created", "timestamp"}:
+                    merged[right_k] = right_v
+                else:
+                    merged[right_k] += right_v
+            else:
+                msg = (
+                    f"Additional kwargs key {right_k} already exists in left dict and "
+                    f"value has unsupported type {type(merged[right_k])}."
+                )
+                raise TypeError(msg)
+    return merged
+
+
+def merge_lists(left: list[Any] | None, *others: list[Any] | None) -> list[Any] | None:
+    """Add many lists, handling `None`.
+
+    Args:
+        left: The first list to merge.
+        others: The other lists to merge.
+
+    Returns:
+        The merged list.
+    """
+    merged = left.copy() if left is not None else None
+    for other in others:
+        if other is None:
+            continue
+        if merged is None:
+            merged = other.copy()
+        else:
+            for e in other:
+                if (
+                    isinstance(e, dict)
+                    and "index" in e
+                    and (
+                        isinstance(e["index"], int)
+                        or (
+                            isinstance(e["index"], str) and e["index"].startswith("lc_")
+                        )
+                    )
+                ):
+                    to_merge = [
+                        i
+                        for i, e_left in enumerate(merged)
+                        if (
+                            isinstance(e_left, dict)
+                            and "index" in e_left
+                            and e_left["index"] == e["index"]  # index matches
+                            and (  # IDs not inconsistent
+                                e_left.get("id") in {None, ""}
+                                or e.get("id") in {None, ""}
+                                or e_left.get("id") == e.get("id")
+                            )
+                        )
+                    ]
+                    if to_merge:
+                        # TODO: Remove this once merge_dict is updated with special
+                        # handling for 'type'.
+                        if (left_type := merged[to_merge[0]].get("type")) and (
+                            e.get("type") == "non_standard" and "value" in e
+                        ):
+                            if left_type != "non_standard":
+                                # standard + non_standard
+                                new_e: dict[str, Any] = {
+                                    "extras": {
+                                        k: v
+                                        for k, v in e["value"].items()
+                                        if k != "type"
+                                    }
+                                }
+                            else:
+                                # non_standard + non_standard
+                                new_e = {
+                                    "value": {
+                                        k: v
+                                        for k, v in e["value"].items()
+                                        if k != "type"
+                                    }
+                                }
+                                if "index" in e:
+                                    new_e["index"] = e["index"]
+                        else:
+                            new_e = (
+                                {k: v for k, v in e.items() if k != "type"}
+                                if "type" in e
+                                else e
+                            )
+                        merged[to_merge[0]] = merge_dicts(merged[to_merge[0]], new_e)
+                    else:
+                        merged.append(e)
+                else:
+                    merged.append(e)
+    return merged
+
+
+def merge_obj(left: Any, right: Any) -> Any:
+    """Merge two objects.
+
+    It handles specific scenarios where a key exists in both dictionaries but has a
+    value of `None` in `'left'`. In such cases, the method uses the value from `'right'`
+    for that key in the merged dictionary.
+
+    Args:
+        left: The first object to merge.
+        right: The other object to merge.
+
+    Returns:
+        The merged object.
+
+    Raises:
+        TypeError: If the key exists in both dictionaries but has a different type.
+        ValueError: If the two objects cannot be merged.
+    """
+    if left is None or right is None:
+        return left if left is not None else right
+    if type(left) is not type(right):
+        msg = (
+            f"left and right are of different types. Left type:  {type(left)}. Right "
+            f"type: {type(right)}."
+        )
+        raise TypeError(msg)
+    if isinstance(left, str):
+        return left + right
+    if isinstance(left, dict):
+        return merge_dicts(left, right)
+    if isinstance(left, list):
+        return merge_lists(left, right)
+    if left == right:
+        return left
+    msg = (
+        f"Unable to merge {left=} and {right=}. Both must be of type str, dict, or "
+        f"list, or else be two equal objects."
+    )
+    raise ValueError(msg)