OpenAI compat: request stream usage
This commit is contained in:
parent
04e99a1264
commit
fad2cc7b65
2 changed files with 96 additions and 5 deletions
|
|
@ -257,6 +257,9 @@ func (p *Provider) ChatStream(
|
||||||
|
|
||||||
requestBody := p.buildRequestBody(messages, tools, model, options)
|
requestBody := p.buildRequestBody(messages, tools, model, options)
|
||||||
requestBody["stream"] = true
|
requestBody["stream"] = true
|
||||||
|
if supportsStreamingUsage(p.apiBase) {
|
||||||
|
requestBody["stream_options"] = map[string]any{"include_usage": true}
|
||||||
|
}
|
||||||
|
|
||||||
jsonData, err := json.Marshal(requestBody)
|
jsonData, err := json.Marshal(requestBody)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
|
|
@ -479,11 +482,10 @@ func isNativeSearchHost(apiBase string) bool {
|
||||||
return host == "api.openai.com" || strings.HasSuffix(host, ".openai.azure.com")
|
return host == "api.openai.com" || strings.HasSuffix(host, ".openai.azure.com")
|
||||||
}
|
}
|
||||||
|
|
||||||
// supportsPromptCacheKey reports whether the given API base is known to
|
// isOpenAINativeBaseURL reports whether the given API base is the OpenAI or
|
||||||
// support the prompt_cache_key request field. Currently only OpenAI's own
|
// Azure OpenAI native endpoint family. We reuse it for request fields that are
|
||||||
// API and Azure OpenAI support this. All other OpenAI-compatible providers
|
// known to work only on those native hosts.
|
||||||
// (Mistral, Gemini, DeepSeek, Groq, etc.) reject unknown fields with 422 errors.
|
func isOpenAINativeBaseURL(apiBase string) bool {
|
||||||
func supportsPromptCacheKey(apiBase string) bool {
|
|
||||||
u, err := url.Parse(apiBase)
|
u, err := url.Parse(apiBase)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return false
|
return false
|
||||||
|
|
@ -491,3 +493,11 @@ func supportsPromptCacheKey(apiBase string) bool {
|
||||||
host := u.Hostname()
|
host := u.Hostname()
|
||||||
return host == "api.openai.com" || strings.HasSuffix(host, ".openai.azure.com")
|
return host == "api.openai.com" || strings.HasSuffix(host, ".openai.azure.com")
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func supportsPromptCacheKey(apiBase string) bool {
|
||||||
|
return isOpenAINativeBaseURL(apiBase)
|
||||||
|
}
|
||||||
|
|
||||||
|
func supportsStreamingUsage(apiBase string) bool {
|
||||||
|
return isOpenAINativeBaseURL(apiBase)
|
||||||
|
}
|
||||||
|
|
|
||||||
|
|
@ -991,6 +991,54 @@ func chatWithCacheKey(t *testing.T, apiBase string) map[string]any {
|
||||||
return requestBody
|
return requestBody
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func chatStreamWithRequestBody(t *testing.T, apiBase string) map[string]any {
|
||||||
|
t.Helper()
|
||||||
|
var requestBody map[string]any
|
||||||
|
|
||||||
|
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
if err := json.NewDecoder(r.Body).Decode(&requestBody); err != nil {
|
||||||
|
http.Error(w, err.Error(), http.StatusBadRequest)
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
w.Header().Set("Content-Type", "text/event-stream")
|
||||||
|
_, _ = io.WriteString(
|
||||||
|
w,
|
||||||
|
"data: {\"choices\":[{\"delta\":{\"content\":\"ok\"},\"finish_reason\":\"stop\"}]}\n\n",
|
||||||
|
)
|
||||||
|
if flusher, ok := w.(http.Flusher); ok {
|
||||||
|
flusher.Flush()
|
||||||
|
}
|
||||||
|
_, _ = io.WriteString(w, "data: [DONE]\n\n")
|
||||||
|
if flusher, ok := w.(http.Flusher); ok {
|
||||||
|
flusher.Flush()
|
||||||
|
}
|
||||||
|
}))
|
||||||
|
defer server.Close()
|
||||||
|
|
||||||
|
p := NewProvider("key", server.URL, "")
|
||||||
|
p.apiBase = apiBase
|
||||||
|
p.httpClient = &http.Client{
|
||||||
|
Transport: roundTripperFunc(func(r *http.Request) (*http.Response, error) {
|
||||||
|
r.URL, _ = url.Parse(server.URL + r.URL.Path)
|
||||||
|
return http.DefaultTransport.RoundTrip(r)
|
||||||
|
}),
|
||||||
|
}
|
||||||
|
|
||||||
|
_, err := p.ChatStream(
|
||||||
|
t.Context(),
|
||||||
|
[]Message{{Role: "user", Content: "hi"}},
|
||||||
|
nil,
|
||||||
|
"test-model",
|
||||||
|
nil,
|
||||||
|
nil,
|
||||||
|
)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("ChatStream() error = %v", err)
|
||||||
|
}
|
||||||
|
return requestBody
|
||||||
|
}
|
||||||
|
|
||||||
func TestProviderChat_PromptCacheKeySentToOpenAI(t *testing.T) {
|
func TestProviderChat_PromptCacheKeySentToOpenAI(t *testing.T) {
|
||||||
body := chatWithCacheKey(t, "https://api.openai.com/v1")
|
body := chatWithCacheKey(t, "https://api.openai.com/v1")
|
||||||
if body["prompt_cache_key"] != "agent-main" {
|
if body["prompt_cache_key"] != "agent-main" {
|
||||||
|
|
@ -998,6 +1046,39 @@ func TestProviderChat_PromptCacheKeySentToOpenAI(t *testing.T) {
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func TestProviderChatStream_IncludesUsageForOpenAI(t *testing.T) {
|
||||||
|
body := chatStreamWithRequestBody(t, "https://api.openai.com/v1")
|
||||||
|
streamOptions, ok := body["stream_options"].(map[string]any)
|
||||||
|
if !ok {
|
||||||
|
t.Fatalf("stream_options = %T, want map[string]any", body["stream_options"])
|
||||||
|
}
|
||||||
|
if got := streamOptions["include_usage"]; got != true {
|
||||||
|
t.Fatalf("stream_options.include_usage = %v, want true", got)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestProviderChatStream_OmitsUsageForNonOpenAI(t *testing.T) {
|
||||||
|
tests := []struct {
|
||||||
|
name string
|
||||||
|
apiBase string
|
||||||
|
}{
|
||||||
|
{"mistral", "https://api.mistral.ai/v1"},
|
||||||
|
{"gemini", "https://generativelanguage.googleapis.com/v1beta"},
|
||||||
|
{"deepseek", "https://api.deepseek.com/v1"},
|
||||||
|
{"groq", "https://api.groq.com/openai/v1"},
|
||||||
|
{"minimax", "https://api.minimaxi.com/v1"},
|
||||||
|
}
|
||||||
|
|
||||||
|
for _, tt := range tests {
|
||||||
|
t.Run(tt.name, func(t *testing.T) {
|
||||||
|
body := chatStreamWithRequestBody(t, tt.apiBase)
|
||||||
|
if _, exists := body["stream_options"]; exists {
|
||||||
|
t.Fatalf("stream_options should NOT be sent to %s, but was included in request", tt.name)
|
||||||
|
}
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
func TestProviderChat_PromptCacheKeyOmittedForNonOpenAI(t *testing.T) {
|
func TestProviderChat_PromptCacheKeyOmittedForNonOpenAI(t *testing.T) {
|
||||||
tests := []struct {
|
tests := []struct {
|
||||||
name string
|
name string
|
||||||
|
|
|
||||||
Loading…
Add table
Reference in a new issue