Improve content-length handling (#115)

ref: See #114

* Improve content-length handling
- Content length was not always being sent
- Add tests for content-length

proxy/proxymanager_test.go

@@ -8,6 +8,7 @@ import (
 	"mime/multipart"
 	"net/http"
 	"net/http/httptest"
+	"strconv"
 	"sync"
 	"testing"
 	"time"
@@ -165,7 +166,9 @@ func TestProxyManager_SwapMultiProcessParallelRequests(t *testing.T) {
 
 			mu.Lock()
 
-			results[key] = w.Body.String()
+			var response map[string]string
+			assert.NoError(t, json.Unmarshal(w.Body.Bytes(), &response))
+			results[key] = response["responseMessage"]
 			mu.Unlock()
 		}(key)
 
@@ -442,6 +445,7 @@ func TestProxyManager_AudioTranscriptionHandler(t *testing.T) {
 	assert.NoError(t, err)
 	assert.Equal(t, "TheExpectedModel", response["model"])
 	assert.Equal(t, response["text"], fmt.Sprintf("The length of the file is %d bytes", contentLength)) // matches simple-responder
+	assert.Equal(t, strconv.Itoa(370+contentLength), response["h_content_length"])
 }
 
 // Test useModelName in configuration sends overrides what is sent to upstream
@@ -592,3 +596,27 @@ func TestProxyManager_Upstream(t *testing.T) {
 	assert.Equal(t, http.StatusOK, rec.Code)
 	assert.Equal(t, "model1", rec.Body.String())
 }
+
+func TestProxyManager_ChatContentLength(t *testing.T) {
+	config := AddDefaultGroupToConfig(Config{
+		HealthCheckTimeout: 15,
+		Models: map[string]ModelConfig{
+			"model1": getTestSimpleResponderConfig("model1"),
+		},
+		LogLevel: "error",
+	})
+
+	proxy := New(config)
+	defer proxy.StopProcesses()
+
+	reqBody := fmt.Sprintf(`{"model":"%s", "x": "this is just some content to push the length out a bit"}`, "model1")
+	req := httptest.NewRequest("POST", "/v1/chat/completions", bytes.NewBufferString(reqBody))
+	w := httptest.NewRecorder()
+
+	proxy.HandlerFunc(w, req)
+	assert.Equal(t, http.StatusOK, w.Code)
+	var response map[string]string
+	assert.NoError(t, json.Unmarshal(w.Body.Bytes(), &response))
+	assert.Equal(t, "81", response["h_content_length"])
+	assert.Equal(t, "model1", response["responseMessage"])
+}