Apply suggestions from code review

Co-authored-by: Jordan Stephens <jordan@stephens.io>

Author: nullfunc

src/api/routers/test_vertex.py

@@ -53,23 +53,23 @@ def test_from_anthropic_to_openai_response():
     assert result["choices"][0]["finish_reason"] == "stop"
     assert result["usage"]["prompt_tokens"] == 5
 
-def test_get_gcp_target_env(monkeypatch):
+def test_get_proxy_target_env(monkeypatch):
     monkeypatch.setenv("PROXY_TARGET", "https://custom-proxy")
-    result = vertex.get_gcp_target("any-model", "/v1/chat/completions")
+    result = vertex.get_proxy_target("any-model", "/v1/chat/completions")
     assert result == "https://custom-proxy"
 
-def test_get_gcp_target_known_chat(monkeypatch):
+def test_get_proxy_target_known_chat(monkeypatch):
     monkeypatch.delenv("PROXY_TARGET", raising=False)
     model = vertex.known_chat_models[0]
     path = "/v1/chat/completions"
-    result = vertex.get_gcp_target(model, path)
+    result = vertex.get_proxy_target(model, path)
     assert "endpoints/openapi/chat/completions" in result
 
-def test_get_gcp_target_raw_predict(monkeypatch):
+def test_get_proxy_target_raw_predict(monkeypatch):
     monkeypatch.delenv("PROXY_TARGET", raising=False)
     model = "unknown-model"
     path = "/v1/other"
-    result = vertex.get_gcp_target(model, path)
+    result = vertex.get_proxy_target(model, path)
     assert ":rawPredict" in result
 
 @patch("api.routers.vertex.get_access_token", return_value="dummy-token")
@@ -84,7 +84,7 @@ def test_get_header_removes_hop_headers(mock_token, dummy_request):
     })
     model = "test-model"
     path = "/v1/chat/completions"
-    with patch("api.routers.vertex.get_gcp_target", return_value="http://target"):
+    with patch("api.routers.vertex.get_proxy_target", return_value="http://target"):
         target_url, headers = vertex.get_header(model, req, path)
     assert target_url == "http://target"
     assert "Host" not in headers

src/api/routers/vertex.py

@@ -73,20 +73,20 @@ def get_access_token():
     credentials.refresh(auth_request)
     return credentials.token
 
-def get_gcp_target(model, path):
+def get_proxy_target(model, path):
     """
     Check if the environment variable is set to use GCP.
     """
     if os.getenv("PROXY_TARGET"):
         return os.getenv("PROXY_TARGET")
-    elif model in known_chat_models and "chat/completions" in path:
+    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"
 
 def get_header(model, request, path):
     path_no_prefix = f"/{path.lstrip('/')}".removeprefix(API_ROUTE_PREFIX)
-    target_url = get_gcp_target(model, path_no_prefix)
+    target_url = get_proxy_target(model, path_no_prefix)
 
     # remove hop-by-hop headers
     headers = {
@@ -168,7 +168,7 @@ async def handle_proxy(request: Request, path: str):
                 headers=headers,
                 content=json.dumps(content_json),
                 params=request.query_params,
-                timeout=30.0,
+                timeout=5.0,
             )
 
         content = response.content