From e1c44fe0881ed714d91942d9751c9505e8eaccc4 Mon Sep 17 00:00:00 2001
From: Harshit28j <Harshitjain0562@gmail.com>
Date: Wed, 11 Mar 2026 14:17:03 +0530
Subject: [PATCH] Fix Greptile round 2 review issues

- DB-only Vantage config now registers VantageLogger at startup so
  background job actually schedules (was silently skipping)
- Add missing `except HTTPException: raise` in /vantage/init endpoint
- Use consistent batch filenames (always .partN suffix)
- Document Tags schema change from pl.Object to pl.String

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
---
 .../focus/destinations/vantage_destination.py |  2 +-
 litellm/integrations/focus/schema.py          |  3 +++
 litellm/proxy/proxy_server.py                 | 24 ++++++++++++++++++-
 .../proxy/spend_tracking/vantage_endpoints.py |  2 ++
 4 files changed, 29 insertions(+), 2 deletions(-)

diff --git a/litellm/integrations/focus/destinations/vantage_destination.py b/litellm/integrations/focus/destinations/vantage_destination.py
index 32ba6df8ea8..e8d9532c41c 100644
--- a/litellm/integrations/focus/destinations/vantage_destination.py
+++ b/litellm/integrations/focus/destinations/vantage_destination.py
@@ -108,7 +108,7 @@ class FocusVantageDestination(FocusDestination):
             if len(batch_csv) > VANTAGE_MAX_BYTES_PER_UPLOAD:
                 await self._upload_size_limited(header, batch_lines, filename, batch_num)
             else:
-                batch_filename = f"{filename}.part{batch_num}" if batch_num > 0 else filename
+                batch_filename = f"{filename}.part{batch_num}"
                 await self._upload_csv(batch_csv, batch_filename)
             batch_num += 1
 
diff --git a/litellm/integrations/focus/schema.py b/litellm/integrations/focus/schema.py
index 6ad725f367d..127e8ee03ad 100644
--- a/litellm/integrations/focus/schema.py
+++ b/litellm/integrations/focus/schema.py
@@ -43,6 +43,9 @@ FOCUS_NORMALIZED_SCHEMA = pl.Schema(
         ("SubAccountId", pl.String),
         ("SubAccountName", pl.String),
         ("SubAccountType", pl.String),
+        # Changed from pl.Object to pl.String to hold JSON metadata
+        # (team_id, user_id, etc.) needed by Vantage Token Allocation.
+        # Previously Tags was always None so no existing data is lost.
         ("Tags", pl.String),
     ]
 )
diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py
index 91196ce5b9a..9bd4cd7f140 100644
--- a/litellm/proxy/proxy_server.py
+++ b/litellm/proxy/proxy_server.py
@@ -6132,9 +6132,31 @@ class ProxyStartupEvent:
         # Vantage Background Job
         ########################################################
         from litellm.integrations.vantage.vantage_logger import VantageLogger
-        from litellm.proxy.spend_tracking.vantage_endpoints import is_vantage_setup
+        from litellm.proxy.spend_tracking.vantage_endpoints import (
+            _get_vantage_settings,
+            is_vantage_setup,
+            is_vantage_setup_in_config,
+            is_vantage_setup_in_db,
+        )
 
         if await is_vantage_setup():
+            # If configured via DB but not in config.yaml callbacks,
+            # instantiate and register a VantageLogger so the scheduler
+            # can find it.
+            if not is_vantage_setup_in_config() and await is_vantage_setup_in_db():
+                try:
+                    db_settings = await _get_vantage_settings()
+                    if db_settings:
+                        vantage_logger = VantageLogger(
+                            api_key=db_settings.get("api_key"),
+                            integration_token=db_settings.get("integration_token"),
+                            base_url=db_settings.get("base_url"),
+                        )
+                        litellm.callbacks.append(vantage_logger)  # type: ignore[arg-type]
+                except Exception as e:
+                    verbose_proxy_logger.warning(
+                        "Failed to register VantageLogger from DB settings: %s", e
+                    )
             await VantageLogger.init_vantage_background_job(scheduler=scheduler)
 
         ########################################################
diff --git a/litellm/proxy/spend_tracking/vantage_endpoints.py b/litellm/proxy/spend_tracking/vantage_endpoints.py
index 47838158b96..d0e01fe325f 100644
--- a/litellm/proxy/spend_tracking/vantage_endpoints.py
+++ b/litellm/proxy/spend_tracking/vantage_endpoints.py
@@ -304,6 +304,8 @@ async def init_vantage_settings(
             message="Vantage settings initialized successfully", status="success"
         )
 
+    except HTTPException:
+        raise
     except Exception as e:
         verbose_proxy_logger.error(
             f"Error initializing Vantage settings: {str(e)}"