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

mindsdb/api/a2a/agent.py

@@ -0,0 +1,308 @@
+import json
+from typing import Any, AsyncIterable, Dict, List, Iterator
+import requests
+import logging
+
+from mindsdb.api.a2a.constants import DEFAULT_STREAM_TIMEOUT
+
+logger = logging.getLogger(__name__)
+
+
+class MindsDBAgent:
+    """An agent that communicates with MindsDB over HTTP following the A2A protocol."""
+
+    # Supported content-types according to A2A spec. We include both the
+    # mime-type form and the simple "text" token so that clients using either
+    # convention succeed.
+    SUPPORTED_CONTENT_TYPES = ["text", "text/plain", "application/json"]
+
+    def __init__(
+        self,
+        agent_name="my_agent",
+        project_name="mindsdb",
+        host="localhost",
+        port=47334,
+    ):
+        self.agent_name = agent_name
+        self.project_name = project_name
+        self.host = host
+        self.port = port
+        self.base_url = f"http://{host}:{port}"
+        self.agent_url = (
+            f"{self.base_url}/api/projects/{project_name}/agents/{agent_name}"
+        )
+        self.sql_url = f"{self.base_url}/api/sql/query"
+        logger.info(f"Initialized MindsDB agent connector to {self.base_url}")
+
+    def invoke(self, query, session_id) -> Dict[str, Any]:
+        """Send a query to the MindsDB agent using SQL API."""
+        try:
+            # Escape single quotes in the query for SQL
+            escaped_query = query.replace("'", "''")
+
+            # Build the SQL query to the agent
+            sql_query = f"SELECT * FROM {self.project_name}.{self.agent_name} WHERE question = '{escaped_query}'"
+
+            # Log request for debugging
+            logger.info(f"Sending SQL query to MindsDB: {sql_query[:100]}...")
+
+            # Send the request to MindsDB SQL API
+            response = requests.post(self.sql_url, json={"query": sql_query})
+            response.raise_for_status()
+
+            # Process the response
+            data = response.json()
+
+            # Log the response for debugging
+            logger.debug(f"Received response from MindsDB: {json.dumps(data)[:200]}...")
+
+            if "data" in data and len(data["data"]) > 0:
+                # The result should be in the first row
+                result_row = data["data"][0]
+
+                # Find the response column (might be 'response', 'answer', 'result', etc.)
+                # Try common column names or just return all content
+                for column in ["response", "result", "answer", "completion", "output"]:
+                    if column in result_row:
+                        content = result_row[column]
+                        logger.info(
+                            f"Found result in column '{column}': {content[:100]}..."
+                        )
+                        return {
+                            "content": content,
+                            "parts": [{"type": "text", "text": content}],
+                        }
+
+                # If no specific column found, return the whole row as JSON
+                logger.info("No specific result column found, returning full row")
+                content = json.dumps(result_row, indent=2)
+
+                # Return structured data only if it is a dictionary (A2A `data` part
+                # must itself be a JSON object).  In some cases MindsDB may return a
+                # list (for instance a list of rows or records).  If that happens we
+                # downgrade it to plain-text to avoid schema-validation errors on the
+                # A2A side.
+
+                parts: List[dict] = [{"type": "text", "text": content}]
+
+                if isinstance(result_row, dict):
+                    parts.append(
+                        {
+                            "type": "data",
+                            "data": result_row,
+                            "metadata": {"subtype": "json"},
+                        }
+                    )
+
+                return {
+                    "content": content,
+                    "parts": parts,
+                }
+            else:
+                error_msg = "Error: No data returned from MindsDB"
+                logger.error(error_msg)
+                return {
+                    "content": error_msg,
+                    "parts": [{"type": "text", "text": error_msg}],
+                }
+
+        except requests.exceptions.RequestException as e:
+            error_msg = f"Error connecting to MindsDB: {str(e)}"
+            logger.error(error_msg)
+            return {
+                "content": error_msg,
+                "parts": [{"type": "text", "text": error_msg}],
+            }
+
+        except Exception as e:
+            error_msg = f"Error: {str(e)}"
+            logger.error(error_msg)
+            return {
+                "content": error_msg,
+                "parts": [{"type": "text", "text": error_msg}],
+            }
+
+    def streaming_invoke(
+        self, messages: List[dict], timeout: int = DEFAULT_STREAM_TIMEOUT
+    ) -> Iterator[Dict[str, Any]]:
+        """Stream responses from the MindsDB agent using the direct API endpoint.
+
+        Args:
+            messages: List of message dictionaries, each containing 'question' and optionally 'answer'.
+                Example: [{'question': 'what is the average rental price for a three bedroom?', 'answer': None}]
+            timeout: Request timeout in seconds (default: 300)
+
+        Returns:
+            Iterator yielding chunks of the streaming response.
+        """
+        try:
+            # Construct the URL for the streaming completions endpoint
+            url = f"{self.base_url}/api/projects/{self.project_name}/agents/{self.agent_name}/completions/stream"
+
+            # Log request for debugging
+            logger.info(
+                f"Sending streaming request to MindsDB agent: {self.agent_name}"
+            )
+            logger.debug(f"Request messages: {json.dumps(messages)[:200]}...")
+
+            # Send the request to MindsDB streaming API with timeout
+            stream = requests.post(
+                url, json={"messages": messages}, stream=True, timeout=timeout
+            )
+            stream.raise_for_status()
+
+            # Process the streaming response directly
+            for line in stream.iter_lines():
+                if line:
+                    # Parse each non-empty line
+                    try:
+                        line = line.decode("utf-8")
+                        if line.startswith("data: "):
+                            # Extract the JSON data from the line that starts with 'data: '
+                            data = line[6:]  # Remove 'data: ' prefix
+                            try:
+                                chunk = json.loads(data)
+                                # Pass through the chunk with minimal modifications
+                                yield chunk
+                            except json.JSONDecodeError as e:
+                                logger.warning(
+                                    f"Failed to parse JSON from line: {data}. Error: {str(e)}"
+                                )
+                                # Yield error information but continue processing
+                                yield {
+                                    "error": f"JSON parse error: {str(e)}",
+                                    "data": data,
+                                    "is_task_complete": False,
+                                    "parts": [
+                                        {
+                                            "type": "text",
+                                            "text": f"Error parsing response: {str(e)}",
+                                        }
+                                    ],
+                                    "metadata": {},
+                                }
+                        else:
+                            # Log other lines for debugging
+                            logger.debug(f"Received non-data line: {line}")
+
+                            # If it looks like a raw text response (not SSE format), wrap it
+                            if not line.startswith("event:") and not line.startswith(
+                                ":"
+                            ):
+                                yield {"content": line, "is_task_complete": False}
+                    except UnicodeDecodeError as e:
+                        logger.warning(f"Failed to decode line: {str(e)}")
+                        # Continue processing despite decode errors
+
+        except requests.exceptions.Timeout as e:
+            error_msg = f"Request timed out after {timeout} seconds: {str(e)}"
+            logger.error(error_msg)
+            yield {
+                "content": error_msg,
+                "parts": [{"type": "text", "text": error_msg}],
+                "is_task_complete": True,
+                "error": "timeout",
+                "metadata": {"error": True},
+            }
+
+        except requests.exceptions.ChunkedEncodingError as e:
+            error_msg = f"Stream was interrupted: {str(e)}"
+            logger.error(error_msg)
+            yield {
+                "content": error_msg,
+                "parts": [{"type": "text", "text": error_msg}],
+                "is_task_complete": True,
+                "error": "stream_interrupted",
+                "metadata": {"error": True},
+            }
+
+        except requests.exceptions.ConnectionError as e:
+            error_msg = f"Connection error: {str(e)}"
+            logger.error(error_msg)
+            yield {
+                "content": error_msg,
+                "parts": [{"type": "text", "text": error_msg}],
+                "is_task_complete": True,
+                "error": "connection_error",
+                "metadata": {"error": True},
+            }
+
+        except requests.exceptions.RequestException as e:
+            error_msg = f"Error connecting to MindsDB streaming API: {str(e)}"
+            logger.error(error_msg)
+            yield {
+                "content": error_msg,
+                "parts": [{"type": "text", "text": error_msg}],
+                "is_task_complete": True,
+                "error": "request_error",
+                "metadata": {"error": True},
+            }
+
+        except Exception as e:
+            error_msg = f"Error in streaming: {str(e)}"
+            logger.error(error_msg)
+            yield {
+                "content": error_msg,
+                "parts": [{"type": "text", "text": error_msg}],
+                "is_task_complete": True,
+                "error": "unknown_error",
+                "metadata": {"error": True},
+            }
+
+        # Send a final completion message
+        yield {"is_task_complete": True, "metadata": {"complete": True}}
+
+    async def stream(self, query, session_id) -> AsyncIterable[Dict[str, Any]]:
+        """Stream responses from the MindsDB agent (uses streaming API endpoint)."""
+        try:
+            logger.info(f"Using streaming API for query: {query[:100]}...")
+
+            # Format the query into the message structure expected by streaming_invoke
+            messages = [{"question": query, "answer": None}]
+
+            # Use the streaming_invoke method to get real streaming responses
+            streaming_response = self.streaming_invoke(messages)
+
+            # Yield all chunks directly from the streaming response
+            for chunk in streaming_response:
+                # Only add required fields if they don't exist
+                # This preserves the original structure as much as possible
+                if "is_task_complete" not in chunk:
+                    chunk["is_task_complete"] = False
+
+                if "metadata" not in chunk:
+                    chunk["metadata"] = {}
+
+                # Ensure parts exist, but try to preserve original content
+                if "parts" not in chunk:
+                    # If content exists, create a part from it
+                    if "content" in chunk:
+                        chunk["parts"] = [{"type": "text", "text": chunk["content"]}]
+                    # If output exists, create a part from it
+                    elif "output" in chunk:
+                        chunk["parts"] = [{"type": "text", "text": chunk["output"]}]
+                    # If actions exist, create empty parts
+                    elif "actions" in chunk or "steps" in chunk or "messages" in chunk:
+                        # These chunks have their own format, just add empty parts
+                        chunk["parts"] = []
+                    else:
+                        # Skip chunks with no content
+                        continue
+
+                yield chunk
+
+        except Exception as e:
+            logger.error(f"Error in streaming: {str(e)}")
+            yield {
+                "is_task_complete": True,
+                "parts": [
+                    {
+                        "type": "text",
+                        "text": f"Error: {str(e)}",
+                    }
+                ],
+                "metadata": {
+                    "type": "reasoning",
+                    "subtype": "error",
+                },
+            }

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

@@ -0,0 +1,4 @@
+from .server import A2AServer
+from .task_manager import TaskManager, InMemoryTaskManager
+
+__all__ = ["A2AServer", "TaskManager", "InMemoryTaskManager"]

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

@@ -0,0 +1,164 @@
+from starlette.applications import Starlette
+from starlette.middleware.cors import CORSMiddleware
+from starlette.responses import JSONResponse
+from sse_starlette.sse import EventSourceResponse
+from starlette.requests import Request
+from ...common.types import (
+    A2ARequest,
+    JSONRPCResponse,
+    InvalidRequestError,
+    JSONParseError,
+    GetTaskRequest,
+    CancelTaskRequest,
+    SendTaskRequest,
+    SetTaskPushNotificationRequest,
+    GetTaskPushNotificationRequest,
+    InternalError,
+    AgentCard,
+    TaskResubscriptionRequest,
+    SendTaskStreamingRequest,
+)
+from pydantic import ValidationError
+import json
+from typing import AsyncIterable, Any
+from ...common.server.task_manager import TaskManager
+
+import logging
+
+logger = logging.getLogger(__name__)
+
+
+class A2AServer:
+    def __init__(
+        self,
+        host="0.0.0.0",
+        port=5000,
+        endpoint="/",
+        agent_card: AgentCard = None,
+        task_manager: TaskManager = None,
+    ):
+        self.host = host
+        self.port = port
+        self.endpoint = endpoint
+        self.task_manager = task_manager
+        self.agent_card = agent_card
+        self.app = Starlette()
+        self.app.add_route(self.endpoint, self._process_request, methods=["POST"])
+        self.app.add_route(
+            "/.well-known/agent.json", self._get_agent_card, methods=["GET"]
+        )
+        # TODO: Remove this when we have a proper CORS policy
+        self.app.add_middleware(
+            CORSMiddleware,
+            allow_origins=["*"],
+            allow_credentials=True,
+            allow_methods=["*"],
+            allow_headers=["*"],
+        )
+
+    def start(self):
+        if self.agent_card is None:
+            raise ValueError("agent_card is not defined")
+
+        if self.task_manager is None:
+            raise ValueError("request_handler is not defined")
+
+        import uvicorn
+
+        # Configure uvicorn with optimized settings for streaming
+        uvicorn.run(
+            self.app,
+            host=self.host,
+            port=self.port,
+            http="h11",
+            timeout_keep_alive=65,
+            log_level="info"
+        )
+
+    def _get_agent_card(self, request: Request) -> JSONResponse:
+        return JSONResponse(self.agent_card.model_dump(exclude_none=True))
+
+    async def _process_request(self, request: Request):
+        try:
+            body = await request.json()
+            json_rpc_request = A2ARequest.validate_python(body)
+
+            if isinstance(json_rpc_request, GetTaskRequest):
+                result = await self.task_manager.on_get_task(json_rpc_request)
+            elif isinstance(json_rpc_request, SendTaskRequest):
+                result = await self.task_manager.on_send_task(json_rpc_request)
+            elif isinstance(json_rpc_request, SendTaskStreamingRequest):
+                # Don't await the async generator, just pass it to _create_response
+                result = self.task_manager.on_send_task_subscribe(
+                    json_rpc_request
+                )
+            elif isinstance(json_rpc_request, CancelTaskRequest):
+                result = await self.task_manager.on_cancel_task(json_rpc_request)
+            elif isinstance(json_rpc_request, SetTaskPushNotificationRequest):
+                result = await self.task_manager.on_set_task_push_notification(
+                    json_rpc_request
+                )
+            elif isinstance(json_rpc_request, GetTaskPushNotificationRequest):
+                result = await self.task_manager.on_get_task_push_notification(
+                    json_rpc_request
+                )
+            elif isinstance(json_rpc_request, TaskResubscriptionRequest):
+                result = await self.task_manager.on_resubscribe_to_task(
+                    json_rpc_request
+                )
+            else:
+                logger.warning(f"Unexpected request type: {type(json_rpc_request)}")
+                raise ValueError(f"Unexpected request type: {type(request)}")
+
+            return self._create_response(result)
+
+        except Exception as e:
+            return self._handle_exception(e)
+
+    def _handle_exception(self, e: Exception) -> JSONResponse:
+        if isinstance(e, json.decoder.JSONDecodeError):
+            json_rpc_error = JSONParseError()
+        elif isinstance(e, ValidationError):
+            json_rpc_error = InvalidRequestError(data=json.loads(e.json()))
+        else:
+            logger.error(f"Unhandled exception: {e}")
+            json_rpc_error = InternalError()
+
+        response = JSONRPCResponse(id=None, error=json_rpc_error)
+        return JSONResponse(response.model_dump(exclude_none=True), status_code=400)
+
+    def _create_response(self, result: Any) -> JSONResponse | EventSourceResponse:
+        if isinstance(result, AsyncIterable):
+
+            async def event_generator(result) -> AsyncIterable[dict[str, str]]:
+                async for item in result:
+                    # Send the data event with immediate flush directive
+                    yield {
+                        "data": item.model_dump_json(exclude_none=True),
+                        "event": "message",
+                        "id": str(id(item)),  # Add a unique ID for each event
+                    }
+                    # Add an empty comment event to force flush
+                    yield {
+                        "comment": " ",  # Empty comment event to force flush
+                    }
+
+            # Create EventSourceResponse with complete headers for browser compatibility
+            return EventSourceResponse(
+                event_generator(result),
+                # Complete set of headers needed for browser streaming
+                headers={
+                    "Cache-Control": "no-cache, no-transform",
+                    "X-Accel-Buffering": "no",
+                    "Connection": "keep-alive",
+                    "Content-Type": "text/event-stream",
+                    "Transfer-Encoding": "chunked"
+                },
+                # Explicitly set media_type
+                media_type="text/event-stream"
+            )
+        elif isinstance(result, JSONRPCResponse):
+            return JSONResponse(result.model_dump(exclude_none=True))
+        else:
+            logger.error(f"Unexpected result type: {type(result)}")
+            raise ValueError(f"Unexpected result type: {type(result)}")