mirror of
https://github.com/Wei-Shaw/sub2api.git
synced 2026-10-07 15:38:22 +08:00
复用请求级图片生成意图判断
This commit is contained in:
@@ -56,6 +56,13 @@ func openAIModelMappedBody(body []byte, mapped bool, mappedModel string, replace
|
||||
return replace(body, mappedModel)
|
||||
}
|
||||
|
||||
func seedOpenAIForwardImageIntentHint(c *gin.Context, channelMapped bool, imageIntent bool) {
|
||||
if channelMapped {
|
||||
return
|
||||
}
|
||||
service.SetOpenAIImageIntentHint(c, imageIntent)
|
||||
}
|
||||
|
||||
func newOpenAIModelMappedBodyCache(body []byte, replace openAIModelBodyReplaceFunc) func(bool, string) []byte {
|
||||
replacedBodies := make(map[string][]byte)
|
||||
return func(mapped bool, mappedModel string) []byte {
|
||||
@@ -282,6 +289,7 @@ func (h *OpenAIGatewayHandler) Responses(c *gin.Context) {
|
||||
// 解析渠道级模型映射
|
||||
channelMapping, _ := h.gatewayService.ResolveChannelMappingAndRestrict(c.Request.Context(), apiKey.GroupID, reqModel)
|
||||
forwardBody := openAIModelMappedBody(body, channelMapping.Mapped, channelMapping.MappedModel, h.gatewayService.ReplaceModelInBody)
|
||||
seedOpenAIForwardImageIntentHint(c, channelMapping.Mapped, imageIntent)
|
||||
|
||||
// 提前校验 function_call_output 是否具备可关联上下文,避免上游 400。
|
||||
if !h.validateFunctionCallOutputRequest(c, body, reqLog) {
|
||||
|
||||
@@ -0,0 +1,44 @@
|
||||
package handler
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestSeedOpenAIForwardImageIntentHint(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
tests := []struct {
|
||||
name string
|
||||
channelMapped bool
|
||||
imageIntent bool
|
||||
wantHint bool
|
||||
}{
|
||||
{name: "seed true", imageIntent: true, wantHint: true},
|
||||
{name: "seed false", imageIntent: false, wantHint: true},
|
||||
{name: "mapped body stays unknown", channelMapped: true, imageIntent: true},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
c := &gin.Context{}
|
||||
service.SetOpenAIClientTransport(c, service.OpenAIClientTransportHTTP)
|
||||
|
||||
seedOpenAIForwardImageIntentHint(c, tt.channelMapped, tt.imageIntent)
|
||||
|
||||
var hintValues []bool
|
||||
for _, value := range c.Keys {
|
||||
if hint, ok := value.(bool); ok {
|
||||
hintValues = append(hintValues, hint)
|
||||
}
|
||||
}
|
||||
if !tt.wantHint {
|
||||
require.Empty(t, hintValues)
|
||||
return
|
||||
}
|
||||
require.Equal(t, []bool{tt.imageIntent}, hintValues)
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -20,6 +20,7 @@ import (
|
||||
// Forward forwards request to OpenAI API
|
||||
func (s *OpenAIGatewayService) Forward(ctx context.Context, c *gin.Context, account *Account, body []byte) (*OpenAIForwardResult, error) {
|
||||
startTime := time.Now()
|
||||
canonicalImageIntentBody := body
|
||||
|
||||
restrictionResult := s.detectCodexClientRestriction(c, account, body)
|
||||
apiKeyID := getAPIKeyIDFromContext(c)
|
||||
@@ -107,6 +108,7 @@ func (s *OpenAIGatewayService) Forward(ctx context.Context, c *gin.Context, acco
|
||||
return nil, errors.New("openai ws v1 is temporarily unsupported; use ws v2")
|
||||
}
|
||||
if passthroughEnabled {
|
||||
attemptImageIntentInvalidated := false
|
||||
if isCodexCLI && codexImageGenerationExplicitToolPolicy == codexImageGenerationExplicitToolPolicyStrip {
|
||||
strippedBody, changed, stripErr := stripOpenAIImageGenerationToolsFromRawPayload(body)
|
||||
if stripErr != nil {
|
||||
@@ -115,6 +117,7 @@ func (s *OpenAIGatewayService) Forward(ctx context.Context, c *gin.Context, acco
|
||||
if changed {
|
||||
body = strippedBody
|
||||
originalBody = strippedBody
|
||||
attemptImageIntentInvalidated = true
|
||||
logger.LegacyPrintf("service.openai_gateway", "[OpenAI] Stripped /responses image_generation tool for Codex client by account policy")
|
||||
}
|
||||
}
|
||||
@@ -123,7 +126,18 @@ func (s *OpenAIGatewayService) Forward(ctx context.Context, c *gin.Context, acco
|
||||
reasoningEffort := extractOpenAIReasoningEffortFromBody(body, mappedModel)
|
||||
// 国产模型默认 effort 补充:也要用 mappedModel 判定是否是 passback-required 上游。
|
||||
reasoningEffort = ApplyThinkingEnabledFallback(reasoningEffort, body, mappedModel)
|
||||
return s.forwardOpenAIPassthrough(ctx, c, account, originalBody, reqModel, reasoningEffort, reqStream, startTime)
|
||||
return s.forwardOpenAIPassthrough(
|
||||
ctx,
|
||||
c,
|
||||
account,
|
||||
originalBody,
|
||||
canonicalImageIntentBody,
|
||||
reqModel,
|
||||
attemptImageIntentInvalidated,
|
||||
reasoningEffort,
|
||||
reqStream,
|
||||
startTime,
|
||||
)
|
||||
}
|
||||
|
||||
bodyModified := false
|
||||
@@ -188,6 +202,7 @@ func (s *OpenAIGatewayService) Forward(ctx context.Context, c *gin.Context, acco
|
||||
codexImageGenerationExplicitToolPolicy != codexImageGenerationExplicitToolPolicyStrip &&
|
||||
s.isCodexImageGenerationBridgeEnabled(ctx, account, apiKey)
|
||||
var imageIntent bool
|
||||
canonicalImageIntent := resolveOpenAIImageIntentHint(c, reqModel, canonicalImageIntentBody, IsImageGenerationIntent)
|
||||
if isCodexCLI && codexImageGenerationExplicitToolPolicy == codexImageGenerationExplicitToolPolicyStrip {
|
||||
decoded, decodeErr := ensureReqBody()
|
||||
if decodeErr != nil {
|
||||
@@ -199,7 +214,7 @@ func (s *OpenAIGatewayService) Forward(ctx context.Context, c *gin.Context, acco
|
||||
}
|
||||
imageIntent = IsImageGenerationIntentMap(openAIResponsesEndpoint, reqModel, decoded)
|
||||
} else {
|
||||
imageIntent = IsImageGenerationIntent(openAIResponsesEndpoint, reqModel, body)
|
||||
imageIntent = canonicalImageIntent
|
||||
}
|
||||
if imageIntent && !imageGenerationAllowed {
|
||||
MarkOpsClientBusinessLimited(c, OpsClientBusinessLimitedReasonLocalFeatureGate)
|
||||
|
||||
@@ -29,7 +29,9 @@ func (s *OpenAIGatewayService) forwardOpenAIPassthrough(
|
||||
c *gin.Context,
|
||||
account *Account,
|
||||
body []byte,
|
||||
canonicalImageIntentBody []byte,
|
||||
reqModel string,
|
||||
attemptImageIntentInvalidated bool,
|
||||
reasoningEffort *string,
|
||||
reqStream bool,
|
||||
startTime time.Time,
|
||||
@@ -44,6 +46,7 @@ func (s *OpenAIGatewayService) forwardOpenAIPassthrough(
|
||||
}
|
||||
body = nextBody
|
||||
upstreamPassthroughModel = compactMappedModel
|
||||
attemptImageIntentInvalidated = true
|
||||
}
|
||||
}
|
||||
|
||||
@@ -100,7 +103,15 @@ func (s *OpenAIGatewayService) forwardOpenAIPassthrough(
|
||||
body = updatedBody
|
||||
|
||||
apiKey := getAPIKeyFromContext(c)
|
||||
imageIntent := IsImageGenerationIntent(openAIResponsesEndpoint, reqModel, body)
|
||||
imageIntent := resolveOpenAIPassthroughImageIntent(
|
||||
c,
|
||||
reqModel,
|
||||
canonicalImageIntentBody,
|
||||
policyModel,
|
||||
body,
|
||||
attemptImageIntentInvalidated,
|
||||
IsImageGenerationIntent,
|
||||
)
|
||||
if imageIntent && !GroupAllowsImageGeneration(apiKeyGroup(apiKey)) {
|
||||
MarkOpsClientBusinessLimited(c, OpsClientBusinessLimitedReasonLocalFeatureGate)
|
||||
c.JSON(http.StatusForbidden, gin.H{
|
||||
|
||||
@@ -221,6 +221,24 @@ func TestOpenAIGatewayService_Forward_MappedImageModelUsesImageGate(t *testing.T
|
||||
require.Nil(t, result)
|
||||
require.Nil(t, upstream.lastReq)
|
||||
require.Equal(t, http.StatusForbidden, rec.Code)
|
||||
cached, known := getOpenAIImageIntentHint(c)
|
||||
require.True(t, known)
|
||||
require.False(t, cached)
|
||||
|
||||
textAccount := *account
|
||||
textAccount.ID = 4
|
||||
textAccount.Credentials = map[string]any{
|
||||
"api_key": "sk-test",
|
||||
"base_url": "https://example.com",
|
||||
}
|
||||
result, err = svc.Forward(context.Background(), c, &textAccount, body)
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, result)
|
||||
require.NotNil(t, upstream.lastReq)
|
||||
require.Len(t, upstream.bodies, 1)
|
||||
cached, known = getOpenAIImageIntentHint(c)
|
||||
require.True(t, known)
|
||||
require.False(t, cached)
|
||||
}
|
||||
|
||||
func TestOpenAIGatewayService_Forward_TextResponsesSetsBillingModelToMappedModel(t *testing.T) {
|
||||
|
||||
@@ -242,6 +242,7 @@ func TestOpenAIGatewayServiceForward_AccountPolicyStripsImageNamespaceTools(t *t
|
||||
}
|
||||
svc := newOpenAIImageGenerationControlTestService(upstream)
|
||||
c, _ := newOpenAIImageGenerationControlTestContext(false, "codex_cli_rs/0.144.1")
|
||||
SetOpenAIClientTransport(c, OpenAIClientTransportHTTP)
|
||||
account := newOpenAIImageGenerationControlTestAccount()
|
||||
account.Extra = map[string]any{
|
||||
featureKeyCodexImageGenerationExplicitToolPolicy: codexImageGenerationExplicitToolPolicyStrip,
|
||||
@@ -274,6 +275,9 @@ func TestOpenAIGatewayServiceForward_AccountPolicyStripsImageNamespaceTools(t *t
|
||||
require.True(t, gjson.GetBytes(upstream.lastBody, `tools.#(name=="shell")`).Exists())
|
||||
require.True(t, gjson.GetBytes(upstream.lastBody, `tools.#(name=="code_tools")`).Exists())
|
||||
require.Equal(t, "write code", gjson.GetBytes(upstream.lastBody, "input.0.content.0.text").String())
|
||||
cached, known := getOpenAIImageIntentHint(c)
|
||||
require.True(t, known)
|
||||
require.True(t, cached)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
@@ -0,0 +1,57 @@
|
||||
package service
|
||||
|
||||
import "github.com/gin-gonic/gin"
|
||||
|
||||
const openAIImageIntentHintContextKey = "openai_image_intent_hint"
|
||||
|
||||
type openAIImageIntentClassifier func(endpoint string, requestedModel string, body []byte) bool
|
||||
|
||||
// SetOpenAIImageIntentHint records the canonical request body's image intent.
|
||||
func SetOpenAIImageIntentHint(c *gin.Context, imageIntent bool) {
|
||||
if c == nil || GetOpenAIClientTransport(c) != OpenAIClientTransportHTTP {
|
||||
return
|
||||
}
|
||||
c.Set(openAIImageIntentHintContextKey, imageIntent)
|
||||
}
|
||||
|
||||
func getOpenAIImageIntentHint(c *gin.Context) (imageIntent bool, known bool) {
|
||||
if c == nil || GetOpenAIClientTransport(c) != OpenAIClientTransportHTTP {
|
||||
return false, false
|
||||
}
|
||||
value, ok := c.Get(openAIImageIntentHintContextKey)
|
||||
if !ok {
|
||||
return false, false
|
||||
}
|
||||
imageIntent, ok = value.(bool)
|
||||
return imageIntent, ok
|
||||
}
|
||||
|
||||
func resolveOpenAIImageIntentHint(
|
||||
c *gin.Context,
|
||||
requestedModel string,
|
||||
canonicalBody []byte,
|
||||
classify openAIImageIntentClassifier,
|
||||
) bool {
|
||||
if imageIntent, known := getOpenAIImageIntentHint(c); known {
|
||||
return imageIntent
|
||||
}
|
||||
imageIntent := classify(openAIResponsesEndpoint, requestedModel, canonicalBody)
|
||||
SetOpenAIImageIntentHint(c, imageIntent)
|
||||
return imageIntent
|
||||
}
|
||||
|
||||
func resolveOpenAIPassthroughImageIntent(
|
||||
c *gin.Context,
|
||||
canonicalRequestedModel string,
|
||||
canonicalBody []byte,
|
||||
attemptRequestedModel string,
|
||||
attemptBody []byte,
|
||||
attemptInvalidated bool,
|
||||
classify openAIImageIntentClassifier,
|
||||
) bool {
|
||||
imageIntent := resolveOpenAIImageIntentHint(c, canonicalRequestedModel, canonicalBody, classify)
|
||||
if attemptInvalidated {
|
||||
imageIntent = classify(openAIResponsesEndpoint, attemptRequestedModel, attemptBody)
|
||||
}
|
||||
return imageIntent
|
||||
}
|
||||
@@ -0,0 +1,338 @@
|
||||
package service
|
||||
|
||||
import (
|
||||
"context"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/require"
|
||||
"github.com/tidwall/gjson"
|
||||
)
|
||||
|
||||
func newOpenAIImageIntentHintTestContext(transport OpenAIClientTransport) *gin.Context {
|
||||
c := &gin.Context{}
|
||||
SetOpenAIClientTransport(c, transport)
|
||||
return c
|
||||
}
|
||||
|
||||
func countingOpenAIImageIntentClassifier(calls *atomic.Int64) openAIImageIntentClassifier {
|
||||
return func(endpoint string, requestedModel string, body []byte) bool {
|
||||
calls.Add(1)
|
||||
return IsImageGenerationIntent(endpoint, requestedModel, body)
|
||||
}
|
||||
}
|
||||
|
||||
func TestResolveOpenAIImageIntentHintCachesTrueAndFalse(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
tests := []struct {
|
||||
name string
|
||||
body []byte
|
||||
want bool
|
||||
}{
|
||||
{name: "true", body: []byte(`{"model":"gpt-5.4","tools":[{"type":"image_generation"}]}`), want: true},
|
||||
{name: "false is known", body: []byte(`{"model":"gpt-5.4","input":"write code"}`), want: false},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
c := newOpenAIImageIntentHintTestContext(OpenAIClientTransportHTTP)
|
||||
var calls atomic.Int64
|
||||
classify := countingOpenAIImageIntentClassifier(&calls)
|
||||
|
||||
require.Equal(t, tt.want, resolveOpenAIImageIntentHint(c, "gpt-5.4", tt.body, classify))
|
||||
require.Equal(t, tt.want, resolveOpenAIImageIntentHint(c, "gpt-5.4", tt.body, classify))
|
||||
require.Equal(t, int64(1), calls.Load())
|
||||
cached, known := getOpenAIImageIntentHint(c)
|
||||
require.True(t, known)
|
||||
require.Equal(t, tt.want, cached)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestResolveOpenAIImageIntentHintUsesHandlerSeed(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
for _, seeded := range []bool{false, true} {
|
||||
c := newOpenAIImageIntentHintTestContext(OpenAIClientTransportHTTP)
|
||||
SetOpenAIImageIntentHint(c, seeded)
|
||||
var calls atomic.Int64
|
||||
|
||||
got := resolveOpenAIImageIntentHint(c, "gpt-5.4", []byte(`{"model":"gpt-5.4"}`), countingOpenAIImageIntentClassifier(&calls))
|
||||
|
||||
require.Equal(t, seeded, got)
|
||||
require.Zero(t, calls.Load())
|
||||
}
|
||||
}
|
||||
|
||||
func TestResolveOpenAIPassthroughImageIntentReusesCanonicalAcrossFailover(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
c := newOpenAIImageIntentHintTestContext(OpenAIClientTransportHTTP)
|
||||
body := []byte(`{"model":"gpt-5.4","input":"write code"}`)
|
||||
var calls atomic.Int64
|
||||
classify := countingOpenAIImageIntentClassifier(&calls)
|
||||
|
||||
for range 3 {
|
||||
require.False(t, resolveOpenAIPassthroughImageIntent(c, "gpt-5.4", body, "gpt-5.4", body, false, classify))
|
||||
}
|
||||
require.Equal(t, int64(1), calls.Load())
|
||||
}
|
||||
|
||||
func TestResolveOpenAIPassthroughImageIntentKeepsCompactMappingAttemptLocal(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
t.Run("text to image", func(t *testing.T) {
|
||||
c := newOpenAIImageIntentHintTestContext(OpenAIClientTransportHTTP)
|
||||
body := []byte(`{"model":"draw-alias","input":"draw"}`)
|
||||
compactBody := []byte(`{"model":"gpt-image-2","input":"draw"}`)
|
||||
var calls atomic.Int64
|
||||
classify := countingOpenAIImageIntentClassifier(&calls)
|
||||
|
||||
require.True(t, resolveOpenAIPassthroughImageIntent(c, "draw-alias", body, "gpt-image-2", compactBody, true, classify))
|
||||
cached, known := getOpenAIImageIntentHint(c)
|
||||
require.True(t, known)
|
||||
require.False(t, cached)
|
||||
|
||||
require.False(t, resolveOpenAIPassthroughImageIntent(c, "draw-alias", body, "draw-alias", body, false, classify))
|
||||
require.Equal(t, int64(2), calls.Load())
|
||||
cached, known = getOpenAIImageIntentHint(c)
|
||||
require.True(t, known)
|
||||
require.False(t, cached)
|
||||
})
|
||||
|
||||
t.Run("image to text", func(t *testing.T) {
|
||||
c := newOpenAIImageIntentHintTestContext(OpenAIClientTransportHTTP)
|
||||
body := []byte(`{"model":"gpt-image-2","input":"draw"}`)
|
||||
compactBody := []byte(`{"model":"gpt-5.4","input":"draw"}`)
|
||||
var calls atomic.Int64
|
||||
classify := countingOpenAIImageIntentClassifier(&calls)
|
||||
|
||||
require.False(t, resolveOpenAIPassthroughImageIntent(c, "gpt-image-2", body, "gpt-5.4", compactBody, true, classify))
|
||||
cached, known := getOpenAIImageIntentHint(c)
|
||||
require.True(t, known)
|
||||
require.True(t, cached)
|
||||
|
||||
require.True(t, resolveOpenAIPassthroughImageIntent(c, "gpt-image-2", body, "gpt-image-2", body, false, classify))
|
||||
require.Equal(t, int64(2), calls.Load())
|
||||
})
|
||||
}
|
||||
|
||||
func TestResolveOpenAIPassthroughImageIntentInvalidationDoesNotPolluteCanonical(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
c := newOpenAIImageIntentHintTestContext(OpenAIClientTransportHTTP)
|
||||
canonicalBody := []byte(`{"model":"gpt-5.4","tools":[{"type":"image_generation"}]}`)
|
||||
strippedBody := []byte(`{"model":"gpt-5.4","tools":[]}`)
|
||||
var calls atomic.Int64
|
||||
classify := countingOpenAIImageIntentClassifier(&calls)
|
||||
|
||||
require.False(t, resolveOpenAIPassthroughImageIntent(c, "gpt-5.4", canonicalBody, "gpt-5.4", strippedBody, true, classify))
|
||||
require.Equal(t, int64(2), calls.Load(), "unknown canonical and invalidated attempt are classified independently")
|
||||
cached, known := getOpenAIImageIntentHint(c)
|
||||
require.True(t, known)
|
||||
require.True(t, cached)
|
||||
|
||||
require.True(t, resolveOpenAIPassthroughImageIntent(c, "gpt-5.4", canonicalBody, "gpt-5.4", canonicalBody, false, classify))
|
||||
require.Equal(t, int64(2), calls.Load())
|
||||
}
|
||||
|
||||
func TestResolveOpenAIPassthroughImageIntentMappedBodyStartsUnknownThenSeedsCanonical(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
c := newOpenAIImageIntentHintTestContext(OpenAIClientTransportHTTP)
|
||||
canonicalBody := []byte(`{"model":"gpt-image-2","input":"draw"}`)
|
||||
strippedAttemptBody := []byte(`{"model":"gpt-5.4","input":"draw"}`)
|
||||
_, known := getOpenAIImageIntentHint(c)
|
||||
require.False(t, known)
|
||||
var calls atomic.Int64
|
||||
classify := countingOpenAIImageIntentClassifier(&calls)
|
||||
|
||||
require.False(t, resolveOpenAIPassthroughImageIntent(c, "gpt-image-2", canonicalBody, "gpt-5.4", strippedAttemptBody, true, classify))
|
||||
require.Equal(t, int64(2), calls.Load())
|
||||
cached, known := getOpenAIImageIntentHint(c)
|
||||
require.True(t, known)
|
||||
require.True(t, cached)
|
||||
}
|
||||
|
||||
func TestResolveOpenAIPassthroughImageIntentReusesAcrossInvariantMutations(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
tests := []struct {
|
||||
name string
|
||||
canonicalBody []byte
|
||||
attemptBody []byte
|
||||
want bool
|
||||
}{
|
||||
{
|
||||
name: "oauth sanitize fast policy and reasoning",
|
||||
canonicalBody: []byte(`{"model":"gpt-5.4","input":[{"type":"input_image","image_url":"data:image/png;base64,"}],"service_tier":"fast","reasoning":{"effort":"minimal"}}`),
|
||||
attemptBody: []byte(`{"model":"gpt-5.4","input":[],"service_tier":"priority","reasoning":{"effort":"none"},"store":false,"stream":true}`),
|
||||
want: false,
|
||||
},
|
||||
{
|
||||
name: "namespace flatten",
|
||||
canonicalBody: []byte(`{"model":"gpt-5.4","tools":[{"type":"namespace","name":"code_tools"}]}`),
|
||||
attemptBody: []byte(`{"model":"gpt-5.4","tools":[{"type":"function","name":"code_tools.run"}]}`),
|
||||
want: false,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
c := newOpenAIImageIntentHintTestContext(OpenAIClientTransportHTTP)
|
||||
var calls atomic.Int64
|
||||
classify := countingOpenAIImageIntentClassifier(&calls)
|
||||
|
||||
require.Equal(t, tt.want, resolveOpenAIPassthroughImageIntent(c, "gpt-5.4", tt.canonicalBody, "gpt-5.4", tt.attemptBody, false, classify))
|
||||
require.Equal(t, tt.want, resolveOpenAIPassthroughImageIntent(c, "gpt-5.4", tt.canonicalBody, "gpt-5.4", tt.attemptBody, false, classify))
|
||||
require.Equal(t, int64(1), calls.Load())
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestOpenAIGatewayServicePassthroughCompactImageIntentIsAttemptLocal(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
tests := []struct {
|
||||
name string
|
||||
canonicalModel string
|
||||
compactModel string
|
||||
wantRejected bool
|
||||
wantCanonical bool
|
||||
}{
|
||||
{
|
||||
name: "text to image rejects",
|
||||
canonicalModel: "gpt-5.4",
|
||||
compactModel: "gpt-image-2",
|
||||
wantRejected: true,
|
||||
},
|
||||
{
|
||||
name: "image to text reaches upstream",
|
||||
canonicalModel: "gpt-image-2",
|
||||
compactModel: "gpt-5.4",
|
||||
wantCanonical: true,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
upstream := &httpUpstreamRecorder{resp: &http.Response{
|
||||
StatusCode: http.StatusOK,
|
||||
Header: http.Header{"Content-Type": []string{"application/json"}},
|
||||
Body: io.NopCloser(strings.NewReader(`{"id":"resp_compact","model":"` + tt.compactModel + `","usage":{"input_tokens":1,"output_tokens":1}}`)),
|
||||
}}
|
||||
svc := newOpenAIImageGenerationControlTestService(upstream)
|
||||
c, recorder := newOpenAIImageGenerationControlTestContext(false, "unit-test-agent/1.0")
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/openai/v1/responses/compact", nil)
|
||||
SetOpenAIClientTransport(c, OpenAIClientTransportHTTP)
|
||||
account := newOpenAIImageGenerationControlTestAccount()
|
||||
account.Extra = map[string]any{"openai_passthrough": true}
|
||||
account.Credentials = map[string]any{
|
||||
"api_key": "sk-test",
|
||||
"compact_model_mapping": map[string]any{
|
||||
tt.canonicalModel: tt.compactModel,
|
||||
},
|
||||
}
|
||||
body := []byte(`{"model":"` + tt.canonicalModel + `","stream":false,"input":"draw"}`)
|
||||
|
||||
result, err := svc.Forward(context.Background(), c, account, body)
|
||||
|
||||
cached, known := getOpenAIImageIntentHint(c)
|
||||
require.True(t, known)
|
||||
require.Equal(t, tt.wantCanonical, cached)
|
||||
if tt.wantRejected {
|
||||
require.Error(t, err)
|
||||
require.Nil(t, result)
|
||||
require.Equal(t, http.StatusForbidden, recorder.Code)
|
||||
require.Nil(t, upstream.lastReq)
|
||||
return
|
||||
}
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, result)
|
||||
require.NotNil(t, upstream.lastReq)
|
||||
require.Equal(t, tt.compactModel, gjson.GetBytes(upstream.lastBody, "model").String())
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestResolveOpenAIImageIntentHintExcludesWebSocketAndUnknownTransport(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
for _, transport := range []OpenAIClientTransport{OpenAIClientTransportWS, OpenAIClientTransportUnknown} {
|
||||
c := newOpenAIImageIntentHintTestContext(transport)
|
||||
var calls atomic.Int64
|
||||
classify := countingOpenAIImageIntentClassifier(&calls)
|
||||
body := []byte(`{"model":"gpt-5.4","input":"write code"}`)
|
||||
|
||||
require.False(t, resolveOpenAIImageIntentHint(c, "gpt-5.4", body, classify))
|
||||
require.False(t, resolveOpenAIImageIntentHint(c, "gpt-5.4", body, classify))
|
||||
require.Equal(t, int64(2), calls.Load())
|
||||
_, known := getOpenAIImageIntentHint(c)
|
||||
require.False(t, known)
|
||||
}
|
||||
}
|
||||
|
||||
func TestResolveOpenAIImageIntentHintConcurrentRequestsAreIsolated(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
const requests = 32
|
||||
var calls atomic.Int64
|
||||
classify := countingOpenAIImageIntentClassifier(&calls)
|
||||
var wg sync.WaitGroup
|
||||
results := make([][2]bool, requests)
|
||||
|
||||
for i := range requests {
|
||||
wg.Add(1)
|
||||
go func(index int, image bool) {
|
||||
defer wg.Done()
|
||||
c := newOpenAIImageIntentHintTestContext(OpenAIClientTransportHTTP)
|
||||
body := []byte(`{"model":"gpt-5.4","input":"write code"}`)
|
||||
if image {
|
||||
body = []byte(`{"model":"gpt-5.4","tools":[{"type":"image_generation"}]}`)
|
||||
}
|
||||
results[index][0] = resolveOpenAIImageIntentHint(c, "gpt-5.4", body, classify)
|
||||
results[index][1] = resolveOpenAIImageIntentHint(c, "gpt-5.4", body, classify)
|
||||
}(i, i%2 == 0)
|
||||
}
|
||||
wg.Wait()
|
||||
for i, result := range results {
|
||||
require.Equal(t, i%2 == 0, result[0])
|
||||
require.Equal(t, result[0], result[1])
|
||||
}
|
||||
require.Equal(t, int64(requests), calls.Load())
|
||||
}
|
||||
|
||||
var openAIImageIntentHintBenchmarkSink bool
|
||||
|
||||
func BenchmarkOpenAIPassthroughImageIntentHintLargeBody(b *testing.B) {
|
||||
body := []byte(`{"model":"gpt-5.4","input":"` + strings.Repeat("x", 4<<20) + `"}`)
|
||||
const attempts = 4
|
||||
|
||||
b.Run("scan_each_attempt", func(b *testing.B) {
|
||||
c := newOpenAIImageIntentHintTestContext(OpenAIClientTransportHTTP)
|
||||
b.ReportAllocs()
|
||||
calls := 0
|
||||
for range b.N {
|
||||
c.Set(openAIImageIntentHintContextKey, struct{}{})
|
||||
for range attempts {
|
||||
calls++
|
||||
openAIImageIntentHintBenchmarkSink = IsImageGenerationIntent(openAIResponsesEndpoint, "gpt-5.4", body)
|
||||
}
|
||||
}
|
||||
b.ReportMetric(float64(calls)/float64(b.N), "classifier_calls/op")
|
||||
})
|
||||
|
||||
b.Run("request_scoped_hint", func(b *testing.B) {
|
||||
c := newOpenAIImageIntentHintTestContext(OpenAIClientTransportHTTP)
|
||||
b.ReportAllocs()
|
||||
calls := 0
|
||||
classify := func(endpoint string, requestedModel string, candidate []byte) bool {
|
||||
calls++
|
||||
return IsImageGenerationIntent(endpoint, requestedModel, candidate)
|
||||
}
|
||||
for range b.N {
|
||||
c.Set(openAIImageIntentHintContextKey, struct{}{})
|
||||
for range attempts {
|
||||
openAIImageIntentHintBenchmarkSink = resolveOpenAIPassthroughImageIntent(c, "gpt-5.4", body, "gpt-5.4", body, false, classify)
|
||||
}
|
||||
}
|
||||
b.ReportMetric(float64(calls)/float64(b.N), "classifier_calls/op")
|
||||
})
|
||||
}
|
||||
Reference in New Issue
Block a user