add test for model rewrite

inference-gateway/src/routes/forward.rs

@@ -54,7 +54,11 @@ pub async fn forward_request(
 
     // log request duration
     let start = std::time::Instant::now();
-    let resp = client.post(new_url).json(&body).send().await?;
+    let resp = client
+        .post(new_url)
+        .json(&rewrite_request.body)
+        .send()
+        .await?;
     let duration = start.elapsed().as_millis() as i32;
     if resp.status().is_success() {
         let llm_resp = resp.json::<serde_json::Value>().await?;

inference-gateway/tests/integration_test.rs

@@ -161,3 +161,34 @@ async fn test_unavailable_model() {
     let resp = test::call_service(&app, req).await;
     assert!(resp.status().is_client_error());
 }
+
+#[ignore]
+#[actix_web::test]
+async fn test_model_rewrite() {
+    let model = "facebook/davinci";
+    std::env::set_var("MODEL_REWRITES", format!("{model}:facebook/opt-125m"));
+
+    let app = common::get_test_app(false).await;
+
+    let mut rng = rand::thread_rng();
+    let rnd = rng.gen_range(0..100000);
+    let instance = format!("MY-TEST-INSTANCE-{}", rnd);
+    let payload = serde_json::json!({
+        "model": model,
+        "messages": [{"role": "user", "content": "the quick brown fox..."}]
+    });
+    let req = test::TestRequest::post()
+        .uri("/v1/chat/completions")
+        .insert_header(("X-TEMBO-ORG", "MY-TEST-ORG"))
+        .insert_header(("X-TEMBO-INSTANCE", instance.clone()))
+        .insert_header((header::CONTENT_TYPE, "application/json"))
+        .set_payload(payload.to_string())
+        .to_request();
+
+    let resp = test::call_service(&app, req).await;
+    assert!(resp.status().is_success());
+
+    let body: serde_json::Value = test::read_body_json(resp).await;
+    let return_model = body.get("model").unwrap().as_str().unwrap();
+    assert_eq!(return_model, "facebook/opt-125m");
+}