refactoring to increase re-usability of code

Author: priyankinfinnov

src/openlayer/lib/integrations/portkey_tracer.py

@@ -185,12 +185,12 @@ def stream_chunks(
             if i == 0:
                 first_token_time = time.time()
                 # Try to detect provider at first chunk
-                provider = detect_provider_from_chunk(chunk, client, model_name)
+                provider = detect_provider(chunk, client, model_name)
             if i > 0:
                 num_of_completion_tokens = i + 1
 
             # Extract usage from chunk if available
-            chunk_usage = extract_usage_from_chunk(chunk)
+            chunk_usage = extract_usage(chunk)
             if any(v is not None for v in chunk_usage.values()):
                 latest_usage_data = chunk_usage
 
@@ -317,9 +317,9 @@ def handle_non_streaming_create(
         output_data = parse_non_streaming_output_data(response)
 
         # Usage (if provided by upstream provider via Portkey)
-        usage_data = extract_usage_from_response(response)
+        usage_data = extract_usage(response)
         model_name = getattr(response, "model", kwargs.get("model", "unknown"))
-        provider = detect_provider_from_response(response, client, model_name)
+        provider = detect_provider(response, client, model_name)
         extra_metadata = extract_portkey_unit_metadata(response, model_name)
         cost = extra_metadata.get("cost", None)
 
@@ -530,36 +530,39 @@ def extract_portkey_unit_metadata(unit: Any, model_name: str) -> Dict[str, Any]:
     return metadata
 
 
-def extract_usage_from_response(response: Any) -> Dict[str, Optional[int]]:
-    """Extract usage from a non-streaming response."""
-    try:
-        if hasattr(response, "usage") and response.usage is not None:
-            usage = response.usage
-            return {
-                "total_tokens": getattr(usage, "total_tokens", None),
-                "prompt_tokens": getattr(usage, "prompt_tokens", None),
-                "completion_tokens": getattr(usage, "completion_tokens", None),
-            }
-    except Exception:
-        pass
-    return {"total_tokens": None, "prompt_tokens": None, "completion_tokens": None}
-
-
-def extract_usage_from_chunk(chunk: Any) -> Dict[str, Optional[int]]:
-    """Extract usage from a streaming chunk if present."""
+def extract_usage(obj: Any) -> Dict[str, Optional[int]]:
+    """Extract usage from a response or chunk object.
+    
+    This function attempts to extract token usage information from various
+    locations where it might be stored, including:
+    - Direct `usage` attribute
+    - `_hidden_params` (for streaming chunks)
+    - `model_dump()` dictionary (for streaming chunks)
+    
+    Parameters
+    ----------
+    obj : Any
+        The response or chunk object to extract usage from.
+    
+    Returns
+    -------
+    Dict[str, Optional[int]]
+        Dictionary with keys: total_tokens, prompt_tokens, completion_tokens.
+        Values are None if usage information is not found.
+    """
     try:
-        # Check for usage attribute
-        if hasattr(chunk, "usage") and chunk.usage is not None:
-            usage = chunk.usage
+        # Check for direct usage attribute (works for both response and chunk)
+        if hasattr(obj, "usage") and obj.usage is not None:
+            usage = obj.usage
             return {
                 "total_tokens": getattr(usage, "total_tokens", None),
                 "prompt_tokens": getattr(usage, "prompt_tokens", None),
                 "completion_tokens": getattr(usage, "completion_tokens", None),
             }
 
-        # Check for usage in _hidden_params (if SDK stores it there)
-        if hasattr(chunk, "_hidden_params"):
-            hidden_params = chunk._hidden_params
+        # Check for usage in _hidden_params (primarily for streaming chunks)
+        if hasattr(obj, "_hidden_params"):
+            hidden_params = obj._hidden_params
             # Check if usage is a direct attribute
             if hasattr(hidden_params, "usage") and hidden_params.usage is not None:
                 usage = hidden_params.usage
@@ -578,11 +581,11 @@ def extract_usage_from_chunk(chunk: Any) -> Dict[str, Optional[int]]:
                         "completion_tokens": usage.get("completion_tokens", None),
                     }
 
-        # Check if chunk model dump has usage
-        if hasattr(chunk, "model_dump"):
-            chunk_dict = chunk.model_dump()
-            if _supports_membership_check(chunk_dict) and "usage" in chunk_dict and chunk_dict["usage"]:
-                usage = chunk_dict["usage"]
+        # Check if object model dump has usage (primarily for streaming chunks)
+        if hasattr(obj, "model_dump"):
+            obj_dict = obj.model_dump()
+            if _supports_membership_check(obj_dict) and "usage" in obj_dict and obj_dict["usage"]:
+                usage = obj_dict["usage"]
                 return {
                     "total_tokens": usage.get("total_tokens", None),
                     "prompt_tokens": usage.get("prompt_tokens", None),
@@ -672,49 +675,54 @@ def calculate_streaming_usage_and_cost(
         return None, None, None, None
 
 
-def detect_provider_from_response(response: Any, client: "Portkey", model_name: str) -> str:
-    """Detect provider for non-streaming responses."""
-    # First: check Portkey headers on the client (authoritative)
-    provider = _provider_from_portkey_headers(client)
-    if provider:
-        return provider
-    # Next: check response metadata if any
+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.
+    Returns None if no provider is found.
+    """
     try:
-        # Some SDKs attach response headers/metadata
-        if hasattr(response, "response_metadata") and _is_dict_like(response.response_metadata):
-            if "provider" in response.response_metadata:
-                return response.response_metadata["provider"]
-        if hasattr(response, "_response_headers"):
-            headers = getattr(response, "_response_headers")
+        # 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
-    # Fallback to model name heuristics
-    return detect_provider_from_model_name(model_name)
+    return None
 
 
-def detect_provider_from_chunk(chunk: Any, client: "Portkey", model_name: str) -> str:
-    """Detect provider for streaming chunks."""
-    # First: check Portkey headers on the client
+def detect_provider(obj: Any, client: "Portkey", model_name: str) -> str:
+    """Detect provider from a response or chunk object.
+    
+    Parameters
+    ----------
+    obj : Any
+        The response or chunk object to extract provider information from.
+    client : Portkey
+        The Portkey client instance.
+    model_name : str
+        The model name to use as a fallback for provider detection.
+    
+    Returns
+    -------
+    str
+        The detected provider name.
+    """
+    # First: check Portkey headers on the client (authoritative)
     provider = _provider_from_portkey_headers(client)
     if provider:
         return provider
-    # Next: see if chunk exposes any metadata
-    try:
-        if hasattr(chunk, "response_metadata") and _is_dict_like(chunk.response_metadata):
-            if "provider" in chunk.response_metadata:
-                return chunk.response_metadata["provider"]
-        if hasattr(chunk, "_response_headers"):
-            headers = getattr(chunk, "_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
+    # Next: check object metadata if any
+    provider = _extract_provider_from_object(obj)
+    if provider:
+        return provider
     # Fallback to model name heuristics
     return detect_provider_from_model_name(model_name)
 

tests/test_portkey_integration.py

@@ -197,32 +197,32 @@ def items(self):
 
     def test_extract_usage_from_response(self) -> None:
         """Usage extraction should read OpenAI-style usage objects."""
-        from openlayer.lib.integrations.portkey_tracer import extract_usage_from_response
+        from openlayer.lib.integrations.portkey_tracer import extract_usage
 
         usage = SimpleNamespace(total_tokens=50, prompt_tokens=20, completion_tokens=30)
         response = SimpleNamespace(usage=usage)
 
-        assert extract_usage_from_response(response) == {
+        assert extract_usage(response) == {
             "total_tokens": 50,
             "prompt_tokens": 20,
             "completion_tokens": 30,
         }
 
         response_no_usage = SimpleNamespace()
-        assert extract_usage_from_response(response_no_usage) == {
+        assert extract_usage(response_no_usage) == {
             "total_tokens": None,
             "prompt_tokens": None,
             "completion_tokens": None,
         }
 
     def test_extract_usage_from_chunk(self) -> None:
         """Usage data should be derived from multiple potential chunk attributes."""
-        from openlayer.lib.integrations.portkey_tracer import extract_usage_from_chunk
+        from openlayer.lib.integrations.portkey_tracer import extract_usage
 
         chunk_direct = SimpleNamespace(
             usage=SimpleNamespace(total_tokens=120, prompt_tokens=40, completion_tokens=80)
         )
-        assert extract_usage_from_chunk(chunk_direct) == {
+        assert extract_usage(chunk_direct) == {
             "total_tokens": 120,
             "prompt_tokens": 40,
             "completion_tokens": 80,
@@ -231,7 +231,7 @@ def test_extract_usage_from_chunk(self) -> None:
         chunk_hidden = SimpleNamespace(
             _hidden_params={"usage": {"total_tokens": 30, "prompt_tokens": 10, "completion_tokens": 20}}
         )
-        assert extract_usage_from_chunk(chunk_hidden) == {
+        assert extract_usage(chunk_hidden) == {
             "total_tokens": 30,
             "prompt_tokens": 10,
             "completion_tokens": 20,
@@ -241,7 +241,7 @@ class ChunkWithModelDump:  # pylint: disable=too-few-public-methods
             def model_dump(self) -> Dict[str, Any]:
                 return {"usage": {"total_tokens": 12, "prompt_tokens": 5, "completion_tokens": 7}}
 
-        assert extract_usage_from_chunk(ChunkWithModelDump()) == {
+        assert extract_usage(ChunkWithModelDump()) == {
             "total_tokens": 12,
             "prompt_tokens": 5,
             "completion_tokens": 7,
@@ -292,31 +292,31 @@ def test_calculate_streaming_usage_and_cost_fallback_estimation(self) -> None:
 
     def test_detect_provider_from_response_prefers_headers(self) -> None:
         """Provider detection should prioritize Portkey headers."""
-        from openlayer.lib.integrations.portkey_tracer import detect_provider_from_response
+        from openlayer.lib.integrations.portkey_tracer import detect_provider
 
         client = SimpleNamespace()
         response = SimpleNamespace()
 
         with patch(
             "openlayer.lib.integrations.portkey_tracer._provider_from_portkey_headers", return_value="header-provider"
         ):
-            assert detect_provider_from_response(response, client, "gpt-4") == "header-provider"
+            assert detect_provider(response, client, "gpt-4") == "header-provider"
 
     def test_detect_provider_from_chunk_prefers_headers(self) -> None:
         """Provider detection from chunk should prioritize header-derived values."""
-        from openlayer.lib.integrations.portkey_tracer import detect_provider_from_chunk
+        from openlayer.lib.integrations.portkey_tracer import detect_provider
 
         client = SimpleNamespace()
         chunk = SimpleNamespace()
 
         with patch(
             "openlayer.lib.integrations.portkey_tracer._provider_from_portkey_headers", return_value="header-provider"
         ):
-            assert detect_provider_from_chunk(chunk, client, "gpt-4") == "header-provider"
+            assert detect_provider(chunk, client, "gpt-4") == "header-provider"
 
     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_from_response
+        from openlayer.lib.integrations.portkey_tracer import detect_provider
 
         client = SimpleNamespace(headers={"x-portkey-provider": "openai"})
         response = SimpleNamespace(
@@ -327,11 +327,11 @@ def test_detect_provider_from_response_fallback(self) -> None:
         with patch(
             "openlayer.lib.integrations.portkey_tracer._provider_from_portkey_headers", return_value=None
         ):
-            assert detect_provider_from_response(response, client, "mistral-7b") == "anthropic"
+            assert detect_provider(response, client, "mistral-7b") == "anthropic"
 
     def test_detect_provider_from_chunk_fallback(self) -> None:
         """Chunk provider detection should fall back gracefully."""
-        from openlayer.lib.integrations.portkey_tracer import detect_provider_from_chunk
+        from openlayer.lib.integrations.portkey_tracer import detect_provider
 
         chunk = SimpleNamespace(
             response_metadata={"provider": "cohere"},
@@ -342,7 +342,7 @@ def test_detect_provider_from_chunk_fallback(self) -> None:
         with patch(
             "openlayer.lib.integrations.portkey_tracer._provider_from_portkey_headers", return_value=None
         ):
-            assert detect_provider_from_chunk(chunk, client, "command-r") == "cohere"
+            assert detect_provider(chunk, client, "command-r") == "cohere"
 
     def test_provider_from_portkey_headers(self) -> None:
         """Header helper should identify provider values on the client."""
@@ -496,11 +496,11 @@ def test_stream_chunks_traces_completion(self) -> None:
         with patch(
             "openlayer.lib.integrations.portkey_tracer.add_to_trace", autospec=True
         ) as mock_add_to_trace, patch(
-            "openlayer.lib.integrations.portkey_tracer.extract_usage_from_chunk", autospec=True
+            "openlayer.lib.integrations.portkey_tracer.extract_usage", autospec=True
         ) as mock_usage, patch(
             "openlayer.lib.integrations.portkey_tracer.extract_portkey_unit_metadata", autospec=True
         ) as mock_unit_metadata, patch(
-            "openlayer.lib.integrations.portkey_tracer.detect_provider_from_chunk", autospec=True
+            "openlayer.lib.integrations.portkey_tracer.detect_provider", autospec=True
         ) as mock_detect_provider, patch(
             "openlayer.lib.integrations.portkey_tracer.get_delta_from_chunk", autospec=True
         ) as mock_delta, patch(
@@ -552,10 +552,10 @@ def test_handle_non_streaming_create_traces_completion(self) -> None:
         with patch(
             "openlayer.lib.integrations.portkey_tracer.parse_non_streaming_output_data", return_value="output"
         ), patch(
-            "openlayer.lib.integrations.portkey_tracer.extract_usage_from_response",
+            "openlayer.lib.integrations.portkey_tracer.extract_usage",
             return_value={"total_tokens": 10, "prompt_tokens": 4, "completion_tokens": 6},
         ), patch(
-            "openlayer.lib.integrations.portkey_tracer.detect_provider_from_response", return_value="OpenAI"
+            "openlayer.lib.integrations.portkey_tracer.detect_provider", return_value="OpenAI"
         ), patch(
             "openlayer.lib.integrations.portkey_tracer.extract_portkey_unit_metadata",
             return_value={"cost": 0.25},