206 lines
8.4 KiB
Go
206 lines
8.4 KiB
Go
package openai
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"testing"
|
|
|
|
edgeservice "iop/apps/edge/internal/service"
|
|
"iop/packages/go/config"
|
|
)
|
|
|
|
func newTestRequestContext(t *testing.T, route routeDispatch, rawBody []byte) openAIRequestContext {
|
|
t.Helper()
|
|
return openAIRequestContext{
|
|
r: httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil),
|
|
route: route,
|
|
rawBody: rawBody,
|
|
callerMetadata: map[string]string{},
|
|
workspace: "ws-1",
|
|
estimate: 11,
|
|
contextClass: "normal",
|
|
endpoint: usageEndpointChatCompletions,
|
|
}
|
|
}
|
|
|
|
// dispatchContextFixture returns a server with strict output enabled, a
|
|
// decoded request, and the caller's raw body.
|
|
func dispatchContextFixture(t *testing.T) (*Server, chatCompletionRequest, []byte) {
|
|
t.Helper()
|
|
srv := NewServer(config.EdgeOpenAIConf{
|
|
Adapter: "ollama",
|
|
Target: "llama-fixed",
|
|
StrictOutput: true,
|
|
}, &fakeRunService{}, nil)
|
|
req := chatCompletionRequest{
|
|
Model: "llama3",
|
|
Messages: []chatMessage{{Role: "user", Content: "hello"}},
|
|
}
|
|
return srv, req, []byte(`{"model":"llama3"}`)
|
|
}
|
|
|
|
// TestNewChatDispatchContextDirectRoute pins that the direct route resolves the
|
|
// adapter, target, and strict-output contract up front.
|
|
func TestNewChatDispatchContextDirectRoute(t *testing.T) {
|
|
srv, req, rawBody := dispatchContextFixture(t)
|
|
|
|
t.Run("direct route resolves adapter, target, and strict instruction", func(t *testing.T) {
|
|
route := routeDispatch{Adapter: "ollama", Target: "llama-fixed", SessionID: "cli"}
|
|
requestCtx := newTestRequestContext(t, route, rawBody)
|
|
basePrompt := promptFromMessages(req.Messages)
|
|
outputPolicy := srv.resolveOutputPolicy(basePrompt)
|
|
|
|
dc := srv.newChatDispatchContext(requestCtx, req, basePrompt, outputPolicy)
|
|
|
|
if dc.submitReq.Adapter != "ollama" || dc.submitReq.Target != "llama-fixed" {
|
|
t.Fatalf("expected the direct route to resolve adapter/target, got %+v", dc.submitReq)
|
|
}
|
|
if dc.submitReq.ProviderPool {
|
|
t.Fatal("expected the direct route not to be marked provider-pool")
|
|
}
|
|
if instruction := strictOutputContractInstruction(outputPolicy); instruction != "" && dc.prompt == basePrompt {
|
|
t.Fatal("expected the strict-output instruction to be prepended on the direct route")
|
|
}
|
|
if dc.submitReq.Workspace != "ws-1" {
|
|
t.Fatalf("expected the workspace to survive into the run request, got %q", dc.submitReq.Workspace)
|
|
}
|
|
})
|
|
|
|
t.Run("run metadata carries the estimate and context class", func(t *testing.T) {
|
|
route := routeDispatch{Adapter: "ollama", Target: "llama-fixed"}
|
|
requestCtx := newTestRequestContext(t, route, rawBody)
|
|
basePrompt := promptFromMessages(req.Messages)
|
|
|
|
dc := srv.newChatDispatchContext(requestCtx, req, basePrompt, srv.resolveOutputPolicy(basePrompt))
|
|
|
|
if dc.runMetadata["context_class"] != dc.contextClass {
|
|
t.Fatalf("expected metadata context_class %q, got %q", dc.contextClass, dc.runMetadata["context_class"])
|
|
}
|
|
if dc.runMetadata["openai_model"] != "llama3" {
|
|
t.Fatalf("expected metadata openai_model to be set, got %q", dc.runMetadata["openai_model"])
|
|
}
|
|
if dc.submitReq.EstimatedInputTokens != dc.estimate {
|
|
t.Fatalf("expected the run request estimate %d to match the context, got %d", dc.estimate, dc.submitReq.EstimatedInputTokens)
|
|
}
|
|
})
|
|
}
|
|
|
|
// TestNewChatDispatchContextProviderPoolRoute pins that the provider-pool route
|
|
// defers adapter/target and the strict-output contract to post-admission
|
|
// dispatch, and preserves the caller's raw body for passthrough.
|
|
func TestNewChatDispatchContextProviderPoolRoute(t *testing.T) {
|
|
srv, req, rawBody := dispatchContextFixture(t)
|
|
|
|
t.Run("provider-pool route defers adapter, target, and strict instruction", func(t *testing.T) {
|
|
route := routeDispatch{SessionID: "cli", ProviderPool: true}
|
|
requestCtx := newTestRequestContext(t, route, rawBody)
|
|
basePrompt := promptFromMessages(req.Messages)
|
|
outputPolicy := srv.resolveOutputPolicy(basePrompt)
|
|
|
|
dc := srv.newChatDispatchContext(requestCtx, req, basePrompt, outputPolicy)
|
|
|
|
if dc.submitReq.Adapter != "" || dc.submitReq.Target != "" {
|
|
t.Fatalf("expected the pool route to leave adapter/target to admission, got %+v", dc.submitReq)
|
|
}
|
|
if !dc.submitReq.ProviderPool {
|
|
t.Fatal("expected the pool route to be marked provider-pool")
|
|
}
|
|
if dc.submitReq.ModelGroupKey != "llama3" {
|
|
t.Fatalf("expected the model group key to reach admission, got %q", dc.submitReq.ModelGroupKey)
|
|
}
|
|
// Strict output is a normalized-dispatch contract; baking it into the
|
|
// pool prompt would leak it into raw tunnel passthrough (SDD D02).
|
|
if dc.prompt != basePrompt {
|
|
t.Fatalf("expected the pool prompt to stay unaugmented, got %q", dc.prompt)
|
|
}
|
|
// The pool estimate is the ingress estimate: the prompt was not
|
|
// re-augmented, so it must not be recomputed.
|
|
if dc.estimate != requestCtx.estimate {
|
|
t.Fatalf("expected the pool route to keep the ingress estimate %d, got %d", requestCtx.estimate, dc.estimate)
|
|
}
|
|
})
|
|
|
|
t.Run("raw body is preserved separately from the run request", func(t *testing.T) {
|
|
route := routeDispatch{SessionID: "cli", ProviderPool: true}
|
|
requestCtx := newTestRequestContext(t, route, rawBody)
|
|
basePrompt := promptFromMessages(req.Messages)
|
|
|
|
dc := srv.newChatDispatchContext(requestCtx, req, basePrompt, srv.resolveOutputPolicy(basePrompt))
|
|
|
|
if string(dc.rawBody) != string(rawBody) {
|
|
t.Fatalf("expected the caller's raw body to be preserved verbatim, got %s", dc.rawBody)
|
|
}
|
|
})
|
|
|
|
}
|
|
|
|
// TestChatDispatchContextWithRetrySubmitDoesNotMutate pins that deriving a
|
|
// retry-bound context leaves the original context untouched.
|
|
func TestChatDispatchContextWithRetrySubmitDoesNotMutate(t *testing.T) {
|
|
dc := newTestChatDispatchContext(httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil), chatCompletionRequest{Model: "llama3"})
|
|
dc.retrySubmit = nil
|
|
|
|
derived := dc.withRetrySubmit(func(ctx context.Context, req edgeservice.SubmitRunRequest) (any, error) {
|
|
return nil, nil
|
|
})
|
|
|
|
if dc.retrySubmit != nil {
|
|
t.Fatal("expected the source context to stay unmodified")
|
|
}
|
|
if derived.retrySubmit == nil {
|
|
t.Fatal("expected the derived context to carry the retry function")
|
|
}
|
|
if derived.req.Model != dc.req.Model {
|
|
t.Fatal("expected the derived context to keep the request")
|
|
}
|
|
}
|
|
|
|
// TestDispatchContextsDeriveMetadataWithoutMutatingIngress pins the immutable
|
|
// boundary between ingress identity metadata and path-specific run metadata.
|
|
// Chat and Responses add different execution fields, but neither may change
|
|
// the metadata snapshot a raw provider-tunnel branch could still read.
|
|
func TestDispatchContextsDeriveMetadataWithoutMutatingIngress(t *testing.T) {
|
|
srv, chatReq, rawBody := dispatchContextFixture(t)
|
|
chatIngress := newTestRequestContext(t, routeDispatch{Adapter: "ollama", Target: "llama-fixed"}, rawBody)
|
|
chatIngress.callerMetadata = map[string]string{"principal_id": "principal-1"}
|
|
|
|
chatDC := srv.newChatDispatchContext(
|
|
chatIngress,
|
|
chatReq,
|
|
promptFromMessages(chatReq.Messages),
|
|
srv.resolveOutputPolicy(promptFromMessages(chatReq.Messages)),
|
|
)
|
|
if _, ok := chatIngress.callerMetadata["openai_model"]; ok {
|
|
t.Fatal("expected chat ingress metadata to remain unchanged")
|
|
}
|
|
if chatDC.runMetadata["principal_id"] != "principal-1" || chatDC.runMetadata["openai_model"] != "llama3" {
|
|
t.Fatalf("expected derived chat metadata, got %#v", chatDC.runMetadata)
|
|
}
|
|
|
|
responsesIngress := &responsesRequestContext{
|
|
openAIRequestContext: openAIRequestContext{
|
|
r: httptest.NewRequest(http.MethodPost, "/v1/responses", nil),
|
|
route: routeDispatch{Adapter: "ollama", Target: "llama-fixed", SessionID: "default"},
|
|
rawBody: []byte(`{"model":"llama3","input":"hello"}`),
|
|
callerMetadata: map[string]string{"principal_id": "principal-1"},
|
|
workspace: "ws-1",
|
|
endpoint: usageEndpointResponses,
|
|
},
|
|
envelope: responsesEnvelope{Model: "llama3"},
|
|
}
|
|
responsesDC, err := srv.newResponsesDispatchContext(responsesIngress, responsesRequest{
|
|
Model: "llama3",
|
|
Input: json.RawMessage(`"hello"`),
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("newResponsesDispatchContext() error = %v", err)
|
|
}
|
|
if _, ok := responsesIngress.callerMetadata["openai_model"]; ok {
|
|
t.Fatal("expected responses ingress metadata to remain unchanged")
|
|
}
|
|
if responsesDC.submitReq.Metadata["principal_id"] != "principal-1" || responsesDC.submitReq.Metadata["openai_model"] != "llama3" {
|
|
t.Fatalf("expected derived responses metadata, got %#v", responsesDC.submitReq.Metadata)
|
|
}
|
|
}
|