mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-08 03:08:45 +00:00
Implement pass-through route registration and removal to prevent duplicates and memory leaks
This commit is contained in:
parent
758fe45401
commit
fd7808e0f3
1 changed files with 60 additions and 0 deletions
|
|
@ -57,6 +57,9 @@ router = APIRouter()
|
|||
|
||||
pass_through_endpoint_logging = PassThroughEndpointLogging()
|
||||
|
||||
# Global registry to track registered pass-through routes and prevent memory leaks
|
||||
_registered_pass_through_routes: Dict[str, Dict[str, str]] = {}
|
||||
|
||||
|
||||
def get_response_body(response: httpx.Response) -> Optional[dict]:
|
||||
try:
|
||||
|
|
@ -952,8 +955,19 @@ class InitPassThroughEndpointHelpers:
|
|||
merge_query_params: Optional[bool],
|
||||
dependencies: Optional[List],
|
||||
cost_per_request: Optional[float],
|
||||
endpoint_id: str,
|
||||
):
|
||||
"""Add exact path route for pass-through endpoint"""
|
||||
route_key = f"{endpoint_id}:exact:{path}"
|
||||
|
||||
# Check if this exact route is already registered
|
||||
if route_key in _registered_pass_through_routes:
|
||||
verbose_proxy_logger.debug(
|
||||
"Skipping duplicate exact pass through endpoint: %s (already registered)",
|
||||
path,
|
||||
)
|
||||
return
|
||||
|
||||
verbose_proxy_logger.debug(
|
||||
"adding exact pass through endpoint: %s, dependencies: %s",
|
||||
path,
|
||||
|
|
@ -974,6 +988,13 @@ class InitPassThroughEndpointHelpers:
|
|||
methods=["GET", "POST", "PUT", "DELETE", "PATCH"],
|
||||
dependencies=dependencies,
|
||||
)
|
||||
|
||||
# Register the route to prevent duplicates
|
||||
_registered_pass_through_routes[route_key] = {
|
||||
"endpoint_id": endpoint_id,
|
||||
"path": path,
|
||||
"type": "exact"
|
||||
}
|
||||
|
||||
@staticmethod
|
||||
def add_subpath_route(
|
||||
|
|
@ -985,9 +1006,20 @@ class InitPassThroughEndpointHelpers:
|
|||
merge_query_params: Optional[bool],
|
||||
dependencies: Optional[List],
|
||||
cost_per_request: Optional[float],
|
||||
endpoint_id: str,
|
||||
):
|
||||
"""Add wildcard route for sub-paths"""
|
||||
wildcard_path = f"{path}/{{subpath:path}}"
|
||||
route_key = f"{endpoint_id}:subpath:{path}"
|
||||
|
||||
# Check if this subpath route is already registered
|
||||
if route_key in _registered_pass_through_routes:
|
||||
verbose_proxy_logger.debug(
|
||||
"Skipping duplicate wildcard pass through endpoint: %s (already registered)",
|
||||
wildcard_path,
|
||||
)
|
||||
return
|
||||
|
||||
verbose_proxy_logger.debug(
|
||||
"adding wildcard pass through endpoint: %s, dependencies: %s",
|
||||
wildcard_path,
|
||||
|
|
@ -1009,6 +1041,26 @@ class InitPassThroughEndpointHelpers:
|
|||
methods=["GET", "POST", "PUT", "DELETE", "PATCH"],
|
||||
dependencies=dependencies,
|
||||
)
|
||||
|
||||
# Register the route to prevent duplicates
|
||||
_registered_pass_through_routes[route_key] = {
|
||||
"endpoint_id": endpoint_id,
|
||||
"path": path,
|
||||
"type": "subpath"
|
||||
}
|
||||
|
||||
@staticmethod
|
||||
def remove_endpoint_routes(endpoint_id: str):
|
||||
"""Remove all routes for a specific endpoint ID from the registry"""
|
||||
keys_to_remove = [
|
||||
key for key, value in _registered_pass_through_routes.items()
|
||||
if value["endpoint_id"] == endpoint_id
|
||||
]
|
||||
for key in keys_to_remove:
|
||||
del _registered_pass_through_routes[key]
|
||||
verbose_proxy_logger.debug(
|
||||
"Removed pass-through route from registry: %s", key
|
||||
)
|
||||
|
||||
|
||||
async def initialize_pass_through_endpoints(
|
||||
|
|
@ -1072,6 +1124,7 @@ async def initialize_pass_through_endpoints(
|
|||
merge_query_params=_merge_query_params,
|
||||
dependencies=_dependencies,
|
||||
cost_per_request=endpoint.get("cost_per_request", None),
|
||||
endpoint_id=endpoint.get("id"),
|
||||
)
|
||||
|
||||
# Add wildcard route for sub-paths
|
||||
|
|
@ -1085,6 +1138,7 @@ async def initialize_pass_through_endpoints(
|
|||
merge_query_params=_merge_query_params,
|
||||
dependencies=_dependencies,
|
||||
cost_per_request=endpoint.get("cost_per_request", None),
|
||||
endpoint_id=endpoint.get("id"),
|
||||
)
|
||||
|
||||
verbose_proxy_logger.debug("Added new pass through endpoint: %s (ID: %s)", _path, endpoint.get("id"))
|
||||
|
|
@ -1229,6 +1283,9 @@ async def update_pass_through_endpoints(
|
|||
# Update the list
|
||||
pass_through_endpoint_data[endpoint_index] = endpoint_dict
|
||||
|
||||
# Remove old routes from registry before they get re-registered
|
||||
InitPassThroughEndpointHelpers.remove_endpoint_routes(endpoint_id)
|
||||
|
||||
## Update db
|
||||
updated_data = ConfigFieldUpdate(
|
||||
field_name="pass_through_endpoints",
|
||||
|
|
@ -1369,6 +1426,9 @@ async def delete_pass_through_endpoints(
|
|||
pass_through_endpoint_data.pop(endpoint_index)
|
||||
response_obj = found_endpoint
|
||||
|
||||
# Remove routes from registry
|
||||
InitPassThroughEndpointHelpers.remove_endpoint_routes(endpoint_id)
|
||||
|
||||
## Update db
|
||||
updated_data = ConfigFieldUpdate(
|
||||
field_name="pass_through_endpoints",
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue