diff --git a/plugins/wasm-go/extensions/ai-proxy/README.md b/plugins/wasm-go/extensions/ai-proxy/README.md index a83bee186..a5599f935 100644 --- a/plugins/wasm-go/extensions/ai-proxy/README.md +++ b/plugins/wasm-go/extensions/ai-proxy/README.md @@ -177,6 +177,10 @@ Grok 所对应的 `type` 为 `grok`。它并无特有的配置字段。 OpenRouter 所对应的 `type` 为 `openrouter`。它并无特有的配置字段。 +#### Fireworks AI + +Fireworks AI 所对应的 `type` 为 `fireworks`。它并无特有的配置字段。 + #### 文心一言（Baidu） 文心一言所对应的 `type` 为 `baidu`。它并无特有的配置字段。 @@ -1018,6 +1022,63 @@ provider: } ``` +### 使用 OpenAI 协议代理 Fireworks AI 服务 + +**配置信息** + +```yaml +provider: + type: fireworks + apiTokens: + - "YOUR_FIREWORKS_API_TOKEN" + modelMapping: + "gpt-4": "accounts/fireworks/models/llama-v3p1-70b-instruct" + "gpt-3.5-turbo": "accounts/fireworks/models/llama-v3p1-8b-instruct" + "*": "accounts/fireworks/models/llama-v3p1-8b-instruct" +``` + +**请求示例** + +```json +{ + "model": "gpt-4", + "messages": [ + { + "role": "user", + "content": "你好，你是谁？" + } + ], + "temperature": 0.7, + "max_tokens": 100 +} +``` + +**响应示例** + +```json +{ + "id": "fw-123456789", + "object": "chat.completion", + "created": 1699123456, + "model": "accounts/fireworks/models/llama-v3p1-70b-instruct", + "choices": [ + { + "index": 0, + "message": { + "role": "assistant", + "content": "你好！我是一个由 Fireworks AI 提供的人工智能助手，基于 Llama 3.1 模型。我可以帮助回答问题、进行对话和提供各种信息。有什么我可以帮助你的吗？" + }, + "finish_reason": "stop" + } + ], + "usage": { + "prompt_tokens": 15, + "completion_tokens": 45, + "total_tokens": 60 + } +} +``` + ### 使用自动协议兼容功能 插件现在支持自动协议检测，可以同时处理 OpenAI 和 Claude 两种协议格式的请求。 @@ -1982,6 +2043,7 @@ provider: } } ``` + ### 使用 OpenAI 协议代理 NVIDIA Triton Interference Server 服务 **配置信息** @@ -2011,6 +2073,7 @@ providers: "stream": false } ``` + **响应示例** ```json diff --git a/plugins/wasm-go/extensions/ai-proxy/README_EN.md b/plugins/wasm-go/extensions/ai-proxy/README_EN.md index e924fe2e3..4031686ff 100644 --- a/plugins/wasm-go/extensions/ai-proxy/README_EN.md +++ b/plugins/wasm-go/extensions/ai-proxy/README_EN.md @@ -148,6 +148,10 @@ For Grok, the corresponding `type` is `grok`. It has no unique configuration fie For OpenRouter, the corresponding `type` is `openrouter`. It has no unique configuration fields. +#### Fireworks AI + +For Fireworks AI, the corresponding `type` is `fireworks`. It has no unique configuration fields. + #### ERNIE Bot For ERNIE Bot, the corresponding `type` is `baidu`. It has no unique configuration fields. @@ -955,6 +959,63 @@ provider: } ``` +### Using OpenAI Protocol Proxy for Fireworks AI Service + +**Configuration Information** + +```yaml +provider: + type: fireworks + apiTokens: + - "YOUR_FIREWORKS_API_TOKEN" + modelMapping: + "gpt-4": "accounts/fireworks/models/llama-v3p1-70b-instruct" + "gpt-3.5-turbo": "accounts/fireworks/models/llama-v3p1-8b-instruct" + "*": "accounts/fireworks/models/llama-v3p1-8b-instruct" +``` + +**Request Example** + +```json +{ + "model": "gpt-4", + "messages": [ + { + "role": "user", + "content": "Hello, who are you?" + } + ], + "temperature": 0.7, + "max_tokens": 100 +} +``` + +**Response Example** + +```json +{ + "id": "fw-123456789", + "object": "chat.completion", + "created": 1699123456, + "model": "accounts/fireworks/models/llama-v3p1-70b-instruct", + "choices": [ + { + "index": 0, + "message": { + "role": "assistant", + "content": "Hello! I am an AI assistant powered by Fireworks AI, based on the Llama 3.1 model. I can help answer questions, engage in conversations, and provide various information. How can I assist you today?" + }, + "finish_reason": "stop" + } + ], + "usage": { + "prompt_tokens": 15, + "completion_tokens": 38, + "total_tokens": 53 + } +} +``` + ### Using Auto Protocol Compatibility The plugin now supports automatic protocol detection, capable of handling both OpenAI and Claude protocol format requests simultaneously. diff --git a/plugins/wasm-go/extensions/ai-proxy/main_test.go b/plugins/wasm-go/extensions/ai-proxy/main_test.go index 43a4b2d7a..2fb38fb68 100644 --- a/plugins/wasm-go/extensions/ai-proxy/main_test.go +++ b/plugins/wasm-go/extensions/ai-proxy/main_test.go @@ -103,3 +103,9 @@ func TestAzure(t *testing.T) { test.RunAzureOnHttpResponseHeadersTests(t) test.RunAzureOnHttpResponseBodyTests(t) } + +func TestFireworks(t *testing.T) { + test.RunFireworksParseConfigTests(t) + test.RunFireworksOnHttpRequestHeadersTests(t) + test.RunFireworksOnHttpRequestBodyTests(t) +} diff --git a/plugins/wasm-go/extensions/ai-proxy/provider/fireworks.go b/plugins/wasm-go/extensions/ai-proxy/provider/fireworks.go new file mode 100644 index 000000000..3d6d3a8de --- /dev/null +++ b/plugins/wasm-go/extensions/ai-proxy/provider/fireworks.go @@ -0,0 +1,83 @@ +package provider + +import ( + "errors" + "net/http" + "strings" + + "github.com/alibaba/higress/plugins/wasm-go/extensions/ai-proxy/util" + "github.com/higress-group/proxy-wasm-go-sdk/proxywasm/types" + "github.com/higress-group/wasm-go/pkg/wrapper" +) + +// fireworksProvider is the provider for Fireworks AI service. + +const ( + fireworksDomain = "api.fireworks.ai" +) + +type fireworksProviderInitializer struct{} + +func (f *fireworksProviderInitializer) ValidateConfig(config *ProviderConfig) error { + if config.apiTokens == nil || len(config.apiTokens) == 0 { + return errors.New("no apiToken found in provider config") + } + return nil +} + +func (f *fireworksProviderInitializer) DefaultCapabilities() map[string]string { + return map[string]string{ + string(ApiNameChatCompletion): PathOpenAIChatCompletions, + string(ApiNameCompletion): PathOpenAICompletions, + string(ApiNameModels): PathOpenAIModels, + } +} + +func (f *fireworksProviderInitializer) CreateProvider(config ProviderConfig) (Provider, error) { + config.setDefaultCapabilities(f.DefaultCapabilities()) + return &fireworksProvider{ + config: config, + contextCache: createContextCache(&config), + }, nil +} + +type fireworksProvider struct { + config ProviderConfig + contextCache *contextCache +} + +func (f *fireworksProvider) GetProviderType() string { + return providerTypeFireworks +} + +func (f *fireworksProvider) OnRequestHeaders(ctx wrapper.HttpContext, apiName ApiName) error { + f.config.handleRequestHeaders(f, ctx, apiName) + return nil +} + +func (f *fireworksProvider) OnRequestBody(ctx wrapper.HttpContext, apiName ApiName, body []byte) (types.Action, error) { + if !f.config.isSupportedAPI(apiName) { + return types.ActionContinue, errUnsupportedApiName + } + return f.config.handleRequestBody(f, f.contextCache, ctx, apiName, body) +} + +func (f *fireworksProvider) TransformRequestHeaders(ctx wrapper.HttpContext, apiName ApiName, headers http.Header) { + util.OverwriteRequestPathHeaderByCapability(headers, string(apiName), f.config.capabilities) + util.OverwriteRequestHostHeader(headers, fireworksDomain) + util.OverwriteRequestAuthorizationHeader(headers, "Bearer "+f.config.GetApiTokenInUse(ctx)) + headers.Del("Content-Length") +} + +func (f *fireworksProvider) GetApiName(path string) ApiName { + if strings.Contains(path, PathOpenAIChatCompletions) { + return ApiNameChatCompletion + } + if strings.Contains(path, PathOpenAICompletions) { + return ApiNameCompletion + } + if strings.Contains(path, PathOpenAIModels) { + return ApiNameModels + } + return "" +} diff --git a/plugins/wasm-go/extensions/ai-proxy/provider/provider.go b/plugins/wasm-go/extensions/ai-proxy/provider/provider.go index 726bce521..aecfb71ff 100644 --- a/plugins/wasm-go/extensions/ai-proxy/provider/provider.go +++ b/plugins/wasm-go/extensions/ai-proxy/provider/provider.go @@ -134,6 +134,7 @@ const ( providerTypeTriton = "triton" providerTypeOpenRouter = "openrouter" providerTypeLongcat = "longcat" + providerTypeFireworks = "fireworks" protocolOpenAI = "openai" protocolOriginal = "original" @@ -215,6 +216,7 @@ var ( providerTypeTriton: &tritonProviderInitializer{}, providerTypeOpenRouter: &openrouterProviderInitializer{}, providerTypeLongcat: &longcatProviderInitializer{}, + providerTypeFireworks: &fireworksProviderInitializer{}, } ) diff --git a/plugins/wasm-go/extensions/ai-proxy/test/fireworks.go b/plugins/wasm-go/extensions/ai-proxy/test/fireworks.go new file mode 100644 index 000000000..99f9fc8b1 --- /dev/null +++ b/plugins/wasm-go/extensions/ai-proxy/test/fireworks.go @@ -0,0 +1,354 @@ +package test + +import ( + "encoding/json" + "strings" + "testing" + + "github.com/higress-group/proxy-wasm-go-sdk/proxywasm/types" + "github.com/higress-group/wasm-go/pkg/test" + "github.com/stretchr/testify/require" +) + +// 测试配置：基本 Fireworks 配置 +var basicFireworksConfig = func() json.RawMessage { + data, _ := json.Marshal(map[string]interface{}{ + "provider": map[string]interface{}{ + "type": "fireworks", + "apiTokens": []string{"fw-test123456789"}, + "modelMapping": map[string]string{ + "*": "accounts/fireworks/models/llama-v3p1-8b-instruct", + }, + }, + }) + return data +}() + +// 测试配置：Fireworks 多模型配置 +var fireworksMultiModelConfig = func() json.RawMessage { + data, _ := json.Marshal(map[string]interface{}{ + "provider": map[string]interface{}{ + "type": "fireworks", + "apiTokens": []string{"fw-multi-model"}, + "modelMapping": map[string]string{ + "gpt-4": "accounts/fireworks/models/llama-v3p1-70b-instruct", + "gpt-3.5-turbo": "accounts/fireworks/models/llama-v3p1-8b-instruct", + "*": "accounts/fireworks/models/llama-v3p1-8b-instruct", + }, + }, + }) + return data +}() + +// 测试配置：无效 Fireworks 配置（缺少 apiToken） +var invalidFireworksConfig = func() json.RawMessage { + data, _ := json.Marshal(map[string]interface{}{ + "provider": map[string]interface{}{ + "type": "fireworks", + "apiTokens": []string{}, + "modelMapping": map[string]string{}, + }, + }) + return data +}() + +// 测试配置：完整 Fireworks 配置 +var completeFireworksConfig = func() json.RawMessage { + data, _ := json.Marshal(map[string]interface{}{ + "provider": map[string]interface{}{ + "type": "fireworks", + "apiTokens": []string{"fw-complete-test"}, + "modelMapping": map[string]string{ + "gpt-4": "accounts/fireworks/models/llama-v3p1-70b-instruct", + "gpt-3.5-turbo": "accounts/fireworks/models/llama-v3p1-8b-instruct", + "*": "accounts/fireworks/models/llama-v3p1-8b-instruct", + }, + }, + }) + return data +}() + +// RunFireworksParseConfigTests 测试 Fireworks 配置解析 +func RunFireworksParseConfigTests(t *testing.T) { + test.RunGoTest(t, func(t *testing.T) { + // 测试基本 Fireworks 配置解析 + t.Run("basic fireworks config", func(t *testing.T) { + host, status := test.NewTestHost(basicFireworksConfig) + defer host.Reset() + require.Equal(t, types.OnPluginStartStatusOK, status) + + config, err := host.GetMatchConfig() + require.NoError(t, err) + require.NotNil(t, config) + }) + + // 测试 Fireworks 多模型配置解析 + t.Run("fireworks multi model config", func(t *testing.T) { + host, status := test.NewTestHost(fireworksMultiModelConfig) + defer host.Reset() + require.Equal(t, types.OnPluginStartStatusOK, status) + + config, err := host.GetMatchConfig() + require.NoError(t, err) + require.NotNil(t, config) + }) + + // 测试无效 Fireworks 配置（缺少 apiToken） + t.Run("invalid fireworks config - missing apiToken", func(t *testing.T) { + host, status := test.NewTestHost(invalidFireworksConfig) + defer host.Reset() + require.Equal(t, types.OnPluginStartStatusFailed, status) + }) + + // 测试完整 Fireworks 配置解析 + t.Run("fireworks complete config", func(t *testing.T) { + host, status := test.NewTestHost(completeFireworksConfig) + defer host.Reset() + require.Equal(t, types.OnPluginStartStatusOK, status) + + config, err := host.GetMatchConfig() + require.NoError(t, err) + require.NotNil(t, config) + }) + }) +} + +// RunFireworksOnHttpRequestHeadersTests 测试 Fireworks 请求头处理 +func RunFireworksOnHttpRequestHeadersTests(t *testing.T) { + test.RunTest(t, func(t *testing.T) { + // 测试 Fireworks 聊天完成请求头处理 + t.Run("fireworks chat completion request headers", func(t *testing.T) { + host, status := test.NewTestHost(basicFireworksConfig) + defer host.Reset() + require.Equal(t, types.OnPluginStartStatusOK, status) + + // 设置请求头 + action := host.CallOnHttpRequestHeaders([][2]string{ + {":authority", "example.com"}, + {":path", "/v1/chat/completions"}, + {":method", "POST"}, + {"Content-Type", "application/json"}, + }) + + // 应该返回 HeaderStopIteration，因为需要处理请求体 + require.Equal(t, types.HeaderStopIteration, action) + + // 验证请求头是否被正确处理 + requestHeaders := host.GetRequestHeaders() + require.NotNil(t, requestHeaders) + + // 验证 Host 是否被改为 Fireworks 域名 + hostValue, hasHost := test.GetHeaderValue(requestHeaders, ":authority") + require.True(t, hasHost, "Host header should exist") + require.Equal(t, "api.fireworks.ai", hostValue, "Host should be changed to Fireworks domain") + + // 验证 Authorization 是否被设置 + authValue, hasAuth := test.GetHeaderValue(requestHeaders, "Authorization") + require.True(t, hasAuth, "Authorization header should exist") + require.Contains(t, authValue, "Bearer fw-test123456789", "Authorization should contain Fireworks API token with Bearer prefix") + + // 验证 Path 保持 OpenAI 兼容格式 + pathValue, hasPath := test.GetHeaderValue(requestHeaders, ":path") + require.True(t, hasPath, "Path header should exist") + require.Equal(t, "/v1/chat/completions", pathValue, "Path should remain OpenAI compatible") + + // 检查是否有相关的处理日志 + debugLogs := host.GetDebugLogs() + hasFireworksLogs := false + for _, log := range debugLogs { + if strings.Contains(log, "fireworks") || strings.Contains(log, "ai-proxy") { + hasFireworksLogs = true + break + } + } + require.True(t, hasFireworksLogs, "Should have Fireworks or ai-proxy processing logs") + }) + + // 测试 Fireworks 文本完成请求头处理 + t.Run("fireworks completion request headers", func(t *testing.T) { + host, status := test.NewTestHost(basicFireworksConfig) + defer host.Reset() + require.Equal(t, types.OnPluginStartStatusOK, status) + + // 设置请求头 + action := host.CallOnHttpRequestHeaders([][2]string{ + {":authority", "example.com"}, + {":path", "/v1/completions"}, + {":method", "POST"}, + {"Content-Type", "application/json"}, + }) + + require.Equal(t, types.HeaderStopIteration, action) + + // 验证请求头处理 + requestHeaders := host.GetRequestHeaders() + require.NotNil(t, requestHeaders) + + // 验证 Host 转换 + hostValue, hasHost := test.GetHeaderValue(requestHeaders, ":authority") + require.True(t, hasHost) + require.Equal(t, "api.fireworks.ai", hostValue) + + // 验证 Path 保持 OpenAI 兼容格式 + pathValue, hasPath := test.GetHeaderValue(requestHeaders, ":path") + require.True(t, hasPath) + require.Equal(t, "/v1/completions", pathValue, "Path should remain OpenAI compatible for completions") + + // 验证 Authorization 设置 + authValue, hasAuth := test.GetHeaderValue(requestHeaders, "Authorization") + require.True(t, hasAuth, "Authorization header should exist for completions") + require.Contains(t, authValue, "Bearer fw-test123456789", "Authorization should contain Fireworks API token") + }) + + // 测试 Fireworks 模型列表请求头处理 + t.Run("fireworks models request headers", func(t *testing.T) { + host, status := test.NewTestHost(basicFireworksConfig) + defer host.Reset() + require.Equal(t, types.OnPluginStartStatusOK, status) + + // 设置请求头 + action := host.CallOnHttpRequestHeaders([][2]string{ + {":authority", "example.com"}, + {":path", "/v1/models"}, + {":method", "GET"}, + }) + + require.Equal(t, types.ActionContinue, action) + + // 验证请求头处理 + requestHeaders := host.GetRequestHeaders() + require.NotNil(t, requestHeaders) + + // 验证 Host 转换 + hostValue, hasHost := test.GetHeaderValue(requestHeaders, ":authority") + require.True(t, hasHost) + require.Equal(t, "api.fireworks.ai", hostValue) + + // 验证 Path 保持 OpenAI 兼容格式 + pathValue, hasPath := test.GetHeaderValue(requestHeaders, ":path") + require.True(t, hasPath) + require.Equal(t, "/v1/models", pathValue, "Path should remain OpenAI compatible for models") + + // 验证 Authorization 设置 + authValue, hasAuth := test.GetHeaderValue(requestHeaders, "Authorization") + require.True(t, hasAuth, "Authorization header should exist for models") + require.Contains(t, authValue, "Bearer fw-test123456789", "Authorization should contain Fireworks API token") + }) + }) +} + +// RunFireworksOnHttpRequestBodyTests 测试 Fireworks 请求体处理 +func RunFireworksOnHttpRequestBodyTests(t *testing.T) { + test.RunTest(t, func(t *testing.T) { + // 测试 Fireworks 聊天完成请求体处理 + t.Run("fireworks chat completion request body", func(t *testing.T) { + host, status := test.NewTestHost(basicFireworksConfig) + defer host.Reset() + require.Equal(t, types.OnPluginStartStatusOK, status) + + // 设置请求头 + host.CallOnHttpRequestHeaders([][2]string{ + {":authority", "example.com"}, + {":path", "/v1/chat/completions"}, + {":method", "POST"}, + {"Content-Type", "application/json"}, + }) + + // 测试请求体 + requestBody := `{ + "model": "gpt-3.5-turbo", + "messages": [ + {"role": "user", "content": "Hello, world!"} + ], + "stream": false + }` + + action := host.CallOnHttpRequestBody([]byte(requestBody)) + require.Equal(t, types.ActionContinue, action) + + // 验证请求体被正确处理 + actualRequestBody := host.GetRequestBody() + require.NotNil(t, actualRequestBody) + + // 验证模型映射 + require.Contains(t, string(actualRequestBody), "accounts/fireworks/models/llama-v3p1-8b-instruct", + "Model should be mapped to Fireworks model") + require.Contains(t, string(actualRequestBody), "Hello, world!", + "Request content should be preserved") + }) + + // 测试 Fireworks 流式聊天完成请求体处理 + t.Run("fireworks streaming chat completion request body", func(t *testing.T) { + host, status := test.NewTestHost(fireworksMultiModelConfig) + defer host.Reset() + require.Equal(t, types.OnPluginStartStatusOK, status) + + // 设置请求头 + host.CallOnHttpRequestHeaders([][2]string{ + {":authority", "example.com"}, + {":path", "/v1/chat/completions"}, + {":method", "POST"}, + {"Content-Type", "application/json"}, + }) + + // 测试流式请求体 + requestBody := `{ + "model": "gpt-4", + "messages": [ + {"role": "user", "content": "Write a poem about AI"} + ], + "stream": true + }` + + action := host.CallOnHttpRequestBody([]byte(requestBody)) + require.Equal(t, types.ActionContinue, action) + + // 验证请求体被正确处理 + actualRequestBody := host.GetRequestBody() + require.NotNil(t, actualRequestBody) + + // 验证模型映射（gpt-4 应该映射到 70b 模型） + require.Contains(t, string(actualRequestBody), "accounts/fireworks/models/llama-v3p1-70b-instruct", + "GPT-4 should be mapped to Fireworks 70b model") + require.Contains(t, string(actualRequestBody), "Write a poem about AI", + "Request content should be preserved") + require.Contains(t, string(actualRequestBody), `"stream": true`, + "Stream flag should be preserved") + }) + + // 测试 Fireworks 文本完成请求体处理 + t.Run("fireworks completion request body", func(t *testing.T) { + host, status := test.NewTestHost(basicFireworksConfig) + defer host.Reset() + require.Equal(t, types.OnPluginStartStatusOK, status) + + // 设置请求头 + host.CallOnHttpRequestHeaders([][2]string{ + {":authority", "example.com"}, + {":path", "/v1/completions"}, + {":method", "POST"}, + {"Content-Type", "application/json"}, + }) + + // 测试完成请求体 + requestBody := `{ + "model": "gpt-3.5-turbo", + "prompt": "The future of AI is", + "max_tokens": 100 + }` + + action := host.CallOnHttpRequestBody([]byte(requestBody)) + require.Equal(t, types.ActionContinue, action) + + // 验证请求体被正确处理 + actualRequestBody := host.GetRequestBody() + require.NotNil(t, actualRequestBody) + + // 验证模型映射 + require.Contains(t, string(actualRequestBody), "accounts/fireworks/models/llama-v3p1-8b-instruct", + "Model should be mapped to Fireworks model") + require.Contains(t, string(actualRequestBody), "The future of AI is", + "Prompt should be preserved") + }) + }) +}