MindsDB 25.7.1.0__py3-none-any.whl → 25.7.2.0__py3-none-any.whl

mindsdb/__about__.py

@@ -1,6 +1,6 @@
 __title__ = "MindsDB"
 __package_name__ = "mindsdb"
-__version__ = "25.7.1.0"
+__version__ = "25.7.2.0"
 __description__ = "MindsDB's AI SQL Server enables developers to build AI tools that need access to real-time data to perform their tasks"
 __email__ = "jorge@mindsdb.com"
 __author__ = "MindsDB Inc"

mindsdb/__main__.py

@@ -1,4 +1,5 @@
 import gc
+
 gc.disable()
 import os
 import sys
@@ -105,9 +106,7 @@ class TrunkProcessData:
         self._restarts_time.append(current_time_seconds)
         if self.max_restart_interval_seconds > 0:
             self._restarts_time = [
-                x
-                for x in self._restarts_time
-                if x >= (current_time_seconds - self.max_restart_interval_seconds)
+                x for x in self._restarts_time if x >= (current_time_seconds - self.max_restart_interval_seconds)
             ]
         if len(self._restarts_time) > self.max_restart_count:
             return False
@@ -124,16 +123,11 @@ class TrunkProcessData:
         if config.is_cloud:
             return False
         if sys.platform in ("linux", "darwin"):
-            return (
-                self.restart_on_failure
-                and self.process.exitcode == -signal.SIGKILL.value
-            )
+            return self.restart_on_failure and self.process.exitcode == -signal.SIGKILL.value
         else:
             if self.max_restart_count == 0:
                 # to prevent infinity restarts, max_restart_count should be > 0
-                logger.warning(
-                    "In the current OS, it is not possible to use `max_restart_count=0`"
-                )
+                logger.warning("In the current OS, it is not possible to use `max_restart_count=0`")
                 return False
             return self.restart_on_failure
 
@@ -174,24 +168,18 @@ def set_error_model_status_by_pids(unexisting_pids: List[int]):
         db.session.query(db.Predictor)
         .filter(
             db.Predictor.deleted_at.is_(None),
-            db.Predictor.status.not_in(
-                [db.PREDICTOR_STATUS.COMPLETE, db.PREDICTOR_STATUS.ERROR]
-            ),
+            db.Predictor.status.not_in([db.PREDICTOR_STATUS.COMPLETE, db.PREDICTOR_STATUS.ERROR]),
         )
         .all()
     )
     for predictor_record in predictor_records:
-        predictor_process_id = (predictor_record.training_metadata or {}).get(
-            "process_id"
-        )
+        predictor_process_id = (predictor_record.training_metadata or {}).get("process_id")
         if predictor_process_id in unexisting_pids:
             predictor_record.status = db.PREDICTOR_STATUS.ERROR
             if isinstance(predictor_record.data, dict) is False:
                 predictor_record.data = {}
             if "error" not in predictor_record.data:
-                predictor_record.data["error"] = (
-                    "The training process was terminated for unknown reasons"
-                )
+                predictor_record.data["error"] = "The training process was terminated for unknown reasons"
                 flag_modified(predictor_record, "data")
             db.session.commit()
 
@@ -204,9 +192,7 @@ def set_error_model_status_for_unfinished():
         db.session.query(db.Predictor)
         .filter(
             db.Predictor.deleted_at.is_(None),
-            db.Predictor.status.not_in(
-                [db.PREDICTOR_STATUS.COMPLETE, db.PREDICTOR_STATUS.ERROR]
-            ),
+            db.Predictor.status.not_in([db.PREDICTOR_STATUS.COMPLETE, db.PREDICTOR_STATUS.ERROR]),
         )
         .all()
     )
@@ -234,11 +220,7 @@ def create_permanent_integrations():
     NOTE: this is intentional to avoid importing integration_controller
     """
     integration_name = "files"
-    existing = (
-        db.session.query(db.Integration)
-        .filter_by(name=integration_name, company_id=None)
-        .first()
-    )
+    existing = db.session.query(db.Integration).filter_by(name=integration_name, company_id=None).first()
     if existing is None:
         integration_record = db.Integration(
             name=integration_name,
@@ -250,9 +232,7 @@ def create_permanent_integrations():
         try:
             db.session.commit()
         except Exception as e:
-            logger.error(
-                f"Failed to commit permanent integration {integration_name}: {e}"
-            )
+            logger.error(f"Failed to commit permanent integration {integration_name}: {e}")
             db.session.rollback()
 
 
@@ -278,9 +258,7 @@ def validate_default_project() -> None:
             func.lower(db.Project.name) == func.lower(new_default_project_name),
         ).first()
         if existing_project is None:
-            logger.critical(
-                f"A project with the name '{new_default_project_name}' does not exist"
-            )
+            logger.critical(f"A project with the name '{new_default_project_name}' does not exist")
             sys.exit(1)
 
         existing_project.metadata_ = {"is_default": True}
@@ -293,9 +271,7 @@ def validate_default_project() -> None:
             func.lower(db.Project.name) == func.lower(new_default_project_name),
         ).first()
         if existing_project is not None:
-            logger.critical(
-                f"A project with the name '{new_default_project_name}' already exists"
-            )
+            logger.critical(f"A project with the name '{new_default_project_name}' already exists")
             sys.exit(1)
         current_default_project.name = new_default_project_name
         db.session.commit()
@@ -317,9 +293,7 @@ def start_process(trunc_process_data: TrunkProcessData) -> None:
         )
         trunc_process_data.process.start()
     except Exception as e:
-        logger.error(
-            f"Failed to start {trunc_process_data.name} API with exception {e}\n{traceback.format_exc()}"
-        )
+        logger.error(f"Failed to start {trunc_process_data.name} API with exception {e}\n{traceback.format_exc()}")
         close_api_gracefully(trunc_processes_struct)
         raise e
 
@@ -329,8 +303,7 @@ if __name__ == "__main__":
     # warn if less than 1Gb of free RAM
     if psutil.virtual_memory().available < (1 << 30):
         logger.warning(
-            "The system is running low on memory. "
-            + "This may impact the stability and performance of the program."
+            "The system is running low on memory. " + "This may impact the stability and performance of the program."
         )
 
     ctx.set_default()
@@ -440,94 +413,86 @@ if __name__ == "__main__":
     mysql_api_config = config.get("api", {}).get("mysql", {})
     mcp_api_config = config.get("api", {}).get("mcp", {})
     litellm_api_config = config.get("api", {}).get("litellm", {})
-    a2a_api_config = config.get("a2a", {})
+    a2a_api_config = config.get("api", {}).get("a2a", {})
     trunc_processes_struct = {
         TrunkProcessEnum.HTTP: TrunkProcessData(
             name=TrunkProcessEnum.HTTP.value,
             entrypoint=start_http,
-            port=http_api_config['port'],
+            port=http_api_config["port"],
             args=(config.cmd_args.verbose, config.cmd_args.no_studio),
-            restart_on_failure=http_api_config.get('restart_on_failure', False),
-            max_restart_count=http_api_config.get('max_restart_count', TrunkProcessData.max_restart_count),
+            restart_on_failure=http_api_config.get("restart_on_failure", False),
+            max_restart_count=http_api_config.get("max_restart_count", TrunkProcessData.max_restart_count),
             max_restart_interval_seconds=http_api_config.get(
-                'max_restart_interval_seconds', TrunkProcessData.max_restart_interval_seconds
-            )
+                "max_restart_interval_seconds", TrunkProcessData.max_restart_interval_seconds
+            ),
         ),
         TrunkProcessEnum.MYSQL: TrunkProcessData(
             name=TrunkProcessEnum.MYSQL.value,
             entrypoint=start_mysql,
-            port=mysql_api_config['port'],
+            port=mysql_api_config["port"],
             args=(config.cmd_args.verbose,),
-            restart_on_failure=mysql_api_config.get('restart_on_failure', False),
-            max_restart_count=mysql_api_config.get('max_restart_count', TrunkProcessData.max_restart_count),
+            restart_on_failure=mysql_api_config.get("restart_on_failure", False),
+            max_restart_count=mysql_api_config.get("max_restart_count", TrunkProcessData.max_restart_count),
             max_restart_interval_seconds=mysql_api_config.get(
-                'max_restart_interval_seconds', TrunkProcessData.max_restart_interval_seconds
-            )
+                "max_restart_interval_seconds", TrunkProcessData.max_restart_interval_seconds
+            ),
         ),
         TrunkProcessEnum.MONGODB: TrunkProcessData(
             name=TrunkProcessEnum.MONGODB.value,
             entrypoint=start_mongo,
-            port=config['api']['mongodb']['port'],
-            args=(config.cmd_args.verbose,)
+            port=config["api"]["mongodb"]["port"],
+            args=(config.cmd_args.verbose,),
         ),
         TrunkProcessEnum.POSTGRES: TrunkProcessData(
             name=TrunkProcessEnum.POSTGRES.value,
             entrypoint=start_postgres,
-            port=config['api']['postgres']['port'],
-            args=(config.cmd_args.verbose,)
+            port=config["api"]["postgres"]["port"],
+            args=(config.cmd_args.verbose,),
         ),
         TrunkProcessEnum.JOBS: TrunkProcessData(
-            name=TrunkProcessEnum.JOBS.value,
-            entrypoint=start_scheduler,
-            args=(config.cmd_args.verbose,)
+            name=TrunkProcessEnum.JOBS.value, entrypoint=start_scheduler, args=(config.cmd_args.verbose,)
         ),
         TrunkProcessEnum.TASKS: TrunkProcessData(
-            name=TrunkProcessEnum.TASKS.value,
-            entrypoint=start_tasks,
-            args=(config.cmd_args.verbose,)
+            name=TrunkProcessEnum.TASKS.value, entrypoint=start_tasks, args=(config.cmd_args.verbose,)
         ),
         TrunkProcessEnum.ML_TASK_QUEUE: TrunkProcessData(
-            name=TrunkProcessEnum.ML_TASK_QUEUE.value,
-            entrypoint=start_ml_task_queue,
-            args=(config.cmd_args.verbose,)
+            name=TrunkProcessEnum.ML_TASK_QUEUE.value, entrypoint=start_ml_task_queue, args=(config.cmd_args.verbose,)
         ),
         TrunkProcessEnum.MCP: TrunkProcessData(
             name=TrunkProcessEnum.MCP.value,
             entrypoint=start_mcp,
-            port=mcp_api_config.get('port', 47337),
+            port=mcp_api_config.get("port", 47337),
             args=(config.cmd_args.verbose,),
-            need_to_run=mcp_api_config.get('need_to_run', False),
-            restart_on_failure=mcp_api_config.get('restart_on_failure', False),
-            max_restart_count=mcp_api_config.get('max_restart_count', TrunkProcessData.max_restart_count),
+            need_to_run=mcp_api_config.get("need_to_run", False),
+            restart_on_failure=mcp_api_config.get("restart_on_failure", False),
+            max_restart_count=mcp_api_config.get("max_restart_count", TrunkProcessData.max_restart_count),
             max_restart_interval_seconds=mcp_api_config.get(
-                'max_restart_interval_seconds', TrunkProcessData.max_restart_interval_seconds
-            )
+                "max_restart_interval_seconds", TrunkProcessData.max_restart_interval_seconds
+            ),
         ),
         TrunkProcessEnum.LITELLM: TrunkProcessData(
             name=TrunkProcessEnum.LITELLM.value,
             entrypoint=start_litellm,
-            port=litellm_api_config.get('port', 8000),
+            port=litellm_api_config.get("port", 8000),
             args=(config.cmd_args.verbose,),
-            restart_on_failure=litellm_api_config.get('restart_on_failure', False),
-            max_restart_count=litellm_api_config.get('max_restart_count', TrunkProcessData.max_restart_count),
+            restart_on_failure=litellm_api_config.get("restart_on_failure", False),
+            max_restart_count=litellm_api_config.get("max_restart_count", TrunkProcessData.max_restart_count),
             max_restart_interval_seconds=litellm_api_config.get(
-                'max_restart_interval_seconds', TrunkProcessData.max_restart_interval_seconds
-            )
+                "max_restart_interval_seconds", TrunkProcessData.max_restart_interval_seconds
+            ),
         ),
         TrunkProcessEnum.A2A: TrunkProcessData(
             name=TrunkProcessEnum.A2A.value,
             entrypoint=start_a2a,
-            port=a2a_api_config.get('port', 8001),
+            port=a2a_api_config.get("port", 8001),
             args=(config.cmd_args.verbose,),
-            need_to_run=a2a_api_config.get('enabled', False),
-            restart_on_failure=a2a_api_config.get('restart_on_failure', True),
-            max_restart_count=a2a_api_config.get(
-                'max_restart_count', TrunkProcessData.max_restart_count
-            ),
+            need_to_run=a2a_api_config.get("enabled", False),
+            restart_on_failure=a2a_api_config.get("restart_on_failure", True),
+            max_restart_count=a2a_api_config.get("max_restart_count", TrunkProcessData.max_restart_count),
             max_restart_interval_seconds=a2a_api_config.get(
-                'max_restart_interval_seconds', TrunkProcessData.max_restart_interval_seconds
-            )
-        )
+                "max_restart_interval_seconds", TrunkProcessData.max_restart_interval_seconds
+            ),
+        ),
     }
 
     for api_enum in api_arr:
@@ -546,10 +511,7 @@ if __name__ == "__main__":
         trunc_processes_struct[TrunkProcessEnum.ML_TASK_QUEUE].need_to_run = True
 
     for trunc_process_data in trunc_processes_struct.values():
-        if (
-            trunc_process_data.started is True
-            or trunc_process_data.need_to_run is False
-        ):
+        if trunc_process_data.started is True or trunc_process_data.need_to_run is False:
             continue
         start_process(trunc_process_data)
 
@@ -572,8 +534,7 @@ if __name__ == "__main__":
                 trunc_process_data.port,
             )
             for trunc_process_data in trunc_processes_struct.values()
-            if trunc_process_data.port is not None
-            and trunc_process_data.need_to_run is True
+            if trunc_process_data.port is not None and trunc_process_data.need_to_run is True
         ]
         for future in asyncio.as_completed(futures):
             api_name, port, started = await future
@@ -596,9 +557,7 @@ if __name__ == "__main__":
             finally:
                 if trunc_process_data.should_restart:
                     if trunc_process_data.request_restart_attempt():
-                        logger.warning(
-                            f"{trunc_process_data.name} API: stopped unexpectedly, restarting"
-                        )
+                        logger.warning(f"{trunc_process_data.name} API: stopped unexpectedly, restarting")
                         trunc_process_data.process = None
                         if trunc_process_data.name == TrunkProcessEnum.HTTP.value:
                             # do not open GUI on HTTP API restart

mindsdb/api/a2a/agent.py

@@ -1,8 +1,9 @@
 import json
-from typing import Any, AsyncIterable, Dict, List, Iterator
+from typing import Any, AsyncIterable, Dict, List
 import requests
 import logging
-
+import httpx
+from mindsdb.api.a2a.utils import to_serializable
 from mindsdb.api.a2a.constants import DEFAULT_STREAM_TIMEOUT
 
 logger = logging.getLogger(__name__)
@@ -11,9 +12,6 @@ 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__(
@@ -35,31 +33,15 @@ class MindsDBAgent:
     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]
@@ -68,19 +50,9 @@ class MindsDBAgent:
                             "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}]
-
+                parts = [{"type": "text", "text": content}]
                 if isinstance(result_row, dict):
                     parts.append(
                         {
@@ -89,7 +61,6 @@ class MindsDBAgent:
                             "metadata": {"subtype": "json"},
                         }
                     )
-
                 return {
                     "content": content,
                     "parts": parts,
@@ -101,7 +72,6 @@ class MindsDBAgent:
                     "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)
@@ -109,7 +79,6 @@ class MindsDBAgent:
                 "content": error_msg,
                 "parts": [{"type": "text", "text": error_msg}],
             }
-
         except Exception as e:
             error_msg = f"Error: {str(e)}"
             logger.error(error_msg)
@@ -118,203 +87,58 @@ class MindsDBAgent:
                 "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 streaming_invoke(self, messages, timeout=DEFAULT_STREAM_TIMEOUT):
+        url = f"{self.base_url}/api/projects/{self.project_name}/agents/{self.agent_name}/completions/stream"
+        logger.info(f"Sending streaming request to MindsDB agent: {self.agent_name}")
+        async with httpx.AsyncClient(timeout=timeout) as client:
+            async with client.stream("POST", url, json={"messages": to_serializable(messages)}) as response:
+                response.raise_for_status()
+                async for line in response.aiter_lines():
+                    if not line.strip():
+                        continue
+                    # Only process actual SSE data lines
+                    if line.startswith("data:"):
+                        payload = line[len("data:") :].strip()
+                        try:
+                            yield json.loads(payload)
+                        except Exception as e:
+                            logger.error(f"Failed to parse SSE JSON payload: {e}; line: {payload}")
+                    # Ignore comments or control lines
+                # Signal the end of the stream
+                yield {"is_task_complete": True}
 
     async def stream(
         self,
         query: str,
         session_id: str,
         history: List[dict] | None = None,
+        timeout: int = DEFAULT_STREAM_TIMEOUT,
     ) -> AsyncIterable[Dict[str, Any]]:
-        """Stream responses from the MindsDB agent (uses streaming API endpoint).
-
-        Args:
-            query: The current query to send to the agent.
-            session_id: Unique identifier for the conversation session.
-            history: Optional list of previous messages in the conversation.
-
-        Returns:
-            AsyncIterable yielding chunks of the streaming response.
-        """
+        """Stream responses from the MindsDB agent (uses streaming API endpoint)."""
         try:
             logger.info(f"Using streaming API for query: {query[:100]}...")
-
-            # Format history into the expected format
             formatted_messages = []
             if history:
                 for msg in history:
-                    # Convert Message object to dict if needed
                     msg_dict = msg.dict() if hasattr(msg, "dict") else msg
                     role = msg_dict.get("role", "user")
-
-                    # Extract text from parts
                     text = ""
                     for part in msg_dict.get("parts", []):
                         if part.get("type") == "text":
                             text = part.get("text", "")
                             break
-
                     if text:
                         if role == "user":
                             formatted_messages.append({"question": text, "answer": None})
                         elif role == "assistant" and formatted_messages:
-                            # Add the answer to the last question
                             formatted_messages[-1]["answer"] = text
-
-            # Add the current query to the messages
             formatted_messages.append({"question": query, "answer": None})
-
             logger.debug(f"Formatted messages for agent: {formatted_messages}")
-
-            # Use the streaming_invoke method to get real streaming responses
-            streaming_response = self.streaming_invoke(formatted_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
-
+            streaming_response = self.streaming_invoke(formatted_messages, timeout=timeout)
+            async for chunk in streaming_response:
+                content_value = chunk.get("text") or chunk.get("output") or json.dumps(chunk)
+                wrapped_chunk = {"is_task_complete": False, "content": content_value, "metadata": {}}
+                yield wrapped_chunk
         except Exception as e:
             logger.error(f"Error in streaming: {str(e)}")
             yield {