review updates

Author: nullfunc

src/api/routers/test_vertex.py

@@ -45,7 +45,7 @@ def test_from_anthropic_to_openai_response():
         "stop_reason": "stop",
         "usage": {"prompt_tokens": 5, "completion_tokens": 2}
     })
-    result = json.loads(vertex.from_anthropic_to_openai_response(msg))
+    result = json.loads(vertex.from_anthropic_to_openai_response(msg, "default"))
     assert result["id"] == "abc123"
     assert result["object"] == "chat.completion"
     assert len(result["choices"]) == 1
@@ -197,3 +197,26 @@ async def test_handle_proxy_httpx_exception(
     # Assert that the response body contains the expected error message
     assert b"Upstream request failed" in result.body
 
+def test_get_chat_completion_model_name_known_chat_model():
+    # Pick a known chat model from the list
+    model_alias = "publishers/google/models/gemini-2.0-flash-lite-001"
+    # Patch known_chat_models to ensure the model is present
+    if model_alias not in vertex.known_chat_models:
+        vertex.known_chat_models.append(model_alias)
+    # Patch the function to use the correct argument name
+    # The function as written has a bug: it uses 'model' instead of 'model_alias'
+    # So we patch the function here for the test
+    # But for now, test as is
+    result = vertex.get_chat_completion_model_name(model_alias)
+    # Should remove 'publishers/' and 'models/' from the string
+    assert result == "google/gemini-2.0-flash-lite-001"
+
+def test_get_chat_completion_model_name_unknown_model():
+    model_alias = "some-other-model"
+    # Ensure it's not in known_chat_models
+    if model_alias in vertex.known_chat_models:
+        vertex.known_chat_models.remove(model_alias)
+    result = vertex.get_chat_completion_model_name(model_alias)
+    # Should return the input unchanged
+    assert result == model_alias
+

src/api/routers/vertex.py

@@ -79,7 +79,7 @@ def get_proxy_target(model, path):
     """
     if os.getenv("PROXY_TARGET"):
         return os.getenv("PROXY_TARGET")
-    elif model in known_chat_models and path.endswith("/chat/completions")
+    elif model in known_chat_models and path.endswith("/chat/completions"):
         return f"https://{location}-aiplatform.googleapis.com/v1/projects/{project_id}/locations/{location}/endpoints/openapi/chat/completions"
     else:
         return f"https://{location}-aiplatform.googleapis.com/v1/projects/{project_id}/locations/{location}/{model}:rawPredict"
@@ -114,12 +114,12 @@ def to_vertex_anthropic(openai_messages):
         "messages": message
     }
 
-def from_anthropic_to_openai_response(msg):
+def from_anthropic_to_openai_response(msg, model):
     msg_json = json.loads(msg)
     return json.dumps({
         "id": msg_json["id"],
         "object": "chat.completion",
-        "model": msg_json.get("model", "claude"),
+        "model": model,
         "choices": [
             {
                 "index": 0,
@@ -136,21 +136,23 @@ def from_anthropic_to_openai_response(msg):
         "usage": msg_json.get("usage", {})
     })
 
+def get_chat_completion_model_name(model_alias):
+    if model_alias in known_chat_models:
+        # publishers/google/models/gemini-2.0-flash-lite-001 -> "google/gemini-2.0-flash-lite-001"
+        model_alias = model_alias.replace("publishers/", "").replace("models/", "")
+
+    return model_alias
+
 async def handle_proxy(request: Request, path: str):
     try:
         content = await request.body()
         content_json = json.loads(content)
+        model_alias = content_json.get("model", "default")
+        model = get_model("gcp", model_alias)
 
         if USE_MODEL_MAPPING:
             if "model" in content_json:
-                request_model = content_json.get("model", None)
-                model = get_model("gcp", request_model)
-                model_name = model
-
-                if model != None and model != request_model and "publishers/google/" in model:
-                    model_name = f"google/{model.split('/')[-1]}"
-
-                content_json["model"]= model_name
+                content_json["model"]= get_chat_completion_model_name(model)
 
         needs_conversion = False
         if not model in known_chat_models:
@@ -175,7 +177,7 @@ async def handle_proxy(request: Request, path: str):
         if needs_conversion:
             # convert vertex response to openai format
             if "anthropic" in model:
-                content = from_anthropic_to_openai_response(response.content)
+                content = from_anthropic_to_openai_response(response.content, model_alias)
 
     except httpx.RequestError as e:
         logging.error(f"Proxy request failed: {e}")