Refine custom endpoint API SSE streaming and configurable A2A RPC endpoints (#216)

Author: pitt-liang

src/agentscope_runtime/engine/app/agent_app.py

@@ -1,12 +1,9 @@
 # -*- coding: utf-8 -*-
-import asyncio
 import logging
 import types
-from contextlib import asynccontextmanager
-from typing import Optional, Any, Callable, List
+from typing import Optional, Callable, List
 
 import uvicorn
-from fastapi import FastAPI
 from pydantic import BaseModel
 
 from .base_app import BaseApp
@@ -43,6 +40,7 @@ def __init__(
         broker_url: Optional[str] = None,
         backend_url: Optional[str] = None,
         runner: Optional[Runner] = None,
+        enable_embedded_worker: bool = False,
         **kwargs,
     ):
         """
@@ -61,6 +59,7 @@ def __init__(
         self.after_finish = after_finish
         self.broker_url = broker_url
         self.backend_url = backend_url
+        self.enable_embedded_worker = enable_embedded_worker
 
         self._runner = runner
         self.custom_endpoints = []  # Store custom endpoints
@@ -79,33 +78,16 @@ def __init__(
         response_protocol = ResponseAPIDefaultAdapter()
         self.protocol_adapters = [a2a_protocol, response_protocol]
 
-        @asynccontextmanager
-        async def lifespan(app: FastAPI) -> Any:
-            """Manage the application lifespan."""
-            if hasattr(self, "before_start") and self.before_start:
-                if asyncio.iscoroutinefunction(self.before_start):
-                    await self.before_start(app, **getattr(self, "kwargs", {}))
-                else:
-                    self.before_start(app, **getattr(self, "kwargs", {}))
-            yield
-            if hasattr(self, "after_finish") and self.after_finish:
-                if asyncio.iscoroutinefunction(self.after_finish):
-                    await self.after_finish(app, **getattr(self, "kwargs", {}))
-                else:
-                    self.after_finish(app, **getattr(self, "kwargs", {}))
-
-        kwargs = {
+        self._app_kwargs = {
             "title": "Agent Service",
             "version": __version__,
             "description": "Production-ready Agent Service API",
-            "lifespan": lifespan,
             **kwargs,
         }
 
         super().__init__(
             broker_url=broker_url,
             backend_url=backend_url,
-            **kwargs,
         )
 
         # Store custom endpoints and tasks for deployment
@@ -167,7 +149,6 @@ def run(
         self,
         host="0.0.0.0",
         port=8090,
-        embed_task_processor=False,
         **kwargs,
     ):
         """
@@ -176,7 +157,6 @@ def run(
         Args:
             host: Host to bind to
             port: Port to bind to
-            embed_task_processor: Whether to embed task processor
             **kwargs: Additional keyword arguments
         """
         # Build runner
@@ -186,24 +166,7 @@ def run(
             logger.info(
                 "[AgentApp] Starting AgentApp with FastAPIAppFactory...",
             )
-
-            # Create FastAPI application using the factory
-            fastapi_app = FastAPIAppFactory.create_app(
-                runner=self._runner,
-                endpoint_path=self.endpoint_path,
-                request_model=self.request_model,
-                response_type=self.response_type,
-                stream=self.stream,
-                before_start=self.before_start,
-                after_finish=self.after_finish,
-                mode=DeploymentMode.DAEMON_THREAD,
-                protocol_adapters=self.protocol_adapters,
-                custom_endpoints=self.custom_endpoints,
-                broker_url=self.broker_url,
-                backend_url=self.backend_url,
-                enable_embedded_worker=embed_task_processor,
-                **kwargs,
-            )
+            fastapi_app = self.get_fastapi_app(**kwargs)
 
             logger.info(f"[AgentApp] Starting server on {host}:{port}")
 
@@ -220,6 +183,29 @@ def run(
             logger.error(f"[AgentApp] Error while running: {e}")
             raise
 
+    def get_fastapi_app(self, **kwargs):
+        """Get the FastAPI application"""
+
+        self._build_runner()
+
+        return FastAPIAppFactory.create_app(
+            runner=self._runner,
+            endpoint_path=self.endpoint_path,
+            request_model=self.request_model,
+            response_type=self.response_type,
+            stream=self.stream,
+            before_start=self.before_start,
+            after_finish=self.after_finish,
+            mode=DeploymentMode.DAEMON_THREAD,
+            protocol_adapters=self.protocol_adapters,
+            custom_endpoints=self.custom_endpoints,
+            broker_url=self.broker_url,
+            backend_url=self.backend_url,
+            enable_embedded_worker=self.enable_embedded_worker,
+            app_kwargs=self._app_kwargs,
+            **kwargs,
+        )
+
     async def deploy(self, deployer: DeployManager, **kwargs):
         """Deploy the agent app with custom endpoints support"""
         # Pass custom endpoints and tasks to the deployer

src/agentscope_runtime/engine/app/base_app.py

@@ -1,38 +1,29 @@
 # -*- coding: utf-8 -*-
 import inspect
 import logging
-import threading
 from typing import Callable, Optional
 
-import uvicorn
-from fastapi import FastAPI, Request
-from fastapi.responses import StreamingResponse
+from fastapi import Request
 
 from .celery_mixin import CeleryMixin
 
 logger = logging.getLogger(__name__)
 
 
-class BaseApp(FastAPI, CeleryMixin):
+class BaseApp(CeleryMixin):
     """
-    BaseApp extends FastAPI and integrates with Celery
-    for asynchronous background task execution.
+    BaseApp integrates Celery for asynchronous background task execution,
+    and provides FastAPI-like routing for task endpoints.
     """
 
     def __init__(
         self,
         broker_url: Optional[str] = None,
         backend_url: Optional[str] = None,
-        **kwargs,
     ):
         # Initialize CeleryMixin
         CeleryMixin.__init__(self, broker_url, backend_url)
 
-        self.server = None
-
-        # Initialize FastAPI
-        FastAPI.__init__(self, **kwargs)
-
     def task(self, path: str, queue: str = "celery"):
         """
         Register an asynchronous task endpoint.
@@ -74,108 +65,3 @@ async def get_task(task_id: str):
             return func
 
         return decorator
-
-    def endpoint(self, path: str):
-        """
-        Unified POST endpoint decorator.
-        Pure FastAPI functionality, independent of Celery.
-        Supports:
-          - Sync functions
-          - Async functions (coroutines)
-          - Sync/async generator functions (streaming responses)
-        """
-
-        def decorator(func: Callable):
-            is_async_gen = inspect.isasyncgenfunction(func)
-            is_sync_gen = inspect.isgeneratorfunction(func)
-
-            if is_async_gen or is_sync_gen:
-                # Handle streaming responses
-                async def _stream_generator(request: Request):
-                    if is_async_gen:
-                        async for chunk in func(request):
-                            yield chunk
-                    else:
-                        for chunk in func(request):
-                            yield chunk
-
-                @self.post(path)
-                async def _wrapped(request: Request):
-                    return StreamingResponse(
-                        _stream_generator(request),
-                        media_type="text/plain",
-                    )
-
-            else:
-                # Handle regular responses
-                @self.post(path)
-                async def _wrapped(request: Request):
-                    if inspect.iscoroutinefunction(func):
-                        return await func(request)
-                    else:
-                        return func(request)
-
-            return func
-
-        return decorator
-
-    def run(
-        self,
-        host="0.0.0.0",
-        port=8090,
-        embed_task_processor=False,
-        **kwargs,
-    ):
-        """
-        Run FastAPI with uvicorn.
-        """
-        if embed_task_processor:
-            if self.celery_app is None:
-                logger.warning(
-                    "[AgentApp] Celery is not configured. "
-                    "Cannot run embedded worker.",
-                )
-            else:
-                logger.warning(
-                    "[AgentApp] embed_task_processor=True: Running "
-                    "task_processor in embedded thread mode. This is "
-                    "intended for development/debug purposes only. In "
-                    "production, run Celery worker in a separate process!",
-                )
-
-                queues = self._registered_queues or {"celery"}
-                queue_list = ",".join(sorted(queues))
-
-                def start_celery_worker():
-                    logger.info(
-                        f"[AgentApp] Embedded worker listening "
-                        f"queues: {queue_list}",
-                    )
-                    self.celery_app.worker_main(
-                        [
-                            "worker",
-                            "--loglevel=INFO",
-                            "-Q",
-                            queue_list,
-                        ],
-                    )
-
-                threading.Thread(
-                    target=start_celery_worker,
-                    daemon=True,
-                ).start()
-                logger.info(
-                    "[AgentApp] Embedded task processor started in background "
-                    "thread (DEV mode).",
-                )
-
-        # TODO: Add CLI to main entrypoint to control run/deploy
-
-        config = uvicorn.Config(
-            app=self,
-            host=host,
-            port=port,
-            **kwargs,
-        )
-        self.server = uvicorn.Server(config)
-        self.server.run()

src/agentscope_runtime/engine/deployers/adapter/__init__.py

@@ -0,0 +1,2 @@
+# -*- coding: utf-8 -*-
+from .a2a import A2AFastAPIDefaultAdapter

src/agentscope_runtime/engine/deployers/adapter/a2a/a2a_protocol_adapter.py

@@ -1,4 +1,5 @@
 # -*- coding: utf-8 -*-
+import posixpath
 from typing import Callable
 
 from a2a.server.apps import A2AFastAPIApplication
@@ -9,12 +10,16 @@
 from .a2a_agent_adapter import A2AExecutor
 from ..protocol_adapter import ProtocolAdapter
 
+A2A_JSON_RPC_URL = "/a2a"
+
 
 class A2AFastAPIDefaultAdapter(ProtocolAdapter):
     def __init__(self, agent_name, agent_description, **kwargs):
         super().__init__(**kwargs)
         self._agent_name = agent_name
         self._agent_description = agent_description
+        self._json_rpc_path = kwargs.get("json_rpc_path", A2A_JSON_RPC_URL)
+        self._base_url = kwargs.get("base_url")
 
     def add_endpoint(self, app, func: Callable, **kwargs):
         request_handler = DefaultRequestHandler(
@@ -32,7 +37,14 @@ def add_endpoint(self, app, func: Callable, **kwargs):
             http_handler=request_handler,
         )
 
-        server.add_routes_to_app(app)
+        server.add_routes_to_app(app, rpc_url=self._json_rpc_path)
+
+    def _get_json_rpc_url(self) -> str:
+        base = self._base_url or "http://127.0.0.1:8000"
+        return posixpath.join(
+            base.rstrip("/"),
+            self._json_rpc_path.lstrip("/"),
+        )
 
     def get_agent_card(
         self,
@@ -62,6 +74,6 @@ def get_agent_card(
             description=agent_description,
             default_input_modes=["text"],
             default_output_modes=["text"],
-            url="http://127.0.0.1:8090/",
+            url=self._get_json_rpc_url(),
             version="1.0.0",
         )

src/agentscope_runtime/engine/deployers/adapter/responses/__init__.py

@@ -0,0 +1,2 @@
+# -*- coding: utf-8 -*-
+from .response_api_protocol_adapter import ResponseAPIDefaultAdapter