fixing headers

Author: priyankinfinnov

src/openlayer/lib/integrations/portkey_tracer.py

@@ -194,7 +194,7 @@ def stream_chunks(
             if any(v is not None for v in chunk_usage.values()):
                 latest_usage_data = chunk_usage
 
-            # Update metadata from latest chunk (headers/cost/etc.)
+            # Update metadata from latest chunk (headers/etc.)
             chunk_metadata = extract_portkey_unit_metadata(chunk, model_name)
             if chunk_metadata:
                 latest_chunk_metadata.update(chunk_metadata)
@@ -487,7 +487,7 @@ def extract_portkey_metadata(client: "Portkey") -> Dict[str, Any]:
 
 
 def extract_portkey_unit_metadata(unit: Any, model_name: str) -> Dict[str, Any]:
-    """Extract metadata from a response or chunk unit (headers, cost, ids)."""
+    """Extract metadata from a response or chunk unit (headers, ids)."""
     metadata: Dict[str, Any] = {}
     try:
         # Extract system fingerprint if available (OpenAI-compatible)
@@ -512,16 +512,12 @@ def extract_portkey_unit_metadata(unit: Any, model_name: str) -> Dict[str, Any]:
             lower = {k.lower(): v for k, v in headers.items()}
             if "x-portkey-trace-id" in lower:
                 metadata["portkey_trace_id"] = lower["x-portkey-trace-id"]
-            if "x-portkey-provider" in lower:
-                metadata["provider"] = lower["x-portkey-provider"]
             if "x-portkey-cache-status" in lower:
                 metadata["portkey_cache_status"] = lower["x-portkey-cache-status"]
-            # Cost if surfaced by gateway
-            if "x-portkey-cost" in lower:
-                try:
-                    metadata["cost"] = float(lower["x-portkey-cost"])
-                except Exception:
-                    pass
+            if "x-portkey-retry-attempt-count" in lower:
+                metadata["portkey_retry_attempt_count"] = lower["x-portkey-retry-attempt-count"]
+            if "x-portkey-last-used-option-index" in lower:
+                metadata["portkey_last_used_option_index"] = lower["x-portkey-last-used-option-index"]
     except Exception:
         pass
     # Attach model for convenience
@@ -662,7 +658,7 @@ def calculate_streaming_usage_and_cost(
 
         total_tokens = (prompt_tokens or 0) + (completion_tokens or 0)
 
-        # Cost from headers if present; otherwise simple heuristic for some models
+        # Cost from metadata if present; otherwise simple heuristic for some models
         cost = latest_chunk_metadata.get("cost")
         if cost is None and total_tokens and model_name:
             ml = model_name.lower()
@@ -678,21 +674,14 @@ def calculate_streaming_usage_and_cost(
 def _extract_provider_from_object(obj: Any) -> Optional[str]:
     """Extract provider from a response or chunk object.
     
-    Checks response_metadata and _response_headers for provider information.
+    Checks response_metadata for provider information.
     Returns None if no provider is found.
     """
     try:
         # Check response_metadata
         if hasattr(obj, "response_metadata") and _is_dict_like(obj.response_metadata):
             if "provider" in obj.response_metadata:
                 return obj.response_metadata["provider"]
-        # Check _response_headers
-        if hasattr(obj, "_response_headers"):
-            headers = getattr(obj, "_response_headers")
-            if _is_dict_like(headers):
-                for k, v in headers.items():
-                    if isinstance(k, str) and k.lower() == "x-portkey-provider" and v:
-                        return str(v)
     except Exception:
         pass
     return None

tests/test_portkey_integration.py

@@ -145,16 +145,16 @@ def test_extract_portkey_metadata(self) -> None:
         assert "Authorization" not in metadata["portkeyHeaders"]
 
     def test_extract_portkey_unit_metadata(self) -> None:
-        """Unit metadata should capture headers, cost, and provider hints."""
+        """Unit metadata should capture headers and retry/option index hints."""
         from openlayer.lib.integrations.portkey_tracer import extract_portkey_unit_metadata
 
         unit = SimpleNamespace(
             system_fingerprint="fingerprint-123",
             _response_headers={
                 "x-portkey-trace-id": "trace-1",
-                "x-portkey-provider": "anthropic",
                 "x-portkey-cache-status": "HIT",
-                "x-portkey-cost": "0.45",
+                "x-portkey-retry-attempt-count": "2",
+                "x-portkey-last-used-option-index": "config.targets[1]",
                 "content-type": "application/json",
             },
         )
@@ -163,9 +163,9 @@ def test_extract_portkey_unit_metadata(self) -> None:
 
         assert metadata["system_fingerprint"] == "fingerprint-123"
         assert metadata["portkey_trace_id"] == "trace-1"
-        assert metadata["provider"] == "anthropic"
         assert metadata["portkey_cache_status"] == "HIT"
-        assert metadata["cost"] == pytest.approx(0.45)
+        assert metadata["portkey_retry_attempt_count"] == "2"
+        assert metadata["portkey_last_used_option_index"] == "config.targets[1]"
         assert metadata["portkey_model"] == "claude-3-opus"
         assert metadata["response_headers"]["content-type"] == "application/json"
 
@@ -178,8 +178,9 @@ class DictLikeHeaders:
             def __init__(self):
                 self._data = {
                     "x-portkey-trace-id": "trace-2",
-                    "x-portkey-provider": "openai",
                     "x-portkey-cache-status": "MISS",
+                    "x-portkey-retry-attempt-count": "3",
+                    "x-portkey-last-used-option-index": "config.targets[1]",
                 }
 
             def items(self):
@@ -192,8 +193,9 @@ def items(self):
         metadata = extract_portkey_unit_metadata(unit, "gpt-4")
 
         assert metadata["portkey_trace_id"] == "trace-2"
-        assert metadata["provider"] == "openai"
         assert metadata["portkey_cache_status"] == "MISS"
+        assert metadata["portkey_retry_attempt_count"] == "3"
+        assert metadata["portkey_last_used_option_index"] == "config.targets[1]"
 
     def test_extract_usage_from_response(self) -> None:
         """Usage extraction should read OpenAI-style usage objects."""
@@ -318,9 +320,8 @@ def test_detect_provider_from_response_fallback(self) -> None:
         """Provider detection should fall back to response metadata or model name."""
         from openlayer.lib.integrations.portkey_tracer import detect_provider
 
-        client = SimpleNamespace(headers={"x-portkey-provider": "openai"})
+        client = SimpleNamespace()
         response = SimpleNamespace(
-            _response_headers={"X-Portkey-Provider": "anthropic"},
             response_metadata={"provider": "anthropic"},
         )
 
@@ -335,7 +336,6 @@ def test_detect_provider_from_chunk_fallback(self) -> None:
 
         chunk = SimpleNamespace(
             response_metadata={"provider": "cohere"},
-            _response_headers={"X-Portkey-Provider": "cohere"},
         )
         client = SimpleNamespace()