feat: add custom http headers to openai related api backends (#1174)

* feat: add custom http headers to openai related api backends Signed-off-by: Aris Boutselis <arisboutselis08@gmail.com> * ci: add custom headers test Signed-off-by: Aris Boutselis <arisboutselis08@gmail.com> * add error handling Signed-off-by: Aris Boutselis <arisboutselis08@gmail.com> * chore(deps): update docker/setup-buildx-action digest to 4fd8129 (#1173) Signed-off-by: renovate[bot] <29139614+renovate[bot]@users.noreply.github.com> Co-authored-by: renovate[bot] <29139614+renovate[bot]@users.noreply.github.com> Signed-off-by: Aris Boutselis <arisboutselis08@gmail.com> * fix(deps): update module buf.build/gen/go/k8sgpt-ai/k8sgpt/grpc-ecosystem/gateway/v2 to v2.20.0-20240406062209-1cc152efbf5c.1 (#1147) Signed-off-by: renovate[bot] <29139614+renovate[bot]@users.noreply.github.com> Co-authored-by: renovate[bot] <29139614+renovate[bot]@users.noreply.github.com> Signed-off-by: Aris Boutselis <arisboutselis08@gmail.com> * chore(deps): update anchore/sbom-action action to v0.16.0 (#1146) Signed-off-by: renovate[bot] <29139614+renovate[bot]@users.noreply.github.com> Co-authored-by: renovate[bot] <29139614+renovate[bot]@users.noreply.github.com> Co-authored-by: Alex Jones <alexsimonjones@gmail.com> Signed-off-by: Aris Boutselis <arisboutselis08@gmail.com> * Update README.md Signed-off-by: Aris Boutselis <arisboutselis08@gmail.com> --------- Signed-off-by: Aris Boutselis <arisboutselis08@gmail.com> Signed-off-by: renovate[bot] <29139614+renovate[bot]@users.noreply.github.com> Co-authored-by: renovate[bot] <29139614+renovate[bot]@users.noreply.github.com> Co-authored-by: Alex Jones <alexsimonjones@gmail.com>
2025-09-06 17:51:28 +00:00 · 2024-07-10 09:59:08 +01:00
parent fef853966f
commit 02e754ed59
8 changed files with 211 additions and 26 deletions
--- a/pkg/ai/openai_header_transport_test.go
+++ b/pkg/ai/openai_header_transport_test.go
@@ -0,0 +1,106 @@
+package ai
+
+import (
+	"context"
+	"net/http"
+	"net/http/httptest"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+// Mock configuration
+type mockConfig struct {
+	baseURL string
+}
+
+func (m *mockConfig) GetPassword() string {
+	return ""
+}
+
+func (m *mockConfig) GetOrganizationId() string {
+	return ""
+}
+
+func (m *mockConfig) GetProxyEndpoint() string {
+	return ""
+}
+
+func (m *mockConfig) GetBaseURL() string {
+	return m.baseURL
+}
+
+func (m *mockConfig) GetCustomHeaders() []http.Header {
+	return []http.Header{
+		{"X-Custom-Header-1": []string{"Value1"}},
+		{"X-Custom-Header-2": []string{"Value2"}},
+		{"X-Custom-Header-2": []string{"Value3"}}, // Testing multiple values for the same header
+	}
+}
+
+func (m *mockConfig) GetModel() string {
+	return ""
+}
+
+func (m *mockConfig) GetTemperature() float32 {
+	return 0.0
+}
+
+func (m *mockConfig) GetTopP() float32 {
+	return 0.0
+}
+func (m *mockConfig) GetCompartmentId() string {
+	return ""
+}
+
+func (m *mockConfig) GetTopK() int32 {
+	return 0.0
+}
+
+func (m *mockConfig) GetMaxTokens() int {
+	return 0
+}
+
+func (m *mockConfig) GetEndpointName() string {
+	return ""
+}
+func (m *mockConfig) GetEngine() string {
+	return ""
+}
+
+func (m *mockConfig) GetProviderId() string {
+	return ""
+}
+
+func (m *mockConfig) GetProviderRegion() string {
+	return ""
+}
+
+func TestOpenAIClient_CustomHeaders(t *testing.T) {
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		assert.Equal(t, "Value1", r.Header.Get("X-Custom-Header-1"))
+		assert.ElementsMatch(t, []string{"Value2", "Value3"}, r.Header["X-Custom-Header-2"])
+		w.WriteHeader(http.StatusOK)
+		// Mock response for openai completion
+		mockResponse := `{"choices": [{"message": {"content": "test"}}]}`
+		n, err := w.Write([]byte(mockResponse))
+		if err != nil {
+			t.Fatalf("error writing response: %v", err)
+		}
+		if n != len(mockResponse) {
+			t.Fatalf("expected to write %d bytes but wrote %d bytes", len(mockResponse), n)
+		}
+	}))
+	defer server.Close()
+
+	config := &mockConfig{baseURL: server.URL}
+
+	client := &OpenAIClient{}
+	err := client.Configure(config)
+	assert.NoError(t, err)
+
+	// Make a completion request to trigger the headers
+	ctx := context.Background()
+	_, err = client.GetCompletion(ctx, "foo prompt")
+	assert.NoError(t, err)
+}