unique-toolkit 2026.30.0.dev18__tar.gz → 2026.30.0.dev20__tar.gz

{unique_toolkit-2026.30.0.dev18 → unique_toolkit-2026.30.0.dev20}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: unique-toolkit
-Version: 2026.30.0.dev18
+Version: 2026.30.0.dev20
 Summary: 
 Author: Cedric Klinkert, Martin Fadler, Sadique Sheik, Fabian Schläpfer, Pascal Hauri, Ejiroghene Oki-Peter, Rami Azouz, Ricardo Visini, Sanzio Monti
 Author-email: Cedric Klinkert <cedric.klinkert@unique.ch>, Martin Fadler <martin.fadler@unique.ch>, Sadique Sheik <sadique@unique.ch>, Fabian Schläpfer <fabian@unique.ch>, Pascal Hauri <pascal@unique.ch>, Ejiroghene Oki-Peter <oki-peter.ejiroghene@unique.ai>, Rami Azouz <rami.ext@unique.ch>, Ricardo Visini <ricardo.visini@unique.ch>, Sanzio Monti <sanzio.monti@unique.ai>
@@ -16,7 +16,7 @@ Requires-Dist: sseclient>=0.0.27,<0.0.28
 Requires-Dist: openai>=1.105.0,<3
 Requires-Dist: platformdirs>=4.0.0,<5
 Requires-Dist: pillow>=12.1.1
-Requires-Dist: unique-sdk>=2026.30.0.dev8,<2026.30.0rc0
+Requires-Dist: unique-sdk>=2026.30.0.dev9,<2026.30.0rc0
 Requires-Dist: jambo>=0.1.2,<0.2
 Requires-Dist: jinja2>=3.1.6,<4
 Requires-Dist: docxtpl>=0.20.1,<0.21

{unique_toolkit-2026.30.0.dev18 → unique_toolkit-2026.30.0.dev20}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "unique_toolkit"
-version = "2026.30.0.dev18"
+version = "2026.30.0.dev20"
 description = ""
 readme = "README.md"
 requires-python = ">=3.12"
@@ -28,7 +28,7 @@ dependencies = [
     "openai>=1.105.0,<3",
     "platformdirs>=4.0.0,<5",
     "pillow>=12.1.1",
-    "unique-sdk>=2026.30.0.dev8,<2026.30.0rc0",
+    "unique-sdk>=2026.30.0.dev9,<2026.30.0rc0",
     "jambo>=0.1.2,<0.2",
     "jinja2>=3.1.6,<4",
     "docxtpl>=0.20.1,<0.21",

{unique_toolkit-2026.30.0.dev18 → unique_toolkit-2026.30.0.dev20}/unique_toolkit/agentic/debug_info_manager/debug_info_manager.py

@@ -2,6 +2,9 @@ from typing import Any, Required, TypedDict
 
 from unique_toolkit.agentic.tools.openai_builtin import OpenAICodeInterpreterTool
 from unique_toolkit.agentic.tools.openai_builtin.base import OpenAIBuiltInToolName
+from unique_toolkit.agentic.tools.openai_builtin.code_interpreter.postprocessors.generated_files import (
+    ArtifactsDebugInfo,
+)
 from unique_toolkit.agentic.tools.schemas import ToolCallResponse
 from unique_toolkit.agentic.tools.tool_manager import ToolManager
 from unique_toolkit.language_model.schemas import (
@@ -55,7 +58,7 @@ class Analytics(TypedDict):
 
 class DebugInfoManager:
     def __init__(self):
-        self.debug_info = {"tools": []}
+        self.debug_info: dict[str, Any] = {"tools": []}
 
     def extract_tool_debug_info(
         self,
@@ -106,6 +109,7 @@ class DebugInfoManager:
         answer_length: int = 0,
         loop_iteration_count: int = 0,
         total_time_to_answer_ms: int | None = None,
+        artifacts: ArtifactsDebugInfo | None = None,
     ) -> None:
         """Add a stable 'analytics' snapshot for downstream ROI/usage reporting.
 
@@ -121,6 +125,10 @@ class DebugInfoManager:
             _to_analytics_tool_entry(tool, tool_display_names)
             for tool in self.debug_info.get("tools", [])
         ]
+        # Artifacts are recorded at the postprocessor seam (Code Interpreter's
+        # DisplayCodeInterpreterFilesPostProcessor) into debug_info["artifacts"].
+        # Absent when no Code Interpreter ran this turn → both fields stay None,
+        # preserving the always-present, null-when-unknown contract.
         analytics = Analytics(
             tools_used=analytics_tools,
             tool_call_count=len(analytics_tools),
@@ -140,10 +148,8 @@ class DebugInfoManager:
             subagent_names_used=_unique_tool_names(analytics_tools, "is_sub_agent"),
             mcp_tool_names_used=_unique_tool_names(analytics_tools, "is_mcp"),
             total_time_to_answer_ms=total_time_to_answer_ms,
-            # Reserved placeholders — schema present so consumers can rely on the
-            # keys, populated in a follow-up step (artifacts instrumentation).
-            artifacts_created_count=None,
-            artifacts_created_filetype=None,
+            artifacts_created_count=artifacts["count"] if artifacts else None,
+            artifacts_created_filetype=artifacts["filetypes"] if artifacts else None,
         )
         self.add("analytics", analytics)
 

{unique_toolkit-2026.30.0.dev18 → unique_toolkit-2026.30.0.dev20}/unique_toolkit/agentic/evaluation/context_relevancy/service.py

@@ -32,6 +32,9 @@ from unique_toolkit.language_model.infos import (
     LanguageModelInfo,
     ModelCapabilities,
 )
+from unique_toolkit.language_model.invocation_stats import (
+    LanguageModelInvocationStats,
+)
 from unique_toolkit.language_model.schemas import (
     LanguageModelMessages,
     LanguageModelSystemMessage,
@@ -145,6 +148,10 @@ class ContextRelevancyEvaluator:
             # Handle regular output
             return await self._handle_regular_output(input, config)
 
+        except EvaluatorException:
+            # Already carries whatever invocation_stats were captured before
+            # the failure -- don't re-wrap it and lose them.
+            raise
         except Exception as e:
             error_message = (
                 "Unknown error occurred during context relevancy metric analysis"
@@ -172,21 +179,49 @@ class ContextRelevancyEvaluator:
             other_options=config.additional_llm_options,
         )
 
+        # Captured as soon as the (billable) LLM call returns, so a validation
+        # failure below can still report it instead of losing it.
+        invocation_stats = (
+            [
+                LanguageModelInvocationStats.from_usage(
+                    config.language_model.name,
+                    result.usage,
+                    source="context_relevancy",
+                )
+            ]
+            if result.usage is not None
+            else []
+        )
+
         try:
             result_content = EvaluationSchemaStructuredOutput.model_validate(
                 result.choices[0].message.parsed
             )
+            evaluation_result = parse_eval_metric_result_structured_output(
+                result_content, EvaluationMetricName.CONTEXT_RELEVANCY
+            )
         except ValidationError as e:
             error_message = "Error occurred during structured output validation of the context relevancy evaluation."
             raise EvaluatorException(
                 error_message=error_message,
                 user_message=error_message,
                 exception=e,
+                invocation_stats=invocation_stats,
+            )
+        except Exception as e:
+            # Catches anything else reading the response (e.g. an empty
+            # `choices` list) -- must still carry invocation_stats, not just
+            # the validation-specific failure above.
+            error_message = "Error occurred while reading the structured output response for context relevancy evaluation."
+            raise EvaluatorException(
+                error_message=f"{error_message}: {e}",
+                user_message=error_message,
+                exception=e,
+                invocation_stats=invocation_stats,
             )
 
-        return parse_eval_metric_result_structured_output(
-            result_content, EvaluationMetricName.CONTEXT_RELEVANCY
-        )
+        evaluation_result.invocation_stats = invocation_stats
+        return evaluation_result
 
     async def _handle_regular_output(
         self,
@@ -201,17 +236,52 @@ class ContextRelevancyEvaluator:
             other_options=config.additional_llm_options,
         )
 
-        result_content = result.choices[0].message.content
-        if not result_content or not isinstance(result_content, str):
-            error_message = "Context relevancy evaluation did not return a result."
+        # Captured as soon as the (billable) LLM call returns, so a failure
+        # below can still report it instead of losing it.
+        invocation_stats = (
+            [
+                LanguageModelInvocationStats.from_usage(
+                    config.language_model.name,
+                    result.usage,
+                    source="context_relevancy",
+                )
+            ]
+            if result.usage is not None
+            else []
+        )
+
+        try:
+            result_content = result.choices[0].message.content
+            if not result_content or not isinstance(result_content, str):
+                error_message = "Context relevancy evaluation did not return a result."
+                raise EvaluatorException(
+                    error_message=error_message,
+                    user_message=error_message,
+                    invocation_stats=invocation_stats,
+                )
+            evaluation_result = parse_eval_metric_result(
+                result_content, EvaluationMetricName.CONTEXT_RELEVANCY
+            )
+        except EvaluatorException as e:
+            # Covers both the raise above and parse_eval_metric_result's own
+            # EvaluatorException (which has no invocation_stats, since it has
+            # no usage context) -- attach ours if it didn't already carry
+            # stats, instead of letting it propagate stats-less.
+            if not e.invocation_stats:
+                e.invocation_stats = invocation_stats
+            raise
+        except Exception as e:
+            # Catches anything else reading the response (e.g. an empty
+            # `choices` list).
+            error_message = "Error occurred while reading the context relevancy evaluation response."
             raise EvaluatorException(
-                error_message=error_message,
+                error_message=f"{error_message}: {e}",
                 user_message=error_message,
+                exception=e,
+                invocation_stats=invocation_stats,
             )
-
-        return parse_eval_metric_result(
-            result_content, EvaluationMetricName.CONTEXT_RELEVANCY
-        )
+        evaluation_result.invocation_stats = invocation_stats
+        return evaluation_result
 
     def _compose_msgs(
         self,

{unique_toolkit-2026.30.0.dev18 → unique_toolkit-2026.30.0.dev20}/unique_toolkit/agentic/evaluation/evaluation_manager.py

@@ -7,6 +7,7 @@ from unique_toolkit._common.execution import (
     Result,
     SafeTaskExecutor,
 )
+from unique_toolkit.agentic.evaluation.exception import EvaluatorException
 from unique_toolkit.agentic.evaluation.schemas import (
     EvaluationAssessmentMessage,
     EvaluationMetricName,
@@ -17,6 +18,9 @@ from unique_toolkit.chat.schemas import (
     ChatMessageAssessmentType,
 )
 from unique_toolkit.chat.service import ChatService
+from unique_toolkit.language_model.invocation_stats import (
+    LanguageModelInvocationStats,
+)
 from unique_toolkit.language_model.schemas import (
     LanguageModelStreamResponse,
 )
@@ -88,6 +92,7 @@ class EvaluationManager:
         self._evaluations: dict[EvaluationMetricName, Evaluation] = {}
         self._evaluation_passed: bool = True
         self._execution_times: dict[str, float] = {}
+        self._invocation_stats: list[LanguageModelInvocationStats] = []
 
     def add_evaluation(self, evaluation: Evaluation):
         self._evaluations[evaluation.get_name()] = evaluation
@@ -102,6 +107,7 @@ class EvaluationManager:
         assistant_message_id: str,
     ) -> list[EvaluationMetricResult]:
         self._execution_times = {}
+        self._invocation_stats = []
 
         task_executor = SafeTaskExecutor(
             logger=self._logger,
@@ -125,6 +131,7 @@ class EvaluationManager:
             )
             if not unpacked_evaluation_result.is_positive:
                 self._evaluation_passed = False
+            self._invocation_stats.extend(unpacked_evaluation_result.invocation_stats)
             evaluation_results_unpacked.append(unpacked_evaluation_result)
 
         for evaluation_name, evaluation_result in zip(
@@ -141,6 +148,10 @@ class EvaluationManager:
     def get_execution_times(self) -> dict[str, float]:
         return self._execution_times.copy()
 
+    def get_invocation_stats(self) -> list[LanguageModelInvocationStats]:
+        """Per-LLM-call stats from every evaluation that made its own LLM call."""
+        return list(self._invocation_stats)
+
     async def execute_evaluation_call(
         self,
         evaluation_name: EvaluationMetricName,
@@ -178,12 +189,18 @@ class EvaluationManager:
         evaluation_name: EvaluationMetricName,
     ) -> EvaluationMetricResult:
         if not result.success:
+            invocation_stats = (
+                result.exception.invocation_stats
+                if isinstance(result.exception, EvaluatorException)
+                else []
+            )
             return EvaluationMetricResult(
                 name=evaluation_name,
                 is_positive=True,
                 value="RED",
                 reason=str(result.exception),
                 error=Exception("Evaluation result is not successful"),
+                invocation_stats=invocation_stats,
             )
         return result.unpack()
 

unique_toolkit-2026.30.0.dev20/unique_toolkit/agentic/evaluation/exception.py

@@ -0,0 +1,23 @@
+from unique_toolkit._common.exception import CommonException
+from unique_toolkit.language_model.invocation_stats import LanguageModelInvocationStats
+
+
+class EvaluatorException(CommonException):
+    def __init__(
+        self,
+        user_message: str,
+        error_message: str,
+        exception: Exception | None = None,
+        invocation_stats: list[LanguageModelInvocationStats] | None = None,
+    ) -> None:
+        super().__init__(
+            user_message=user_message,
+            error_message=error_message,
+            exception=exception,
+        )
+        # Carries whatever usage was already captured before the failure, so a
+        # successful (billable) LLM call isn't silently dropped from the report
+        # just because something after it (parsing, validation) raised.
+        self.invocation_stats: list[LanguageModelInvocationStats] = (
+            invocation_stats or []
+        )

{unique_toolkit-2026.30.0.dev18 → unique_toolkit-2026.30.0.dev20}/unique_toolkit/agentic/evaluation/hallucination/utils.py

@@ -18,6 +18,9 @@ from unique_toolkit.content import ContentReference
 from unique_toolkit.content.schemas import ContentChunk
 from unique_toolkit.language_model.builder import MessagesBuilder
 from unique_toolkit.language_model.infos import ModelCapabilities
+from unique_toolkit.language_model.invocation_stats import (
+    LanguageModelInvocationStats,
+)
 from unique_toolkit.language_model.schemas import (
     LanguageModelMessages,
     LanguageModelStreamResponse,
@@ -73,29 +76,58 @@ async def check_hallucination(
 
     try:
         msgs = _get_msgs(input, config)
-
         result = await LanguageModelService.complete_async_util(
             company_id=company_id, user_id=user_id, messages=msgs, model_name=model_name
         )
+    except Exception as e:
+        error_message = "Error occurred during hallucination metric analysis"
+        raise EvaluatorException(
+            error_message=f"{error_message}: {e}",
+            user_message=error_message,
+            exception=e,
+        )
+
+    # Captured as soon as the (billable) LLM call returns, so a failure below
+    # (empty content, parsing) can still report it instead of losing it.
+    invocation_stats = (
+        [
+            LanguageModelInvocationStats.from_usage(
+                model_name, result.usage, source="hallucination"
+            )
+        ]
+        if result.usage is not None
+        else []
+    )
+
+    try:
         result_content = result.choices[0].message.content
         if not result_content or not isinstance(result_content, str):
             error_message = "Hallucination evaluation did not return a text result."
             raise EvaluatorException(
                 error_message=error_message,
                 user_message=error_message,
+                invocation_stats=invocation_stats,
             )
-        result = parse_eval_metric_result(
+        eval_result = parse_eval_metric_result(
             result_content,
             EvaluationMetricName.HALLUCINATION,
         )
-
-        return result
+        eval_result.invocation_stats = invocation_stats
+        return eval_result
+    except EvaluatorException as e:
+        # parse_eval_metric_result raises its own EvaluatorException (with no
+        # invocation_stats, since it has no usage context) -- attach ours if
+        # it didn't already carry stats, instead of re-raising it stats-less.
+        if not e.invocation_stats:
+            e.invocation_stats = invocation_stats
+        raise
     except Exception as e:
         error_message = "Error occurred during hallucination metric analysis"
         raise EvaluatorException(
             error_message=f"{error_message}: {e}",
             user_message=error_message,
             exception=e,
+            invocation_stats=invocation_stats,
         )
 
 

{unique_toolkit-2026.30.0.dev18 → unique_toolkit-2026.30.0.dev20}/unique_toolkit/agentic/evaluation/schemas.py

@@ -10,6 +10,7 @@ from unique_toolkit.chat.schemas import (
     ChatMessageAssessmentStatus,
     ChatMessageAssessmentType,
 )
+from unique_toolkit.language_model.invocation_stats import LanguageModelInvocationStats
 
 
 class EvaluationMetricName(StrEnum):
@@ -102,6 +103,7 @@ class EvaluationMetricResult(BaseModel):
     user_info: Optional[str] = None
     error: Exception | None = None
     fact_list: list[str] = Field(default_factory=list[str])
+    invocation_stats: list[LanguageModelInvocationStats] = Field(default_factory=list)
 
 
 class EvaluationAssessmentMessage(BaseModel):

{unique_toolkit-2026.30.0.dev18 → unique_toolkit-2026.30.0.dev20}/unique_toolkit/agentic/evaluation/tests/test_context_relevancy_service.py

@@ -16,6 +16,7 @@ from unique_toolkit.agentic.evaluation.schemas import (
 from unique_toolkit.language_model.schemas import (
     LanguageModelAssistantMessage,
     LanguageModelCompletionChoice,
+    LanguageModelTokenUsage,
 )
 from unique_toolkit.language_model.service import LanguageModelResponse
 
@@ -271,3 +272,73 @@ async def test_analyze__raises_evaluator_exception__with_unknown_error(
         assert "Unknown error occurred during context relevancy metric analysis" in str(
             exc_info.value
         )
+
+
+@pytest.mark.ai
+@pytest.mark.asyncio
+async def test_analyze__preserves_invocation_stats__on_regular_output_read_error(
+    context_relevancy_evaluator: MagicMock,
+    sample_evaluation_input: EvaluationMetricInput,
+    basic_evaluation_config: EvaluationMetricConfig,
+) -> None:
+    """
+    Purpose: A billed LLM call that returns no choices must still surface its
+    usage, instead of losing it to the generic error path.
+    Why this matters: `result.choices[0]` access is not otherwise guarded; an
+    empty `choices` list must not silently drop invocation_stats.
+    """
+    mock_result: LanguageModelResponse = LanguageModelResponse(
+        choices=[],
+        usage=LanguageModelTokenUsage(
+            completion_tokens=1, prompt_tokens=2, total_tokens=3
+        ),
+    )
+
+    with patch.object(
+        context_relevancy_evaluator.language_model_service,
+        "complete_async",
+        return_value=mock_result,
+    ):
+        with pytest.raises(EvaluatorException) as exc_info:
+            await context_relevancy_evaluator.analyze(
+                sample_evaluation_input, basic_evaluation_config
+            )
+
+    assert len(exc_info.value.invocation_stats) == 1
+    assert exc_info.value.invocation_stats[0].token_usage.total_tokens == 3
+
+
+@pytest.mark.ai
+@pytest.mark.asyncio
+async def test_analyze__preserves_invocation_stats__on_structured_output_read_error(
+    context_relevancy_evaluator: MagicMock,
+    sample_evaluation_input: EvaluationMetricInput,
+    structured_evaluation_config: EvaluationMetricConfig,
+) -> None:
+    """
+    Purpose: Same as the regular-output case, for the structured-output path.
+    """
+    mock_result: LanguageModelResponse = LanguageModelResponse(
+        choices=[],
+        usage=LanguageModelTokenUsage(
+            completion_tokens=1, prompt_tokens=2, total_tokens=3
+        ),
+    )
+    structured_output_schema: type[EvaluationSchemaStructuredOutput] = (
+        EvaluationSchemaStructuredOutput
+    )
+
+    with patch.object(
+        context_relevancy_evaluator.language_model_service,
+        "complete_async",
+        return_value=mock_result,
+    ):
+        with pytest.raises(EvaluatorException) as exc_info:
+            await context_relevancy_evaluator.analyze(
+                sample_evaluation_input,
+                structured_evaluation_config,
+                structured_output_schema,
+            )
+
+    assert len(exc_info.value.invocation_stats) == 1
+    assert exc_info.value.invocation_stats[0].token_usage.total_tokens == 3

{unique_toolkit-2026.30.0.dev18 → unique_toolkit-2026.30.0.dev20}/unique_toolkit/agentic/loop_runner/__init__.py

@@ -6,6 +6,7 @@ from unique_toolkit.agentic.loop_runner._iteration_handler_utils import (
 from unique_toolkit.agentic.loop_runner.base import (
     LoopIterationRunner,
     ResponsesLoopIterationRunner,
+    SupportsInvocationStats,
 )
 from unique_toolkit.agentic.loop_runner.middleware import (
     PlanningConfig,
@@ -42,4 +43,5 @@ __all__ = [
     "QWEN_MAX_LOOP_ITERATIONS",
     "ResponsesLoopIterationRunner",
     "ResponsesBasicLoopIterationRunner",
+    "SupportsInvocationStats",
 ]

{unique_toolkit-2026.30.0.dev18 → unique_toolkit-2026.30.0.dev20}/unique_toolkit/agentic/loop_runner/_iteration_handler_utils.py

@@ -6,6 +6,7 @@ from unique_toolkit.agentic.loop_runner.base import (
     _LoopIterationRunnerKwargs,
 )
 from unique_toolkit.chat.functions import LanguageModelStreamResponse
+from unique_toolkit.language_model.schemas import LanguageModelTokenUsage
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -103,5 +104,6 @@ async def run_forced_tools_iteration(
     response = responses[0]
     response.tool_calls = tool_calls if len(tool_calls) > 0 else None
     response.message.references = references
+    response.usage = LanguageModelTokenUsage.sum_usages(r.usage for r in responses)
 
     return response

{unique_toolkit-2026.30.0.dev18 → unique_toolkit-2026.30.0.dev20}/unique_toolkit/agentic/loop_runner/_responses_iteration_handler_utils.py

@@ -7,7 +7,10 @@ from unique_toolkit.agentic.loop_runner._responses_stream_handler_utils import (
 from unique_toolkit.agentic.loop_runner.base import (
     _ResponsesLoopIterationRunnerKwargs,
 )
-from unique_toolkit.language_model.schemas import ResponsesLanguageModelStreamResponse
+from unique_toolkit.language_model.schemas import (
+    LanguageModelTokenUsage,
+    ResponsesLanguageModelStreamResponse,
+)
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -59,5 +62,6 @@ async def handle_responses_forced_tools_iteration(
     response = responses[0]
     response.tool_calls = tool_calls if len(tool_calls) > 0 else None
     response.message.references = references
+    response.usage = LanguageModelTokenUsage.sum_usages(r.usage for r in responses)
 
     return response

{unique_toolkit-2026.30.0.dev18 → unique_toolkit-2026.30.0.dev20}/unique_toolkit/agentic/loop_runner/base.py

@@ -1,4 +1,4 @@
-from typing import Any, Protocol, Required, Unpack
+from typing import Any, Protocol, Required, Unpack, runtime_checkable
 
 from openai.types.chat import ChatCompletionNamedToolChoiceParam
 from openai.types.responses import (
@@ -15,6 +15,9 @@ from unique_toolkit.chat.functions import LanguageModelStreamResponse
 from unique_toolkit.chat.service import LanguageModelMessages
 from unique_toolkit.content import ContentChunk
 from unique_toolkit.language_model.infos import LanguageModelInfo
+from unique_toolkit.language_model.invocation_stats import (
+    LanguageModelInvocationStats,
+)
 from unique_toolkit.language_model.schemas import (
     ResponsesLanguageModelStreamResponse,
     ResponsesMessageInput,
@@ -25,6 +28,17 @@ from unique_toolkit.protocols.support import (
 )
 
 
+@runtime_checkable
+class SupportsInvocationStats(Protocol):
+    """Optional capability for a loop iteration runner (or middleware wrapping
+    one) that makes its own extra LLM call(s) beyond the main iteration
+    response — e.g. a planning step. `unique_ai.py` checks for this via
+    `isinstance` and, if present, folds the returned stats into the turn's
+    accumulator; runners that don't make extra calls need not implement it."""
+
+    def get_invocation_stats(self) -> list[LanguageModelInvocationStats]: ...
+
+
 class _LoopIterationRunnerKwargs(TypedDict, total=False):
     iteration_index: Required[int]
     streaming_handler: Required[SupportCompleteWithReferences]

{unique_toolkit-2026.30.0.dev18 → unique_toolkit-2026.30.0.dev20}/unique_toolkit/agentic/loop_runner/middleware/planning/planning.py

@@ -37,9 +37,13 @@ from unique_toolkit.language_model import (
     LanguageModelAssistantMessage,
     LanguageModelUserMessage,
 )
+from unique_toolkit.language_model.invocation_stats import (
+    LanguageModelInvocationStats,
+)
 from unique_toolkit.language_model.schemas import (
     LanguageModelMessageOptions,
     LanguageModelMessages,
+    LanguageModelTokenUsage,
     ResponsesLanguageModelStreamResponse,
 )
 
@@ -136,6 +140,12 @@ class PlanningMiddleware(LoopIterationRunner):
         self._loop_runner = loop_runner
         self._history_manager = history_manager
         self._llm_service = llm_service
+        self._invocation_stats: list[LanguageModelInvocationStats] = []
+
+    def get_invocation_stats(self) -> list[LanguageModelInvocationStats]:
+        """Return and clear the planning-step stats captured since the last call."""
+        stats, self._invocation_stats = self._invocation_stats, []
+        return stats
 
     @failsafe_async(failure_return_value=None, logger=_LOGGER)
     async def _run_plan_step(
@@ -155,6 +165,13 @@ class PlanningMiddleware(LoopIterationRunner):
             other_options=other_options,
         )
 
+        if response.usage is not None:
+            self._invocation_stats.append(
+                LanguageModelInvocationStats.from_usage(
+                    kwargs["model"].name, response.usage, source="planning"
+                )
+            )
+
         if response.choices[0].message.parsed is None:
             _LOGGER.info("Error parsing planning response")
             return None
@@ -196,6 +213,12 @@ class ResponsesPlanningMiddleware(ResponsesLoopIterationRunner):
         self._loop_runner = loop_runner
         self._history_manager = history_manager
         self._openai_client = openai_client
+        self._invocation_stats: list[LanguageModelInvocationStats] = []
+
+    def get_invocation_stats(self) -> list[LanguageModelInvocationStats]:
+        """Return and clear the planning-step stats captured since the last call."""
+        stats, self._invocation_stats = self._invocation_stats, []
+        return stats
 
     def _build_forwarded_options(
         self, kwargs: _ResponsesLoopIterationRunnerKwargs
@@ -260,6 +283,15 @@ class ResponsesPlanningMiddleware(ResponsesLoopIterationRunner):
             **forwarded_options,
         )
 
+        if response.usage is not None:
+            self._invocation_stats.append(
+                LanguageModelInvocationStats.from_usage(
+                    model_name,
+                    LanguageModelTokenUsage.model_validate(response.usage),
+                    source="planning",
+                )
+            )
+
         output_text = _get_first_tool_call_arguments(response)
         if not output_text:
             _LOGGER.info("Empty planning response")