CI token usage

Author: mateusz834

splunklib/ai/engines/langchain.py

@@ -120,6 +120,8 @@
 LC_AgentMiddleware = Langchain_AgentMiddleware[Any, "InvokeContext", Any]
 LC_ModelRequest = Langchain_ModelRequest["InvokeContext"]
 
+total_token_usage: int = 0
+
 # Set to True to enable debugging mode.
 _DEBUG = False
 
@@ -291,7 +293,6 @@ async def awrap_model_call(
                 request: LC_ModelRequest,
                 handler: Callable[[LC_ModelRequest], Awaitable[LC_ModelCallResult]],
             ) -> LC_ModelCallResult:
-
                 agent_thread_ids: dict[str, set[str]] = {}
 
                 # Update the subagent schema definitions to include all thread_ids that the
@@ -498,6 +499,9 @@ async def awrap_model_call(
                 print("LLM CALL", request)
                 try:
                     resp = await handler(request)
+                except LC_StructuredOutputError as e:
+                    print("LLM FAILURE", e, e.ai_message)
+                    raise
                 except Exception as e:
                     print("LLM FAILURE", e)
                     raise
@@ -528,6 +532,45 @@ async def awrap_tool_call(
         if _DEBUG:
             lc_middleware.append(_DEBUGMiddleware())
 
+        class _TOKENUsage(LC_AgentMiddleware):
+            @override
+            async def awrap_model_call(
+                self,
+                request: LC_ModelRequest,
+                handler: Callable[[LC_ModelRequest], Awaitable[LC_ModelCallResult]],
+            ) -> LC_ModelCallResult:
+                global total_token_usage
+
+                def _extract_tokens(resp: LC_ModelCallResult) -> int:
+                    ai_message = resp
+                    if isinstance(ai_message, LC_ExtendedModelResponse):
+                        ai_message = ai_message.model_response
+                    if isinstance(ai_message, LC_ModelResponse):
+                        ai_message = next(
+                            (
+                                m
+                                for m in ai_message.result
+                                if isinstance(m, LC_AIMessage)
+                            ),
+                            None,
+                        )
+                    if ai_message is not None and ai_message.usage_metadata:
+                        return ai_message.usage_metadata.get("total_tokens", 0)
+                    return 0
+
+                try:
+                    resp = await handler(request)
+                    total_token_usage += _extract_tokens(resp)
+                    return resp
+                except LC_StructuredOutputError as e:
+                    if e.ai_message.usage_metadata:
+                        total_token_usage += e.ai_message.usage_metadata.get(
+                            "total_tokens", 0
+                        )
+                    raise
+
+        lc_middleware.append(_TOKENUsage())
+
         response_format = None
         if agent.output_schema is not None:
             if _supports_provider_strategy(model_impl):

tests/ai_testlib.py

@@ -1,15 +1,18 @@
 from typing import override
+import splunklib.ai.engines.langchain as langchain_engine
 from splunklib.ai.model import PredefinedModel
 from tests.ai_test_model import InternalAIModel, TestLLMSettings, create_model
 from tests.testlib import SDKTestCase
 
 
 class AITestCase(SDKTestCase):
     _model: PredefinedModel | None = None
+    _token_usage_before: int = 0
 
     @override
     def setUp(self) -> None:
         super().setUp()
+        self._token_usage_before = langchain_engine.total_token_usage
 
         # Our tests don't expect this app to be installed, if needed it is
         # installed on demand.
@@ -18,6 +21,12 @@ def setUp(self) -> None:
                 app.delete()
                 self.restart_splunk()
 
+    @override
+    def tearDown(self) -> None:
+        tokens_used = langchain_engine.total_token_usage - self._token_usage_before
+        print(f"\n[token usage] {self.id()}: {tokens_used} tokens")
+        super().tearDown()
+
     @property
     def test_llm_settings(self) -> TestLLMSettings:
         client_id: str = self.opts.kwargs["internal_ai_client_id"]