Feat/improve proxy strategy

pyproject.toml

@@ -42,7 +42,7 @@ dev = [
 framex = "framex.cli:framex"
 
 [project.optional-dependencies]
-ray = ["ray[serve]==2.53.0"]
+ray = ["ray[serve]==2.54.0"]
 release = ["poethepoet>=0.36.0", "python-semantic-release>=10.2.0"]
 
 

src/framex/driver/application.py

@@ -100,6 +100,7 @@ async def _on_start(deployment: Any) -> None:
             settings.auth.oauth.redirect_uri,
             oauth_callback,
             methods=["GET"],
+            include_in_schema=False,
         )
 
     @application.get(DOCS_URL, include_in_schema=False)

src/framex/driver/ingress.py

@@ -1,3 +1,4 @@
+import os
 import re
 from collections.abc import Callable
 from enum import Enum
@@ -26,6 +27,14 @@ async def health() -> str:
     return "ok"
 
 
+@app.get("/version")
+async def version() -> str:
+
+    from framex.config import settings
+
+    return settings.server.reversion or os.getenv("REVERSION") or "unknown"
+
+
 @api_ingress(app=app, name=BACKEND_NAME)
 class APIIngress:
     def __init__(self, deployments: list[Any], plugin_apis: list["PluginApi"]) -> None:
@@ -65,6 +74,7 @@ def register_route(
         direct_output: bool = False,
         tags: list[str | Enum] | None = None,
         auth_keys: list[str] | None = None,
+        include_in_schema: bool = True,
     ) -> bool:
         from framex.log import logger
 
@@ -131,6 +141,7 @@ def _verify_api_key(request: Request, api_key: str | None = Depends(api_key_head
                 tags=tags,
                 response_class=StreamingResponse if stream else JSONResponse,
                 dependencies=dependencies,
+                include_in_schema=include_in_schema,
             )
             methods_str = ",".join(m.upper() for m in methods)
             short_path = shorten_str(path)

src/framex/plugins/proxy/__init__.py

@@ -43,11 +43,11 @@ def __init__(self, **kwargs: Any) -> None:
 
     @override
     async def on_start(self) -> None:
-        if not settings.proxy_urls:  # pragma: no cover
+        if not settings.proxy_url_list:  # pragma: no cover
             logger.opt(colors=True).warning("<y>No url provided, skipping proxy plugin</y>")
             return
 
-        for url in settings.proxy_urls:
+        for url in settings.proxy_url_list:
             await self._parse_openai_docs(url)
 
         if settings.proxy_functions:
@@ -84,7 +84,7 @@ async def _parse_openai_docs(self, url: str) -> None:
         components = openapi_data.get("components", {}).get("schemas", {})
         for path, details in paths.items():
             # Check if the path is legal!
-            if not settings.is_white_url(path):
+            if not settings.is_white_url(url, path):
                 logger.opt(colors=True).warning(f"Proxy api(<y>{path}</y>) not in white_list, skipping...")
                 continue
 
@@ -142,7 +142,7 @@ async def _parse_openai_docs(self, url: str) -> None:
                     handle=handle,
                     stream=is_stream,
                     direct_output=True,
-                    tags=[__plugin_meta__.name],
+                    tags=[f"{__plugin_meta__.name}({url})"],
                 )
 
                 # Proxy api to map
@@ -169,6 +169,7 @@ async def register_proxy_func_route(
             stream=False,
             direct_output=False,
             tags=[__plugin_meta__.name],
+            include_in_schema=False,
         )
 
     async def _proxy_func_route(self, model: ProxyFuncHttpBody) -> Any:

src/framex/plugins/proxy/config.py

@@ -6,33 +6,64 @@
 from framex.plugin import get_plugin_config
 
 
+class ProxyUrlRuleConfig(BaseModel):
+    enable: list[str] = Field(default_factory=list)
+    disable: list[str] = Field(default_factory=list)
+
+
 class ProxyPluginConfig(BaseModel):
-    proxy_urls: list[str] = Field(default_factory=list)
+    proxy_urls: list[str] | dict[str, ProxyUrlRuleConfig] = Field(default_factory=list)
     force_stream_apis: list[str] = Field(default_factory=list)
     timeout: int = 600
     ingress_config: dict[str, Any] = Field(default_factory=lambda: {"max_ongoing_requests": 60})
 
-    white_list: list[str] = Field(default_factory=list)
+    white_list: list[str] = Field(default=["/*"])
 
     auth: AuthConfig = Field(default_factory=AuthConfig)
 
     proxy_functions: dict[str, list[str]] = Field(default_factory=dict)
 
-    def is_white_url(self, url: str) -> bool:
-        """Check if a URL is protected by any auth_urls rule."""
-        if self.white_list == []:  # pragma: no cover
-            return True
-        for rule in self.white_list:
-            if rule == url:
+    def is_white_url(self, base_url: str, path: str) -> bool:
+        """
+        Return True if the path is allowed by white rules.
+        Base_url specific rules take precedence.
+        """
+
+        if isinstance(self.proxy_urls, dict):
+            config = self.proxy_urls.get(base_url)
+            if config is not None:
+                if self._match_rules(config.disable, path):
+                    return False
+
+                if self._match_rules(config.enable, path):
+                    return True
+
+        # fallback global rule
+        return self._match_rules(self.white_list, path)
+
+    @staticmethod
+    def _match_rules(white_list: list[str], path: str) -> bool:
+        for rule in white_list:
+            if rule == "/*":
                 return True
-            if rule.endswith("/*") and url.startswith(rule[:-1]):
+            if rule == path:
+                return True
+            if rule.endswith("/*") and path.startswith(rule[:-1]):
                 return True
         return False
 
+    @property
+    def proxy_url_list(self) -> list[str]:
+        if isinstance(self.proxy_urls, list):
+            return self.proxy_urls
+        if isinstance(self.proxy_urls, dict):
+            return list(self.proxy_urls.keys())
+        raise TypeError("Invalid proxy_urls type")  # pragma: no cover
+
     @model_validator(mode="after")
     def validate_proxy_functions(self) -> Self:
         for url in self.proxy_functions:
-            if url not in self.proxy_urls:  # pragma: no cover
+            if url not in self.proxy_url_list:  # pragma: no cover
                 raise ValueError(f"proxy_functions url '{url}' is not covered by any proxy_urls rule")
         return self
 

tests/api/test_health.py → tests/api/test_base.py

@@ -2,6 +2,12 @@
 
 
 def test_get_health(client: TestClient):
-    """Test the version endpoint."""
+    """Test the health endpoint."""
     r = client.get("/health").json()
     assert r == "ok"
+
+
+def test_get_version(client: TestClient):
+    """Test the version endpoint."""
+    r = client.get("/version")
+    assert r.status_code == 200

tests/test_config.py

@@ -46,3 +46,17 @@ def test_oauth_config_does_not_override_custom_urls():
     assert cfg.authorization_url == "https://custom.auth"
     assert cfg.token_url == "https://custom.token"  # noqa: S105
     assert cfg.user_info_url == "https://custom.user"
+
+
+def test_proxy_config():
+    from framex.plugins.proxy.config import ProxyPluginConfig, ProxyUrlRuleConfig
+
+    proxy_config = ProxyPluginConfig(
+        proxy_urls={
+            "http://localhost:10000": ProxyUrlRuleConfig(enable=["/*"], disable=["/health"]),
+            "http://localhost:10001": ProxyUrlRuleConfig(enable=["/*"]),
+        },
+    )
+    assert not proxy_config.is_white_url("http://localhost:10000", "/health")
+    assert proxy_config.is_white_url("http://localhost:10000", "/echo")
+    assert proxy_config.is_white_url("http://localhost:10001", "/health")