agentscope-runtime 0.1.5b1__py3-none-any.whl → 0.2.0b1__py3-none-any.whl

agentscope_runtime/engine/deployers/utils/service_utils/fastapi_factory.py

@@ -1,15 +1,17 @@
 # -*- coding: utf-8 -*-
-# pylint:disable=too-many-branches, unused-argument
+# pylint:disable=too-many-branches, unused-argument, too-many-return-statements
 
 
 import asyncio
+import inspect
 import json
 from contextlib import asynccontextmanager
-from typing import Optional, Callable, Type, Any
+from typing import Optional, Callable, Type, Any, List, Dict
 
 from fastapi import FastAPI, Request
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.responses import StreamingResponse, JSONResponse
+from pydantic import BaseModel
 
 from .service_config import ServicesConfig, DEFAULT_SERVICES_CONFIG
 from .service_factory import ServiceFactory
@@ -17,6 +19,13 @@ from ..deployment_modes import DeploymentMode
 from ...adapter.protocol_adapter import ProtocolAdapter
 
 
+async def error_stream(e):
+    yield (
+        f"data: "
+        f"{json.dumps({'error': f'Request parsing error: {str(e)}'})}\n\n"
+    )
+
+
 class FastAPIAppFactory:
     """Factory for creating FastAPI applications with unified architecture."""
 
@@ -33,6 +42,13 @@ class FastAPIAppFactory:
         mode: DeploymentMode = DeploymentMode.DAEMON_THREAD,
         services_config: Optional[ServicesConfig] = None,
         protocol_adapters: Optional[list[ProtocolAdapter]] = None,
+        custom_endpoints: Optional[
+            List[Dict]
+        ] = None,  # New parameter for custom endpoints
+        # Celery parameters
+        broker_url: Optional[str] = None,
+        backend_url: Optional[str] = None,
+        enable_embedded_worker: bool = False,
         **kwargs: Any,
     ) -> FastAPI:
         """Create a FastAPI application with unified architecture.
@@ -49,6 +65,10 @@ class FastAPIAppFactory:
             mode: Deployment mode
             services_config: Services configuration
             protocol_adapters: Protocol adapters
+            custom_endpoints: List of custom endpoint configurations
+            broker_url: Celery broker URL
+            backend_url: Celery backend URL
+            enable_embedded_worker: Whether to run embedded Celery worker
             **kwargs: Additional keyword arguments
 
         Returns:
@@ -58,6 +78,20 @@ class FastAPIAppFactory:
         if services_config is None:
             services_config = DEFAULT_SERVICES_CONFIG
 
+        # Initialize Celery mixin if broker and backend URLs are provided
+        celery_mixin = None
+        if broker_url and backend_url:
+            try:
+                from ....app.celery_mixin import CeleryMixin
+
+                celery_mixin = CeleryMixin(
+                    broker_url=broker_url,
+                    backend_url=backend_url,
+                )
+            except ImportError:
+                # CeleryMixin not available, will use fallback task processing
+                celery_mixin = None
+
         # Create lifespan manager
         @asynccontextmanager
         async def lifespan(app: FastAPI):
@@ -93,6 +127,15 @@ class FastAPIAppFactory:
         app.state.external_runner = runner
         app.state.endpoint_path = endpoint_path
         app.state.protocol_adapters = protocol_adapters  # Store for later use
+        app.state.custom_endpoints = (
+            custom_endpoints or []
+        )  # Store custom endpoints
+
+        # Store Celery configuration
+        app.state.celery_mixin = celery_mixin
+        app.state.broker_url = broker_url
+        app.state.backend_url = backend_url
+        app.state.enable_embedded_worker = enable_embedded_worker
 
         # Add middleware
         FastAPIAppFactory._add_middleware(app, mode)
@@ -168,6 +211,49 @@ class FastAPIAppFactory:
                 for protocol_adapter in app.state.protocol_adapters:
                     protocol_adapter.add_endpoint(app=app, func=effective_func)
 
+        # Add custom endpoints after runner is available
+        if (
+            hasattr(app.state, "custom_endpoints")
+            and app.state.custom_endpoints
+        ):
+            FastAPIAppFactory._add_custom_endpoints(app)
+
+        # Start embedded Celery worker if enabled
+        if (
+            hasattr(app.state, "enable_embedded_worker")
+            and app.state.enable_embedded_worker
+            and hasattr(app.state, "celery_mixin")
+            and app.state.celery_mixin
+        ):
+            # Start Celery worker in background thread
+            import threading
+
+            def start_celery_worker():
+                try:
+                    celery_mixin = app.state.celery_mixin
+                    # Get registered queues or use default
+                    queues = (
+                        list(celery_mixin.get_registered_queues())
+                        if celery_mixin.get_registered_queues()
+                        else ["celery"]
+                    )
+                    celery_mixin.run_task_processor(
+                        loglevel="INFO",
+                        concurrency=1,
+                        queues=queues,
+                    )
+                except Exception as e:
+                    import logging
+
+                    logger = logging.getLogger(__name__)
+                    logger.error(f"Failed to start Celery worker: {e}")
+
+            worker_thread = threading.Thread(
+                target=start_celery_worker,
+                daemon=True,
+            )
+            worker_thread.start()
+
     @staticmethod
     async def _handle_shutdown(
         app: FastAPI,
@@ -188,9 +274,10 @@ class FastAPIAppFactory:
             and not app.state.runner_managed_externally
         ):
             runner = app.state.runner
-            if runner and hasattr(runner, "context_manager"):
+            if runner:
                 try:
-                    await runner.context_manager.__aexit__(None, None, None)
+                    # Clean up runner
+                    await runner.__aexit__(None, None, None)
                 except Exception as e:
                     print(f"Warning: Error during runner cleanup: {e}")
 
@@ -211,15 +298,15 @@ class FastAPIAppFactory:
             memory_service=services["memory"],
         )
 
-        # Initialize context manager
-        await context_manager.__aenter__()
-
         # Create runner (agent will be set later)
         runner = Runner(
             agent=None,  # Will be set by the specific deployment
             context_manager=context_manager,
         )
 
+        # Initialize runner
+        await runner.__aenter__()
+
         return runner
 
     @staticmethod
@@ -502,3 +589,476 @@ class FastAPIAppFactory:
         if hasattr(app.state, "runner"):
             return app.state.runner
         return None
+
+    @staticmethod
+    def _create_parameter_wrapper(handler: Callable):
+        """Create a wrapper that handles parameter parsing based on function
+        signature.
+
+        This method inspects the handler function's parameters and creates
+        appropriate wrappers to parse request data into the expected
+        parameter types.
+        """
+        try:
+            sig = inspect.signature(handler)
+            params = list(sig.parameters.values())
+
+            if not params:
+                # No parameters, call function directly
+                return handler
+
+            # Get the first parameter (assuming single parameter for now)
+            first_param = params[0]
+            param_annotation = first_param.annotation
+
+            # If no annotation or annotation is Request, pass Request directly
+            if param_annotation in [inspect.Parameter.empty, Request]:
+                return handler
+
+            # Check if the annotation is a Pydantic model
+            if isinstance(param_annotation, type) and issubclass(
+                param_annotation,
+                BaseModel,
+            ):
+                # Create wrapper that parses JSON to Pydantic model
+                if inspect.iscoroutinefunction(handler):
+
+                    async def async_pydantic_wrapper(request: Request):
+                        try:
+                            body = await request.json()
+                            parsed_param = param_annotation(**body)
+                            return await handler(parsed_param)
+                        except Exception as e:
+                            return JSONResponse(
+                                status_code=422,
+                                content={
+                                    "detail": f"Request parsing error: "
+                                    f"{str(e)}",
+                                },
+                            )
+
+                    return async_pydantic_wrapper
+                else:
+
+                    async def sync_pydantic_wrapper(request: Request):
+                        try:
+                            body = await request.json()
+                            parsed_param = param_annotation(**body)
+                            return handler(parsed_param)
+                        except Exception as e:
+                            return JSONResponse(
+                                status_code=422,
+                                content={
+                                    "detail": f"Request parsing error: "
+                                    f"{str(e)}",
+                                },
+                            )
+
+                    return sync_pydantic_wrapper
+
+            # For other types, fall back to original behavior
+            return handler
+
+        except Exception:
+            # If anything goes wrong with introspection, fall back to
+            # original behavior
+            return handler
+
+    @staticmethod
+    def _create_streaming_parameter_wrapper(
+        handler: Callable,
+        is_async_gen: bool = False,
+    ):
+        """Create a wrapper for streaming handlers that handles parameter
+        parsing."""
+        try:
+            sig = inspect.signature(handler)
+            params = list(sig.parameters.values())
+            no_params = False
+            param_annotation = None
+
+            if not params:
+                no_params = True
+            else:
+                # Get the first parameter
+                first_param = params[0]
+                param_annotation = first_param.annotation
+
+                # If no annotation or annotation is Request, goto no params
+                # logic
+                if param_annotation in [inspect.Parameter.empty, Request]:
+                    no_params = True
+
+            if no_params:
+                if is_async_gen:
+
+                    async def async_no_param_wrapper():
+                        async def generate():
+                            async for chunk in handler():
+                                yield str(chunk)
+
+                        return StreamingResponse(
+                            generate(),
+                            media_type="text/plain",
+                        )
+
+                    return async_no_param_wrapper
+                else:
+
+                    async def sync_no_param_wrapper():
+                        def generate():
+                            for chunk in handler():
+                                yield str(chunk)
+
+                        return StreamingResponse(
+                            generate(),
+                            media_type="text/plain",
+                        )
+
+                    return sync_no_param_wrapper
+
+            # Check if the annotation is a Pydantic model
+            if isinstance(param_annotation, type) and issubclass(
+                param_annotation,
+                BaseModel,
+            ):
+                if is_async_gen:
+
+                    async def async_stream_pydantic_wrapper(
+                        request: Request,
+                    ):
+                        try:
+                            body = await request.json()
+                            parsed_param = param_annotation(**body)
+
+                            async def generate():
+                                async for chunk in handler(parsed_param):
+                                    yield str(chunk)
+
+                            return StreamingResponse(
+                                generate(),
+                                media_type="text/plain",
+                            )
+                        except Exception as e:
+                            return StreamingResponse(
+                                error_stream(e),
+                                media_type="text/event-stream",
+                            )
+
+                    return async_stream_pydantic_wrapper
+                else:
+
+                    async def sync_stream_pydantic_wrapper(
+                        request: Request,
+                    ):
+                        try:
+                            body = await request.json()
+                            parsed_param = param_annotation(**body)
+
+                            def generate():
+                                for chunk in handler(parsed_param):
+                                    yield str(chunk)
+
+                            return StreamingResponse(
+                                generate(),
+                                media_type="text/plain",
+                            )
+                        except Exception as e:
+                            return JSONResponse(
+                                status_code=422,
+                                content={
+                                    "detail": f"Request parsing error:"
+                                    f" {str(e)}",
+                                },
+                            )
+
+                    return sync_stream_pydantic_wrapper
+
+            return handler
+
+        except Exception:
+            return handler
+
+    @staticmethod
+    def _add_custom_endpoints(app: FastAPI):
+        """Add all custom endpoints to the FastAPI application."""
+        if (
+            not hasattr(app.state, "custom_endpoints")
+            or not app.state.custom_endpoints
+        ):
+            return
+
+        for endpoint in app.state.custom_endpoints:
+            FastAPIAppFactory._register_single_custom_endpoint(
+                app,
+                endpoint["path"],
+                endpoint["handler"],
+                endpoint["methods"],
+                endpoint,  # Pass the full endpoint config
+            )
+
+    @staticmethod
+    def _register_single_custom_endpoint(
+        app: FastAPI,
+        path: str,
+        handler: Callable,
+        methods: List[str],
+        endpoint_config: Dict = None,
+    ):
+        """Register a single custom endpoint with proper async/sync
+        handling."""
+
+        for method in methods:
+            # Check if this is a task endpoint
+            if endpoint_config and endpoint_config.get("task_type"):
+                # Create task endpoint with proper execution logic
+                task_handler = FastAPIAppFactory._create_task_handler(
+                    app,
+                    handler,
+                    endpoint_config.get("queue", "default"),
+                )
+                app.add_api_route(path, task_handler, methods=[method])
+
+                # Add task status endpoint - align with BaseApp pattern
+                status_path = f"{path}/{{task_id}}"
+                status_handler = FastAPIAppFactory._create_task_status_handler(
+                    app,
+                )
+                app.add_api_route(
+                    status_path,
+                    status_handler,
+                    methods=["GET"],
+                )
+
+            else:
+                # Regular endpoint handling with automatic parameter parsing
+                # Check in the correct order: async gen > sync gen > async &
+                # sync
+                if inspect.isasyncgenfunction(handler):
+                    # Async generator -> Streaming response with parameter
+                    # parsing
+                    wrapped_handler = (
+                        FastAPIAppFactory._create_streaming_parameter_wrapper(
+                            handler,
+                            is_async_gen=True,
+                        )
+                    )
+
+                    app.add_api_route(
+                        path,
+                        wrapped_handler,
+                        methods=[method],
+                    )
+                elif inspect.isgeneratorfunction(handler):
+                    # Sync generator -> Streaming response with parameter
+                    # parsing
+                    wrapped_handler = (
+                        FastAPIAppFactory._create_streaming_parameter_wrapper(
+                            handler,
+                            is_async_gen=False,
+                        )
+                    )
+                    app.add_api_route(
+                        path,
+                        wrapped_handler,
+                        methods=[method],
+                    )
+                else:
+                    # Sync function -> Async wrapper with parameter parsing
+                    wrapped_handler = (
+                        FastAPIAppFactory._create_parameter_wrapper(handler)
+                    )
+                    app.add_api_route(path, wrapped_handler, methods=[method])
+
+    @staticmethod
+    def _create_task_handler(app: FastAPI, task_func: Callable, queue: str):
+        """Create a task handler that executes functions asynchronously."""
+
+        async def task_endpoint(request: dict):
+            try:
+                import uuid
+
+                # Generate task ID
+                task_id = str(uuid.uuid4())
+
+                # Check if Celery is available
+                if (
+                    hasattr(app.state, "celery_mixin")
+                    and app.state.celery_mixin
+                ):
+                    # Use Celery for task processing
+                    celery_mixin = app.state.celery_mixin
+
+                    # Register the task function if not already registered
+                    if not hasattr(task_func, "celery_task"):
+                        celery_task = celery_mixin.register_celery_task(
+                            task_func,
+                            queue,
+                        )
+                        task_func.celery_task = celery_task
+
+                    # Submit task to Celery
+                    result = celery_mixin.submit_task(task_func, request)
+
+                    return {
+                        "task_id": result.id,
+                        "status": "submitted",
+                        "queue": queue,
+                        "message": f"Task {result.id} submitted to Celery "
+                        f"queue {queue}",
+                    }
+
+                else:
+                    # Fallback to in-memory task processing
+                    import time
+
+                    # Initialize task storage if not exists
+                    if not hasattr(app.state, "active_tasks"):
+                        app.state.active_tasks = {}
+
+                    # Create task info for tracking
+                    task_info = {
+                        "task_id": task_id,
+                        "status": "submitted",
+                        "queue": queue,
+                        "submitted_at": time.time(),
+                        "request": request,
+                    }
+                    app.state.active_tasks[task_id] = task_info
+
+                    # Execute task asynchronously in background
+                    asyncio.create_task(
+                        FastAPIAppFactory._execute_background_task(
+                            app,
+                            task_id,
+                            task_func,
+                            request,
+                            queue,
+                        ),
+                    )
+
+                    return {
+                        "task_id": task_id,
+                        "status": "submitted",
+                        "queue": queue,
+                        "message": f"Task {task_id} submitted to queue "
+                        f"{queue}",
+                    }
+
+            except Exception as e:
+                return {
+                    "error": str(e),
+                    "type": "task",
+                    "queue": queue,
+                    "status": "failed",
+                }
+
+        return task_endpoint
+
+    @staticmethod
+    async def _execute_background_task(
+        app: FastAPI,
+        task_id: str,
+        func: Callable,
+        request: dict,
+        queue: str,
+    ):
+        """Execute task in background and update status."""
+        try:
+            import time
+            import concurrent.futures
+
+            # Update status to running
+            if (
+                hasattr(app.state, "active_tasks")
+                and task_id in app.state.active_tasks
+            ):
+                app.state.active_tasks[task_id].update(
+                    {
+                        "status": "running",
+                        "started_at": time.time(),
+                    },
+                )
+
+            # Execute the actual task function
+            if asyncio.iscoroutinefunction(func):
+                result = await func(request)
+            else:
+                # Run sync function in thread pool to avoid blocking
+                with concurrent.futures.ThreadPoolExecutor() as executor:
+                    result = await asyncio.get_event_loop().run_in_executor(
+                        executor,
+                        func,
+                        request,
+                    )
+
+            # Update status to completed
+            if (
+                hasattr(app.state, "active_tasks")
+                and task_id in app.state.active_tasks
+            ):
+                app.state.active_tasks[task_id].update(
+                    {
+                        "status": "completed",
+                        "result": result,
+                        "completed_at": time.time(),
+                    },
+                )
+
+        except Exception as e:
+            # Update status to failed
+            if (
+                hasattr(app.state, "active_tasks")
+                and task_id in app.state.active_tasks
+            ):
+                app.state.active_tasks[task_id].update(
+                    {
+                        "status": "failed",
+                        "error": str(e),
+                        "failed_at": time.time(),
+                    },
+                )
+
+    @staticmethod
+    def _create_task_status_handler(app: FastAPI):
+        """Create a handler for checking task status."""
+
+        async def task_status_handler(task_id: str):
+            if not task_id:
+                return {"error": "task_id required"}
+
+            # Check if Celery is available
+            if hasattr(app.state, "celery_mixin") and app.state.celery_mixin:
+                # Use Celery for task status checking
+                celery_mixin = app.state.celery_mixin
+                return celery_mixin.get_task_status(task_id)
+
+            else:
+                # Fallback to in-memory task status checking
+                if (
+                    not hasattr(app.state, "active_tasks")
+                    or task_id not in app.state.active_tasks
+                ):
+                    return {"error": f"Task {task_id} not found"}
+
+                task_info = app.state.active_tasks[task_id]
+                task_status = task_info.get("status", "unknown")
+
+                # Align with BaseApp.get_task logic - map internal status to
+                # external status format
+                if task_status in ["submitted", "running"]:
+                    return {"status": "pending", "result": None}
+                elif task_status == "completed":
+                    return {
+                        "status": "finished",
+                        "result": task_info.get("result"),
+                    }
+                elif task_status == "failed":
+                    return {
+                        "status": "error",
+                        "result": task_info.get("error", "Unknown error"),
+                    }
+                else:
+                    return {"status": task_status, "result": None}
+
+        return task_status_handler