switched to fib backoff

2026-07-30 09:53:36 +03:00
parent 1cc6ddb2e4
commit d49339d770
4 changed files with 519 additions and 5 deletions
@@ -779,5 +779,216 @@ func TestMultipleTokensKeyFile(t *testing.T) {
  }
 }

+func TestRequestMessagesAutodetectionAndContentArrays(t *testing.T) {
+  var receivedBody map[string]any
+
+  server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+    var b map[string]any
+    _ = json.NewDecoder(r.Body).Decode(&b)
+    receivedBody = b
+    w.Header().Set("Content-Type", "application/json")
+    w.WriteHeader(http.StatusOK)
+    _, _ = w.Write([]byte(`{"choices":[{"message":{"content":"OK"}}]} `))
+  }))
+  defer server.Close()
+
+  configs := []ModelConfig{
+    {Model: "test-model", Key: "test-key", Endpoint: server.URL},
+  }
+  cm := &ConfigManager{
+    configs:      configs,
+    uniqueModels: []string{"test-model"},
+  }
+  handler := handleChatCompletions(cm, nil)
+
+  t.Run("Pi agent format - array content", func(t *testing.T) {
+    reqObj := map[string]any{
+      "model": "test-model",
+      "messages": []any{
+        map[string]any{
+          "role": "user",
+          "content": []any{
+            map[string]any{"type": "text", "text": "Hello from Pi agent"},
+          },
+        },
+      },
+    }
+    reqBytes, _ := json.Marshal(reqObj)
+    req := httptest.NewRequest("POST", "/v1/chat/completions", bytes.NewReader(reqBytes))
+    req.Header.Set("Content-Type", "application/json")
+    w := httptest.NewRecorder()
+
+    handler.ServeHTTP(w, req)
+    if w.Result().StatusCode != http.StatusOK {
+      t.Fatalf("Expected status 200, got %d", w.Result().StatusCode)
+    }
+
+    msgs, ok := receivedBody["messages"].([]any)
+    if !ok || len(msgs) != 1 {
+      t.Fatalf("Expected 1 message in forwarded body, got %+v", receivedBody)
+    }
+    userMsg := msgs[0].(map[string]any)
+    if userMsg["role"] != "user" || userMsg["content"] != "Hello from Pi agent" {
+      t.Errorf("Expected content 'Hello from Pi agent', got %+v", userMsg)
+    }
+  })
+
+  t.Run("Top-level system prompt", func(t *testing.T) {
+    reqObj := map[string]any{
+      "model":  "test-model",
+      "system": "You are a helpful coding assistant",
+      "messages": []any{
+        map[string]any{
+          "role":    "user",
+          "content": "Write a test",
+        },
+      },
+    }
+    reqBytes, _ := json.Marshal(reqObj)
+    req := httptest.NewRequest("POST", "/v1/chat/completions", bytes.NewReader(reqBytes))
+    req.Header.Set("Content-Type", "application/json")
+    w := httptest.NewRecorder()
+
+    handler.ServeHTTP(w, req)
+
+    msgs, ok := receivedBody["messages"].([]any)
+    if !ok || len(msgs) != 2 {
+      t.Fatalf("Expected 2 messages in forwarded body, got %+v", receivedBody)
+    }
+    sysMsg := msgs[0].(map[string]any)
+    userMsg := msgs[1].(map[string]any)
+    if sysMsg["role"] != "system" || sysMsg["content"] != "You are a helpful coding assistant" {
+      t.Errorf("Unexpected system message: %+v", sysMsg)
+    }
+    if userMsg["role"] != "user" || userMsg["content"] != "Write a test" {
+      t.Errorf("Unexpected user message: %+v", userMsg)
+    }
+    if _, exists := receivedBody["system"]; exists {
+      t.Errorf("Expected top-level system field to be removed")
+    }
+  })
+
+  t.Run("Prompt field autodetection", func(t *testing.T) {
+    reqObj := map[string]any{
+      "model":  "test-model",
+      "prompt": "Explain recursion",
+    }
+    reqBytes, _ := json.Marshal(reqObj)
+    req := httptest.NewRequest("POST", "/v1/chat/completions", bytes.NewReader(reqBytes))
+    req.Header.Set("Content-Type", "application/json")
+    w := httptest.NewRecorder()
+
+    handler.ServeHTTP(w, req)
+
+    msgs, ok := receivedBody["messages"].([]any)
+    if !ok || len(msgs) != 1 {
+      t.Fatalf("Expected 1 message in forwarded body, got %+v", receivedBody)
+    }
+    userMsg := msgs[0].(map[string]any)
+    if userMsg["role"] != "user" || userMsg["content"] != "Explain recursion" {
+      t.Errorf("Unexpected user message: %+v", userMsg)
+    }
+  })
+
+  t.Run("Multimodal array content preserved", func(t *testing.T) {
+    reqObj := map[string]any{
+      "model": "test-model",
+      "messages": []any{
+        map[string]any{
+          "role": "user",
+          "content": []any{
+            map[string]any{"type": "text", "text": "Describe this image"},
+            map[string]any{"type": "image_url", "image_url": map[string]any{"url": "http://example.com/img.png"}},
+          },
+        },
+      },
+    }
+    reqBytes, _ := json.Marshal(reqObj)
+    req := httptest.NewRequest("POST", "/v1/chat/completions", bytes.NewReader(reqBytes))
+    req.Header.Set("Content-Type", "application/json")
+    w := httptest.NewRecorder()
+
+    handler.ServeHTTP(w, req)
+
+    msgs, ok := receivedBody["messages"].([]any)
+    if !ok || len(msgs) != 1 {
+      t.Fatalf("Expected 1 message in forwarded body, got %+v", receivedBody)
+    }
+    userMsg := msgs[0].(map[string]any)
+    cntArray, isArray := userMsg["content"].([]any)
+    if !isArray || len(cntArray) != 2 {
+      t.Fatalf("Expected multimodal array content to be preserved, got %+v", userMsg["content"])
+    }
+  })
+}
+
+func TestFibonacciBackoff(t *testing.T) {
+  expectedSeq := []struct {
+    n        int
+    expected int64
+  }{
+    {0, 0},
+    {1, 1},
+    {2, 1},
+    {3, 2},
+    {4, 3},
+    {5, 5},
+    {6, 8},
+    {7, 13},
+    {8, 21},
+    {9, 34},
+    {10, 55},
+  }
+
+  for _, tc := range expectedSeq {
+    got := fibonacci(tc.n)
+    if got != tc.expected {
+      t.Errorf("fibonacci(%d) = %d; want %d", tc.n, got, tc.expected)
+    }
+  }
+
+  // Verify failover works with zero-delay in handleChatCompletions
+  serverA := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+    w.WriteHeader(http.StatusInternalServerError)
+  }))
+  defer serverA.Close()
+
+  serverB := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+    w.Header().Set("Content-Type", "application/json")
+    w.WriteHeader(http.StatusOK)
+    _, _ = w.Write([]byte(`{"choices":[{"message":{"content":"Success from B"}}]} `))
+  }))
+  defer serverB.Close()
+
+  configs := []ModelConfig{
+    {Model: "m1", Key: "k1", Endpoint: serverA.URL},
+    {Model: "m1", Key: "k2", Endpoint: serverB.URL},
+  }
+  cm := &ConfigManager{
+    configs:      configs,
+    uniqueModels: []string{"m1"},
+  }
+  handler := handleChatCompletions(cm, nil, 0) // 0 delay for fast test
+
+  reqBodyObj := map[string]any{
+    "model": "m1",
+    "messages": []map[string]string{
+      {"role": "user", "content": "test"},
+    },
+  }
+  reqBytes, _ := json.Marshal(reqBodyObj)
+
+  req := httptest.NewRequest("POST", "/v1/chat/completions", bytes.NewReader(reqBytes))
+  req.Header.Set("Content-Type", "application/json")
+  w := httptest.NewRecorder()
+
+  handler.ServeHTTP(w, req)
+  if w.Result().StatusCode != http.StatusOK {
+    t.Fatalf("Expected status 200 after Fibonacci retry, got %d", w.Result().StatusCode)
+  }
+}
+
+
+