Fix request model creation and MCP execution for pipeline wrappers with only run_api_async implemented (#125)

* remove unused imports

* Fix requset_model creation for pipelines with only run_api_async implemented

* run correctly pipeline as MCP tools when they have only run_api_async implemented

Author: mpangrazzi

src/hayhooks/cli/mcp.py

@@ -1,4 +1,3 @@
-import asyncio
 import typer
 import uvicorn
 import sys

src/hayhooks/server/utils/deploy_utils.py

@@ -423,10 +423,28 @@ def add_pipeline_to_registry(
     clog.debug("Running setup()")
     pipeline_wrapper.setup()
 
-    docstring = docstring_parser.parse(inspect.getdoc(pipeline_wrapper.run_api) or "")
+    # Determine which run_api method to use for creating request model (prefer async if available)
+    if pipeline_wrapper._is_run_api_async_implemented:
+        run_method_to_inspect = pipeline_wrapper.run_api_async
+        clog.debug("Using `run_api_async` for metadata creation.")
+    elif pipeline_wrapper._is_run_api_implemented:
+        run_method_to_inspect = pipeline_wrapper.run_api
+        clog.debug("Using `run_api` for metadata creation.")
+    else:
+        # If neither run_api nor run_api_async is implemented, skip creating request model
+        run_method_to_inspect = None
+        clog.debug("No run_api method implemented, skipping request model creation.")
+
+    if run_method_to_inspect:
+        docstring = docstring_parser.parse(inspect.getdoc(run_method_to_inspect) or "")
+        request_model = create_request_model_from_callable(run_method_to_inspect, f'{pipeline_name}Run', docstring)
+    else:
+        docstring = docstring_parser.Docstring()
+        request_model = None
+
     metadata = {
         "description": docstring.short_description or "",
-        "request_model": create_request_model_from_callable(pipeline_wrapper.run_api, f'{pipeline_name}Run', docstring),
+        "request_model": request_model,
         "skip_mcp": pipeline_wrapper.skip_mcp,
     }
 

src/hayhooks/server/utils/mcp_utils.py

@@ -1,4 +1,5 @@
 import asyncio
+import traceback
 from enum import Enum
 from pathlib import Path
 from typing import TYPE_CHECKING, List, Union
@@ -13,8 +14,10 @@
 )
 from hayhooks.settings import settings
 from hayhooks.server.pipelines import registry
+from hayhooks.server.pipelines.registry import PipelineType
 from haystack.lazy_imports import LazyImport
 from hayhooks.server.routers.deploy import PipelineFilesRequest
+from fastapi.concurrency import run_in_threadpool
 
 
 with LazyImport("Run 'pip install \"mcp\"' to install MCP.") as mcp_import:
@@ -133,12 +136,21 @@ async def run_pipeline_as_tool(
     mcp_import.check()
 
     log.debug(f"Calling pipeline as tool '{name}' with arguments: {arguments}")
-    pipeline_wrapper: Union[BasePipelineWrapper, None] = registry.get(name)
+    pipeline: Union[PipelineType, None] = registry.get(name)
 
-    if not pipeline_wrapper:
+    if not pipeline:
         raise ValueError(f"Pipeline '{name}' not found")
 
-    result = await asyncio.to_thread(pipeline_wrapper.run_api, **arguments)
+    # Only BasePipelineWrapper instances support run_api/run_api_async methods
+    if not isinstance(pipeline, BasePipelineWrapper):
+        raise ValueError(f"Pipeline '{name}' is not a BasePipelineWrapper and cannot be used as an MCP tool")
+
+    # Use the same async/sync pattern as in deploy_utils.py
+    if pipeline._is_run_api_async_implemented:
+        result = await pipeline.run_api_async(**arguments)
+    else:
+        result = await run_in_threadpool(pipeline.run_api, **arguments)
+
     log.trace(f"Pipeline '{name}' returned result: {result}")
 
     return [TextContent(text=result, type="text")]
@@ -206,6 +218,8 @@ async def call_tool(name: str, arguments: dict) -> List[TextContent | ImageConte
 
         except Exception as exc:
             msg = f"General unhandled error in call_tool for tool '{name}': {exc}"
+            if settings.show_tracebacks:
+                msg += f"\n{traceback.format_exc()}"
             log.error(msg)
             raise Exception(msg) from exc
 

tests/test_deploy_utils.py

@@ -15,6 +15,7 @@
     create_pipeline_wrapper_instance,
     deploy_pipeline_files,
     undeploy_pipeline,
+    add_pipeline_to_registry,
 )
 from hayhooks.server.exceptions import (
     PipelineFilesError,
@@ -378,3 +379,29 @@ def test_undeploy_pipeline_without_app(test_settings):
 
     # 4. Assert pipeline files are deleted
     assert not pipeline_dir.exists()
+
+
+def test_add_pipeline_to_registry_with_async_run_api():
+    pipeline_name = "async_question_answer"
+    pipeline_wrapper_path = Path("tests/test_files/files/async_question_answer/pipeline_wrapper.py")
+    pipeline_yml_path = Path("tests/test_files/files/async_question_answer/question_answer.yml")
+    files = {
+        "pipeline_wrapper.py": pipeline_wrapper_path.read_text(),
+        "question_answer.yml": pipeline_yml_path.read_text(),
+    }
+
+    pipeline_wrapper = add_pipeline_to_registry(pipeline_name=pipeline_name, files=files, save_files=False)
+    assert registry.get(pipeline_name) == pipeline_wrapper
+
+    metadata = registry.get_metadata(pipeline_name)
+    assert metadata is not None
+    assert "request_model" in metadata
+    assert metadata["request_model"] is not None
+
+    assert pipeline_wrapper._is_run_api_async_implemented is True
+    assert pipeline_wrapper._is_run_api_implemented is False
+
+    request_model = metadata["request_model"]
+    schema = request_model.model_json_schema()
+    assert "question" in schema["properties"]
+    assert schema["properties"]["question"]["type"] == "string"

tests/test_files/files/async_question_answer/pipeline_wrapper.py

@@ -21,32 +21,15 @@ def setup(self) -> None:
     async def run_api_async(self, question: str) -> str:
         log.trace(f"Running pipeline with question: {question}")
 
-        result = await self.pipeline.run_async(
-            {
-                "prompt_builder": {
-                    "template": [
-                        ChatMessage.from_system(SYSTEM_MESSAGE),
-                        ChatMessage.from_user(question),
-                    ]
-                }
-            }
-        )
-        return result["llm"]["replies"][0].text
+        return "This is a mock response from the pipeline"
 
     async def run_chat_completion_async(self, model: str, messages: List[dict], body: dict) -> AsyncGenerator:
         log.trace(f"Running pipeline with model: {model}, messages: {messages}, body: {body}")
 
-        question = get_last_user_message(messages)
-        log.trace(f"Question: {question}")
-
-        return async_streaming_generator(
-            pipeline=self.pipeline,
-            pipeline_run_args={
-                "prompt_builder": {
-                    "template": [
-                        ChatMessage.from_system(SYSTEM_MESSAGE),
-                        ChatMessage.from_user(question),
-                    ]
-                },
-            },
-        )
+        mock_response = "This is a mock response from the pipeline"
+
+        async def mock_generator():
+            for word in mock_response.split():
+                yield word + " "
+
+        return mock_generator()

tests/test_it_mcp_server.py

@@ -43,6 +43,19 @@ def deploy_chat_with_website_mcp_pipeline():
     return pipeline_name
 
 
+@pytest.fixture
+def deploy_async_question_answer_mcp_pipeline():
+    pipeline_name = "async_question_answer"
+    pipeline_wrapper_path = Path("tests/test_files/files/async_question_answer/pipeline_wrapper.py")
+    pipeline_yml_path = Path("tests/test_files/files/async_question_answer/question_answer.yml")
+    files = {
+        "pipeline_wrapper.py": pipeline_wrapper_path.read_text(),
+        "question_answer.yml": pipeline_yml_path.read_text(),
+    }
+    add_pipeline_to_registry(pipeline_name=pipeline_name, files=files)
+    return pipeline_name
+
+
 @pytest.fixture
 def mcp_server_instance() -> "Server":
     return create_mcp_server()
@@ -105,6 +118,19 @@ async def test_call_pipeline_as_tool(mcp_server_instance, deploy_chat_with_websi
         assert result.content == [TextContent(type="text", text="This is a mock response from the pipeline")]
 
 
+@pytest.mark.asyncio
+async def test_call_async_pipeline_as_tool(mcp_server_instance, deploy_async_question_answer_mcp_pipeline):
+    async with client_session(mcp_server_instance) as client:
+        result = await client.call_tool(
+            deploy_async_question_answer_mcp_pipeline, {"question": "What is the capital of France?"}
+        )
+
+        assert isinstance(result, CallToolResult)
+
+        # In the deployed pipeline, the response is mocked
+        assert result.content == [TextContent(type="text", text="This is a mock response from the pipeline")]
+
+
 @pytest.mark.asyncio
 async def test_call_pipeline_as_tool_with_invalid_arguments(mcp_server_instance, deploy_chat_with_website_mcp_pipeline):
     async with client_session(mcp_server_instance) as client: