autobyteus 1.2.1__py3-none-any.whl → 1.2.3__py3-none-any.whl

autobyteus/rpc/hosting.py

@@ -1,244 +0,0 @@
-# file: autobyteus/autobyteus/rpc/hosting.py
-import asyncio
-import logging
-import signal
-from typing import Optional, Dict, Union # Added Dict, Union
-
-from autobyteus.agent.agent import Agent
-from autobyteus.rpc.server import AgentServerEndpoint
-from autobyteus.rpc.config import AgentServerConfig
-from autobyteus.rpc.transport_type import TransportType
-
-logger = logging.getLogger(__name__)
-
-async def serve_agent_stdio(
-    agent: Agent,
-    stop_event: asyncio.Event,
-    server_id_for_config: Optional[str] = None
-) -> None:
-    """
-    Hosts a single Agent instance over STDIO within the current process.
-    This function will run until the stop_event is set.
-    The script calling this function will dedicate its sys.stdin/sys.stdout
-    to RPC communication.
-
-    Args:
-        agent: The Agent instance to serve.
-        stop_event: An asyncio.Event that signals when to stop the server.
-        server_id_for_config: Optional string to use as server_id in the internal
-                              AgentServerConfig. Defaults to f"embedded_stdio_{agent.agent_id}".
-    """
-    if not isinstance(agent, Agent):
-        raise TypeError("agent must be an instance of autobyteus.agent.agent.Agent")
-    if not isinstance(stop_event, asyncio.Event):
-        raise TypeError("stop_event must be an asyncio.Event")
-
-    actual_server_id = server_id_for_config or f"embedded_stdio_{agent.agent_id}"
-    
-    stdio_config = AgentServerConfig(
-        server_id=actual_server_id,
-        transport_type=TransportType.STDIO,
-        stdio_command=["in-process-stdio-server"] 
-    )
-    logger.info(f"Configuring in-process STDIO server for agent '{agent.agent_id}' with config_id '{actual_server_id}'.")
-
-    server_endpoint = AgentServerEndpoint(agent) # Pass single agent
-
-    try:
-        if not agent.is_running:
-            logger.info(f"Starting agent '{agent.agent_id}' for STDIO hosting...")
-            agent.start() 
-
-        logger.info(f"Starting STDIO server endpoint for agent '{agent.agent_id}'...")
-        await server_endpoint.start(stdio_config)
-        
-        logger.info(f"Agent '{agent.agent_id}' now being served via STDIO. Waiting for stop signal...")
-        await stop_event.wait()
-        logger.info(f"Stop signal received for STDIO server of agent '{agent.agent_id}'.")
-
-    except Exception as e:
-        logger.error(f"Error during STDIO hosting for agent '{agent.agent_id}': {e}", exc_info=True)
-    finally:
-        logger.info(f"Shutting down STDIO server for agent '{agent.agent_id}'...")
-        if server_endpoint.is_running:
-            await server_endpoint.stop()
-        
-        if agent.is_running: 
-            logger.info(f"Stopping agent '{agent.agent_id}' after STDIO hosting...")
-            await agent.stop()
-        logger.info(f"STDIO hosting for agent '{agent.agent_id}' has shut down.")
-
-
-async def serve_single_agent_http_sse(
-    agent: Agent,
-    host: str,
-    port: int,
-    stop_event: asyncio.Event,
-    server_id_prefix: str = "embedded_http_single",
-    request_endpoint: str = "/rpc",
-    events_endpoint_base: str = "/events" # Base path for events
-) -> None:
-    """
-    Hosts a single Agent instance over HTTP/SSE within the current process.
-    The agent will be accessible using its own agent_id as the key on the server
-    (e.g., for SSE events: /events/{agent.agent_id}).
-
-    Args:
-        agent: The Agent instance to serve.
-        host: The hostname or IP address to bind the HTTP server to.
-        port: The port number to bind the HTTP server to.
-        stop_event: An asyncio.Event that signals when to stop the server.
-        server_id_prefix: Prefix for the server_id in the internal AgentServerConfig.
-        request_endpoint: The URL path for RPC requests.
-        events_endpoint_base: The base URL path for SSE event streams (agent_id will be appended).
-    """
-    # This function now essentially calls serve_multiple_agents_http_sse
-    # The agent's own ID will be used as the key for routing on the server.
-    agents_to_serve = {agent.agent_id: agent}
-    await serve_multiple_agents_http_sse(
-        agents=agents_to_serve,
-        host=host,
-        port=port,
-        stop_event=stop_event,
-        server_id_prefix=server_id_prefix,
-        request_endpoint=request_endpoint,
-        events_endpoint_base=events_endpoint_base
-    )
-
-
-async def serve_multiple_agents_http_sse(
-    agents: Dict[str, Agent], # Key: server-routable ID, Value: Agent instance
-    host: str,
-    port: int,
-    stop_event: asyncio.Event,
-    server_id_prefix: str = "embedded_http_multi",
-    request_endpoint: str = "/rpc", # Single RPC endpoint for the gateway
-    events_endpoint_base: str = "/events" # Base for /events/{agent_id_on_server}
-) -> None:
-    """
-    Hosts multiple Agent instances over a single HTTP/SSE server within the current process.
-    This function will run until the stop_event is set.
-
-    Args:
-        agents: A dictionary where keys are server-routable string IDs for each agent,
-                and values are the Agent instances to serve.
-        host: The hostname or IP address to bind the HTTP server to.
-        port: The port number to bind the HTTP server to.
-        stop_event: An asyncio.Event that signals when to stop the server.
-        server_id_prefix: Prefix for the server_id in the internal AgentServerConfig for the gateway.
-        request_endpoint: The URL path for all RPC requests (target_agent_id must be in payload).
-        events_endpoint_base: The base URL path for SSE event streams. Clients subscribe to
-                              {events_endpoint_base}/{server_routable_agent_id}.
-    """
-    if not isinstance(agents, dict) or not agents:
-        raise TypeError("agents must be a non-empty dictionary of [str, Agent] instances.")
-    if not all(isinstance(k, str) and isinstance(v, Agent) for k, v in agents.items()):
-        raise TypeError("agents dictionary keys must be strings and values must be Agent instances.")
-    if not isinstance(stop_event, asyncio.Event):
-        raise TypeError("stop_event must be an asyncio.Event")
-
-    # server_id for the gateway server config
-    gateway_server_id = f"{server_id_prefix}_gateway" 
-    
-    sse_config = AgentServerConfig(
-        server_id=gateway_server_id,
-        transport_type=TransportType.SSE,
-        sse_base_url=f"http://{host}:{port}", # type: ignore
-        sse_request_endpoint=request_endpoint,
-        sse_events_endpoint=events_endpoint_base # SseServerHandler will append /{agent_id_on_server}
-    )
-    routable_agent_ids = list(agents.keys())
-    logger.info(f"Configuring in-process HTTP/SSE gateway server at {sse_config.sse_base_url} "
-                f"(config_id: '{gateway_server_id}') to serve agents: {routable_agent_ids}.")
-
-    server_endpoint = AgentServerEndpoint(agents) # Pass the dictionary of agents
-
-    try:
-        for agent_id_on_server, agent_instance in agents.items():
-            if not agent_instance.is_running:
-                logger.info(f"Starting agent '{agent_instance.agent_id}' (server key: '{agent_id_on_server}') for HTTP/SSE gateway hosting...")
-                agent_instance.start()
-
-        logger.info(f"Starting HTTP/SSE gateway server endpoint for agents: {routable_agent_ids}...")
-        await server_endpoint.start(sse_config)
-        
-        logger.info(f"Multi-agent gateway now being served via HTTP/SSE on http://{host}:{port}. Waiting for stop signal...")
-        await stop_event.wait()
-        logger.info(f"Stop signal received for multi-agent HTTP/SSE gateway server.")
-
-    except Exception as e:
-        logger.error(f"Error during multi-agent HTTP/SSE hosting: {e}", exc_info=True)
-    finally:
-        logger.info(f"Shutting down multi-agent HTTP/SSE gateway server...")
-        if server_endpoint.is_running:
-            await server_endpoint.stop()
-        
-        for agent_id_on_server, agent_instance in agents.items():
-            if agent_instance.is_running:
-                logger.info(f"Stopping agent '{agent_instance.agent_id}' (server key: '{agent_id_on_server}')...")
-                await agent_instance.stop()
-        logger.info(f"Multi-agent HTTP/SSE hosting has shut down.")
-
-
-# Renamed old serve_agent_http_sse to serve_single_agent_http_sse
-# Kept the old signature for backward compatibility if any tests used it,
-# but it now delegates to the multi-agent version.
-async def serve_agent_http_sse(
-    agent: Agent,
-    host: str,
-    port: int,
-    stop_event: asyncio.Event,
-    server_id_prefix: str = "embedded_http", # Note: prefix change if desired
-    request_endpoint: str = "/rpc",
-    events_endpoint_base: str = "/events" 
-) -> None:
-     await serve_single_agent_http_sse(
-         agent, host, port, stop_event, server_id_prefix, request_endpoint, events_endpoint_base
-     )
-
-
-if __name__ == "__main__": # pragma: no cover
-    logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
-    from autobyteus.agent.registry.agent_definition import AgentDefinition
-    from autobyteus.agent.registry.agent_registry import default_agent_registry, default_definition_registry
-    from autobyteus.llm.models import LLMModel
-    try: import autobyteus.agent.input_processor 
-    except ImportError: print("Warning: input_processor module not found.", file=sys.stderr)
-
-    def create_example_agent(name: str, role: str) -> Optional[Agent]:
-        agent_def = default_definition_registry.get(name, role)
-        if not agent_def:
-            try:
-                agent_def = AgentDefinition(name=name, role=role, description=f"{name} agent",
-                                            system_prompt=f"You are {name}.", tool_names=[],
-                                            input_processor_names=["PassthroughInputProcessor"])
-            except Exception as e_def:
-                logger.error(f"Failed to create def for {name}: {e_def}"); return None
-        
-        try: chosen_llm_model = next(iter(LLMModel)) if LLMModel else None
-        except StopIteration: chosen_llm_model = None
-        if not chosen_llm_model:
-            class MockLLMModel(str, Enum): DUMMY = "dummy_model" # type: ignore
-            LLMModel = MockLLMModel # type: ignore
-            chosen_llm_model = LLMModel.DUMMY
-        
-        if chosen_llm_model:
-            try: return default_agent_registry.create_agent(definition=agent_def, llm_model=chosen_llm_model)
-            except Exception as e_create: logger.error(f"Failed to create agent {name}: {e_create}", exc_info=True)
-        return None
-
-    agent1 = create_example_agent("AgentAlpha", "analyzer")
-    agent2 = create_example_agent("AgentBeta", "writer")
-
-    if agent1 and agent2:
-        agents_for_server = {
-            "alpha": agent1, # Routable ID "alpha" maps to AgentAlpha
-            "beta": agent2   # Routable ID "beta" maps to AgentBeta
-        }
-        # Example: Run the multi-agent SSE server
-        asyncio.run(serve_multiple_agents_http_sse(agents_for_server, "127.0.0.1", 8888, asyncio.Event()))
-        # To test, connect to http://127.0.0.1:8888/events/alpha and /events/beta for SSE
-        # Send POST to http://127.0.0.1:8888/rpc with {"target_agent_id": "alpha", ...}
-    else:
-        logger.error("Could not create example agents for multi-agent hosting demo.")
-

autobyteus/rpc/protocol.py

@@ -1,244 +0,0 @@
-# file: autobyteus/autobyteus/rpc/protocol.py
-import logging
-from enum import Enum
-from typing import Optional, Dict, Any, List, Union
-from pydantic import BaseModel, Field, validator
-import uuid
-
-logger = logging.getLogger(__name__)
-
-class MessageType(str, Enum):
-    """Type of the ProtocolMessage."""
-    REQUEST = "request"
-    RESPONSE = "response"
-    ERROR = "error"
-    EVENT = "event"  # For server-pushed events, e.g., via SSE
-
-class RequestType(str, Enum):
-    """Specific method/type for a 'request' message."""
-    DISCOVER_CAPABILITIES = "discover_capabilities"
-    INVOKE_METHOD = "invoke_method"
-    REQUEST_STREAM_DOWNLOAD = "request_stream_download" # New request type for stream download
-    # Other specific control messages can be added here
-
-class ResponseType(str, Enum):
-    """Specific type for a 'response' message, often mirroring RequestType."""
-    CAPABILITIES_RESPONSE = "capabilities_response"
-    METHOD_RESULT = "method_result"
-    ACKNOWLEDGEMENT = "acknowledgement" # Generic ack
-    STREAM_DOWNLOAD_READY = "stream_download_ready" # New response type for stream download
-
-class EventType(str, Enum):
-    """Specific type for an 'event' message (server-pushed)."""
-    AGENT_OUTPUT_CHUNK = "agent_output_chunk"
-    AGENT_FINAL_MESSAGE = "agent_final_message"
-    AGENT_STATUS_UPDATE = "agent_status_update"
-    TOOL_LOG_ENTRY = "tool_log_entry"
-    # Other event types
-
-class ErrorCode(Enum):
-    """Standard error codes for RPC communication."""
-    PARSE_ERROR = -32700
-    INVALID_REQUEST = -32600
-    METHOD_NOT_FOUND = -32601
-    INVALID_PARAMS = -32602
-    INTERNAL_ERROR = -32603
-    # Custom server errors: -32000 to -32099
-    SERVER_ERROR_AGENT_PROCESSING_FAILED = -32000
-    SERVER_ERROR_CAPABILITY_DISCOVERY_FAILED = -32001
-    SERVER_ERROR_UNAUTHORIZED = -32002 # If auth is added
-    SERVER_ERROR_TIMEOUT = -32003
-    SERVER_ERROR_STREAM_PREPARATION_FAILED = -32004 # For stream download errors
-    SERVER_ERROR_STREAM_NOT_FOUND = -32005 # If stream_id is invalid
-
-class ErrorDetails(BaseModel):
-    """Structure for error details within a ProtocolMessage."""
-    code: Union[ErrorCode, int] # Allow standard ErrorCode or custom int
-    message: str
-    data: Optional[Any] = None # Optional additional error data
-
-    @validator('code', pre=True)
-    def _validate_code(cls, v):
-        if isinstance(v, ErrorCode):
-            return v.value # Store the integer value of the enum
-        if isinstance(v, int):
-            return v
-        raise ValueError("Error code must be an ErrorCode enum member or an integer.")
-
-class ProtocolMessage(BaseModel):
-    """
-    Defines the structure for all RPC communications.
-    Based on JSON-RPC 2.0 concepts but adapted for AutoByteUs needs.
-    """
-    # Common fields for all message types
-    type: MessageType = Field(..., description="The main type of the message (request, response, error, event).")
-    id: Optional[str] = Field(default=None, description="Correlation ID for requests/responses. Can be string or number, UUID for simplicity here.")
-
-    # Fields specific to 'request' type
-    method: Optional[Union[RequestType, str]] = Field(default=None, description="Method name for 'request' type (e.g., 'discover_capabilities', 'invoke_method', or a custom agent method name).")
-    params: Optional[Dict[str, Any]] = Field(default=None, description="Parameters for the method in 'request' type.")
-
-    # Fields specific to 'response' type
-    result: Optional[Any] = Field(default=None, description="Payload of a successful 'response'.")
-    response_type: Optional[ResponseType] = Field(default=None, description="Specific type of a 'response', e.g., 'capabilities_response'.")
-
-
-    # Fields specific to 'error' type (when type is MessageType.ERROR)
-    error: Optional[ErrorDetails] = Field(default=None, description="Error details if the message type is 'error'.")
-
-    # Fields specific to 'event' type (server-pushed)
-    event_type: Optional[Union[EventType, str]] = Field(default=None, description="Specific type of 'event' (server-pushed).")
-    payload: Optional[Dict[str, Any]] = Field(default=None, description="Data payload for an 'event'.")
-    
-    # Auto-generated fields (optional, can be added by sender if needed)
-    # timestamp: datetime = Field(default_factory=lambda: datetime.now(timezone.utc))
-
-    class Config:
-        use_enum_values = True # Serialize enums to their values
-        validate_assignment = True
-
-    @validator('id', pre=True, always=True)
-    def set_default_id_for_requests(cls, v, values):
-        # Ensure requests have an ID, generate if not provided
-        if values.get('type') == MessageType.REQUEST and v is None:
-            return str(uuid.uuid4())
-        return v
-    
-    @validator('error')
-    def check_error_for_error_type(cls, v, values):
-        if values.get('type') == MessageType.ERROR and v is None:
-            raise ValueError("Field 'error' is required when message 'type' is ERROR.")
-        if values.get('type') != MessageType.ERROR and v is not None:
-            raise ValueError("Field 'error' should only be present when message 'type' is ERROR.")
-        return v
-
-    @validator('result')
-    def check_result_for_response_type(cls, v, values):
-        if values.get('type') == MessageType.RESPONSE and values.get('error') is not None:
-            raise ValueError("Fields 'result' and 'error' must not coexist in a RESPONSE message.")
-        if values.get('type') == MessageType.RESPONSE and v is None and values.get('error') is None:
-            # Allow responses with null result explicitly if that's valid (e.g. for simple ack)
-            # This validation might be too strict depending on use case.
-            # For now, permit null result for acknowledgements and stream_download_ready (which will be populated by server handler)
-            if values.get('response_type') not in [ResponseType.ACKNOWLEDGEMENT, ResponseType.STREAM_DOWNLOAD_READY]:
-                 logger.debug("Field 'result' is None for a non-acknowledgement/non-stream_download_ready RESPONSE message.")
-        return v
-        
-    @validator('method', 'params')
-    def check_request_fields(cls, v, field, values):
-        if values.get('type') == MessageType.REQUEST and v is None and field.name == 'method':
-            raise ValueError(f"Field '{field.name}' is required when message 'type' is REQUEST.")
-        if values.get('type') != MessageType.REQUEST and v is not None:
-            raise ValueError(f"Field '{field.name}' should only be present when message 'type' is REQUEST.")
-        return v
-
-    @validator('event_type', 'payload')
-    def check_event_fields(cls, v, field, values):
-        if values.get('type') == MessageType.EVENT and v is None and field.name == 'event_type':
-            raise ValueError(f"Field 'event_type' is required when message 'type' is EVENT.")
-        # Payload can be optional even for event type
-        if values.get('type') != MessageType.EVENT and v is not None:
-            raise ValueError(f"Field '{field.name}' should only be present when message 'type' is EVENT.")
-        return v
-
-    def to_json_str(self) -> str:
-        """Serializes the ProtocolMessage to a JSON string."""
-        return self.model_dump_json(exclude_none=True)
-
-    @classmethod
-    def from_json_str(cls, json_str: str) -> 'ProtocolMessage':
-        """Deserializes a ProtocolMessage from a JSON string."""
-        return cls.model_validate_json(json_str)
-
-    # Factory methods for convenience
-    @classmethod
-    def create_request(cls, method: Union[RequestType, str], params: Optional[Dict[str, Any]] = None, id: Optional[str] = None) -> 'ProtocolMessage':
-        return cls(type=MessageType.REQUEST, method=method, params=params or {}, id=id or str(uuid.uuid4()))
-
-    @classmethod
-    def create_response(cls, id: str, result: Any, response_type: Optional[ResponseType] = ResponseType.METHOD_RESULT) -> 'ProtocolMessage':
-        return cls(type=MessageType.RESPONSE, id=id, result=result, response_type=response_type)
-
-    @classmethod
-    def create_error_response(cls, id: Optional[str], code: Union[ErrorCode, int], message: str, data: Optional[Any] = None) -> 'ProtocolMessage':
-        actual_code = code.value if isinstance(code, ErrorCode) else code
-        return cls(type=MessageType.ERROR, id=id, error=ErrorDetails(code=actual_code, message=message, data=data))
-
-    @classmethod
-    def create_event(cls, event_type: Union[EventType, str], payload: Optional[Dict[str, Any]] = None) -> 'ProtocolMessage':
-        return cls(type=MessageType.EVENT, event_type=event_type, payload=payload or {})
-
-
-if __name__ == "__main__": # pragma: no cover
-    # Example Usage
-    logging.basicConfig(level=logging.DEBUG)
-
-    # Request
-    cap_request = ProtocolMessage.create_request(RequestType.DISCOVER_CAPABILITIES)
-    logger.info(f"Capability Request: {cap_request.to_json_str()}")
-
-    invoke_request_params = {"tool_name": "calculator", "args": {"operation": "add", "a": 5, "b": 3}}
-    invoke_request = ProtocolMessage.create_request(RequestType.INVOKE_METHOD, params=invoke_request_params)
-    logger.info(f"Invoke Request: {invoke_request.to_json_str()}")
-    
-    custom_invoke_req = ProtocolMessage.create_request("custom_agent_method", params={"input_data": "hello"})
-    logger.info(f"Custom Invoke Request: {custom_invoke_req.to_json_str()}")
-
-    stream_download_req = ProtocolMessage.create_request(RequestType.REQUEST_STREAM_DOWNLOAD, params={"resource_id": "large_file.dat"})
-    logger.info(f"Stream Download Request: {stream_download_req.to_json_str()}")
-
-    # Response
-    cap_response_result = {"agent_id": "server_agent_001", "capabilities": ["post_inter_agent_message", "get_status"]}
-    cap_response = ProtocolMessage.create_response(id=cap_request.id, result=cap_response_result, response_type=ResponseType.CAPABILITIES_RESPONSE)
-    logger.info(f"Capability Response: {cap_response.to_json_str()}")
-
-    invoke_response_result = {"output": 8}
-    invoke_response = ProtocolMessage.create_response(id=invoke_request.id, result=invoke_response_result)
-    logger.info(f"Invoke Response: {invoke_response.to_json_str()}")
-
-    ack_response = ProtocolMessage.create_response(id=custom_invoke_req.id, result=None, response_type=ResponseType.ACKNOWLEDGEMENT)
-    logger.info(f"Ack Response: {ack_response.to_json_str()}")
-
-    stream_download_ready_result = {"stream_id": "uuid-for-stream", "metadata": {"filename": "large_file.dat", "size": 1024000}}
-    # Note: download_url would be added by SseServerHandler before sending to client
-    stream_download_ready_resp = ProtocolMessage.create_response(id=stream_download_req.id, result=stream_download_ready_result, response_type=ResponseType.STREAM_DOWNLOAD_READY)
-    logger.info(f"Stream Download Ready Response (initial): {stream_download_ready_resp.to_json_str()}")
-
-
-    # Error Response
-    error_resp = ProtocolMessage.create_error_response(id=invoke_request.id, code=ErrorCode.METHOD_NOT_FOUND, message="Method not supported by agent.")
-    logger.info(f"Error Response: {error_resp.to_json_str()}")
-    
-    parse_error_resp = ProtocolMessage.create_error_response(id=None, code=ErrorCode.PARSE_ERROR, message="Invalid JSON received.")
-    logger.info(f"Parse Error Response (no ID): {parse_error_resp.to_json_str()}")
-
-    stream_error_resp = ProtocolMessage.create_error_response(
-        id=stream_download_req.id,
-        code=ErrorCode.SERVER_ERROR_STREAM_PREPARATION_FAILED,
-        message="Agent failed to prepare the stream for resource 'large_file.dat'."
-    )
-    logger.info(f"Stream Preparation Error Response: {stream_error_resp.to_json_str()}")
-
-    # Event
-    status_event = ProtocolMessage.create_event(EventType.AGENT_STATUS_UPDATE, payload={"status": "running"})
-    logger.info(f"Status Event: {status_event.to_json_str()}")
-    
-    # Test validation
-    try:
-        invalid_msg = ProtocolMessage(type=MessageType.REQUEST) # Missing method
-    except ValueError as e:
-        logger.error(f"Validation Error: {e}")
-
-    try:
-        invalid_error_msg = ProtocolMessage(type=MessageType.ERROR, id="123") # Missing error field
-    except ValueError as e:
-        logger.error(f"Validation Error: {e}")
-
-    try:
-        error_details_invalid_code = ErrorDetails(code="INVALID_CODE_STR", message="test")
-    except ValueError as e:
-        logger.error(f"Validation Error (ErrorDetails): {e}")
-    
-    error_details_valid_int_code = ErrorDetails(code=-32050, message="Custom server error")
-    logger.info(f"ErrorDetails with int code: {error_details_valid_int_code.model_dump_json()}")
-    assert error_details_valid_int_code.code == -32050

autobyteus/rpc/server/__init__.py

@@ -1,20 +0,0 @@
-# file: autobyteus/autobyteus/rpc/server/__init__.py
-"""
-Server-side components for the AutoByteUs RPC framework.
-These components enable an AgentRuntime to expose its capabilities remotely.
-"""
-from .base_method_handler import BaseMethodHandler
-from .method_handlers import DiscoverCapabilitiesHandler, InvokeMethodHandler, InitiateStreamDownloadHandler # Added InitiateStreamDownloadHandler
-from .stdio_server_handler import StdioServerHandler
-from .sse_server_handler import SseServerHandler 
-from .agent_server_endpoint import AgentServerEndpoint
-
-__all__ = [
-    "BaseMethodHandler",
-    "DiscoverCapabilitiesHandler",
-    "InvokeMethodHandler",
-    "InitiateStreamDownloadHandler", # Added InitiateStreamDownloadHandler
-    "StdioServerHandler",
-    "SseServerHandler", 
-    "AgentServerEndpoint",
-]