Use extProc's bodyMutation capability to allow setting model args

[tiswanso]

Hi ai-gateway devs,

Great job so far on this project! Very cool to see the progress and implementation details! After experimenting with the project to solve a specific use-case, I am wondering if there's any proposal yet to allow setting model args via an AIServiceBackend config (similar to headerMutation)?

From an analogous functionality at the client SDK level, this follows a similar pattern as model_kwargs: https://python.langchain.com/api_reference/openai/llms/langchain_openai.llms.base.OpenAI.html#langchain_openai.llms.base.OpenAI.model_kwargs

This functionality would be essential to enable the use of unmodified agents with various models they weren't originally implemented with exposing the myriad of exact client SDK options.

Is this type of functionality part of any existing proposed or in-progress work?

Detailed Thoughts/Proposal

Proposed API change

new modelArgMutation field in AIServicesBackend (similar to headerMutation)

apiVersion: aigateway.envoyproxy.io/v1alpha1
kind: AIServiceBackend
metadata:
  name: envoy-ai-gateway-custom
  namespace: default
spec:
  schema:
    name: AzureOpenAI
    version: 2024-08-01-preview
  backendRef:
    name: envoy-ai-gateway-custom
    kind: Backend
    group: gateway.envoyproxy.io
  headerMutation:
    set:
    - name: foo
      value: bar
  modelArgMutation:
    # args similar to model_kwargs in client SDKs: https://python.langchain.com/api_reference/openai/llms/langchain_openai.llms.base.OpenAI.html#langchain_openai.llms.base.OpenAI.model_kwargs
    set:
    - name: user
      value: '{"key": "my-key"}'

Use-case Background

In a use-case I was trying to see if I could get working out of the box with ai-gateway, I need to set a model arg for a specific model backend. The model backend is an corporate internal service wrapping access to a specific openAI compatible model but requires setting the model argument user:. I need to make setting this transparent to existing agents to enable them to use this internal model. I've done this type of thing in the past with hacky envoy Lua filters but was hoping setting model args would already be part of the ai-gateway project :-) -- unfortunately not the case.

However, it was easy to quickly follow the handling of the model name override in the extproc/translator code to see how the RequestBody methods can be enhanced to set model args in the bodyMutation.

Proof-of-concept

Quick PoC diff showing setting a user: arg

diff --git a/internal/extproc/translator/openai_openai.go b/internal/extproc/translator/openai_openai.go
index 1ccbca4..23dd2e2 100644
--- a/internal/extproc/translator/openai_openai.go
+++ b/internal/extproc/translator/openai_openai.go
@@ -74,6 +74,13 @@ func (o *openAIToOpenAITranslatorV1ChatCompletion) RequestBody(original []byte,
                        }},
                },
        }
+       modelArgUserOverride := "{\"key\": \"my-key\"}"
+       req.User = modelArgUserOverride
+       // set the model arg 'user' to be used for the request.
+       newBody, err = sjson.SetBytesOptions(original, "user", modelArgUserOverride, sjsonOptions)
+       if err != nil {
+               return nil, nil, fmt.Errorf("failed to set model arg for 'user': %w", err)
+       }
 
        if forceBodyMutation && len(newBody) == 0 {
                newBody = original

The model arg can be seen when tested with the quickstart envoy-ai-gateway-basic-testupstream mock model... and similar functionality with a hardcoded value in extproc/translator/openai_azureopenai.go.

[mathetake]

Cool, i think i got to understand the motivation and use case. Thank you for initiating the discussion! @tiswanso

I would rather make this called bodyMutation and inside it, i would like to expose JsonPatch (https://datatracker.ietf.org/doc/html/rfc6902) as it's well-defined standard. For example, Envoy Gateway has something called EnvoyPatchPolicy where they use it https://gateway.envoyproxy.io/v1.4/api/extension_types/#jsonpatchoperation

So, i am thinking about the change like the below

diff --git a/api/v1alpha1/ai_service_backend.go b/api/v1alpha1/ai_service_backend.go
index a99450ad..850416ab 100644
--- a/api/v1alpha1/ai_service_backend.go
+++ b/api/v1alpha1/ai_service_backend.go
@@ -6,6 +6,7 @@
 package v1alpha1
 
 import (
+       "github.com/envoyproxy/gateway/api/v1alpha1"
        metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
        gwapiv1 "sigs.k8s.io/gateway-api/apis/v1"
 )
@@ -67,10 +68,19 @@ type AIServiceBackendSpec struct {
        // +optional
        HeaderMutation *HTTPHeaderMutation `json:"headerMutation,omitempty"`
 
+       // BodyMutation defines the mutation of HTTP body that will be applied to the request
+       // before sending it to the backend.
+       // +optional
+       BodyMutation *HTTPBodyMutation `json:"bodyMutation,omitempty"`
+
        // TODO: maybe add backend-level LLMRequestCost configuration that overrides the AIGatewayRoute-level LLMRequestCost.
        //      That may be useful for the backend that has a different cost calculation logic.
 }
 
+type HTTPBodyMutation struct {
+       JSONPatches []v1alpha1.JSONPatchOperation `json:"jsonPatches,omitempty"`
+}
+
 // HTTPHeaderMutation defines the mutation of HTTP headers that will be applied to the request
 type HTTPHeaderMutation struct {
        // Set overwrites/adds the request with the given header (name, value)

wdyt?

[tiswanso]

nice! Yeah, that'd definitely work! It could be very powerful for setting any fields in the json payload. Likely more future-proof than just a list of {key, val} pairs. I'd be a little concerned with API complexity for that field, though.

I think my case would be like:

apiVersion: aigateway.envoyproxy.io/v1alpha1
kind: AIServiceBackend
metadata:
  name: envoy-ai-gateway-custom
  namespace: default
spec:
  schema:
    name: AzureOpenAI
    version: 2024-08-01-preview
  backendRef:
    name: envoy-ai-gateway-custom
    kind: Backend
    group: gateway.envoyproxy.io
  headerMutation:
    set:
    - name: foo
      value: bar
  bodyMutation:
    - op: add
       path: "user"
       value:
          key: mykey

But folks can do some complicated things with jsonpath and from settings that might be tough to maintain.

One thing I was thinking about was possibly needing valueFrom functionality to get values from configmaps/secrets. I don't think the JsonPatch can do that right?

e.g.

      valueFrom:
        configMapKeyRef:
          name: my-configmap
          key: my-config-key

[mathetake]

yeah I think we don't necessarily need to reuse the EG's jsonpath as-is but define our own version including the valueFrom secret/configmap stuff. My main point was not to tie this API with something like "model" but more about having a generic mechanism like jsonpatch since it can be used to insert vendor specific fields as well.

[tiswanso]

Sounds good! Just to add, if the implementation is a lot more complicated with valueFrom semantics I think it's a fine start to leave it out.

[tiswanso]

I think the existence of this feature would allow folks to have a configuration path to create bodyMutations that functionally do the same as something like: #1396

NOTE: my comment is not in any opposition to adding model specific functionality as more first class config. It's to point out the power of the approach for exposing config for custom bodyMutations using jsonpatches as @mathetake describes.

[mathetake]

so the #1396 and this will have a bit of different personas: The initial effort towards "vendor specific fields" is described in https://github.com/envoyproxy/ai-gateway/blob/main/docs/proposals/004-vendor-specific-fields/proposal.md and it focuses on giving clients more control over vendor specific fields. On the other hand, this body mutation thingy will give that ability to the Gateway operators to enforce some Gateway level rule such as vendor specific safe guard configuration, max token, etc. I think both use cases are orthogonal to each other

[mathetake]

anyways i think we can add this body mutation stuff! The application of jsonpatches should happen after the translation is done at the later stage of request body processing