Add TTFT & ITL Metrics for LLM Inference (#28)

Time-to-First-Token (TTFT) and Inter-Token Latency (ITL) metrics for LLM inference observability Metrics: llm_ttft_seconds, llm_itl_seconds, llm_tokens_total Closes #31 #32 #33 --------- Co-authored-by: poimen <[email protected]> Reviewed-on: #28
2026-09-15 08:53:58 +00:00
parent 30e0a83a50
commit d439536ca9
22 changed files with 2224 additions and 199 deletions
@@ -381,3 +381,66 @@ func TestBodySizeCappedDispatch(t *testing.T) {
 		t.Errorf("expected 200 for reasonable body, got %d", resp.StatusCode)
 	}
 }
+
+// TestUpstreamModelRewrite verifies that the model field is rewritten when upstreamModel is set.
+func TestUpstreamModelRewrite(t *testing.T) {
+	var receivedModel string
+
+	upstreamServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		body, _ := io.ReadAll(r.Body)
+		var payload map[string]interface{}
+		json.Unmarshal(body, &payload)
+		receivedModel = payload["model"].(string)
+		w.Header().Set("Content-Type", "application/json")
+		w.WriteHeader(http.StatusOK)
+		fmt.Fprint(w, `{}`)
+	}))
+	defer upstreamServer.Close()
+
+	upstreamAddr := strings.TrimPrefix(upstreamServer.URL, "http://")
+
+	cfg := &config.Config{
+		Models: map[string]*config.ModelUpstream{
+			// Client sends "ornith:35b", upstream expects "qwen2.5:72b-instruct"
+			"ornith:35b": {
+				Name:          "ornith:35b",
+				Address:       upstreamAddr,
+				UpstreamModel: "qwen2.5:72b-instruct",
+			},
+			// No rewrite - upstream model same as client model
+			"reasoning": {
+				Name:    "reasoning",
+				Address: upstreamAddr,
+			},
+		},
+		Routes: make(map[string]*config.Route),
+	}
+
+	handler := New(cfg)
+	defer handler.Close()
+
+	server := httptest.NewServer(handler)
+	defer server.Close()
+
+	// Test 1: Model should be rewritten
+	requestBody := `{"model":"ornith:35b","messages":[{"role":"user","content":"hi"}]}`
+	resp, err := http.Post(server.URL+"/v1/chat/completions", "application/json", strings.NewReader(requestBody))
+	if err != nil {
+		t.Fatalf("request failed: %v", err)
+	}
+	resp.Body.Close()
+
+	if receivedModel != "qwen2.5:72b-instruct" {
+		t.Errorf("expected upstream to receive model 'qwen2.5:72b-instruct', got '%s'", receivedModel)
+	}
+
+	// Test 2: No rewrite when upstreamModel is empty
+	receivedModel = ""
+	requestBody = `{"model":"reasoning","messages":[{"role":"user","content":"hi"}]}`
+	resp, _ = http.Post(server.URL+"/v1/chat/completions", "application/json", strings.NewReader(requestBody))
+	resp.Body.Close()
+
+	if receivedModel != "reasoning" {
+		t.Errorf("expected upstream to receive model 'reasoning', got '%s'", receivedModel)
+	}
+}