MindsDB 25.5.3.0__py3-none-any.whl → 25.5.4.0__py3-none-any.whl

mindsdb/api/a2a/common/server/task_manager.py

@@ -0,0 +1,287 @@
+from abc import ABC, abstractmethod
+from typing import Union, AsyncIterable, List
+from ...common.types import Task
+from ...common.types import (
+    JSONRPCResponse,
+    TaskIdParams,
+    TaskQueryParams,
+    GetTaskRequest,
+    TaskNotFoundError,
+    SendTaskRequest,
+    CancelTaskRequest,
+    TaskNotCancelableError,
+    SetTaskPushNotificationRequest,
+    GetTaskPushNotificationRequest,
+    GetTaskResponse,
+    CancelTaskResponse,
+    SendTaskResponse,
+    SetTaskPushNotificationResponse,
+    GetTaskPushNotificationResponse,
+    TaskSendParams,
+    TaskStatus,
+    TaskState,
+    TaskResubscriptionRequest,
+    SendTaskStreamingRequest,
+    SendTaskStreamingResponse,
+    Artifact,
+    PushNotificationConfig,
+    TaskStatusUpdateEvent,
+    JSONRPCError,
+    TaskPushNotificationConfig,
+    InternalError,
+)
+from ...common.server.utils import new_not_implemented_error
+import asyncio
+import logging
+
+logger = logging.getLogger(__name__)
+
+
+class TaskManager(ABC):
+    @abstractmethod
+    async def on_get_task(self, request: GetTaskRequest) -> GetTaskResponse:
+        pass
+
+    @abstractmethod
+    async def on_cancel_task(self, request: CancelTaskRequest) -> CancelTaskResponse:
+        pass
+
+    @abstractmethod
+    async def on_send_task(self, request: SendTaskRequest) -> SendTaskResponse:
+        pass
+
+    @abstractmethod
+    async def on_send_task_subscribe(
+        self, request: SendTaskStreamingRequest
+    ) -> Union[AsyncIterable[SendTaskStreamingResponse], JSONRPCResponse]:
+        pass
+
+    @abstractmethod
+    async def on_set_task_push_notification(
+        self, request: SetTaskPushNotificationRequest
+    ) -> SetTaskPushNotificationResponse:
+        pass
+
+    @abstractmethod
+    async def on_get_task_push_notification(
+        self, request: GetTaskPushNotificationRequest
+    ) -> GetTaskPushNotificationResponse:
+        pass
+
+    @abstractmethod
+    async def on_resubscribe_to_task(
+        self, request: TaskResubscriptionRequest
+    ) -> Union[AsyncIterable[SendTaskResponse], JSONRPCResponse]:
+        pass
+
+
+class InMemoryTaskManager(TaskManager):
+    def __init__(self):
+        self.tasks: dict[str, Task] = {}
+        self.push_notification_infos: dict[str, PushNotificationConfig] = {}
+        self.lock = asyncio.Lock()
+        self.task_sse_subscribers: dict[str, List[asyncio.Queue]] = {}
+        self.subscriber_lock = asyncio.Lock()
+
+    async def on_get_task(self, request: GetTaskRequest) -> GetTaskResponse:
+        logger.info(f"Getting task {request.params.id}")
+        task_query_params: TaskQueryParams = request.params
+
+        async with self.lock:
+            task = self.tasks.get(task_query_params.id)
+            if task is None:
+                return GetTaskResponse(id=request.id, error=TaskNotFoundError())
+
+            task_result = self.append_task_history(
+                task, task_query_params.historyLength
+            )
+
+        return GetTaskResponse(id=request.id, result=task_result)
+
+    async def on_cancel_task(self, request: CancelTaskRequest) -> CancelTaskResponse:
+        logger.info(f"Cancelling task {request.params.id}")
+        task_id_params: TaskIdParams = request.params
+
+        async with self.lock:
+            task = self.tasks.get(task_id_params.id)
+            if task is None:
+                return CancelTaskResponse(id=request.id, error=TaskNotFoundError())
+
+        return CancelTaskResponse(id=request.id, error=TaskNotCancelableError())
+
+    @abstractmethod
+    async def on_send_task(self, request: SendTaskRequest) -> SendTaskResponse:
+        pass
+
+    @abstractmethod
+    async def on_send_task_subscribe(
+        self, request: SendTaskStreamingRequest
+    ) -> Union[AsyncIterable[SendTaskStreamingResponse], JSONRPCResponse]:
+        pass
+
+    async def set_push_notification_info(
+        self, task_id: str, notification_config: PushNotificationConfig
+    ):
+        async with self.lock:
+            task = self.tasks.get(task_id)
+            if task is None:
+                raise ValueError(f"Task not found for {task_id}")
+
+            self.push_notification_infos[task_id] = notification_config
+
+        return
+
+    async def get_push_notification_info(self, task_id: str) -> PushNotificationConfig:
+        async with self.lock:
+            task = self.tasks.get(task_id)
+            if task is None:
+                raise ValueError(f"Task not found for {task_id}")
+
+            return self.push_notification_infos[task_id]
+
+        return
+
+    async def has_push_notification_info(self, task_id: str) -> bool:
+        async with self.lock:
+            return task_id in self.push_notification_infos
+
+    async def on_set_task_push_notification(
+        self, request: SetTaskPushNotificationRequest
+    ) -> SetTaskPushNotificationResponse:
+        logger.info(f"Setting task push notification {request.params.id}")
+        task_notification_params: TaskPushNotificationConfig = request.params
+
+        try:
+            await self.set_push_notification_info(
+                task_notification_params.id,
+                task_notification_params.pushNotificationConfig,
+            )
+        except Exception as e:
+            logger.error(f"Error while setting push notification info: {e}")
+            return JSONRPCResponse(
+                id=request.id,
+                error=InternalError(
+                    message="An error occurred while setting push notification info"
+                ),
+            )
+
+        return SetTaskPushNotificationResponse(
+            id=request.id, result=task_notification_params
+        )
+
+    async def on_get_task_push_notification(
+        self, request: GetTaskPushNotificationRequest
+    ) -> GetTaskPushNotificationResponse:
+        logger.info(f"Getting task push notification {request.params.id}")
+        task_params: TaskIdParams = request.params
+
+        try:
+            notification_info = await self.get_push_notification_info(task_params.id)
+        except Exception as e:
+            logger.error(f"Error while getting push notification info: {e}")
+            return GetTaskPushNotificationResponse(
+                id=request.id,
+                error=InternalError(
+                    message="An error occurred while getting push notification info"
+                ),
+            )
+
+        return GetTaskPushNotificationResponse(
+            id=request.id,
+            result=TaskPushNotificationConfig(
+                id=task_params.id, pushNotificationConfig=notification_info
+            ),
+        )
+
+    async def upsert_task(self, task_send_params: TaskSendParams) -> Task:
+        logger.info(f"Upserting task {task_send_params.id}")
+        async with self.lock:
+            task = self.tasks.get(task_send_params.id)
+            if task is None:
+                task = Task(
+                    id=task_send_params.id,
+                    sessionId=task_send_params.sessionId,
+                    messages=[task_send_params.message],
+                    status=TaskStatus(state=TaskState.SUBMITTED),
+                    history=[task_send_params.message],
+                )
+                self.tasks[task_send_params.id] = task
+            else:
+                task.history.append(task_send_params.message)
+
+            return task
+
+    async def on_resubscribe_to_task(
+        self, request: TaskResubscriptionRequest
+    ) -> Union[AsyncIterable[SendTaskStreamingResponse], JSONRPCResponse]:
+        return new_not_implemented_error(request.id)
+
+    async def update_store(
+        self, task_id: str, status: TaskStatus, artifacts: list[Artifact]
+    ) -> Task:
+        async with self.lock:
+            try:
+                task = self.tasks[task_id]
+            except KeyError:
+                logger.error(f"Task {task_id} not found for updating the task")
+                raise ValueError(f"Task {task_id} not found")
+
+            task.status = status
+
+            if status.message is not None:
+                task.history.append(status.message)
+
+            if artifacts is not None:
+                if task.artifacts is None:
+                    task.artifacts = []
+                task.artifacts.extend(artifacts)
+
+            return task
+
+    def append_task_history(self, task: Task, historyLength: int | None):
+        new_task = task.model_copy()
+        if historyLength is not None and historyLength > 0:
+            new_task.history = new_task.history[-historyLength:]
+        else:
+            new_task.history = []
+
+        return new_task
+
+    async def setup_sse_consumer(self, task_id: str, is_resubscribe: bool = False):
+        async with self.subscriber_lock:
+            if task_id not in self.task_sse_subscribers:
+                if is_resubscribe:
+                    raise ValueError("Task not found for resubscription")
+                else:
+                    self.task_sse_subscribers[task_id] = []
+
+            sse_event_queue = asyncio.Queue(maxsize=0)  # <=0 is unlimited
+            self.task_sse_subscribers[task_id].append(sse_event_queue)
+            return sse_event_queue
+
+    async def enqueue_events_for_sse(self, task_id, task_update_event):
+        async with self.subscriber_lock:
+            if task_id not in self.task_sse_subscribers:
+                return
+
+            current_subscribers = self.task_sse_subscribers[task_id]
+            for subscriber in current_subscribers:
+                await subscriber.put(task_update_event)
+
+    async def dequeue_events_for_sse(
+        self, request_id, task_id, sse_event_queue: asyncio.Queue
+    ) -> AsyncIterable[SendTaskStreamingResponse] | JSONRPCResponse:
+        try:
+            while True:
+                event = await sse_event_queue.get()
+                if isinstance(event, JSONRPCError):
+                    yield SendTaskStreamingResponse(id=request_id, error=event)
+                    break
+
+                yield SendTaskStreamingResponse(id=request_id, result=event)
+                if isinstance(event, TaskStatusUpdateEvent) and event.final:
+                    break
+        finally:
+            async with self.subscriber_lock:
+                if task_id in self.task_sse_subscribers:
+                    self.task_sse_subscribers[task_id].remove(sse_event_queue)

mindsdb/api/a2a/common/server/utils.py

@@ -0,0 +1,28 @@
+from ...common.types import (
+    JSONRPCResponse,
+    ContentTypeNotSupportedError,
+    UnsupportedOperationError,
+)
+from typing import List
+
+
+def are_modalities_compatible(
+    server_output_modes: List[str], client_output_modes: List[str]
+):
+    """Modalities are compatible if they are both non-empty
+    and there is at least one common element."""
+    if client_output_modes is None or len(client_output_modes) == 0:
+        return True
+
+    if server_output_modes is None or len(server_output_modes) == 0:
+        return True
+
+    return any(x in server_output_modes for x in client_output_modes)
+
+
+def new_incompatible_types_error(request_id):
+    return JSONRPCResponse(id=request_id, error=ContentTypeNotSupportedError())
+
+
+def new_not_implemented_error(request_id):
+    return JSONRPCResponse(id=request_id, error=UnsupportedOperationError())

mindsdb/api/a2a/common/types.py

@@ -0,0 +1,365 @@
+from typing import Union, Any
+from pydantic import BaseModel, Field, TypeAdapter
+from typing import Literal, List, Annotated, Optional
+from datetime import datetime
+from pydantic import model_validator, ConfigDict, field_serializer
+from uuid import uuid4
+from enum import Enum
+from typing_extensions import Self
+
+
+class TaskState(str, Enum):
+    SUBMITTED = "submitted"
+    WORKING = "working"
+    INPUT_REQUIRED = "input-required"
+    COMPLETED = "completed"
+    CANCELED = "canceled"
+    FAILED = "failed"
+    UNKNOWN = "unknown"
+
+
+class TextPart(BaseModel):
+    type: Literal["text"] = "text"
+    text: str
+    metadata: dict[str, Any] | None = None
+
+
+class FileContent(BaseModel):
+    name: str | None = None
+    mimeType: str | None = None
+    bytes: str | None = None
+    uri: str | None = None
+
+    @model_validator(mode="after")
+    def check_content(self) -> Self:
+        if not (self.bytes or self.uri):
+            raise ValueError("Either 'bytes' or 'uri' must be present in the file data")
+        if self.bytes and self.uri:
+            raise ValueError(
+                "Only one of 'bytes' or 'uri' can be present in the file data"
+            )
+        return self
+
+
+class FilePart(BaseModel):
+    type: Literal["file"] = "file"
+    file: FileContent
+    metadata: dict[str, Any] | None = None
+
+
+class DataPart(BaseModel):
+    type: Literal["data"] = "data"
+    data: dict[str, Any]
+    metadata: dict[str, Any] | None = None
+
+
+Part = Annotated[Union[TextPart, FilePart, DataPart], Field(discriminator="type")]
+
+
+class Message(BaseModel):
+    role: Literal["user", "agent"]
+    parts: List[Part]
+    metadata: dict[str, Any] | None = None
+
+
+class TaskStatus(BaseModel):
+    state: TaskState
+    message: Message | None = None
+    timestamp: datetime = Field(default_factory=datetime.now)
+
+    @field_serializer("timestamp")
+    def serialize_dt(self, dt: datetime, _info):
+        return dt.isoformat()
+
+
+class Artifact(BaseModel):
+    name: str | None = None
+    description: str | None = None
+    parts: List[Part]
+    metadata: dict[str, Any] | None = None
+    index: int = 0
+    append: bool | None = None
+    lastChunk: bool | None = None
+
+
+class Task(BaseModel):
+    id: str
+    sessionId: str | None = None
+    status: TaskStatus
+    artifacts: List[Artifact] | None = None
+    history: List[Message] | None = None
+    metadata: dict[str, Any] | None = None
+
+
+class TaskStatusUpdateEvent(BaseModel):
+    id: str
+    status: TaskStatus
+    final: bool = False
+    metadata: dict[str, Any] | None = None
+
+
+class TaskArtifactUpdateEvent(BaseModel):
+    id: str
+    artifact: Artifact
+    metadata: dict[str, Any] | None = None
+
+
+class AuthenticationInfo(BaseModel):
+    model_config = ConfigDict(extra="allow")
+
+    schemes: List[str]
+    credentials: str | None = None
+
+
+class PushNotificationConfig(BaseModel):
+    url: str
+    token: str | None = None
+    authentication: AuthenticationInfo | None = None
+
+
+class TaskIdParams(BaseModel):
+    id: str
+    metadata: dict[str, Any] | None = None
+
+
+class TaskQueryParams(TaskIdParams):
+    historyLength: int | None = None
+
+
+class TaskSendParams(BaseModel):
+    id: str
+    sessionId: str = Field(default_factory=lambda: uuid4().hex)
+    message: Message
+    acceptedOutputModes: Optional[List[str]] = None
+    pushNotification: PushNotificationConfig | None = None
+    historyLength: int | None = None
+    metadata: dict[str, Any] | None = None
+
+
+class TaskPushNotificationConfig(BaseModel):
+    id: str
+    pushNotificationConfig: PushNotificationConfig
+
+
+# RPC Messages
+
+
+class JSONRPCMessage(BaseModel):
+    jsonrpc: Literal["2.0"] = "2.0"
+    id: int | str | None = Field(default_factory=lambda: uuid4().hex)
+
+
+class JSONRPCRequest(JSONRPCMessage):
+    method: str
+    params: dict[str, Any] | None = None
+
+
+class JSONRPCError(BaseModel):
+    code: int
+    message: str
+    data: Any | None = None
+
+
+class JSONRPCResponse(JSONRPCMessage):
+    result: Any | None = None
+    error: JSONRPCError | None = None
+
+
+class SendTaskRequest(JSONRPCRequest):
+    method: Literal["tasks/send"] = "tasks/send"
+    params: TaskSendParams
+
+
+class SendTaskResponse(JSONRPCResponse):
+    result: Task | None = None
+
+
+class SendTaskStreamingRequest(JSONRPCRequest):
+    method: Literal["tasks/sendSubscribe"] = "tasks/sendSubscribe"
+    params: TaskSendParams
+
+
+class SendTaskStreamingResponse(JSONRPCResponse):
+    result: TaskStatusUpdateEvent | TaskArtifactUpdateEvent | None = None
+
+
+class GetTaskRequest(JSONRPCRequest):
+    method: Literal["tasks/get"] = "tasks/get"
+    params: TaskQueryParams
+
+
+class GetTaskResponse(JSONRPCResponse):
+    result: Task | None = None
+
+
+class CancelTaskRequest(JSONRPCRequest):
+    method: Literal["tasks/cancel",] = "tasks/cancel"
+    params: TaskIdParams
+
+
+class CancelTaskResponse(JSONRPCResponse):
+    result: Task | None = None
+
+
+class SetTaskPushNotificationRequest(JSONRPCRequest):
+    method: Literal["tasks/pushNotification/set",] = "tasks/pushNotification/set"
+    params: TaskPushNotificationConfig
+
+
+class SetTaskPushNotificationResponse(JSONRPCResponse):
+    result: TaskPushNotificationConfig | None = None
+
+
+class GetTaskPushNotificationRequest(JSONRPCRequest):
+    method: Literal["tasks/pushNotification/get",] = "tasks/pushNotification/get"
+    params: TaskIdParams
+
+
+class GetTaskPushNotificationResponse(JSONRPCResponse):
+    result: TaskPushNotificationConfig | None = None
+
+
+class TaskResubscriptionRequest(JSONRPCRequest):
+    method: Literal["tasks/resubscribe",] = "tasks/resubscribe"
+    params: TaskIdParams
+
+
+A2ARequest = TypeAdapter(
+    Annotated[
+        Union[
+            SendTaskRequest,
+            GetTaskRequest,
+            CancelTaskRequest,
+            SetTaskPushNotificationRequest,
+            GetTaskPushNotificationRequest,
+            TaskResubscriptionRequest,
+            SendTaskStreamingRequest,
+        ],
+        Field(discriminator="method"),
+    ]
+)
+
+# Error types
+
+
+class JSONParseError(JSONRPCError):
+    code: int = -32700
+    message: str = "Invalid JSON payload"
+    data: Any | None = None
+
+
+class InvalidRequestError(JSONRPCError):
+    code: int = -32600
+    message: str = "Request payload validation error"
+    data: Any | None = None
+
+
+class MethodNotFoundError(JSONRPCError):
+    code: int = -32601
+    message: str = "Method not found"
+    data: None = None
+
+
+class InvalidParamsError(JSONRPCError):
+    code: int = -32602
+    message: str = "Invalid parameters"
+    data: Any | None = None
+
+
+class InternalError(JSONRPCError):
+    code: int = -32603
+    message: str = "Internal error"
+    data: Any | None = None
+
+
+class TaskNotFoundError(JSONRPCError):
+    code: int = -32001
+    message: str = "Task not found"
+    data: None = None
+
+
+class TaskNotCancelableError(JSONRPCError):
+    code: int = -32002
+    message: str = "Task cannot be canceled"
+    data: None = None
+
+
+class PushNotificationNotSupportedError(JSONRPCError):
+    code: int = -32003
+    message: str = "Push Notification is not supported"
+    data: None = None
+
+
+class UnsupportedOperationError(JSONRPCError):
+    code: int = -32004
+    message: str = "This operation is not supported"
+    data: None = None
+
+
+class ContentTypeNotSupportedError(JSONRPCError):
+    code: int = -32005
+    message: str = "Incompatible content types"
+    data: None = None
+
+
+class AgentProvider(BaseModel):
+    organization: str
+    url: str | None = None
+
+
+class AgentCapabilities(BaseModel):
+    streaming: bool = False
+    pushNotifications: bool = False
+    stateTransitionHistory: bool = False
+
+
+class AgentAuthentication(BaseModel):
+    schemes: List[str]
+    credentials: str | None = None
+
+
+class AgentSkill(BaseModel):
+    id: str
+    name: str
+    description: str | None = None
+    tags: List[str] | None = None
+    examples: List[str] | None = None
+    inputModes: List[str] | None = None
+    outputModes: List[str] | None = None
+
+
+class AgentCard(BaseModel):
+    name: str
+    description: str | None = None
+    url: str
+    provider: AgentProvider | None = None
+    version: str
+    documentationUrl: str | None = None
+    capabilities: AgentCapabilities
+    authentication: AgentAuthentication | None = None
+    defaultInputModes: List[str] = ["text"]
+    defaultOutputModes: List[str] = ["text"]
+    skills: List[AgentSkill]
+
+
+class A2AClientError(Exception):
+    pass
+
+
+class A2AClientHTTPError(A2AClientError):
+    def __init__(self, status_code: int, message: str):
+        self.status_code = status_code
+        self.message = message
+        super().__init__(f"HTTP Error {status_code}: {message}")
+
+
+class A2AClientJSONError(A2AClientError):
+    def __init__(self, message: str):
+        self.message = message
+        super().__init__(f"JSON Error: {message}")
+
+
+class MissingAPIKeyError(Exception):
+    """Exception for missing API key."""
+
+    pass

mindsdb/api/a2a/constants.py

@@ -0,0 +1,9 @@
+TEXT_BY_SUBTYPE = {
+    "plan": "To determine the distribution of companies by size, we first need to identify the relevant table and column that contains information about company sizes. Let's start by listing the tables in the database to find the appropriate one.",
+    "query": "Sales_Data_Expert_Demo_Data.public.`call_summaries`, Sales_Data_Expert_Demo_Data.public.`prospects_details`",
+    "curate": "To find the distribution of companies by size, we should examine the `prospects_details` table, as it likely contains information about potential customers, including company size. Let's first check the schema of the `prospects_details` table to identify the relevant column for company size.",
+    "validate": "Table named `Sales_Data_Expert_Demo_Data.public.prospects_details`:\n\nSample with first 3 rows from table Sales_Data_Expert_Demo_Data.public.prospects_details in CSV format (dialect is 'excel'):\n\nid,created_at,first_name,last_name,email,company,job_title,company_size,use_case\r\n1,2025-01-16 13:55:03,Alexandra,Dixon,alexandra.dixon@harveypayneandhudson.com,\"Harvey, Payne and Hudson\",Manager,Medium company (500-999 employees),Looking into predictive analytics for customer behavior\r\n2,2025-01-28 08:01:46,David,Clark,david.clark@boyd-hernandez.com,Boyd-Hernandez,CEO,Medium company (500-999 employees),Looking into predictive analytics for customer behavior\r\n3,2025-02-03 03:37:12,Jason,Mcmillan,jason.mcmillan@stokesandsons.com,Stokes and Sons,Marketing Manager,Small company (100-499 employees),Interested in AI search over my data\r\n\n\nColumn data types: \n`id` : `MYSQL_DATA_TYPE.INT`,\t\n`created_at` : `MYSQL_DATA_TYPE.DATETIME`,\t\n`first_name` : `MYSQL_DATA_TYPE.TEXT`,\t\n`last_name` : `MYSQL_DATA_TYPE.TEXT`,\t\n`email` : `MYSQL_DATA_TYPE.TEXT`,\t\n`company` : `MYSQL_DATA_TYPE.TEXT`,\t\n`job_title` : `MYSQL_DATA_TYPE.TEXT`,\t\n`company_size` : `MYSQL_DATA_TYPE.TEXT`,\t\n`use_case` : `MYSQL_DATA_TYPE.TEXT`\n",
+    "respond": "The `prospects_details` table contains a `company_size` column, which is what we need to determine the distribution of companies by size. Let's first retrieve the distinct values in the `company_size` column to understand the categories available, and then calculate the distribution.",
+}
+
+DEFAULT_STREAM_TIMEOUT = 300