mirror of
https://github.com/icereed/paperless-gpt.git
synced 2025-03-12 12:58:02 -05:00
feat: refactor HTTP client creation to use a custom client with headers
This commit is contained in:
parent
23632f6fcd
commit
51946743af
2 changed files with 40 additions and 26 deletions
44
main.go
44
main.go
|
@ -640,22 +640,10 @@ func createLLM() (llms.Model, error) {
|
|||
return nil, fmt.Errorf("OpenAI API key is not set")
|
||||
}
|
||||
|
||||
// Create custom transport that adds headers
|
||||
customTransport := &headerTransport{
|
||||
transport: http.DefaultTransport,
|
||||
headers: map[string]string{
|
||||
"X-Title": "paperless-gpt",
|
||||
},
|
||||
}
|
||||
|
||||
// Create custom client with the transport
|
||||
httpClient := http.DefaultClient
|
||||
httpClient.Transport = customTransport
|
||||
|
||||
return openai.New(
|
||||
openai.WithModel(llmModel),
|
||||
openai.WithToken(openaiAPIKey),
|
||||
openai.WithHTTPClient(httpClient),
|
||||
openai.WithHTTPClient(createCustomHTTPClient()),
|
||||
)
|
||||
case "ollama":
|
||||
host := os.Getenv("OLLAMA_HOST")
|
||||
|
@ -678,22 +666,10 @@ func createVisionLLM() (llms.Model, error) {
|
|||
return nil, fmt.Errorf("OpenAI API key is not set")
|
||||
}
|
||||
|
||||
// Create custom transport that adds headers
|
||||
customTransport := &headerTransport{
|
||||
transport: http.DefaultTransport,
|
||||
headers: map[string]string{
|
||||
"X-Title": "paperless-gpt",
|
||||
},
|
||||
}
|
||||
|
||||
// Create custom client with the transport
|
||||
httpClient := http.DefaultClient
|
||||
httpClient.Transport = customTransport
|
||||
|
||||
return openai.New(
|
||||
openai.WithModel(visionLlmModel),
|
||||
openai.WithToken(openaiAPIKey),
|
||||
openai.WithHTTPClient(httpClient),
|
||||
openai.WithHTTPClient(createCustomHTTPClient()),
|
||||
)
|
||||
case "ollama":
|
||||
host := os.Getenv("OLLAMA_HOST")
|
||||
|
@ -710,6 +686,22 @@ func createVisionLLM() (llms.Model, error) {
|
|||
}
|
||||
}
|
||||
|
||||
func createCustomHTTPClient() *http.Client {
|
||||
// Create custom transport that adds headers
|
||||
customTransport := &headerTransport{
|
||||
transport: http.DefaultTransport,
|
||||
headers: map[string]string{
|
||||
"X-Title": "paperless-gpt",
|
||||
},
|
||||
}
|
||||
|
||||
// Create custom client with the transport
|
||||
httpClient := http.DefaultClient
|
||||
httpClient.Transport = customTransport
|
||||
|
||||
return httpClient
|
||||
}
|
||||
|
||||
// headerTransport is a custom http.RoundTripper that adds custom headers to requests
|
||||
type headerTransport struct {
|
||||
transport http.RoundTripper
|
||||
|
|
22
main_test.go
22
main_test.go
|
@ -4,6 +4,7 @@ import (
|
|||
"encoding/json"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"slices"
|
||||
"testing"
|
||||
"text/template"
|
||||
|
@ -175,3 +176,24 @@ func TestProcessAutoTagDocuments(t *testing.T) {
|
|||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestCreateCustomHTTPClient(t *testing.T) {
|
||||
// Create a test server
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
// Verify custom header
|
||||
assert.Equal(t, "paperless-gpt", r.Header.Get("X-Title"), "Expected X-Title header")
|
||||
w.WriteHeader(http.StatusOK)
|
||||
}))
|
||||
defer server.Close()
|
||||
|
||||
// Get custom client
|
||||
client := createCustomHTTPClient()
|
||||
require.NotNil(t, client, "HTTP client should not be nil")
|
||||
|
||||
// Make a request
|
||||
resp, err := client.Get(server.URL)
|
||||
require.NoError(t, err, "Request should not fail")
|
||||
defer resp.Body.Close()
|
||||
|
||||
assert.Equal(t, http.StatusOK, resp.StatusCode, "Expected 200 OK response")
|
||||
}
|
||||
|
|
Loading…
Reference in a new issue