deepset-mcp 0.0.4__py3-none-any.whl → 0.0.5rc1__py3-none-any.whl

deepset_mcp/tool_models.py

@@ -0,0 +1,42 @@
+# SPDX-FileCopyrightText: 2025-present deepset GmbH <info@deepset.ai>
+#
+# SPDX-License-Identifier: Apache-2.0
+
+from dataclasses import dataclass, field
+from enum import StrEnum
+from typing import Any
+
+
+class MemoryType(StrEnum):
+    """Configuration for how memory is provided to tools."""
+
+    EXPLORABLE = "explorable"
+    REFERENCEABLE = "referenceable"
+    BOTH = "both"
+    NO_MEMORY = "no_memory"
+
+
+@dataclass
+class ToolConfig:
+    """Configuration for tool registration."""
+
+    needs_client: bool = False
+    needs_workspace: bool = False
+    memory_type: MemoryType = MemoryType.NO_MEMORY
+    custom_args: dict[str, Any] = field(default_factory=dict)
+
+
+class WorkspaceMode(StrEnum):
+    """Configuration for how workspace is provided to tools."""
+
+    STATIC = "static"  # workspace from env, no parameter in tool signature
+    DYNAMIC = "dynamic"  # workspace as required parameter in tool signature
+
+
+@dataclass
+class DeepsetDocsConfig:
+    """Configuration for deepset documentation search tool."""
+
+    pipeline_name: str
+    api_key: str
+    workspace_name: str

deepset_mcp/tool_registry.py

@@ -0,0 +1,208 @@
+# SPDX-FileCopyrightText: 2025-present deepset GmbH <info@deepset.ai>
+#
+# SPDX-License-Identifier: Apache-2.0
+
+from collections.abc import Callable
+from typing import Any
+
+from deepset_mcp.api.client import AsyncDeepsetClient
+from deepset_mcp.config import DEFAULT_CLIENT_HEADER, DOCS_SEARCH_TOOL_NAME
+from deepset_mcp.initialize_embedding_model import get_initialized_model
+from deepset_mcp.tool_models import DeepsetDocsConfig, MemoryType, ToolConfig
+from deepset_mcp.tools.custom_components import (
+    get_latest_custom_component_installation_logs as get_latest_custom_component_installation_logs_tool,
+    list_custom_component_installations as list_custom_component_installations_tool,
+)
+from deepset_mcp.tools.doc_search import search_docs as search_docs_tool
+from deepset_mcp.tools.haystack_service import (
+    get_component_definition as get_component_definition_tool,
+    get_custom_components as get_custom_components_tool,
+    list_component_families as list_component_families_tool,
+    search_component_definition as search_component_definition_tool,
+)
+from deepset_mcp.tools.indexes import (
+    create_index as create_index_tool,
+    deploy_index as deploy_index_tool,
+    get_index as get_index_tool,
+    list_indexes as list_indexes_tool,
+    update_index as update_index_tool,
+)
+from deepset_mcp.tools.object_store import create_get_from_object_store, create_get_slice_from_object_store
+from deepset_mcp.tools.pipeline import (
+    create_pipeline as create_pipeline_tool,
+    deploy_pipeline as deploy_pipeline_tool,
+    get_pipeline as get_pipeline_tool,
+    get_pipeline_logs as get_pipeline_logs_tool,
+    list_pipelines as list_pipelines_tool,
+    search_pipeline as search_pipeline_tool,
+    update_pipeline as update_pipeline_tool,
+    validate_pipeline as validate_pipeline_tool,
+)
+from deepset_mcp.tools.pipeline_template import (
+    get_template as get_pipeline_template_tool,
+    list_templates as list_pipeline_templates_tool,
+    search_templates as search_pipeline_templates_tool,
+)
+from deepset_mcp.tools.secrets import get_secret as get_secret_tool, list_secrets as list_secrets_tool
+from deepset_mcp.tools.workspace import (
+    create_workspace as create_workspace_tool,
+    get_workspace as get_workspace_tool,
+    list_workspaces as list_workspaces_tool,
+)
+
+
+def get_docs_search_tool(config: DeepsetDocsConfig) -> Callable[..., Any]:
+    """Get a docs search tool configured with the provided config."""
+
+    async def search_docs(query: str) -> str:
+        """Search the deepset platform documentation.
+
+        This tool allows you to search through deepset's official documentation to find
+        information about features, API usage, best practices, and troubleshooting guides.
+        Use this when you need to look up specific deepset functionality or help users
+        understand how to use deepset features.
+
+        :param query: The search query to execute against the documentation.
+        :returns: The formatted search results from the documentation.
+        """
+        async with AsyncDeepsetClient(api_key=config.api_key, transport_config=DEFAULT_CLIENT_HEADER) as client:
+            response = await search_docs_tool(
+                client=client,
+                workspace=config.workspace_name,
+                pipeline_name=config.pipeline_name,
+                query=query,
+            )
+        return response
+
+    return search_docs
+
+
+TOOL_REGISTRY: dict[str, tuple[Callable[..., Any], ToolConfig]] = {
+    # Workspace tools
+    "list_pipelines": (
+        list_pipelines_tool,
+        ToolConfig(needs_client=True, needs_workspace=True, memory_type=MemoryType.EXPLORABLE),
+    ),
+    "create_pipeline": (
+        create_pipeline_tool,
+        ToolConfig(
+            needs_client=True,
+            needs_workspace=True,
+            memory_type=MemoryType.BOTH,
+            custom_args={"skip_validation_errors": True},
+        ),
+    ),
+    "update_pipeline": (
+        update_pipeline_tool,
+        ToolConfig(
+            needs_client=True,
+            needs_workspace=True,
+            memory_type=MemoryType.BOTH,
+            custom_args={"skip_validation_errors": True},
+        ),
+    ),
+    "get_pipeline": (
+        get_pipeline_tool,
+        ToolConfig(needs_client=True, needs_workspace=True, memory_type=MemoryType.EXPLORABLE),
+    ),
+    "deploy_pipeline": (
+        deploy_pipeline_tool,
+        ToolConfig(
+            needs_client=True,
+            needs_workspace=True,
+            memory_type=MemoryType.EXPLORABLE,
+            custom_args={"wait_for_deployment": True, "timeout_seconds": 600, "poll_interval": 5},
+        ),
+    ),
+    "validate_pipeline": (
+        validate_pipeline_tool,
+        ToolConfig(needs_client=True, needs_workspace=True, memory_type=MemoryType.BOTH),
+    ),
+    "get_pipeline_logs": (
+        get_pipeline_logs_tool,
+        ToolConfig(needs_client=True, needs_workspace=True, memory_type=MemoryType.EXPLORABLE),
+    ),
+    "search_pipeline": (
+        search_pipeline_tool,
+        ToolConfig(needs_client=True, needs_workspace=True, memory_type=MemoryType.EXPLORABLE),
+    ),
+    "list_indexes": (
+        list_indexes_tool,
+        ToolConfig(needs_client=True, needs_workspace=True, memory_type=MemoryType.EXPLORABLE),
+    ),
+    "get_index": (
+        get_index_tool,
+        ToolConfig(needs_client=True, needs_workspace=True, memory_type=MemoryType.EXPLORABLE),
+    ),
+    "create_index": (
+        create_index_tool,
+        ToolConfig(needs_client=True, needs_workspace=True, memory_type=MemoryType.BOTH),
+    ),
+    "update_index": (
+        update_index_tool,
+        ToolConfig(needs_client=True, needs_workspace=True, memory_type=MemoryType.BOTH),
+    ),
+    "deploy_index": (
+        deploy_index_tool,
+        ToolConfig(needs_client=True, needs_workspace=True, memory_type=MemoryType.EXPLORABLE),
+    ),
+    "list_templates": (
+        list_pipeline_templates_tool,
+        ToolConfig(
+            needs_client=True,
+            needs_workspace=True,
+            memory_type=MemoryType.EXPLORABLE,
+            custom_args={"field": "created_at", "order": "DESC", "limit": 100},
+        ),
+    ),
+    "get_template": (
+        get_pipeline_template_tool,
+        ToolConfig(needs_client=True, needs_workspace=True, memory_type=MemoryType.EXPLORABLE),
+    ),
+    "search_templates": (
+        search_pipeline_templates_tool,
+        ToolConfig(
+            needs_client=True,
+            needs_workspace=True,
+            memory_type=MemoryType.EXPLORABLE,
+            custom_args={"model": get_initialized_model()},
+        ),
+    ),
+    "list_custom_component_installations": (
+        list_custom_component_installations_tool,
+        ToolConfig(needs_client=True, needs_workspace=True, memory_type=MemoryType.EXPLORABLE),
+    ),
+    "get_latest_custom_component_installation_logs": (
+        get_latest_custom_component_installation_logs_tool,
+        ToolConfig(needs_client=True, needs_workspace=True, memory_type=MemoryType.EXPLORABLE),
+    ),
+    # Non-workspace tools
+    "list_component_families": (
+        list_component_families_tool,
+        ToolConfig(needs_client=True, memory_type=MemoryType.EXPLORABLE),
+    ),
+    "get_component_definition": (
+        get_component_definition_tool,
+        ToolConfig(needs_client=True, memory_type=MemoryType.EXPLORABLE),
+    ),
+    "search_component_definitions": (
+        search_component_definition_tool,
+        ToolConfig(
+            needs_client=True, memory_type=MemoryType.EXPLORABLE, custom_args={"model": get_initialized_model()}
+        ),
+    ),
+    "get_custom_components": (
+        get_custom_components_tool,
+        ToolConfig(needs_client=True, memory_type=MemoryType.EXPLORABLE),
+    ),
+    "list_secrets": (list_secrets_tool, ToolConfig(needs_client=True, memory_type=MemoryType.EXPLORABLE)),
+    "get_secret": (get_secret_tool, ToolConfig(needs_client=True, memory_type=MemoryType.EXPLORABLE)),
+    "list_workspaces": (list_workspaces_tool, ToolConfig(needs_client=True, memory_type=MemoryType.EXPLORABLE)),
+    "get_workspace": (get_workspace_tool, ToolConfig(needs_client=True, memory_type=MemoryType.EXPLORABLE)),
+    "create_workspace": (create_workspace_tool, ToolConfig(needs_client=True, memory_type=MemoryType.EXPLORABLE)),
+    "get_from_object_store": (create_get_from_object_store, ToolConfig(memory_type=MemoryType.NO_MEMORY)),
+    "get_slice_from_object_store": (create_get_slice_from_object_store, ToolConfig(memory_type=MemoryType.NO_MEMORY)),
+    DOCS_SEARCH_TOOL_NAME: (get_docs_search_tool, ToolConfig(memory_type=MemoryType.NO_MEMORY)),
+}
+
+ALL_DEEPSET_TOOLS = set(TOOL_REGISTRY.keys())

deepset_mcp/tools/object_store.py

@@ -0,0 +1,49 @@
+# SPDX-FileCopyrightText: 2025-present deepset GmbH <info@deepset.ai>
+#
+# SPDX-License-Identifier: Apache-2.0
+
+from collections.abc import Callable
+from typing import Any
+
+from deepset_mcp.tools.tokonomics import RichExplorer
+
+
+def create_get_from_object_store(explorer: RichExplorer) -> Callable[..., Any]:
+    """Creates the `get_from_object_store` tool."""
+
+    def get_from_object_store(object_id: str, path: str = "") -> str:
+        """Use this tool to fetch an object from the object store.
+
+        You can fetch a specific object by using the object's id (e.g. `@obj_001`).
+        You can also fetch any nested path by using the path-parameter
+            (e.g. `{"object_id": "@obj_001", "path": "user_info.given_name"}`
+            -> returns the content at obj.user_info.given_name).
+
+        :param object_id: The id of the object to fetch in the format `@obj_001`.
+        :param path: The path of the object to fetch in the format of `access.to.attr` or `["access"]["to"]["attr"]`.
+        """
+        return explorer.explore(obj_id=object_id, path=path)
+
+    return get_from_object_store
+
+
+def create_get_slice_from_object_store(explorer: RichExplorer) -> Callable[..., Any]:
+    """Creates the `get_slice_from_object_store` tool."""
+
+    def get_slice_from_object_store(
+        object_id: str,
+        start: int = 0,
+        end: int | None = None,
+        path: str = "",
+    ) -> str:
+        """Extract a slice from a string or list object that is stored in the object store.
+
+        :param object_id: Identifier of the object.
+        :param start: Start index for slicing.
+        :param end: End index for slicing (optional - leave empty to get slice from start to end of sequence).
+        :param path: Navigation path to object to slice (optional).
+        :return: String representation of the slice.
+        """
+        return explorer.slice(obj_id=object_id, start=start, end=end, path=path)
+
+    return get_slice_from_object_store

deepset_mcp/tools/tokonomics/__init__.py

@@ -2,70 +2,13 @@
 #
 # SPDX-License-Identifier: Apache-2.0
 
-"""
-Tokonomics: Explorable and Referenceable Tools for LLM Agents.
-
-=============================================================
-
-A library that provides token-efficient object exploration and reference
-passing capabilities for LLM agents.
-
-Key Features:
-- TTL-based object storage for temporary results
-- Rich object exploration with multiple rendering modes
-- Reference-based parameter passing (@obj_001.path.to.value)
-- Type-safe decorators that preserve function signatures
-- Configurable preview truncation and custom rendering callbacks
-
-Usage:
-------
-
-Basic explorable tool:
-
-    >>> from deepset_mcp.tools.tokonomics import explorable
-    >>>
-    >>> @explorable
-    >>> def get_data():
-    ...     return {"users": [{"name": "Alice", "age": 30}]}
-    >>>
-    >>> result = get_data()
-    >>> print(result)  # Shows rich preview
-    >>> result.obj_id  # "obj_001"
-    >>> result.value   # Original data
-
-Referenceable tool that accepts references:
-
-    >>> from deepset_mcp.tools.tokonomics import referenceable
-    >>>
-    >>> @referenceable
-    >>> def process_users(users: list) -> str:
-    ...     return f"Processed {len(users)} users"
-    >>>
-    >>> # Use with direct data
-    >>> process_users([{"name": "Bob"}])
-    >>>
-    >>> # Use with reference
-    >>> process_users("@obj_001.users")
-
-Exploration utilities:
-
-    >>> from deepset_mcp.tools.tokonomics import explore, search
-    >>>
-    >>> # Explore object structure
-    >>> explore("obj_001", mode="tree")
-    >>>
-    >>> # Search within objects
-    >>> search("obj_001", "Alice")
-"""
-
 from .decorators import explorable, explorable_and_referenceable, referenceable
 from .explorer import RichExplorer
-from .object_store import Explorable, ObjectRef, ObjectStore
+from .object_store import InMemoryBackend, ObjectStore
 
 __all__ = [
     # Core classes
-    "Explorable",
-    "ObjectRef",
+    "InMemoryBackend",
     "ObjectStore",
     "RichExplorer",
     # Decorators
@@ -73,5 +16,3 @@ __all__ = [
     "referenceable",
     "explorable_and_referenceable",
 ]
-
-__version__ = "0.1.0"

deepset_mcp/tools/tokonomics/decorators.py

@@ -19,15 +19,15 @@ from typing import Any, TypeVar, Union, get_args, get_origin
 
 from glom import GlomError, glom
 
-from .explorer import RichExplorer
-from .object_store import Explorable, ObjectRef, ObjectStore
+from deepset_mcp.tools.tokonomics.explorer import RichExplorer
+from deepset_mcp.tools.tokonomics.object_store import ObjectStore
 
 F = TypeVar("F", bound=Callable[..., Any])
 
 
 def _is_reference(value: Any) -> bool:
     """Check if a value is a reference string."""
-    return isinstance(value, str) and ObjectRef.parse(value) is not None
+    return isinstance(value, str) and value.startswith("@") and len(value) > 1
 
 
 def _type_allows_str(annotation: Any) -> bool:
@@ -49,75 +49,54 @@ def _add_str_to_type(annotation: Any) -> Any:
     return annotation | str
 
 
-def _enhance_docstring_for_references(original: str, param_info: dict[str, dict[str, Any]], func_name: str) -> str:
-    """Add reference documentation to function docstring.
+def _enhance_docstring_for_references(original: str, func_name: str) -> str:
+    """Create complete docstring for LLM tool with reference support.
 
     :param original: Original docstring.
-    :param param_info: Parameter modification info.
     :param func_name: Function name for examples.
-    :return: Enhanced docstring.
+    :return: Complete docstring for LLM tool.
     """
     if not original:
-        original = f"{func_name} function with reference support."
+        original = f"{func_name} function."
 
-    # Build the reference section
-    ref_section = [
-        "",
-        "**Reference Support**",
+    enhancement = [
         "",
         "All parameters accept object references in the form ``@obj_id`` or ``@obj_id.path.to.value``.",
         "",
+        "Examples::",
+        "",
+        "    # Direct call with values",
+        f"    {func_name}(data={{'key': 'value'}}, threshold=10)",
+        "",
+        "    # Call with references",
+        f"    {func_name}(data='@obj_123', threshold='@obj_456.config.threshold')",
+        "",
+        "    # Mixed call",
+        f"    {func_name}(data='@obj_123.items', threshold=10)",
     ]
 
-    if param_info:
-        ref_section.append("Parameter types after decoration:")
-        ref_section.append("")
-        for name, info in param_info.items():
-            if info["accepts_str"]:
-                ref_section.append(f"- ``{name}``: {info['original']} (already accepts strings)")
-            else:
-                ref_section.append(f"- ``{name}``: {info['original']} → {info['modified']} (now accepts references)")
-        ref_section.append("")
-
-    ref_section.extend(
-        [
-            "Examples::",
-            "",
-            "    # Direct call with values",
-            f"    {func_name}(data={{'key': 'value'}}, threshold=10)",
-            "",
-            "    # Call with references",
-            f"    {func_name}(data='@obj_001', threshold='@obj_002.config.threshold')",
-            "",
-            "    # Mixed call",
-            f"    {func_name}(data='@obj_001.items', threshold=10)",
-        ]
-    )
-
-    return original.rstrip() + "\n" + "\n".join(ref_section)
+    return original.rstrip() + "\n" + "\n".join(enhancement)
 
 
 def _enhance_docstring_for_explorable(original: str, func_name: str) -> str:
-    """Add explorable documentation to function docstring.
+    """Create complete docstring for LLM tool with output storage.
 
     :param original: Original docstring.
     :param func_name: Function name.
-    :return: Enhanced docstring.
+    :return: Complete docstring for LLM tool.
     """
     if not original:
-        original = f"{func_name} function with stored output."
+        original = f"{func_name} function."
 
-    section = [
+    enhancement = [
         "",
-        "**Output Storage**",
-        "",
-        "The output of this function is automatically stored and can be referenced in other functions.",
-        "The function returns a formatted preview of the result along with an object ID (e.g., ``@obj_123``).",
-        "",
-        "Use the returned object ID to pass this result to other functions that accept references.",
+        "The output is automatically stored and can be referenced in other functions.",
+        "Returns a formatted preview with an object ID (e.g., ``@obj_123``).",
+        "Use the object store tools in combination with the object ID to view nested properties of the object.",
+        "Use the returned object ID to pass this result to other functions.",
     ]
 
-    return original.rstrip() + "\n" + "\n".join(section)
+    return original.rstrip() + "\n" + "\n".join(enhancement)
 
 
 def explorable(
@@ -141,33 +120,36 @@ def explorable(
     ...     return {"processed": data}
     ...
     >>> result = process_data({"input": "value"})
-    >>> # result contains a preview and object ID like "@obj_001"
+    >>> # result contains a preview and object ID like "@obj_123"
     """
 
     def decorator(func: F) -> F:
         if inspect.iscoroutinefunction(func):
 
             @wraps(func)
-            async def async_wrapper(*args: Any, **kwargs: Any) -> Explorable[Any]:
+            async def async_wrapper(*args: Any, **kwargs: Any) -> str:
                 result = await func(*args, **kwargs)
                 obj_id = object_store.put(result)
                 preview = explorer.explore(obj_id)
-                return Explorable(obj_id, result, preview)
+
+                return preview
 
             # Enhance docstring
             async_wrapper.__doc__ = _enhance_docstring_for_explorable(func.__doc__ or "", func.__name__)
+            async_wrapper.__annotations__["return"] = str
             return async_wrapper  # type: ignore[return-value]
         else:
 
             @wraps(func)
-            def sync_wrapper(*args: Any, **kwargs: Any) -> Explorable[Any]:
+            def sync_wrapper(*args: Any, **kwargs: Any) -> str:
                 result = func(*args, **kwargs)
                 obj_id = object_store.put(result)
                 preview = explorer.explore(obj_id)
-                return Explorable(obj_id, result, preview)
+                return preview
 
             # Enhance docstring
             sync_wrapper.__doc__ = _enhance_docstring_for_explorable(func.__doc__ or "", func.__name__)
+            sync_wrapper.__annotations__["return"] = str
             return sync_wrapper  # type: ignore[return-value]
 
     return decorator
@@ -180,7 +162,7 @@ def referenceable(
 ) -> Callable[[F], F]:
     """Decorator factory that enables parameters to accept object references.
 
-    Parameters can accept reference strings like '@obj_001' or '@obj_001.path.to.value'
+    Parameters can accept reference strings like '@obj_id' or '@obj_id.path.to.value'
     which are automatically resolved before calling the function.
 
     :param object_store: The object store instance to use for lookups.
@@ -200,27 +182,25 @@ def referenceable(
     >>> process_data({"a": 1, "b": 2}, 10)
     >>>
     >>> # Call with references
-    >>> process_data("@obj_001", "@obj_002.config.threshold")
+    >>> process_data("@obj_123", "@obj_456.config.threshold")
     """
 
     def resolve_reference(ref_str: str) -> Any:
         """Resolve a reference string to its actual value."""
-        ref = ObjectRef.parse(ref_str)
-        if ref is None:
-            raise ValueError(f"Invalid reference format: {ref_str}")
+        obj_id, path = explorer.parse_reference(ref_str)
 
-        obj = object_store.get(ref.obj_id)
+        obj = object_store.get(obj_id)
         if obj is None:
-            raise ValueError(f"Object @{ref.obj_id} not found or expired")
+            raise ValueError(f"Object @{obj_id} not found or expired")
 
-        if ref.path:
+        if path:
             try:
-                explorer._validate_path(ref.path)
-                return glom(obj, explorer._parse_path(ref.path))
+                explorer._validate_path(path)
+                return glom(obj, explorer._parse_path(path))
             except GlomError as exc:
-                raise ValueError(f"Navigation error at {ref.path}: {exc}") from exc
+                raise ValueError(f"Navigation error at {path}: {exc}") from exc
             except ValueError as exc:
-                raise ValueError(f"Invalid path {ref.path}: {exc}") from exc
+                raise ValueError(f"Invalid path {path}: {exc}") from exc
 
         return obj
 
@@ -289,7 +269,7 @@ def referenceable(
 
             # Update signature and docstring
             async_wrapper.__signature__ = new_sig  # type: ignore[attr-defined]
-            async_wrapper.__doc__ = _enhance_docstring_for_references(func.__doc__ or "", param_info, func.__name__)
+            async_wrapper.__doc__ = _enhance_docstring_for_references(func.__doc__ or "", func.__name__)
             return async_wrapper  # type: ignore[return-value]
         else:
 
@@ -332,7 +312,7 @@ def referenceable(
 
             # Update signature and docstring
             sync_wrapper.__signature__ = new_sig  # type: ignore[attr-defined]
-            sync_wrapper.__doc__ = _enhance_docstring_for_references(func.__doc__ or "", param_info, func.__name__)
+            sync_wrapper.__doc__ = _enhance_docstring_for_references(func.__doc__ or "", func.__name__)
             return sync_wrapper  # type: ignore[return-value]
 
     return decorator
@@ -362,7 +342,7 @@ def explorable_and_referenceable(
     ...     return {**data1, **data2}
     ...
     >>> # Accepts references and returns preview with object ID
-    >>> result = merge_data("@obj_001", {"new": "data"})
+    >>> result = merge_data("@obj_123", {"new": "data"})
     >>> # result contains a preview and can be referenced as "@obj_002"
     """
 
@@ -374,26 +354,19 @@ def explorable_and_referenceable(
 
         # Combine docstrings (remove duplicate function name line)
         if ref_func.__doc__ and exp_func.__doc__:
-            # Take the reference part from ref_func and explorable part from exp_func
-            ref_lines = ref_func.__doc__.split("\n")
             exp_lines = exp_func.__doc__.split("\n")
-
-            # Find where the reference section starts
-            ref_start = next((i for i, line in enumerate(ref_lines) if "**Reference Support**" in line), len(ref_lines))
             # Find where the explorable section starts
-            exp_start = next((i for i, line in enumerate(exp_lines) if "**Output Storage**" in line), 0)
-
-            # Combine: original + reference section + explorable section
-            # Take everything from ref_func including reference section but excluding examples
-            # Find end of reference section (before examples)
-            ref_end = len(ref_lines)
-            for i in range(ref_start, len(ref_lines)):
-                if "Examples::" in ref_lines[i]:
-                    ref_end = i
-                    break
-
-            combined = ref_lines[:ref_end] + exp_lines[exp_start:]
-            exp_func.__doc__ = "\n".join(combined)
+            exp_start = next(
+                (
+                    i
+                    for i, line in enumerate(exp_lines)
+                    if "The output is automatically stored and can be referenced" in line
+                ),
+                0,
+            )
+
+            combined = ref_func.__doc__ + "\n".join(exp_lines[exp_start:])
+            exp_func.__doc__ = combined
 
         return exp_func