OPIK-651 Log usage information for LlamaIndex

sdks/python/src/opik/integrations/llama_index/callback.py

@@ -156,6 +156,10 @@ def on_event_end(
             # Log the output to the span with the matching id
             if event_id in self._map_event_id_to_span_data:
                 span_data = self._map_event_id_to_span_data[event_id]
+
+                model, usage_info = event_parsing_utils.get_usage_data(payload)
+                span_data.update(usage=usage_info, model=model)
+
                 span_data.update(output=span_output).init_end_time()
                 self._opik_client.span(**span_data.__dict__)
 

sdks/python/src/opik/integrations/llama_index/event_parsing_utils.py

@@ -1,6 +1,8 @@
-from typing import Optional, Dict, Any
+from typing import Optional, Dict, Any, Tuple
 from llama_index.core.callbacks import schema as llama_index_schema
 
+from opik.types import UsageDict
+
 
 def get_span_input_from_events(
     event_type: llama_index_schema.CBEventType, payload: Optional[Dict[str, Any]]
@@ -109,3 +111,25 @@ def get_span_output_from_event(
         return {"output": payload_copy}
     else:
         return None
+
+
+def get_usage_data(
+    payload: Optional[Dict[str, Any]],
+) -> Tuple[Optional[str], Optional[UsageDict]]:
+    model = None
+    usage_info = None
+
+    if payload is None:
+        return model, usage_info
+
+    response: Any = payload.get(llama_index_schema.EventPayload.RESPONSE)
+
+    if hasattr(response, "raw"):
+        if hasattr(response.raw, "model"):
+            model = response.raw.model
+        if hasattr(response.raw, "usage"):
+            usage_info = response.raw.usage.model_dump()
+            usage_info.pop("completion_tokens_details", None)
+            usage_info.pop("prompt_tokens_details", None)
+
+    return model, usage_info