unique_deep_research 2026.32.0.dev0__tar.gz → 2026.32.0.dev2__tar.gz

{unique_deep_research-2026.32.0.dev0 → unique_deep_research-2026.32.0.dev2}/PKG-INFO

@@ -1,18 +1,18 @@
 Metadata-Version: 2.3
 Name: unique-deep-research
-Version: 2026.32.0.dev0
+Version: 2026.32.0.dev2
 Summary: Deep Research Tool for complex research tasks
 Author: Martin Fadler, Sadique Sheik, Fabian Schläpfer, Pascal Hauri
 Author-email: Martin Fadler <martin.fadler@unique.ch>, Sadique Sheik <sadique@unique.ch>, Fabian Schläpfer <fabian@unique.ch>, Pascal Hauri <pascal@unique.ch>
 License: Proprietary
 Requires-Dist: pydantic>=2.8.2,<3
 Requires-Dist: typing-extensions>=4.9.0,<5
-Requires-Dist: unique-toolkit>=2026.32.0.dev0,<2026.32.0rc0
+Requires-Dist: unique-toolkit>=2026.32.0.dev9,<2026.32.0rc0
 Requires-Dist: jinja2>=3.1.2,<4
 Requires-Dist: openai>=1.99.0,<3
 Requires-Dist: langgraph>=1.0.0,<2
 Requires-Dist: langchain-core>=1.0.0,<2
-Requires-Dist: unique-web-search>=2026.32.0.dev0,<2026.32.0rc0
+Requires-Dist: unique-web-search>=2026.32.0.dev3,<2026.32.0rc0
 Requires-Dist: httpx>=0.28.0,<0.29
 Requires-Dist: timeout-decorator>=0.5.0,<0.6
 Requires-Dist: beautifulsoup4>=4.12.0,<5

{unique_deep_research-2026.32.0.dev0 → unique_deep_research-2026.32.0.dev2}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "unique_deep_research"
-version = "2026.32.0.dev0"
+version = "2026.32.0.dev2"
 description = "Deep Research Tool for complex research tasks"
 readme = "README.md"
 license = { text = "Proprietary" }
@@ -14,12 +14,12 @@ requires-python = ">=3.12,<4"
 dependencies = [
     "pydantic>=2.8.2,<3",
     "typing-extensions>=4.9.0,<5",
-    "unique-toolkit>=2026.32.0.dev0,<2026.32.0rc0",
+    "unique-toolkit>=2026.32.0.dev9,<2026.32.0rc0",
     "jinja2>=3.1.2,<4",
     "openai>=1.99.0,<3",
     "langgraph>=1.0.0,<2",
     "langchain-core>=1.0.0,<2",
-    "unique-web-search>=2026.32.0.dev0,<2026.32.0rc0",
+    "unique-web-search>=2026.32.0.dev3,<2026.32.0rc0",
     "httpx>=0.28.0,<0.29",
     "timeout-decorator>=0.5.0,<0.6",
     "beautifulsoup4>=4.12.0,<5",

unique_deep_research-2026.32.0.dev2/unique_deep_research/invocation_stats.py

@@ -0,0 +1,84 @@
+"""Run-scoped collection of Deep Research language-model usage."""
+
+import logging
+from collections.abc import Iterator
+from contextlib import contextmanager
+from contextvars import ContextVar, Token
+from typing import Any
+
+from unique_toolkit.language_model.invocation_stats import LanguageModelInvocationStats
+from unique_toolkit.language_model.schemas import LanguageModelTokenUsage
+
+_LOGGER = logging.getLogger(__name__)
+
+_CURRENT_INVOCATION_STATS: ContextVar[list[LanguageModelInvocationStats] | None] = (
+    ContextVar("deep_research_invocation_stats", default=None)
+)
+
+
+@contextmanager
+def invocation_stats_scope() -> Iterator[list[LanguageModelInvocationStats]]:
+    """Collect stats for one Deep Research run and isolate concurrent runs."""
+    invocation_stats: list[LanguageModelInvocationStats] = []
+    token: Token[list[LanguageModelInvocationStats] | None] = (
+        _CURRENT_INVOCATION_STATS.set(invocation_stats)
+    )
+    try:
+        yield invocation_stats
+    finally:
+        _CURRENT_INVOCATION_STATS.reset(token)
+
+
+def record_token_usage(
+    *,
+    model_name: str,
+    usage: Any,
+    source: str,
+) -> None:
+    """Record provider usage when both a run scope and usage are available."""
+    invocation_stats = _CURRENT_INVOCATION_STATS.get()
+    if invocation_stats is None or usage is None:
+        return
+
+    try:
+        token_usage = LanguageModelTokenUsage.model_validate(usage)
+        invocation_stats.append(
+            LanguageModelInvocationStats.from_usage(
+                model_name=model_name,
+                token_usage=token_usage,
+                source=source,
+            )
+        )
+    except Exception:
+        _LOGGER.warning(
+            "Unable to parse Deep Research token usage for %s",
+            source,
+            exc_info=True,
+        )
+
+
+def record_invocation_stats(
+    invocations: list[LanguageModelInvocationStats],
+) -> None:
+    """Add invocation stats collected by a nested Deep Research dependency."""
+    invocation_stats = _CURRENT_INVOCATION_STATS.get()
+    if invocation_stats is not None:
+        invocation_stats.extend(invocations)
+
+
+def record_language_model_response(
+    *,
+    model_name: str,
+    response: Any,
+    source: str,
+) -> None:
+    """Record usage from OpenAI SDK, toolkit, or LangChain responses."""
+    usage = getattr(response, "usage", None)
+    if usage is None:
+        usage = getattr(response, "usage_metadata", None)
+    if usage is None:
+        response_metadata = getattr(response, "response_metadata", None)
+        if isinstance(response_metadata, dict):
+            usage = response_metadata.get("token_usage")
+
+    record_token_usage(model_name=model_name, usage=usage, source=source)

{unique_deep_research-2026.32.0.dev0 → unique_deep_research-2026.32.0.dev2}/unique_deep_research/service.py

@@ -50,6 +50,9 @@ from unique_toolkit.language_model.schemas import (
     LanguageModelMessageRole,
 )
 from unique_toolkit.short_term_memory.service import ShortTermMemoryService
+from unique_web_search.invocation_stats import (
+    invocation_stats_scope as web_search_invocation_stats_scope,
+)
 
 from unique_deep_research.unique_custom.tools import crawl_url, get_today_str
 
@@ -60,6 +63,12 @@ from .config import (
     DeepResearchToolConfig,
     UniqueEngine,
 )
+from .invocation_stats import (
+    invocation_stats_scope,
+    record_invocation_stats,
+    record_language_model_response,
+    record_token_usage,
+)
 from .markdown_utils import (
     export_report_to_docx,
     postprocess_research_result_with_chunks,
@@ -261,39 +270,65 @@ class DeepResearchTool(Tool[DeepResearchToolConfig]):
         )
 
     async def run(self, tool_call: LanguageModelFunction) -> ToolCallResponse:
-        sub = self.chat_service.cancellation.on_cancellation.subscribe(
-            self._on_cancellation
-        )
-        try:
-            return await self._run(tool_call)
-        except Exception as e:
-            self.write_message_log_text_message(
-                "**Research failed for an unknown reason**"
+        with invocation_stats_scope() as invocation_stats:
+            sub = self.chat_service.cancellation.on_cancellation.subscribe(
+                self._on_cancellation
             )
-            if self.is_message_execution():
-                await self._update_execution_status(MessageExecutionUpdateStatus.FAILED)
+            try:
+                response = await self._run(tool_call)
+                response.invocation_stats.extend(invocation_stats)
+                return response
+            except Exception as e:
+                _LOGGER.exception("Deep Research tool run failed")
+
+                # The failure-handling below issues several awaits (status
+                # update, debugInfo read/write, message edit); if any of them
+                # raises it must not escape and take the tokens already spent
+                # this run down with it -- the finally block still returns a
+                # stats-bearing error response so analytics stay accurate.
+                try:
+                    self.write_message_log_text_message(
+                        "**Research failed for an unknown reason**"
+                    )
+                    if self.is_message_execution():
+                        await self._update_execution_status(
+                            MessageExecutionUpdateStatus.FAILED
+                        )
 
-            _LOGGER.exception(f"Deep Research tool run failed: {e}")
+                    # update_debug_info_async replaces the message's debugInfo
+                    # wholesale (see ChatService.update_debug_info_async /
+                    # replace_debug_info), so the existing debugInfo must be read
+                    # and merged in first -- otherwise this write wipes out
+                    # `llm_invocations` persisted from an earlier turn of this same
+                    # research (e.g. the clarification round), and UniqueAI's later
+                    # merge (unique_ai.py) reads back an empty prior list.
+                    existing_debug_info = await self.chat_service.get_debug_info_async()
+                    debug_info = {
+                        **existing_debug_info,
+                        **self._get_tool_debug_info(),
+                        "error": str(e),
+                        "traceback": traceback.format_exc(),
+                    }
+                    await self.chat_service.update_debug_info_async(debug_info)
+                    await self.chat_service.modify_assistant_message_async(
+                        content="Deep Research failed to complete for an unknown reason",
+                        set_completed_at=True,
+                    )
+                except Exception:
+                    _LOGGER.exception(
+                        "Failed to persist Deep Research failure state; "
+                        "returning collected invocation stats anyway"
+                    )
+            finally:
+                sub.cancel()
 
-            debug_info = {
-                **self._get_tool_debug_info(),
-                "error": str(e),
-                "traceback": traceback.format_exc(),
-            }
-            await self.chat_service.update_debug_info_async(debug_info)
-            await self.chat_service.modify_assistant_message_async(
-                content="Deep Research failed to complete for an unknown reason",
-                set_completed_at=True,
+            return ToolCallResponse(
+                id=tool_call.id or "",
+                name=self.name,
+                content="Failed to complete research",
+                error_message="Research process failed or returned empty results",
+                invocation_stats=invocation_stats,
             )
-        finally:
-            sub.cancel()
-
-        return ToolCallResponse(
-            id=tool_call.id or "",
-            name=self.name,
-            content="Failed to complete research",
-            error_message="Research process failed or returned empty results",
-        )
 
     async def _run(self, tool_call: LanguageModelFunction) -> ToolCallResponse:
         _LOGGER.info("Starting Deep Research tool run")
@@ -337,11 +372,25 @@ class DeepResearchTool(Tool[DeepResearchToolConfig]):
             if processed_result == "":
                 raise ValueError("Research returned empty result")
 
-            await self.chat_service.modify_assistant_message_async(
-                set_completed_at=True,
-            )
-
-            await self._update_execution_status(MessageExecutionUpdateStatus.COMPLETED)
+            # The real report is already persisted by run_research() above; these
+            # two calls are trailing bookkeeping (completed-flag, execution status).
+            # If either raises, run()'s except handler would otherwise treat the
+            # whole turn as failed and overwrite the already-correct report with a
+            # generic failure message -- swallow here instead so a bookkeeping
+            # hiccup can't destroy a successful research result.
+            try:
+                await self.chat_service.modify_assistant_message_async(
+                    set_completed_at=True,
+                )
+                await self._update_execution_status(
+                    MessageExecutionUpdateStatus.COMPLETED
+                )
+            except Exception:
+                _LOGGER.warning(
+                    "Failed to finalize Deep Research completion bookkeeping "
+                    "after a successful research result",
+                    exc_info=True,
+                )
 
             return ToolCallResponse(
                 id=tool_call.id or "",
@@ -506,10 +555,18 @@ class DeepResearchTool(Tool[DeepResearchToolConfig]):
             },
         }
 
-        result = await self.chat_service.cancellation.run_with_cancellation(
-            custom_agent.ainvoke(initial_state, config=config),  # type: ignore[arg-type]
-            cancel_result={"final_report": ""},
-        )
+        with web_search_invocation_stats_scope() as web_search_invocation_stats:
+            try:
+                result = await self.chat_service.cancellation.run_with_cancellation(
+                    custom_agent.ainvoke(initial_state, config=config),  # type: ignore[arg-type]
+                    cancel_result={"final_report": ""},
+                )
+            finally:
+                # web_search_invocation_stats is populated live as nested web-search
+                # calls happen during ainvoke(); merge it into the outer Deep Research
+                # scope even if the graph raises, so tokens already spent before the
+                # failure aren't dropped from the failure ToolCallResponse.
+                record_invocation_stats(web_search_invocation_stats)
 
         research_result = result.get("final_report", "")
 
@@ -634,6 +691,11 @@ class DeepResearchTool(Tool[DeepResearchToolConfig]):
                             _LOGGER.info(
                                 f"OpenAI research token usage: {event.response.usage}"
                             )
+                            record_token_usage(
+                                model_name=self.config.engine.research_model.name,
+                                usage=event.response.usage,
+                                source="deep_research.openai_research",
+                            )
                         # Extract the final output with annotations
                         if event.response.output and len(event.response.output) > 0:
                             final_output = event.response.output[-1]
@@ -810,6 +872,11 @@ class DeepResearchTool(Tool[DeepResearchToolConfig]):
             temperature=0.1,
             max_tokens=5000,
         )
+        record_language_model_response(
+            model_name=self.config.engine.large_model.name,
+            response=response,
+            source="deep_research.report_cleanup",
+        )
 
         formatted_result = response.choices[0].message.content
         if formatted_result:
@@ -870,6 +937,11 @@ class DeepResearchTool(Tool[DeepResearchToolConfig]):
             model_name=self.config.engine.small_model.name,
             content_chunks=None,
         )
+        record_language_model_response(
+            model_name=self.config.engine.small_model.name,
+            response=response,
+            source="deep_research.clarification",
+        )
         assert isinstance(response.choices[0].message.content, str), (
             "No clarifying questions generated"
         )
@@ -897,6 +969,11 @@ class DeepResearchTool(Tool[DeepResearchToolConfig]):
             model=self.config.engine.large_model.name,
             messages=chat_messages,
         )
+        record_language_model_response(
+            model_name=self.config.engine.large_model.name,
+            response=research_response,
+            source="deep_research.research_brief",
+        )
 
         research_instructions = research_response.choices[0].message.content
         assert research_instructions, "No research instructions generated"

{unique_deep_research-2026.32.0.dev0 → unique_deep_research-2026.32.0.dev2}/unique_deep_research/unique_custom/agents.py

@@ -198,7 +198,10 @@ async def research_supervisor(
     supervisor_messages = state.get("supervisor_messages", [])
 
     response = await ainvoke_with_token_handling(
-        research_model, supervisor_messages, model_info=engine_config.research_model
+        research_model,
+        supervisor_messages,
+        model_info=engine_config.research_model,
+        source="deep_research.supervisor",
     )
     if should_force_complete and not response.tool_calls:
         _LOGGER.error("Failed to force research_complete tool call")
@@ -347,7 +350,10 @@ async def researcher(
     researcher_messages = state.get("researcher_messages", [])
     messages = [SystemMessage(content=researcher_prompt)] + researcher_messages
     response = await ainvoke_with_token_handling(
-        research_model, messages, model_info=engine_config.research_model
+        research_model,
+        messages,
+        model_info=engine_config.research_model,
+        source="deep_research.researcher",
     )
 
     return Command(
@@ -496,7 +502,10 @@ async def compress_research(
 
     # Use proactive token handling instead of retry logic
     response = await ainvoke_with_token_handling(
-        compression_model, compression_messages, model_info=custom_config.large_model
+        compression_model,
+        compression_messages,
+        model_info=custom_config.large_model,
+        source="deep_research.research_compression",
     )
 
     compressed_research = (
@@ -549,7 +558,10 @@ async def final_report_generation(
         HumanMessage(content=message),
     ]
     raw_report = await ainvoke_with_token_handling(
-        llm, messages, model_info=custom_config.large_model
+        llm,
+        messages,
+        model_info=custom_config.large_model,
+        source="deep_research.final_report",
     )
 
     refinement_prompt = TEMPLATE_ENV.get_template("report_cleanup_prompt.j2").render(
@@ -560,7 +572,10 @@ async def final_report_generation(
         HumanMessage(content=raw_report.content),
     ]
     refined_report = await ainvoke_with_token_handling(
-        llm, messages, model_info=custom_config.large_model
+        llm,
+        messages,
+        model_info=custom_config.large_model,
+        source="deep_research.final_report_refinement",
     )
 
     # Return successful report generation

{unique_deep_research-2026.32.0.dev0 → unique_deep_research-2026.32.0.dev2}/unique_deep_research/unique_custom/utils.py

@@ -29,6 +29,7 @@ from unique_toolkit.content.service import ContentService
 from unique_toolkit.language_model.infos import LanguageModelInfo
 
 from ..config import UniqueEngine
+from ..invocation_stats import record_language_model_response
 from .citation import GlobalCitationManager
 from .state import AgentState, ResearcherState, SupervisorState
 
@@ -533,6 +534,8 @@ async def ainvoke_with_token_handling(
     model: Runnable[List[BaseMessage], Any],
     messages: List[BaseMessage],
     model_info: LanguageModelInfo,
+    *,
+    source: str,
 ) -> Any:
     """
     Invoke model with proactive token filtering and retry logic with exponential backoff.
@@ -543,6 +546,7 @@ async def ainvoke_with_token_handling(
         model: The model to invoke
         messages: List of messages to send
         model_info: Model info for token limits
+        source: Stable identifier for the LLM call site
 
     Returns:
         Model response
@@ -556,7 +560,7 @@ async def ainvoke_with_token_handling(
 
     for attempt in range(max_retries + 1):
         try:
-            return await model.ainvoke(prepared_messages)
+            response = await model.ainvoke(prepared_messages)
         except Exception as e:
             _LOGGER.exception(f"Error invoking model: {e}")
             # Handle token errors by truncating history and retrying
@@ -581,3 +585,10 @@ async def ainvoke_with_token_handling(
                     f"Model invocation failed after {max_retries + 1} attempts: {str(e)}"
                 )
                 raise e
+        else:
+            record_language_model_response(
+                model_name=model_info.name,
+                response=response,
+                source=source,
+            )
+            return response