fix: docstring enhancement with decorators (#169)

mathislucka · web-flow · commit 4d1e6a90ac45 · 2025-07-28T13:59:09.000+02:00
* feat: improve programmatic use

* chore: remove unused param

* chore: simplify combined decorator

* chore: format

* fix: api key pass through

* chore: format
diff --git a/src/deepset_mcp/tool_factory.py b/src/deepset_mcp/tool_factory.py
@@ -227,6 +227,7 @@ def build_tool(
     base_func: Callable[..., Any],
     config: ToolConfig,
     workspace_mode: WorkspaceMode,
+    api_key: str | None = None,
     workspace: str | None = None,
     use_request_context: bool = True,
     base_url: str | None = None,
@@ -240,6 +241,7 @@ def build_tool(
     :param base_func: The base tool function.
     :param config: Tool configuration specifying dependencies and custom arguments.
     :param workspace_mode: How the workspace should be handled.
+    :param api_key: The deepset API key to use.
     :param workspace: The workspace to use when using a static workspace.
     :param use_request_context: Whether to collect the API key from the request context.
     :param base_url: Base URL for the deepset API.
@@ -258,7 +260,9 @@ def build_tool(
     enhanced_func = apply_workspace(enhanced_func, config, workspace_mode, workspace)
 
     # Apply client injection (adds ctx parameter if needed)
-    enhanced_func = apply_client(enhanced_func, config, use_request_context=use_request_context, base_url=base_url)
+    enhanced_func = apply_client(
+        enhanced_func, config, use_request_context=use_request_context, base_url=base_url, api_key=api_key
+    )
 
     # Create final async wrapper if needed
     if not inspect.iscoroutinefunction(enhanced_func):
@@ -354,13 +358,14 @@ def register_tools(
             enhanced_tool = base_func(explorer=explorer)
         else:
             enhanced_tool = build_tool(
-                base_func,
-                config,
-                workspace_mode,
-                workspace,
-                get_api_key_from_authorization_header,
-                base_url,
-                object_store,
+                base_func=base_func,
+                config=config,
+                workspace_mode=workspace_mode,
+                workspace=workspace,
+                use_request_context=get_api_key_from_authorization_header,
+                base_url=base_url,
+                object_store=object_store,
+                api_key=api_key,
             )
 
         mcp_server_instance.add_tool(enhanced_tool, name=tool_name)
diff --git a/src/deepset_mcp/tools/tokonomics/decorators.py b/src/deepset_mcp/tools/tokonomics/decorators.py
@@ -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_123', threshold='@obj_456.config.threshold')",
-            "",
-            "    # Mixed call",
-            f"    {func_name}(data='@obj_123.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 = [
-        "",
-        "**Output Storage**",
+    enhancement = [
         "",
-        "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(
@@ -290,7 +269,7 @@ async def async_wrapper(*args: Any, **kwargs: Any) -> Any:
 
             # 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:
 
@@ -333,7 +312,7 @@ def sync_wrapper(*args: Any, **kwargs: Any) -> Any:
 
             # 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
@@ -375,26 +354,19 @@ def decorator(func: F) -> F:
 
         # 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
 
diff --git a/test/unit/test_tool_factory.py b/test/unit/test_tool_factory.py
@@ -549,7 +549,7 @@ async def sample_func(client: AsyncClientProtocol, workspace: str, a: int, custo
             custom_args={"custom_arg": "injected"},
         )
 
-        result = build_tool(sample_func, config, WorkspaceMode.STATIC, "test-workspace")
+        result = build_tool(sample_func, config, WorkspaceMode.STATIC, workspace="test-workspace")
 
         # Check final signature
         sig = inspect.signature(result)
@@ -578,7 +578,9 @@ async def sample_func(client: AsyncClientProtocol, workspace: str, a: int) -> st
             needs_workspace=True,
         )
 
-        result = build_tool(sample_func, config, WorkspaceMode.STATIC, "test-workspace")
+        result = build_tool(
+            sample_func, config, WorkspaceMode.STATIC, workspace="test-workspace", use_request_context=True
+        )
 
         # Mock the context and use FakeClient
         mock_ctx = MagicMock()
diff --git a/test/unit/tools/tokonomics/test_decorators.py b/test/unit/tools/tokonomics/test_decorators.py
@@ -90,26 +90,24 @@ def test_add_str_to_type_other_type(self) -> None:
     def test_enhance_docstring_for_references(self) -> None:
         """Test docstring enhancement for referenceable functions."""
         original = "Original docstring."
-        param_info = {
-            "data": {"original": "dict", "modified": "dict | str", "accepts_str": False},
-            "text": {"original": "str", "modified": "str", "accepts_str": True},
-        }
 
-        result = _enhance_docstring_for_references(original, param_info, "test_func")
+        result = _enhance_docstring_for_references(original, "test_func")
 
         assert "Original docstring." in result
-        assert "**Reference Support**" in result
         assert "All parameters accept object references" in result
-        assert "``data``: dict → dict | str" in result
-        assert "``text``: str (already accepts strings)" in result
         assert "test_func(data='@obj_123'" in result
+        # Should not contain the old reference support formatting
+        assert "**Reference Support**" not in result
+        assert "dict → dict | str" not in result
 
     def test_enhance_docstring_for_references_empty_original(self) -> None:
         """Test docstring enhancement with empty original docstring."""
-        result = _enhance_docstring_for_references("", {}, "test_func")
+        result = _enhance_docstring_for_references("", "test_func")
 
-        assert "test_func function with reference support." in result
-        assert "**Reference Support**" in result
+        assert "test_func function." in result
+        assert "All parameters accept object references" in result
+        # Should not contain the old reference support formatting
+        assert "**Reference Support**" not in result
 
     def test_enhance_docstring_for_explorable(self) -> None:
         """Test docstring enhancement for explorable functions."""
@@ -118,16 +116,69 @@ def test_enhance_docstring_for_explorable(self) -> None:
         result = _enhance_docstring_for_explorable(original, "test_func")
 
         assert "Original docstring." in result
-        assert "**Output Storage**" in result
         assert "automatically stored and can be referenced" in result
         assert "object ID (e.g., ``@obj_123``)" in result
+        # Should not contain the old output storage formatting
+        assert "**Output Storage**" not in result
 
     def test_enhance_docstring_for_explorable_empty_original(self) -> None:
         """Test docstring enhancement with empty original docstring."""
         result = _enhance_docstring_for_explorable("", "test_func")
 
-        assert "test_func function with stored output." in result
-        assert "**Output Storage**" in result
+        assert "test_func function." in result
+        assert "automatically stored and can be referenced" in result
+        # Should not contain the old output storage formatting
+        assert "**Output Storage**" not in result
+
+    def test_enhance_docstring_for_references_preserves_original(self) -> None:
+        """Test that _enhance_docstring_for_references preserves all original content."""
+        original = """Process data and return results.
+
+        This function processes the input data and returns processed results.
+        It may raise exceptions if the data is invalid.
+
+        :param data: Input data to process
+        :param options: Processing options
+        :return: Processed results
+        :raises ValueError: If data is invalid
+        :raises RuntimeError: If processing fails
+        """
+
+        result = _enhance_docstring_for_references(original, "test_func")
+
+        # Check that all original content is preserved exactly
+        assert "Process data and return results." in result
+        assert "This function processes the input data" in result
+        assert ":param data: Input data to process" in result
+        assert ":param options: Processing options" in result
+        assert ":return: Processed results" in result
+        assert ":raises ValueError: If data is invalid" in result
+        assert ":raises RuntimeError: If processing fails" in result
+
+        # Check that enhancement is added
+        assert "All parameters accept object references" in result
+        assert "test_func(data='@obj_123'" in result
+
+    def test_enhance_docstring_for_explorable_preserves_original(self) -> None:
+        """Test that _enhance_docstring_for_explorable preserves all original content."""
+        original = """Calculate statistics from data.
+
+        :param values: List of numbers
+        :return: Statistical summary dict
+        :raises ValueError: If values is empty
+        """
+
+        result = _enhance_docstring_for_explorable(original, "calc_stats")
+
+        # Check that all original content is preserved exactly
+        assert "Calculate statistics from data." in result
+        assert ":param values: List of numbers" in result
+        assert ":return: Statistical summary dict" in result
+        assert ":raises ValueError: If values is empty" in result
+
+        # Check that enhancement is added
+        assert "automatically stored and can be referenced" in result
+        assert "object ID (e.g., ``@obj_123``)" in result
 
 
 class TestExplorableDecorator:
@@ -205,8 +256,9 @@ def test_func() -> dict:
             return {}
 
         assert "Original docstring." in test_func.__doc__
-        assert "**Output Storage**" in test_func.__doc__
         assert "automatically stored" in test_func.__doc__
+        # Should not contain the old output storage formatting
+        assert "**Output Storage**" not in test_func.__doc__
 
     def test_explorable_with_args(self, store: ObjectStore, explorer: RichExplorer) -> None:
         """Test @explorable decorated function with arguments."""
@@ -390,8 +442,9 @@ def test_func(data: dict) -> str:
             return "result"
 
         assert "Original docstring." in test_func.__doc__
-        assert "**Reference Support**" in test_func.__doc__
         assert "All parameters accept object references" in test_func.__doc__
+        # Should not contain the old reference support formatting
+        assert "**Reference Support**" not in test_func.__doc__
 
 
 class TestExplorableAndReferenceableDecorator:
@@ -465,8 +518,11 @@ def test_func(data: dict) -> dict:
 
         # Should contain both reference and explorable documentation
         assert "Original docstring." in docstring
-        assert "**Reference Support**" in docstring
-        assert "**Output Storage**" in docstring
+        assert "All parameters accept object references" in docstring
+        assert "automatically stored" in docstring
+        # Should not contain the old section formatting
+        assert "**Reference Support**" not in docstring
+        assert "**Output Storage**" not in docstring
 
     def test_combined_decorator_signature(self, store: ObjectStore, explorer: RichExplorer) -> None:
         """Test that combined decorator modifies signature correctly."""
diff --git a/test/unit/tools/tokonomics/test_integration.py b/test/unit/tools/tokonomics/test_integration.py
@@ -438,14 +438,16 @@ def documented_referenceable(stats: dict, precision: int = 2) -> str:
         # Check enhanced docstrings
         explorable_doc = documented_explorable.__doc__
         assert "Process a list of items" in explorable_doc
-        assert "**Output Storage**" in explorable_doc
         assert "automatically stored" in explorable_doc
+        # Should not contain the old output storage formatting
+        assert "**Output Storage**" not in explorable_doc
 
         referenceable_doc = documented_referenceable.__doc__
         assert "Format statistics" in referenceable_doc
-        assert "**Reference Support**" in referenceable_doc
         assert "All parameters accept object references" in referenceable_doc
-        assert "``stats``: <class 'dict'> → dict | str" in referenceable_doc
+        # Should not contain the old reference support formatting
+        assert "**Reference Support**" not in referenceable_doc
+        assert "dict | str" not in referenceable_doc
 
         # Test functionality with enhanced docs
         result1 = documented_explorable([10, 20, 30, 40, 50])
