diff --git a/apiproxy/proxy.go b/apiproxy/proxy.go index d15effd..0c52ccd 100644 --- a/apiproxy/proxy.go +++ b/apiproxy/proxy.go @@ -86,10 +86,13 @@ type ProxyDB interface { } func (h *baseHandle) ServeHTTP(w http.ResponseWriter, r *http.Request) { - if r.URL.Path == "/api/rerank" || r.URL.Path == "/api/rerank/" || - r.URL.Path == "/api/v1/rerank" || r.URL.Path == "/api/v1/rerank/" { + switch r.URL.Path { + case "/api/v2/rerank", "/api/v2/rerank/", "/api/rerank", "/api/rerank/": h.HandleRerank(w, r) return + case "/api/v1/rerank", "/api/v1/rerank/": + http.NotFound(w, r) + return } // Intercept OpenAI-compatible models endpoints and serve locally if strings.HasPrefix(r.URL.Path, "/api/models") || strings.HasPrefix(r.URL.Path, "/api/v1/models") { diff --git a/apiproxy/rerank.go b/apiproxy/rerank.go index 37cc32e..40820b4 100644 --- a/apiproxy/rerank.go +++ b/apiproxy/rerank.go @@ -1,6 +1,7 @@ package apiproxy import ( + "bytes" "encoding/json" "errors" "fmt" @@ -23,11 +24,10 @@ const ( ) type RerankRequest struct { - Model string `json:"model"` - Query string `json:"query"` - Documents []string `json:"documents"` - TopN *int `json:"top_n,omitempty"` - ReturnDocuments *bool `json:"return_documents,omitempty"` + Model string `json:"model"` + Query string `json:"query"` + Documents []string `json:"documents"` + TopN *int `json:"top_n,omitempty"` } type rerankUpstreamResponse struct { @@ -45,24 +45,8 @@ type rerankBilledUnits struct { } type rerankUpstreamResult struct { - Index int `json:"index"` - RelevanceScore float64 `json:"relevance_score"` - Document *RerankDocument `json:"document,omitempty"` -} - -type RerankDocument struct { - Text string `json:"text"` -} - -type rerankResponse struct { - ID string `json:"id,omitempty"` - Results []RerankResult `json:"results"` -} - -type RerankResult struct { - Index int `json:"index"` - RelevanceScore float64 `json:"relevance_score"` - Document *RerankDocument `json:"document"` + Index int `json:"index"` + RelevanceScore float64 `json:"relevance_score"` } func rerankTimeoutFromEnv() time.Duration { @@ -135,13 +119,18 @@ func (h *baseHandle) HandleRerank(w http.ResponseWriter, r *http.Request) { if h.az.RerankMaxRequestBytes > 0 { r.Body = http.MaxBytesReader(w, r.Body, h.az.RerankMaxRequestBytes) } - var request RerankRequest - decoder := json.NewDecoder(r.Body) - if err := decoder.Decode(&request); err != nil { + requestBody, err := io.ReadAll(r.Body) + if err != nil { if isMaxBytesError(err) { writeRerankError(w, http.StatusRequestEntityTooLarge, "rerank request body is too large", "request_too_large") return } + writeRerankError(w, http.StatusBadRequest, "could not read rerank request body", "invalid_request") + return + } + var request RerankRequest + decoder := json.NewDecoder(bytes.NewReader(requestBody)) + if err := decoder.Decode(&request); err != nil { writeRerankError(w, http.StatusBadRequest, "request body must be valid JSON", "invalid_request") return } @@ -152,6 +141,15 @@ func (h *baseHandle) HandleRerank(w http.ResponseWriter, r *http.Request) { writeRerankError(w, http.StatusBadRequest, "request body must contain exactly one JSON value", "invalid_request") return } + var requestFields map[string]json.RawMessage + if err := json.Unmarshal(requestBody, &requestFields); err != nil { + writeRerankError(w, http.StatusBadRequest, "request body must be a JSON object", "invalid_request") + return + } + if _, found := requestFields["return_documents"]; found { + writeRerankError(w, http.StatusBadRequest, "return_documents is not supported by Cohere v2", "invalid_request") + return + } if err := validateRerankRequest(&request); err != nil { writeRerankError(w, http.StatusBadRequest, err.Error(), "invalid_request") return @@ -178,12 +176,7 @@ func (h *baseHandle) HandleRerank(w http.ResponseWriter, r *http.Request) { return } - body, err := json.Marshal(request) - if err != nil { - writeRerankError(w, http.StatusInternalServerError, "could not encode rerank request", "internal_error") - return - } - upstreamRequest, err := http.NewRequestWithContext(r.Context(), http.MethodPost, target.String(), strings.NewReader(string(body))) + upstreamRequest, err := http.NewRequestWithContext(r.Context(), http.MethodPost, target.String(), bytes.NewReader(requestBody)) if err != nil { writeRerankError(w, http.StatusInternalServerError, "could not create rerank request", "internal_error") return @@ -252,11 +245,6 @@ func (h *baseHandle) HandleRerank(w http.ResponseWriter, r *http.Request) { return } - translated := rerankResponse{ - ID: upstream.ID, - Results: make([]RerankResult, 0, len(upstream.Results)), - } - includeDocuments := request.ReturnDocuments != nil && *request.ReturnDocuments for _, result := range upstream.Results { if math.IsNaN(result.RelevanceScore) || math.IsInf(result.RelevanceScore, 0) { writeRerankError(w, http.StatusBadGateway, "rerank upstream returned an invalid relevance score", "upstream_contract_error") @@ -266,18 +254,6 @@ func (h *baseHandle) HandleRerank(w http.ResponseWriter, r *http.Request) { writeRerankError(w, http.StatusBadGateway, fmt.Sprintf("rerank upstream returned invalid document index %d", result.Index), "upstream_error") return } - var document *RerankDocument - if includeDocuments { - document = result.Document - if document == nil { - document = &RerankDocument{Text: request.Documents[result.Index]} - } - } - translated.Results = append(translated.Results, RerankResult{ - Index: result.Index, - RelevanceScore: result.RelevanceScore, - Document: document, - }) } searchUnits := *upstream.Meta.BilledUnits.SearchUnits @@ -298,7 +274,7 @@ func (h *baseHandle) HandleRerank(w http.ResponseWriter, r *http.Request) { } w.Header().Set("Content-Type", "application/json") - if err := json.NewEncoder(w).Encode(translated); err != nil { + if _, err := w.Write(upstreamBody); err != nil { return } } diff --git a/apiproxy/rerank_test.go b/apiproxy/rerank_test.go index 9751709..1fffbb8 100644 --- a/apiproxy/rerank_test.go +++ b/apiproxy/rerank_test.go @@ -1,6 +1,7 @@ package apiproxy import ( + "bytes" "encoding/json" "io" "net/http" @@ -54,25 +55,31 @@ func TestAzureServiceHost_SupportsSharedBaseURL(t *testing.T) { } } -func TestHandleRerank_ForwardsRequestAndNormalizesResponse(t *testing.T) { +func TestHandleRerank_ForwardsV2RequestAndResponse(t *testing.T) { var gotAuthorization string + var gotRequestBody string var gotRequest RerankRequest + upstreamResponse := []byte(`{ + "id": "rerank-test", + "results": [ + {"index": 1, "relevance_score": 0.91, "relevant_spans": [{"start": 0, "end": 12, "score": 0.8}]}, + {"index": 0, "relevance_score": 0.12} + ], + "meta": { + "api_version": {"version": "2"}, + "billed_units": {"search_units": 1}, + "tokens": {"input_tokens": 42} + } + }`) upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { gotAuthorization = r.Header.Get("Authorization") body, _ := io.ReadAll(r.Body) + gotRequestBody = string(body) if err := json.Unmarshal(body, &gotRequest); err != nil { t.Fatalf("failed to decode upstream request: %v", err) } w.Header().Set("Content-Type", "application/json") - _, _ = w.Write([]byte(`{ - "id": "rerank-test", - "results": [ - {"index": 1, "relevance_score": 0.91}, - {"index": 0, "relevance_score": 0.12} - ], - "model": "backend-model", - "meta": {"billed_units": {"search_units": 1}} - }`)) + _, _ = w.Write(upstreamResponse) })) defer upstream.Close() @@ -86,13 +93,8 @@ func TestHandleRerank_ForwardsRequestAndNormalizesResponse(t *testing.T) { return true }, } - req := httptest.NewRequest(http.MethodPost, "http://proxy/api/v1/rerank", strings.NewReader(`{ - "model": "Cohere-rerank-v4.0-fast", - "query": "green tea benefits", - "documents": ["first document", "second document"], - "top_n": 2, - "return_documents": true - }`)) + requestBody := `{"model":"Cohere-rerank-v4.0-fast","query":"green tea benefits","documents":["first document","second document"],"top_n":2,"max_tokens_per_doc":4096,"priority":0,"relevant_spans":{"mode":"enabled","threshold":0.1,"max_spans":2}}` + req := httptest.NewRequest(http.MethodPost, "http://proxy/api/v2/rerank", strings.NewReader(requestBody)) req.Header.Set("Authorization", "Bearer proxy-key") req.Header.Set("Content-Type", "application/json") recorder := httptest.NewRecorder() @@ -111,32 +113,19 @@ func TestHandleRerank_ForwardsRequestAndNormalizesResponse(t *testing.T) { if gotRequest.TopN == nil || *gotRequest.TopN != 2 { t.Fatalf("expected top_n=2, got %v", gotRequest.TopN) } - - var response map[string]interface{} - if err := json.Unmarshal(recorder.Body.Bytes(), &response); err != nil { - t.Fatalf("failed to decode response: %v", err) - } - if response["id"] != "rerank-test" { - t.Fatalf("unexpected response id: %v", response["id"]) - } - if _, ok := response["model"]; ok { - t.Fatal("response should not expose upstream model") + if gotRequestBody != requestBody { + t.Fatalf("expected the original v2 JSON body to be forwarded unchanged; got %s", gotRequestBody) } - if _, ok := response["meta"]; ok { - t.Fatal("response should not expose upstream meta") - } - results := response["results"].([]interface{}) - first := results[0].(map[string]interface{}) - if first["index"] != float64(1) || first["relevance_score"] != 0.91 { - t.Fatalf("unexpected first result: %v", first) - } - if first["document"].(map[string]interface{})["text"] != "second document" { - t.Fatalf("expected document text to be reconstructed: %v", first["document"]) + + if !bytes.Equal(recorder.Body.Bytes(), upstreamResponse) { + t.Fatalf("expected the original Cohere v2 response to be returned unchanged; got %s", recorder.Body.String()) } } -func TestHandleRerank_ReturnsNullDocumentWhenNotRequested(t *testing.T) { +func TestServeHTTP_UnversionedRerankAliasUsesV2Handler(t *testing.T) { + called := false upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + called = true _, _ = w.Write([]byte(`{"id":"rerank-test","results":[{"index":0,"relevance_score":0.5}],"meta":{"billed_units":{"search_units":1}}}`)) })) defer upstream.Close() @@ -149,24 +138,67 @@ func TestHandleRerank_ReturnsNullDocumentWhenNotRequested(t *testing.T) { return true }, } - req := httptest.NewRequest(http.MethodPost, "http://proxy/api/v1/rerank", strings.NewReader(`{"model":"test-model","query":"query","documents":["document"],"return_documents":false}`)) + req := httptest.NewRequest(http.MethodPost, "http://proxy/api/rerank", strings.NewReader(`{"model":"test-model","query":"query","documents":["document"]}`)) recorder := httptest.NewRecorder() - h.HandleRerank(recorder, req) + h.ServeHTTP(recorder, req) - if recorder.Code != http.StatusOK { - t.Fatalf("expected 200, got %d: %s", recorder.Code, recorder.Body.String()) + if recorder.Code != http.StatusOK || !called { + t.Fatalf("expected unversioned alias to use rerank handler, got status=%d called=%v", recorder.Code, called) + } +} + +func TestServeHTTP_RetiredV1RerankRoutesReturnNotFound(t *testing.T) { + for _, path := range []string{"/api/v1/rerank", "/api/v1/rerank/"} { + t.Run(path, func(t *testing.T) { + validated := false + h := &baseHandle{ + az: &AzureConfig{}, + clientTokenValidator: func(http.ResponseWriter, *http.Request) bool { + validated = true + return false + }, + } + req := httptest.NewRequest(http.MethodPost, "http://proxy"+path, strings.NewReader(`{}`)) + recorder := httptest.NewRecorder() + + h.ServeHTTP(recorder, req) + + if recorder.Code != http.StatusNotFound { + t.Fatalf("expected 404 for retired v1 route, got %d", recorder.Code) + } + if validated { + t.Fatal("retired v1 route should not fall through to generic Azure handling") + } + }) } - var response struct { - Results []struct { - Document *RerankDocument `json:"document"` - } `json:"results"` +} + +func TestHandleRerank_RejectsReturnDocuments(t *testing.T) { + called := false + upstream := httptest.NewServer(http.HandlerFunc(func(http.ResponseWriter, *http.Request) { + called = true + })) + defer upstream.Close() + + h := &baseHandle{ + az: &AzureConfig{ApiKey: "azure-key"}, + rerankClient: upstream.Client(), + rerankEndpointOverride: upstream.URL, + clientTokenValidator: func(http.ResponseWriter, *http.Request) bool { + return true + }, } - if err := json.Unmarshal(recorder.Body.Bytes(), &response); err != nil { - t.Fatalf("failed to decode response: %v", err) + req := httptest.NewRequest(http.MethodPost, "http://proxy/api/v2/rerank", strings.NewReader(`{"model":"test-model","query":"query","documents":["document"],"return_documents":false}`)) + recorder := httptest.NewRecorder() + + h.HandleRerank(recorder, req) + + if recorder.Code != http.StatusBadRequest { + t.Fatalf("expected 400 for return_documents, got %d", recorder.Code) } - if len(response.Results) != 1 || response.Results[0].Document != nil { - t.Fatalf("expected document=null, got %+v", response.Results) + if called { + t.Fatal("return_documents request should be rejected before calling upstream") } } @@ -185,7 +217,7 @@ func TestHandleRerank_RejectsInvalidRequestBeforeUpstream(t *testing.T) { return true }, } - req := httptest.NewRequest(http.MethodPost, "http://proxy/api/v1/rerank", strings.NewReader(`{"query":"","documents":[]}`)) + req := httptest.NewRequest(http.MethodPost, "http://proxy/api/v2/rerank", strings.NewReader(`{"query":"","documents":[]}`)) recorder := httptest.NewRecorder() h.HandleRerank(recorder, req) @@ -212,7 +244,7 @@ func TestHandleRerank_RejectsInvalidUpstreamIndex(t *testing.T) { return true }, } - req := httptest.NewRequest(http.MethodPost, "http://proxy/api/v1/rerank", strings.NewReader(`{"model":"test-model","query":"query","documents":["document"]}`)) + req := httptest.NewRequest(http.MethodPost, "http://proxy/api/v2/rerank", strings.NewReader(`{"model":"test-model","query":"query","documents":["document"]}`)) recorder := httptest.NewRecorder() h.HandleRerank(recorder, req) @@ -236,7 +268,7 @@ func TestHandleRerank_RejectsMissingUpstreamID(t *testing.T) { return true }, } - req := httptest.NewRequest(http.MethodPost, "http://proxy/api/v1/rerank", strings.NewReader(`{"model":"test-model","query":"query","documents":["document"]}`)) + req := httptest.NewRequest(http.MethodPost, "http://proxy/api/v2/rerank", strings.NewReader(`{"model":"test-model","query":"query","documents":["document"]}`)) recorder := httptest.NewRecorder() h.HandleRerank(recorder, req) @@ -260,7 +292,7 @@ func TestHandleRerank_RejectsMissingBilledUnits(t *testing.T) { return true }, } - req := httptest.NewRequest(http.MethodPost, "http://proxy/api/v1/rerank", strings.NewReader(`{"model":"test-model","query":"query","documents":["document"]}`)) + req := httptest.NewRequest(http.MethodPost, "http://proxy/api/v2/rerank", strings.NewReader(`{"model":"test-model","query":"query","documents":["document"]}`)) recorder := httptest.NewRecorder() h.HandleRerank(recorder, req) @@ -286,7 +318,7 @@ func TestHandleRerank_ReturnsPersistenceErrorBeforeResponse(t *testing.T) { return true }, } - req := httptest.NewRequest(http.MethodPost, "http://proxy/api/v1/rerank", strings.NewReader(`{"model":"test-model","query":"query","documents":["document"]}`)) + req := httptest.NewRequest(http.MethodPost, "http://proxy/api/v2/rerank", strings.NewReader(`{"model":"test-model","query":"query","documents":["document"]}`)) recorder := httptest.NewRecorder() h.HandleRerank(recorder, req) @@ -315,7 +347,7 @@ func TestHandleRerankPersistsExactSearchUnits(t *testing.T) { return true }, } - req := httptest.NewRequest(http.MethodPost, "http://proxy/api/v1/rerank", strings.NewReader(`{"model":"test-model","query":"query","documents":["document"]}`)) + req := httptest.NewRequest(http.MethodPost, "http://proxy/api/v2/rerank", strings.NewReader(`{"model":"test-model","query":"query","documents":["document"]}`)) recorder := httptest.NewRecorder() h.HandleRerank(recorder, req) @@ -347,7 +379,7 @@ func TestHandleRerank_RejectsNonRerankModel(t *testing.T) { return true }, } - req := httptest.NewRequest(http.MethodPost, "http://proxy/api/v1/rerank", strings.NewReader(`{"model":"test-model","query":"query","documents":["document"]}`)) + req := httptest.NewRequest(http.MethodPost, "http://proxy/api/v2/rerank", strings.NewReader(`{"model":"test-model","query":"query","documents":["document"]}`)) recorder := httptest.NewRecorder() h.HandleRerank(recorder, req) diff --git a/readme.md b/readme.md index 92919de..cc9b1ee 100644 --- a/readme.md +++ b/readme.md @@ -37,7 +37,8 @@ curl "http://localhost:8082/api/chat/completions" \ Reranking uses the same `DEPLOYMENT_NAME`, `BASE_URL`, and `AZURE_API_KEY` configuration as the Azure OpenAI proxy. `BASE_URL` is the shared Azure domain suffix, for example `azure.com`. The proxy exposes reranking at -`/api/v1/rerank`: +`/api/v2/rerank`. The unversioned `/api/rerank` path is also available and uses +the same Cohere v2 contract. The retired `/api/v1/rerank` path returns 404. The Rerank model must be registered in the admin model catalog as type `RERANK`. The current pricing policy for `Cohere-rerank-v4.0-fast` is `$2 per @@ -45,7 +46,7 @@ The Rerank model must be registered in the admin model catalog as type for other Cohere models. ```bash -curl "http://localhost:8082/api/v1/rerank" \ +curl "http://localhost:8082/api/v2/rerank" \ -H "Content-Type: application/json" \ -H "Authorization: Bearer " \ -d '{