mirror of
https://github.com/agentscope-ai/ReMe.git
synced 2026-10-05 02:41:43 +00:00
99 lines
3.6 KiB
Python
99 lines
3.6 KiB
Python
"""HTTP service implementation for ReMe."""
|
|
|
|
import asyncio
|
|
import json
|
|
import os
|
|
import warnings
|
|
from collections.abc import AsyncGenerator
|
|
from contextlib import asynccontextmanager
|
|
from typing import TYPE_CHECKING
|
|
|
|
import uvicorn
|
|
from fastapi import FastAPI
|
|
from fastapi.middleware.cors import CORSMiddleware
|
|
from fastapi.responses import StreamingResponse
|
|
|
|
from .base_service import BaseService
|
|
from ..component_registry import R
|
|
from ..job import BaseJob, StreamJob
|
|
from ...constants import REME_DEFAULT_HOST, REME_DEFAULT_PORT, REME_SERVICE_INFO
|
|
from ...schema import Request, Response
|
|
from ...utils import execute_stream_task
|
|
|
|
if TYPE_CHECKING:
|
|
from ...application import Application
|
|
|
|
|
|
@R.register("http")
|
|
class HttpService(BaseService):
|
|
"""HTTP service: normal jobs -> JSON endpoints, stream jobs -> SSE endpoints."""
|
|
|
|
def __init__(self, host: str = REME_DEFAULT_HOST, port: int = REME_DEFAULT_PORT, **kwargs):
|
|
super().__init__(**kwargs)
|
|
self.host: str = host
|
|
self.port: int = port
|
|
|
|
def _add_job(self, job: BaseJob) -> None:
|
|
async def execute_endpoint(request: Request) -> Response:
|
|
return await job(**request.model_dump(exclude_none=True))
|
|
|
|
self.service.post(path=f"/{job.name}", response_model=Response, description=job.description)(execute_endpoint)
|
|
|
|
def _add_stream_job(self, job: StreamJob) -> None:
|
|
async def execute_stream_endpoint(request: Request) -> StreamingResponse:
|
|
stream_queue = asyncio.Queue()
|
|
task = asyncio.create_task(job(stream_queue=stream_queue, **request.model_dump(exclude_none=True)))
|
|
|
|
async def generate_stream() -> AsyncGenerator[bytes, None]:
|
|
async for chunk in execute_stream_task(
|
|
stream_queue=stream_queue,
|
|
task=task,
|
|
task_name=job.name,
|
|
output_format="bytes",
|
|
):
|
|
assert isinstance(chunk, bytes)
|
|
yield chunk
|
|
|
|
return StreamingResponse(generate_stream(), media_type="text/event-stream")
|
|
|
|
self.service.post(f"/{job.name}")(execute_stream_endpoint)
|
|
|
|
def add_job(self, job: BaseJob) -> None:
|
|
if isinstance(job, StreamJob):
|
|
self._add_stream_job(job)
|
|
else:
|
|
self._add_job(job)
|
|
|
|
def build_service(self, app: "Application") -> None:
|
|
@asynccontextmanager
|
|
async def lifespan(_: FastAPI):
|
|
await app.start()
|
|
service_info = json.dumps({"host": self.host, "port": self.port})
|
|
os.environ[REME_SERVICE_INFO] = service_info
|
|
self.logger.info(f"ReMe Service started: {REME_SERVICE_INFO}={service_info}")
|
|
yield
|
|
await app.close()
|
|
|
|
self.service = FastAPI(title=app.config.app_name, lifespan=lifespan)
|
|
self.service.add_middleware(
|
|
CORSMiddleware, # type: ignore[arg-type]
|
|
allow_origins=["*"],
|
|
allow_credentials=True,
|
|
allow_methods=["*"],
|
|
allow_headers=["*"],
|
|
)
|
|
|
|
def start_service(self, app: "Application") -> None:
|
|
# uvicorn 0.41 still imports websockets.legacy / WebSocketServerProtocol
|
|
# on startup; silence those specific lines since we don't use WebSocket.
|
|
warnings.filterwarnings(
|
|
"ignore",
|
|
category=DeprecationWarning,
|
|
message=r".*websockets\.legacy is deprecated.*",
|
|
)
|
|
warnings.filterwarnings(
|
|
"ignore",
|
|
category=DeprecationWarning,
|
|
message=r".*WebSocketServerProtocol is deprecated.*",
|
|
)
|
|
uvicorn.run(self.service, host=self.host, port=self.port, **self.kwargs)
|