Class OpenAICodec

python/tests/test_integration_codecs.py:26–76

Codec for OpenAI Chat Completions API payloads. Decodes flat dict payloads (messages, model, temperature, etc.) into AnnotatedLLMRequest. Unrecognized keys go to extra for lossless round-trip.

class OpenAICodec(LlmCodec):
    """Codec for OpenAI Chat Completions API payloads.

    Decodes flat dict payloads (messages, model, temperature, etc.) into
    AnnotatedLLMRequest. Unrecognized keys go to extra for lossless
    round-trip.
    """

    _PARAM_KEYS = {"temperature", "max_tokens", "max_completion_tokens", "top_p", "stop"}
    _MODELED_KEYS = {"messages", "model", "tools", "tool_choice"} | _PARAM_KEYS

    def decode(self, request):
        c = request.content
        messages = c.get("messages", [])
        model = c.get("model")
        params = {}
        for k in self._PARAM_KEYS:
            if k in c and c[k] is not None:
                # Normalize max_completion_tokens -> max_tokens
                key = "max_tokens" if k == "max_completion_tokens" else k
                params[key] = c[k]
        extra = {k: v for k, v in c.items() if k not in self._MODELED_KEYS}
        return AnnotatedLLMRequest(
            messages,
            model=model,
            params=params or None,
            tools=c.get("tools"),
            tool_choice=c.get("tool_choice"),
            extra=extra or None,
        )

    def encode(self, annotated, original):
        content = dict(original.content)
        content["messages"] = annotated.messages
        if annotated.model is not None:
            content["model"] = annotated.model
        if annotated.params:
            for k, v in annotated.params.items():
                # Write back to max_completion_tokens when original used that key
                if k == "max_tokens" and "max_completion_tokens" in original.content:
                    content["max_completion_tokens"] = v
                else:
                    content[k] = v
        if annotated.tools is not None:
            content["tools"] = annotated.tools
        if annotated.tool_choice is not None:
            content["tool_choice"] = annotated.tool_choice
        if annotated.extra:
            for k, v in annotated.extra.items():
                content[k] = v
        return LLMRequest(original.headers, content)


class NIMCodec(OpenAICodec):