agentscope-runtime 0.2.0b2__py3-none-any.whl → 1.0.0__py3-none-any.whl

agentscope_runtime/engine/runner.py

@@ -1,12 +1,21 @@
 # -*- coding: utf-8 -*-
+# pylint: disable=not-callable
+import asyncio
+import logging
+import inspect
+import traceback
 import uuid
 from contextlib import AsyncExitStack
-from typing import Optional, List, AsyncGenerator, Any, Union, Dict
-
-from agentscope_runtime.engine.deployers.utils.service_utils import (
-    ServicesConfig,
+from typing import (
+    Optional,
+    List,
+    AsyncGenerator,
+    Any,
+    Union,
+    Dict,
+    AsyncIterator,
 )
-from .agents import Agent
+
 from .deployers import (
     DeployManager,
     LocalDeployManager,
@@ -18,67 +27,93 @@ from .schemas.agent_schemas import (
     RunStatus,
     AgentResponse,
     SequenceNumberGenerator,
+    Error,
 )
-from .schemas.context import Context
-from .services.context_manager import ContextManager
-from .services.environment_manager import EnvironmentManager
 from .tracing import TraceType
 from .tracing.wrapper import trace
 from .tracing.message_util import (
     merge_agent_response,
     get_agent_response_finish_reason,
 )
+from .constant import ALLOWED_FRAMEWORK_TYPES
+
+
+logger = logging.getLogger(__name__)
 
 
 class Runner:
-    def __init__(
-        self,
-        agent: Agent,
-        environment_manager: Optional[EnvironmentManager] = None,
-        context_manager: Optional[ContextManager] = None,
-    ) -> None:
+    def __init__(self) -> None:
         """
-        Initializes a runner as core function.
-        Args:
-            agent: The agent to run.
-            environment_manager: The environment manager
-            context_manager: The context manager
-        """
-        self._agent = agent
-        self._environment_manager = environment_manager
-        self._context_manager = (
-            context_manager or ContextManager()
-        )  # Add default context manager
+        Initializes a runner as core instance.
+        """
+        self.framework_type = None
+
         self._deploy_managers = {}
+        self._health = False
         self._exit_stack = AsyncExitStack()
 
-    async def __aenter__(self) -> "Runner":
+    async def query_handler(self, *args, **kwargs):
         """
-        Initializes the runner and ensures context/environment managers
-        are fully entered so that attributes like compose_session are
-        available.
+        Handle agent query.
         """
-        if self._environment_manager:
-            # enter_async_context returns the "real" object
-            self._environment_manager = (
-                await self._exit_stack.enter_async_context(
-                    self._environment_manager,
-                )
-            )
+        raise NotImplementedError("query_handler not implemented")
 
-        if self._context_manager:
-            self._context_manager = await self._exit_stack.enter_async_context(
-                self._context_manager,
-            )
+    async def init_handler(self, *args, **kwargs):
+        """
+        Init handler.
+        """
 
+    async def shutdown_handler(self, *args, **kwargs):
+        """
+        Shutdown handler.
+        """
+
+    async def start(self):
+        init_fn = getattr(self, "init_handler", None)
+        if callable(init_fn):
+            if inspect.iscoroutinefunction(init_fn):
+                await init_fn()
+            else:
+                init_fn()
+        else:
+            logger.warning("[Runner] init_handler is not callable")
+        self._health = True
         return self
 
-    async def __aexit__(self, exc_type, exc_val, exc_tb):
+    async def stop(self):
+        shutdown_fn = getattr(self, "shutdown_handler", None)
+        try:
+            if callable(shutdown_fn):
+                if inspect.iscoroutinefunction(shutdown_fn):
+                    await shutdown_fn()
+                else:
+                    shutdown_fn()
+        except Exception as e:
+            logger.warning(f"[Runner] Exception in shutdown handler: {e}")
         try:
             await self._exit_stack.aclose()
         except Exception:
             pass
 
+        self._health = False
+
+    async def __aenter__(self) -> "Runner":
+        """
+        Initializes the runner
+        """
+        await self.start()
+        return self
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        await self.stop()
+
+        if hasattr(self, "_deploy_manager") and self._deploy_manager:
+            for deploy_id in self._deploy_manager:
+                await self._deploy_manager[deploy_id].stop()
+        else:
+            # No deploy manager found, nothing to stop
+            pass
+
     async def deploy(
         self,
         deploy_manager: DeployManager = LocalDeployManager(),
@@ -90,7 +125,6 @@ class Runner:
         base_image: str = "python:3.9-slim",
         environment: Optional[Dict[str, str]] = None,
         runtime_config: Optional[Dict] = None,
-        services_config: Optional[Union[ServicesConfig, dict]] = None,
         **kwargs,
     ):
         """
@@ -106,7 +140,6 @@ class Runner:
             base_image: Docker base image (for containerized deployment)
             environment: Environment variables dict
             runtime_config: Runtime configuration dict
-            services_config: Services configuration dict
             **kwargs: Additional arguments passed to deployment manager
         Returns:
             URL of the deployed service
@@ -124,14 +157,34 @@ class Runner:
             base_image=base_image,
             environment=environment,
             runtime_config=runtime_config,
-            services_config=services_config,
             **kwargs,
         )
 
-        # TODO: add redis or other persistant method
+        # TODO: add redis or other persistent method
         self._deploy_managers[deploy_manager.deploy_id] = deploy_result
         return deploy_result
 
+    async def _call_handler_streaming(self, handler, *args, **kwargs):
+        """
+        Call handler and yield results in streaming fashion, async or sync.
+        """
+        result = handler(*args, **kwargs)
+
+        if inspect.isasyncgenfunction(handler):
+            async for item in result:
+                yield item
+
+        elif inspect.isgenerator(result):
+            for item in result:
+                yield item
+
+        elif asyncio.iscoroutine(result):
+            res = await result
+            yield res
+
+        else:
+            yield result
+
     @trace(
         TraceType.AGENT_STEP,
         trace_name="agent_step",
@@ -141,115 +194,103 @@ class Runner:
     async def stream_query(  # pylint:disable=unused-argument
         self,
         request: Union[AgentRequest, dict],
-        user_id: Optional[str] = None,
-        tools: Optional[List] = None,
         **kwargs: Any,
     ) -> AsyncGenerator[Event, None]:
         """
         Streams the agent.
         """
+        if self.framework_type not in ALLOWED_FRAMEWORK_TYPES:
+            raise RuntimeError(
+                f"Framework type '{self.framework_type}' is invalid or not "
+                f"set. Please set `self.framework_type` to one of:"
+                f" {', '.join(ALLOWED_FRAMEWORK_TYPES)}.",
+            )
+
+        if not self._health:
+            raise RuntimeError(
+                "Runner has not been started. "
+                "Please call 'await runner.start()' or use 'async with "
+                "Runner()' before calling 'stream_query'.",
+            )
+
         if isinstance(request, dict):
             request = AgentRequest(**request)
 
+        # Assign session ID
+        request.session_id = request.session_id or str(uuid.uuid4())
+
+        # Assign user ID
+        request.user_id = request.user_id or request.session_id
+
         seq_gen = SequenceNumberGenerator()
 
         # Initial response
-        response = AgentResponse()
+        response = AgentResponse(id=request.id)
+        response.session_id = request.session_id
         yield seq_gen.yield_with_sequence(response)
 
         # Set to in-progress status
         response.in_progress()
         yield seq_gen.yield_with_sequence(response)
 
-        user_id = user_id or str(uuid.uuid4())
-        session_id = request.session_id or str(uuid.uuid4())
-        request_input = request.input
-        session = await self._context_manager.compose_session(
-            user_id=user_id,
-            session_id=session_id,
-        )
+        query_kwargs = {
+            "request": request,
+        }
 
-        context = Context(
-            user_id=session.user_id,
-            session=session,
-            request=request,
-            current_messages=request_input,
-            context_manager=self._context_manager,
-            environment_manager=self._environment_manager,
-            agent=self._agent,
-        )
+        if self.framework_type == "text":
+            from ..adapters.text.stream import adapt_text_stream
 
-        # TODO: Update activate tools into the context (not schema only)
-        tools = tools or getattr(self._agent, "tools", None)
-        if tools:
-            # Lazy import
-            from ..sandbox.tools.utils import setup_tools
-
-            activated_tools, schemas = setup_tools(
-                tools=tools,
-                environment_manager=context.environment_manager,
-                session_id=session.id,
-                user_id=session.user_id,
-                include_schemas=True,
+            stream_adapter = adapt_text_stream
+        elif self.framework_type == "agentscope":
+            from ..adapters.agentscope.stream import (
+                adapt_agentscope_message_stream,
             )
+            from ..adapters.agentscope.message import message_to_agentscope_msg
 
-            # update the context
-            context.activate_tools = activated_tools
+            stream_adapter = adapt_agentscope_message_stream
+            kwargs.update(
+                {"msgs": message_to_agentscope_msg(request.input)},
+            )
+        # TODO: support other frameworks
+        else:
 
-            # convert schema to a function call tool lists
-            # TODO: use pydantic model
-            if hasattr(context.request, "tools") and context.request.tools:
-                context.request.tools.extend(schemas)
+            def identity_stream_adapter(
+                source_stream: AsyncIterator[Any],
+            ) -> AsyncIterator[Any]:
+                return source_stream
 
-        # update message in session
-        await context.context_manager.compose_context(
-            session=context.session,
-            request_input=request_input,
-        )
+            stream_adapter = identity_stream_adapter
 
-        async for event in context.agent.run_async(context):
-            if (
-                event.status == RunStatus.Completed
-                and event.object == "message"
+        try:
+            async for event in stream_adapter(
+                source_stream=self._call_handler_streaming(
+                    self.query_handler,
+                    **query_kwargs,
+                    **kwargs,
+                ),
             ):
-                response.add_new_message(event)
-            yield seq_gen.yield_with_sequence(event)
-
-        await context.context_manager.append(
-            session=context.session,
-            event_output=response.output,
-        )
-        yield seq_gen.yield_with_sequence(response.completed())
-
-    #  TODO: will be added before 2025/11/30
-    # @trace(TraceType.AGENT_STEP)
-    # async def query(  # pylint:disable=unused-argument
-    #     self,
-    #     message: List[dict],
-    #     session_id: Optional[str] = None,
-    #     **kwargs: Any,
-    # ) -> ChatCompletion:
-    #     """
-    #     Streams the agent.
-    #     """
-    #     return self._agent.query(message, session_id)
-
-    # TODO: should be sync method?
-    async def stop(
-        self,
-        deploy_id: str,
-    ) -> None:
-        """
-        Stops the agent service.
-
-        Args:
-            deploy_id: Optional deploy ID (not used for service shutdown)
+                if (
+                    event.status == RunStatus.Completed
+                    and event.object == "message"
+                ):
+                    response.add_new_message(event)
+                yield seq_gen.yield_with_sequence(event)
+        except Exception as e:
+            # TODO: fix code
+            error = Error(
+                code="500",
+                message=f"Error happens in `query_handler`: {e}",
+            )
+            logger.error(f"{error.model_dump()}: {traceback.format_exc()}")
+            yield seq_gen.yield_with_sequence(response.failed(error))
+            return
 
-        Raises:
-            RuntimeError: If stopping fails
-        """
-        if hasattr(self, "_deploy_manager") and self._deploy_manager:
-            await self._deploy_manager[deploy_id].stop()
-        else:
-            # No deploy manager found, nothing to stop
+        # Obtain token usage
+        try:
+            if response.output:
+                response.usage = response.output[-1].usage
+        except IndexError:
+            # Avoid empty message
             pass
+
+        yield seq_gen.yield_with_sequence(response.completed())

agentscope_runtime/engine/schemas/agent_schemas.py

@@ -11,7 +11,7 @@ except ImportError:
     from typing_extensions import Self
 from uuid import uuid4
 
-from pydantic import BaseModel, Field, field_validator
+from pydantic import BaseModel, Field, field_validator, ConfigDict
 from openai.types.chat import ChatCompletionChunk
 
 
@@ -144,10 +144,79 @@ class FunctionCallOutput(BaseModel):
     call_id: str
     """The ID of the tool call."""
 
+    name: Optional[str] = None
+    """The name of the tool call."""
+
     output: str
     """The result of the function."""
 
 
+class McpCall(BaseModel):
+    id: str
+    """The unique ID of the tool call."""
+
+    arguments: str
+    """A JSON string of the arguments passed to the tool."""
+
+    name: str
+    """The name of the tool that was run."""
+
+    server_label: str
+    """The label of the MCP server running the tool."""
+
+    error: Optional[str] = None
+    """The error from the tool call, if any."""
+
+    output: Optional[str] = None
+    """The output from the tool call."""
+
+
+class McpListToolsTool(BaseModel):
+    input_schema: object
+    """The JSON schema describing the tool's input."""
+
+    name: str
+    """The name of the tool."""
+
+    annotations: Optional[object] = None
+    """Additional annotations about the tool."""
+
+    description: Optional[str] = None
+    """The description of the tool."""
+
+
+class McpListTools(BaseModel):
+    id: str
+    """The unique ID of the list."""
+
+    server_label: str
+    """The label of the MCP server."""
+
+    tools: List[McpListToolsTool]
+    """The tools available on the server."""
+
+    error: Optional[str] = None
+    """Error message if the server could not list tools."""
+
+
+class McpApprovalRequest(BaseModel):
+    """
+    mcp approval request
+    """
+
+    id: str
+    """The unique ID of the approval request."""
+
+    arguments: str
+    """A json string of arguments for the tool."""
+
+    name: str
+    """The name of the tool to run."""
+
+    server_label: str
+    """The label of the mcp server making the request."""
+
+
 class Error(BaseModel):
     code: str
     """The error code of the message."""
@@ -393,6 +462,8 @@ class Message(Event):
     usage: Optional[Dict] = None
     """response usage for output"""
 
+    metadata: Optional[Dict] = None
+
     @staticmethod
     def from_openai_message(message: Union[BaseModel, dict]) -> "Message":
         """Create a message object from an openai message."""
@@ -627,10 +698,50 @@ class BaseRequest(BaseModel):
     stream: bool = True
     """If set, partial message deltas will be sent, like in ChatGPT. """
 
+    id: Optional[str] = None
+    """request unique id"""
+
 
 class AgentRequest(BaseRequest):
     """agent request"""
 
+    model_config = ConfigDict(
+        extra="allow",
+        json_schema_extra={
+            "example": {
+                "input": [
+                    {
+                        "role": "user",
+                        "type": "message",
+                        "content": [{"type": "text", "text": "hello"}],
+                    },
+                    {
+                        "type": "message",
+                        "role": "assistant",
+                        "content": [
+                            {
+                                "object": "content",
+                                "type": "text",
+                                "text": "Hello! How can I assist you today?",
+                            },
+                        ],
+                    },
+                    {
+                        "role": "user",
+                        "type": "message",
+                        "content": [
+                            {
+                                "type": "text",
+                                "text": "What is the capital of France?",
+                            },
+                        ],
+                    },
+                ],
+                "session_id": "1764056632961",
+            },
+        },
+    )
+
     model: Optional[str] = None
     """
     model id
@@ -702,8 +813,8 @@ class AgentRequest(BaseRequest):
     session_id: Optional[str] = None
     """conversation id for dialog"""
 
-    response_id: Optional[str] = None
-    """response unique id"""
+    user_id: Optional[str] = None
+    """User id for dialog"""
 
 
 class BaseResponse(Event):

agentscope_runtime/engine/schemas/modelstudio_llm.py

@@ -220,7 +220,8 @@ class RagOptions(BaseModel):
     """Whether to enable hybrid generations."""
 
     @field_validator("prompt_strategy")
-    def prompt_strategy_check(self, value: str) -> str:
+    @classmethod
+    def prompt_strategy_check(cls, value: str) -> str:
         if value:
             value = value.lower()
             if value in ["topk", "top_k"]:
@@ -228,7 +229,8 @@ class RagOptions(BaseModel):
         return value
 
     @field_validator("maximum_allowed_chunk_num")
-    def maximum_allowed_chunk_num_check(self, value: int) -> int:
+    @classmethod
+    def maximum_allowed_chunk_num_check(cls, value: int) -> int:
         if value < int(PARAM_MAXIMUM_ALLOWED_CHUNK_NUM_MIN) or value > int(
             PARAM_MAXIMUM_ALLOWED_CHUNK_NUM_MAX,
         ):

agentscope_runtime/engine/schemas/oai_llm.py

@@ -116,7 +116,7 @@ class OpenAIMessage(BaseModel):
     """The role of the messages author, should be in `user`,`system`,
     'assistant', 'tool'."""
 
-    content: Optional[Union[str, List[ChatCompletionMessage]]] = None
+    content: Optional[Union[List[ChatCompletionMessage], str]] = None
     """The contents of the message.
 
     Can be a string, a list of content parts for multimodal messages.
@@ -137,7 +137,7 @@ class OpenAIMessage(BaseModel):
         Extract the first text content from the message.
 
         :return: First text string found in the content, or None if no text
-        content
+            content.
         """
         if self.content is None:
             return None
@@ -158,7 +158,7 @@ class OpenAIMessage(BaseModel):
         Extract all image content (URLs or base64 data) from the message.
 
         :return: List of image URLs or base64 encoded strings found in the
-        content
+            content.
         """
         images = []
 
@@ -183,7 +183,7 @@ class OpenAIMessage(BaseModel):
         Extract all audio content (URLs or base64 data) from the message.
 
         :return: List of audio URLs or base64 encoded strings found in the
-        content
+            content.
         """
         audios = []
 
@@ -307,34 +307,34 @@ class ResponseFormat(BaseModel):
     """The type of response format being defined.
 
     - `text`: The default response format, which can be either text or any
-    value needed.
+        value needed.
     - `json_object`: Enables JSON mode, which guarantees the message the model
-      generates is valid JSON.
+        generates is valid JSON.
     - `json_schema`: Enables Structured Outputs which guarantees the model will
-      match your supplied JSON schema.
+        match your supplied JSON schema.
     """
 
     json_schema: Optional[JsonSchema] = None
     """The JSON schema for the response format."""
 
-    @model_validator(mode="before")
-    def validate_schema(self, values: dict) -> dict:
-        if not isinstance(values, dict) or "type" not in values:
-            raise ValueError(f"Json schema not valid with type {type(values)}")
-        format_type = values.get("type")
-        json_schema = values.get("json_schema")
-
-        if format_type in ["text", "json_object"] and json_schema is not None:
+    @model_validator(mode="after")
+    def validate_schema(self) -> "ResponseFormat":
+        if (
+            self.type
+            in [
+                "text",
+                "json_object",
+            ]
+            and self.json_schema is not None
+        ):
             raise ValueError(
-                f"Json schema is not allowed for type {format_type}",
+                f"Json schema is not allowed for type {self.type}",
             )
-
-        if format_type == "json_schema":
-            if json_schema is None:
-                raise ValueError(
-                    f"Json schema is required for type {format_type}",
-                )
-        return values
+        if self.type == "json_schema" and self.json_schema is None:
+            raise ValueError(
+                f"Json schema is required for type {self.type}",
+            )
+        return self
 
 
 class ToolChoiceInputFunction(BaseModel):