agentscope-runtime 0.1.6__py3-none-any.whl → 0.2.0__py3-none-any.whl

agentscope_runtime/engine/schemas/realtime.py

@@ -0,0 +1,254 @@
+# -*- coding: utf-8 -*-
+import os
+import uuid
+from enum import Enum
+from typing import Optional, List, Union, Any
+
+from openai.types.chat.chat_completion_chunk import ChoiceDeltaToolCall
+from pydantic import BaseModel, model_validator, Field
+
+
+class AsrVendor(str, Enum):
+    MODELSTUDIO = "modelstudio"
+    AZURE = "azure"
+
+
+class TtsVendor(str, Enum):
+    MODELSTUDIO = "modelstudio"
+    AZURE = "azure"
+
+
+class ModelstudioConnection(BaseModel):
+    base_url: Optional[str] = None
+    api_key: Optional[str] = None
+    workspace_id: Optional[str] = None
+    user_agent: Optional[str] = None
+    data_inspection: Optional[str] = None
+
+
+class TtsConfig(BaseModel):
+    model: Optional[str] = None
+    voice: Optional[str] = None
+    sample_rate: Optional[int] = None
+    format: Optional[str] = None
+    bits_per_sample: Optional[int] = None
+    nb_channels: Optional[int] = None
+    chat_id: Optional[str] = None
+
+
+class ModelstudioTtsConfig(TtsConfig, ModelstudioConnection):
+    model: str = "cosyvoice-v2"
+    voice: str = os.getenv("TTS_VOICE", "longcheng_v2")
+    sample_rate: int = 16000
+    format: Optional[str] = "pcm"
+
+
+class AsrConfig(BaseModel):
+    model: Optional[str] = None
+    language: Optional[str] = None
+    sample_rate: Optional[int] = None
+    format: Optional[str] = None
+    bits_per_sample: Optional[int] = None
+    nb_channels: Optional[int] = None
+    initial_silence_timeout: Optional[int] = None
+    max_end_silence: Optional[int] = None
+    fast_vad_min_duration: Optional[int] = None
+    fast_vad_max_duration: Optional[int] = None
+
+
+class ModelstudioAsrConfig(AsrConfig, ModelstudioConnection):
+    model: Optional[str] = "gummy-realtime-v1"
+    sample_rate: Optional[int] = 16000
+    format: Optional[str] = "pcm"
+    max_end_silence: Optional[int] = 700
+    fast_vad_min_duration: Optional[int] = 200
+    fast_vad_max_duration: Optional[int] = 1100
+
+
+class ModelstudioKnowledgeBaseConfig(BaseModel):
+    index_ids: List[str]
+    workspace_id: str
+    api_key: str
+
+
+class ModelstudioVoiceChatUpstream(BaseModel):
+    dialog_mode: Optional[str] = "duplex"
+    enable_server_vad: Optional[bool] = True
+    modalities: Optional[List[str]] = Field(default_factory=lambda: ["audio"])
+    asr_vendor: Optional[AsrVendor] = AsrVendor(
+        os.getenv("ASR_VENDOR", AsrVendor.MODELSTUDIO.value),
+    )
+    asr_options: Optional[AsrConfig] = AsrConfig()
+
+
+class ModelstudioVoiceChatDownstream(BaseModel):
+    modalities: Optional[List[str]] = Field(
+        default_factory=lambda: ["audio", "text"],
+    )
+    tts_vendor: Optional[TtsVendor] = TtsVendor(
+        os.getenv("TTS_VENDOR", TtsVendor.MODELSTUDIO.value),
+    )
+    tts_options: Optional[TtsConfig] = TtsConfig()
+
+
+class ModelstudioVoiceChatParameters(BaseModel):
+    modelstudio_kb: Optional[ModelstudioKnowledgeBaseConfig] = None
+    enable_tool_call: Optional[bool] = False
+
+
+class ModelstudioVoiceChatInput(BaseModel):
+    dialog_id: Optional[str] = None
+    app_id: Optional[str] = None
+    text: Optional[str] = None
+
+
+class ModelstudioVoiceChatDirective(str, Enum):
+    SESSION_START = "SessionStart"
+    SESSION_STOP = "SessionStop"
+
+
+class ModelstudioVoiceChatEvent(str, Enum):
+    SESSION_STARTED = "SessionStarted"
+    SESSION_STOPPED = "SessionStopped"
+    AUDIO_TRANSCRIPT = "AudioTranscript"
+    RESPONSE_TEXT = "ResponseText"
+    RESPONSE_AUDIO_STARTED = "ResponseAudioStarted"
+    RESPONSE_AUDIO_ENDED = "ResponseAudioEnded"
+
+
+class ModelstudioVoiceChatInPayload(BaseModel):
+    pass
+
+
+class ModelstudioVoiceChatSessionStartPayload(ModelstudioVoiceChatInPayload):
+    session_id: Optional[str] = Field(
+        default_factory=lambda: str(uuid.uuid4()),
+    )
+    upstream: Optional[
+        ModelstudioVoiceChatUpstream
+    ] = ModelstudioVoiceChatUpstream()
+    downstream: Optional[
+        ModelstudioVoiceChatDownstream
+    ] = ModelstudioVoiceChatDownstream()
+    parameters: Optional[
+        ModelstudioVoiceChatParameters
+    ] = ModelstudioVoiceChatParameters()
+
+
+class ModelstudioVoiceChatSessionStopPayload(ModelstudioVoiceChatInPayload):
+    pass
+
+
+class ModelstudioVoiceChatRequest(BaseModel):
+    directive: ModelstudioVoiceChatDirective
+    payload: Union[
+        ModelstudioVoiceChatSessionStartPayload,
+        ModelstudioVoiceChatSessionStopPayload,
+    ]
+
+    @model_validator(mode="wrap")
+    def parse_payload_based_on_directive(
+        self,
+        values: Any,
+        handler: Any,
+    ) -> None:
+        data = values if isinstance(values, dict) else values.model_dump()
+
+        directive = data.get("directive")
+        payload_data = data.get("payload", {})
+
+        if directive == ModelstudioVoiceChatDirective.SESSION_START:
+            data["payload"] = ModelstudioVoiceChatSessionStartPayload(
+                **payload_data,
+            )
+        elif directive == ModelstudioVoiceChatDirective.SESSION_STOP:
+            data["payload"] = ModelstudioVoiceChatSessionStopPayload(
+                **payload_data,
+            )
+        else:
+            raise ValueError(f"Unsupported directive: {directive}")
+
+        return handler(data)
+
+
+class ModelstudioVoiceChatOutPayload(BaseModel):
+    session_id: Optional[str] = None
+
+
+class ModelstudioVoiceChatSessionStartedPayload(
+    ModelstudioVoiceChatOutPayload,
+):
+    pass
+
+
+class ModelstudioVoiceChatSessionStoppedPayload(
+    ModelstudioVoiceChatOutPayload,
+):
+    pass
+
+
+class ModelstudioVoiceChatAudioTranscriptPayload(
+    ModelstudioVoiceChatOutPayload,
+):
+    text: Optional[str] = ""
+    finished: bool
+
+
+class ModelstudioVoiceChatResponseTextPayload(
+    ModelstudioVoiceChatOutPayload,
+):
+    text: Optional[str] = ""
+    tool_calls: Optional[List[ChoiceDeltaToolCall]] = Field(
+        default_factory=list,
+    )
+    finished: bool
+
+
+class ModelstudioVoiceChatResponseAudioStartedPayload(
+    ModelstudioVoiceChatOutPayload,
+):
+    pass
+
+
+class ModelstudioVoiceChatResponseAudioStoppedPayload(
+    ModelstudioVoiceChatOutPayload,
+):
+    pass
+
+
+class ModelstudioVoiceChatResponse(BaseModel):
+    event: Optional[ModelstudioVoiceChatEvent]
+    payload: Union[
+        ModelstudioVoiceChatSessionStartedPayload,
+        ModelstudioVoiceChatSessionStoppedPayload,
+        ModelstudioVoiceChatAudioTranscriptPayload,
+        ModelstudioVoiceChatResponseTextPayload,
+        ModelstudioVoiceChatResponseAudioStartedPayload,
+        ModelstudioVoiceChatResponseAudioStoppedPayload,
+    ]
+
+
+class AzureConnection(BaseModel):
+    key: Optional[str] = None
+    region: Optional[str] = None
+
+
+class AzureAsrConfig(AsrConfig, AzureConnection):
+    sample_rate: Optional[int] = 16000
+    format: Optional[str] = "pcm"
+    bits_per_sample: Optional[int] = 16
+    nb_channels: Optional[int] = 1
+    initial_silence_timeout: Optional[int] = 5000
+    max_end_silence: Optional[int] = 800
+    language: Optional[str] = os.getenv("ASR_LANG", "en-US")
+
+
+class AzureTtsConfig(TtsConfig, AzureConnection):
+    voice: Optional[str] = os.getenv(
+        "TTS_VOICE",
+        "en-US-AvaMultilingualNeural",
+    )
+    sample_rate: Optional[int] = 16000
+    format: Optional[str] = "pcm"
+    bits_per_sample: Optional[int] = 16
+    nb_channels: Optional[int] = 1

agentscope_runtime/engine/services/tablestore_memory_service.py

@@ -51,7 +51,7 @@ class TablestoreMemoryService(MemoryService):
         self,
         tablestore_client: AsyncTablestoreClient,
         search_strategy: SearchStrategy = SearchStrategy.FULL_TEXT,
-        embedding_model: Optional[Embeddings] = DashScopeEmbeddings(),
+        embedding_model: Optional[Embeddings] = None,
         vector_dimension: int = 1536,
         table_name: Optional[str] = "agentscope_runtime_memory",
         search_index_schema: Optional[List[tablestore.FieldSchema]] = (
@@ -63,6 +63,9 @@ class TablestoreMemoryService(MemoryService):
         vector_metric_type: VectorMetricType = VectorMetricType.VM_COSINE,
         **kwargs: Any,
     ):
+        if embedding_model is None:
+            embedding_model = DashScopeEmbeddings()
+
         self._search_strategy = search_strategy
         self._embedding_model = (
             embedding_model  # the parameter is None, don't store vector.

agentscope_runtime/engine/tracing/__init__.py

@@ -3,6 +3,14 @@ from typing import Any, List, Union
 
 from .base import BaseLogHandler, Tracer, TracerHandler
 from .tracing_metric import TraceType
+from .tracing_util import TracingUtil
+from .wrapper import trace
+
+__all__ = [
+    "trace",
+    "TraceType",
+    "TracingUtil",
+]
 
 
 def create_handler(
@@ -15,9 +23,7 @@ def create_handler(
     if "default" in eval_mode:
         handlers.append(BaseLogHandler())
     elif "local_logging" in eval_mode:
-        from .local_logging_handler import (
-            LocalLogHandler,
-        )
+        from .local_logging_handler import LocalLogHandler
 
         handlers.append(LocalLogHandler(**eval_params))
     return handlers

agentscope_runtime/engine/tracing/asyncio_util.py

@@ -0,0 +1,24 @@
+# -*- coding: utf-8 -*-
+from typing import AsyncIterable, AsyncIterator, Tuple, TypeVar
+
+T_co = TypeVar("T_co", covariant=True)
+
+
+async def aenumerate(
+    asequence: AsyncIterable[T_co],
+    start: int = 0,
+) -> AsyncIterator[Tuple[int, T_co]]:
+    """Asynchronously enumerate an async iterator from a given start value.
+
+    Args:
+        asequence (AsyncIterable[T_co]): The async iterable to enumerate.
+        start (int): The starting value for enumeration. Defaults to 0.
+
+    Yields:
+        Tuple[int, T_co]: A tuple containing the index and the item from the
+        async iterable.
+    """
+    n = start
+    async for elem in asequence:
+        yield n, elem
+        n += 1

agentscope_runtime/engine/tracing/base.py

@@ -1,12 +1,12 @@
 # -*- coding: utf-8 -*-
-# pylint: disable=protected-access
 import time
 import traceback
 from abc import ABC, abstractmethod
 from contextlib import contextmanager
 from typing import Any, Dict, List
 
-from .tracing_metric import TraceType
+from opentelemetry.propagate import inject, extract
+from opentelemetry.context.context import Context
 
 
 # Handler Interface
@@ -16,22 +16,23 @@ class TracerHandler(ABC):
     @abstractmethod
     def on_start(
         self,
-        event_type: TraceType,
+        event_name: str,
         payload: Dict[str, Any],
         **kwargs: Any,
     ) -> None:
         """Handle the start of a trace event.
 
         Args:
-            event_type (TraceType): The type of event being traced.
+            event_name (str): The name of event being traced.
             payload (Dict[str, Any]): The payload data for the event.
             **kwargs (Any): Additional keyword arguments.
         """
+        raise NotImplementedError("Subclasses must implement on_start method")
 
     @abstractmethod
     def on_end(
         self,
-        event_type: TraceType,
+        event_name: str,
         start_payload: Dict[str, Any],
         end_payload: Dict[str, Any],
         start_time: float,
@@ -40,12 +41,13 @@ class TracerHandler(ABC):
         """Handle the end of a trace event.
 
         Args:
-            event_type (TraceType): The type of event being traced.
+            event_name (str): The name of event being traced.
             start_payload (Dict[str, Any]): The payload data from event start.
             end_payload (Dict[str, Any]): The payload data from event end.
             start_time (float): The timestamp when the event started.
             **kwargs (Any): Additional keyword arguments.
         """
+        raise NotImplementedError("Subclasses must implement on_end method")
 
     @abstractmethod
     def on_log(self, message: str, **kwargs: Any) -> None:
@@ -55,11 +57,12 @@ class TracerHandler(ABC):
             message (str): The log message.
             **kwargs (Any): Additional keyword arguments.
         """
+        raise NotImplementedError("Subclasses must implement on_log method")
 
     @abstractmethod
     def on_error(
         self,
-        event_type: TraceType,
+        event_name: str,
         start_payload: Dict[str, Any],
         error: Exception,
         start_time: float,
@@ -69,15 +72,17 @@ class TracerHandler(ABC):
         """Handle an error during tracing.
 
         Args:
-            event_type (TraceType): The type of event being traced.
+            event_name (str): The type of event being traced.
             start_payload (Dict[str, Any]): The payload data from event start.
             error (Exception): The exception that occurred.
             start_time (float): The timestamp when the event started.
             traceback_info (str): The traceback information.
             **kwargs (Any): Additional keyword arguments.
         """
+        raise NotImplementedError("Subclasses must implement on_error method")
 
 
+# 新增基础的LogHandler类
 class BaseLogHandler(TracerHandler):
     """Basic log handler implementation using Python's logging module."""
 
@@ -87,22 +92,22 @@ class BaseLogHandler(TracerHandler):
 
     def on_start(
         self,
-        event_type: TraceType,
+        event_name: str,
         payload: Dict[str, Any],
         **kwargs: Any,
     ) -> None:
         """Log the start of a trace event.
 
         Args:
-            event_type (TraceType): The type of event being traced.
+            event_name (str): The name of event being traced.
             payload (Dict[str, Any]): The payload data for the event.
             **kwargs (Any): Additional keyword arguments.
         """
-        self.logger.info(f"Event {event_type} started with payload: {payload}")
+        self.logger.info(f"Event {event_name} started with payload: {payload}")
 
     def on_end(
         self,
-        event_type: TraceType,
+        event_name: str,
         start_payload: Dict[str, Any],
         end_payload: Dict[str, Any],
         start_time: float,
@@ -111,14 +116,14 @@ class BaseLogHandler(TracerHandler):
         """Log the end of a trace event.
 
         Args:
-            event_type (TraceType): The type of event being traced.
+            event_name (str): The name of event being traced.
             start_payload (Dict[str, Any]): The payload data from event start.
             end_payload (Dict[str, Any]): The payload data from event end.
             start_time (float): The timestamp when the event started.
             **kwargs (Any): Additional keyword arguments.
         """
         self.logger.info(
-            f"Event {event_type} ended with start payload: {start_payload}, "
+            f"Event {event_name} ended with start payload: {start_payload}, "
             f"end payload: {end_payload}, duration: "
             f"{time.time() - start_time} seconds, kwargs: {kwargs}",
         )
@@ -130,12 +135,11 @@ class BaseLogHandler(TracerHandler):
             message (str): The log message.
             **kwargs (Any): Additional keyword arguments.
         """
-        if message:
-            self.logger.info(f"Log: {message}")
+        self.logger.info(f"Log: {message}")
 
     def on_error(
         self,
-        event_type: TraceType,
+        event_name: str,
         start_payload: Dict[str, Any],
         error: Exception,
         start_time: float,
@@ -145,7 +149,7 @@ class BaseLogHandler(TracerHandler):
         """Log an error during tracing.
 
         Args:
-            event_type (TraceType): The type of event being traced.
+            event_name (str): The name of event being traced.
             start_payload (Dict[str, Any]): The payload data from event start.
             error (Exception): The exception that occurred.
             start_time (float): The timestamp when the event started.
@@ -153,7 +157,7 @@ class BaseLogHandler(TracerHandler):
             **kwargs (Any): Additional keyword arguments.
         """
         self.logger.error(
-            f"Error in event {event_type} with payload: {start_payload}, "
+            f"Error in event {event_name} with payload: {start_payload}, "
             f"error: {error}, "
             f"traceback: {traceback_info}, duration: "
             f"{time.time() - start_time} seconds, kwargs: {kwargs}",
@@ -184,14 +188,16 @@ class Tracer:
     @contextmanager
     def event(
         self,
-        event_type: TraceType,
+        span: Any,
+        event_name: str,
         payload: Dict[str, Any],
         **kwargs: Any,
     ) -> Any:
         """Create a context manager for tracing an event.
 
         Args:
-            event_type (TraceType): The type of event being traced.
+            span(Any): span of event
+            event_name (str): The name of event being traced.
             payload (Dict[str, Any]): The payload data for the event.
             **kwargs (Any): Additional keyword arguments.
 
@@ -201,28 +207,35 @@ class Tracer:
         start_time = time.time()
 
         for handle in self.handlers:
-            handle.on_start(event_type, payload, **kwargs)
+            handle.on_start(
+                event_name,
+                payload,
+                **kwargs,
+            )
 
         event_context = EventContext(
+            span,
             self.handlers,
-            event_type,
+            event_name,
             start_time,
             payload,
         )
+
         try:
             yield event_context
         except Exception as e:
             traceback_info = traceback.format_exc()
             for handle in self.handlers:
                 handle.on_error(
-                    event_type,
+                    event_name,
                     payload,
                     e,
                     start_time,
                     traceback_info=traceback_info,
                 )
             raise
-        event_context._end(payload)
+
+        event_context.finalize(payload)
 
     def log(self, message: str, **kwargs: Any) -> None:
         """Log a message using all registered handlers.
@@ -240,8 +253,9 @@ class EventContext:
 
     def __init__(
         self,
+        span: Any,
         handlers: List[TracerHandler],
-        event_type: TraceType,
+        event_name: str,
         start_time: float,
         start_payload: Dict[str, Any],
     ) -> None:
@@ -250,12 +264,13 @@ class EventContext:
         Args:
             handlers (List[TracerHandler]): List of handlers to process
             trace  events.
-            event_type (TraceType): The type of event being traced.
+            event_name (str): The name of event being traced.
             start_time (float): The timestamp when the event started.
             start_payload (Dict[str, Any]): The payload data from event start.
         """
+        self.span = span
         self.handlers = handlers
-        self.event_type = event_type
+        self.event_name = event_name
         self.start_time = start_time
         self.start_payload = start_payload
         self.end_payload = {}
@@ -278,12 +293,21 @@ class EventContext:
             message (str): The log message.
             **kwargs (Any): Additional keyword arguments.
         """
-        kwargs["event_type"] = self.event_type
+        kwargs["event_name"] = self.event_name
         kwargs["start_time"] = self.start_time
         kwargs["start_payload"] = self.start_payload
         for handle in self.handlers:
             handle.on_log(message, **kwargs)
 
+    def finalize(self, start_payload: Dict[str, Any] = None) -> None:
+        """Public method to finalize the event.
+
+        Args:
+            start_payload (Dict[str, Any], optional): The payload data from
+            event start.
+        """
+        self._end(start_payload)
+
     def _end(self, start_payload: Dict[str, Any] = None) -> None:
         """Finalize the event by calling on_end for all handlers.
 
@@ -293,17 +317,25 @@ class EventContext:
         """
         for handle in self.handlers:
             handle.on_end(
-                self.event_type,
+                self.event_name,
                 start_payload,
                 self.end_payload,
                 self.start_time,
                 **self.kwargs,
             )
 
-    def set_trace_header(self, trace_header: Dict[str, Any]) -> None:
-        """Set trace header for compatible handlers.
+    def set_attribute(self, key: str, value: Any) -> None:
+        """Set attribute for the current span.
 
         Args:
-            trace_header (Dict[str, Any]): The trace header information.
+            key (str): key of attribute
+            value(str): value of attribute
         """
-        # TODO: Implement this
+
+        self.span.set_attribute(key, value)
+
+    def get_trace_context(self) -> Context:
+        carrier = {}
+        inject(carrier)
+        context = extract(carrier)
+        return context