From 7e3f7e1a34c037bb3f678b996904b771c1a0d9e5 Mon Sep 17 00:00:00 2001 From: "autofix-ci[bot]" <114827586+autofix-ci[bot]@users.noreply.github.com> Date: Wed, 10 Jun 2026 12:28:54 +0000 Subject: [PATCH 1/3] [autofix.ci] apply automated fixes --- src/lfx/src/lfx/_assets/component_index.json | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/lfx/src/lfx/_assets/component_index.json b/src/lfx/src/lfx/_assets/component_index.json index bce86b23c4..9fc99281be 100644 --- a/src/lfx/src/lfx/_assets/component_index.json +++ b/src/lfx/src/lfx/_assets/component_index.json @@ -70596,7 +70596,7 @@ "icon": "file-text", "legacy": false, "metadata": { - "code_hash": "837120b2e497", + "code_hash": "2897652b1a16", "dependencies": { "dependencies": [ { @@ -70807,7 +70807,7 @@ "show": true, "title_case": false, "type": "code", - "value": "import json\nfrom collections.abc import AsyncIterator, Iterator\nfrom pathlib import Path\nfrom typing import Any\n\nimport orjson\nimport pandas as pd\nfrom fastapi import UploadFile\nfrom fastapi.encoders import jsonable_encoder\n\nfrom lfx.custom import Component\nfrom lfx.inputs import SortableListInput\nfrom lfx.inputs.inputs import DataFrameInput\nfrom lfx.io import BoolInput, DropdownInput, SecretStrInput, StrInput\nfrom lfx.schema import Data, DataFrame, Message\nfrom lfx.services.deps import get_settings_service, get_storage_service, session_scope\nfrom lfx.template.field.base import Output\nfrom lfx.utils.validate_cloud import is_astra_cloud_environment\n\n\ndef _get_storage_location_options():\n \"\"\"Get storage location options, filtering out Local if in Astra cloud environment.\"\"\"\n all_options = [{\"name\": \"AWS\", \"icon\": \"Amazon\"}, {\"name\": \"Google Drive\", \"icon\": \"google\"}]\n if is_astra_cloud_environment():\n return all_options\n return [{\"name\": \"Local\", \"icon\": \"hard-drive\"}, *all_options]\n\n\ndef _get_default_storage_location() -> list[dict[str, str]]:\n \"\"\"Return the default storage selection for the component template.\"\"\"\n return [_get_storage_location_options()[0]]\n\n\ndef _is_default_storage(storage_name: str) -> bool:\n \"\"\"Check whether a storage type is the default selection.\"\"\"\n return _get_default_storage_location()[0][\"name\"] == storage_name\n\n\nclass SaveToFileComponent(Component):\n display_name = \"Write File\"\n description = (\n \"Save data to a file. \"\n \"Arguments: 'input' — the content to save (pass a DataFrame directly, or a JSON string \"\n \"for tabular data, or plain text for messages); \"\n \"'file_name' — the name to save as, without extension (e.g. 'report'); \"\n \"'file_format' — output format: 'csv', 'json', 'txt', 'html', 'excel', 'markdown' (optional). \"\n \"Returns a confirmation with the file path or URL.\"\n )\n documentation: str = \"https://docs.langflow.org/write-file\"\n icon = \"file-text\"\n name = \"SaveToFile\"\n\n # File format options for different storage types\n LOCAL_DATA_FORMAT_CHOICES = [\"csv\", \"excel\", \"json\", \"markdown\"]\n LOCAL_MESSAGE_FORMAT_CHOICES = [\"txt\", \"html\", \"json\", \"markdown\"]\n AWS_FORMAT_CHOICES = [\n \"txt\",\n \"json\",\n \"csv\",\n \"xml\",\n \"html\",\n \"md\",\n \"yaml\",\n \"log\",\n \"tsv\",\n \"jsonl\",\n \"parquet\",\n \"xlsx\",\n \"zip\",\n ]\n GDRIVE_FORMAT_CHOICES = [\"txt\", \"html\", \"json\", \"csv\", \"xlsx\", \"slides\", \"docs\", \"jpg\", \"mp3\"]\n\n inputs = [\n SortableListInput(\n name=\"storage_location\",\n display_name=\"Storage Location\",\n placeholder=\"Select Location\",\n info=\"Choose where to save the file.\",\n options=_get_storage_location_options(),\n real_time_refresh=True,\n limit=1,\n value=_get_default_storage_location(),\n advanced=True,\n ),\n # Common inputs\n DataFrameInput(\n name=\"input\",\n display_name=\"File Content\",\n info=(\n \"The content to save. Accepts a DataFrame, Data, or Message object directly. \"\n 'Can also accept a JSON string (e.g. \\'[{\"col1\": \"val1\"}]\\') which will be '\n \"parsed into a DataFrame, or plain text which will be saved as a Message.\"\n ),\n input_types=[\"Data\", \"JSON\", \"DataFrame\", \"Table\", \"Message\"],\n required=True,\n tool_mode=True,\n ),\n StrInput(\n name=\"file_name\",\n display_name=\"File Name\",\n info=\"File name without extension (e.g. 'report'). Extension is added automatically.\",\n required=True,\n show=True,\n tool_mode=True,\n ),\n StrInput(\n name=\"file_format\",\n display_name=\"File Format (Tool)\",\n info=\"Output format: 'csv', 'json', 'txt', 'html', 'excel', 'markdown'. Overrides pre-configured format.\",\n required=False,\n show=False,\n tool_mode=True,\n ),\n BoolInput(\n name=\"append_mode\",\n display_name=\"Append\",\n info=(\n \"Append to file if it exists (only for Local storage with plain text formats). \"\n \"Not supported for cloud storage (AWS/Google Drive).\"\n ),\n value=False,\n show=_is_default_storage(\"Local\"),\n ),\n # Format inputs (dynamic based on storage location)\n DropdownInput(\n name=\"local_format\",\n display_name=\"File Format\",\n options=list(dict.fromkeys(LOCAL_DATA_FORMAT_CHOICES + LOCAL_MESSAGE_FORMAT_CHOICES)),\n info=\"Select the file format for local storage.\",\n value=\"json\",\n show=_is_default_storage(\"Local\"),\n ),\n DropdownInput(\n name=\"aws_format\",\n display_name=\"File Format\",\n options=AWS_FORMAT_CHOICES,\n info=\"Select the file format for AWS S3 storage.\",\n value=\"txt\",\n show=_is_default_storage(\"AWS\"),\n ),\n DropdownInput(\n name=\"gdrive_format\",\n display_name=\"File Format\",\n options=GDRIVE_FORMAT_CHOICES,\n info=\"Select the file format for Google Drive storage.\",\n value=\"txt\",\n show=_is_default_storage(\"Google Drive\"),\n ),\n # AWS S3 specific inputs\n SecretStrInput(\n name=\"aws_access_key_id\",\n display_name=\"AWS Access Key ID\",\n info=\"AWS Access key ID.\",\n show=_is_default_storage(\"AWS\"),\n advanced=not _is_default_storage(\"AWS\"),\n required=True,\n ),\n SecretStrInput(\n name=\"aws_secret_access_key\",\n display_name=\"AWS Secret Key\",\n info=\"AWS Secret Key.\",\n show=_is_default_storage(\"AWS\"),\n advanced=not _is_default_storage(\"AWS\"),\n required=True,\n ),\n StrInput(\n name=\"bucket_name\",\n display_name=\"S3 Bucket Name\",\n info=\"Enter the name of the S3 bucket.\",\n show=_is_default_storage(\"AWS\"),\n advanced=not _is_default_storage(\"AWS\"),\n required=True,\n ),\n StrInput(\n name=\"aws_region\",\n display_name=\"AWS Region\",\n info=\"AWS region (e.g., us-east-1, eu-west-1).\",\n show=_is_default_storage(\"AWS\"),\n advanced=not _is_default_storage(\"AWS\"),\n ),\n StrInput(\n name=\"s3_prefix\",\n display_name=\"S3 Prefix\",\n info=\"Prefix for all files in S3.\",\n show=_is_default_storage(\"AWS\"),\n advanced=not _is_default_storage(\"AWS\"),\n ),\n # Google Drive specific inputs\n SecretStrInput(\n name=\"service_account_key\",\n display_name=\"GCP Credentials Secret Key\",\n info=\"Your Google Cloud Platform service account JSON key as a secret string (complete JSON content).\",\n show=_is_default_storage(\"Google Drive\"),\n advanced=not _is_default_storage(\"Google Drive\"),\n required=True,\n ),\n StrInput(\n name=\"folder_id\",\n display_name=\"Google Drive Folder ID\",\n info=(\n \"The Google Drive folder ID where the file will be uploaded. \"\n \"The folder must be shared with the service account email.\"\n ),\n required=True,\n show=_is_default_storage(\"Google Drive\"),\n advanced=not _is_default_storage(\"Google Drive\"),\n ),\n ]\n\n outputs = [Output(display_name=\"File Path\", name=\"message\", method=\"save_to_file\")]\n\n def update_build_config(self, build_config, field_value, field_name=None):\n \"\"\"Update build configuration to show/hide fields based on storage location selection.\"\"\"\n # Update options dynamically based on cloud environment\n # This ensures options are refreshed when build_config is updated\n if \"storage_location\" in build_config:\n updated_options = _get_storage_location_options()\n build_config[\"storage_location\"][\"options\"] = updated_options\n\n # When tool_mode is toggled, hide storage-specific format dropdowns\n # (the agent uses the unified file_format input instead)\n if field_name == \"tool_mode\":\n format_fields = [\"local_format\", \"aws_format\", \"gdrive_format\"]\n for f_name in format_fields:\n if f_name in build_config:\n build_config[f_name][\"show\"] = not bool(field_value)\n return build_config\n\n if field_name != \"storage_location\":\n return build_config\n\n # Extract selected storage location\n selected = [location[\"name\"] for location in field_value] if isinstance(field_value, list) else []\n\n # Hide all dynamic fields first\n dynamic_fields = [\n \"file_name\", # Common fields (input is always visible)\n \"append_mode\",\n \"local_format\",\n \"aws_format\",\n \"gdrive_format\",\n \"aws_access_key_id\",\n \"aws_secret_access_key\",\n \"bucket_name\",\n \"aws_region\",\n \"s3_prefix\",\n \"service_account_key\",\n \"folder_id\",\n ]\n\n for f_name in dynamic_fields:\n if f_name in build_config:\n build_config[f_name][\"show\"] = False\n\n # Show fields based on selected storage location\n is_tool_mode = build_config.get(\"tools_metadata\", {}).get(\"show\", False)\n\n if len(selected) == 1:\n location = selected[0]\n\n # Show file_name when any storage location is selected\n if \"file_name\" in build_config:\n build_config[\"file_name\"][\"show\"] = True\n\n # Show append_mode only for Local storage (not supported for cloud storage)\n if \"append_mode\" in build_config:\n build_config[\"append_mode\"][\"show\"] = location == \"Local\"\n\n if location == \"Local\":\n if \"local_format\" in build_config:\n build_config[\"local_format\"][\"show\"] = not is_tool_mode\n\n elif location == \"AWS\":\n aws_fields = [\n \"aws_format\",\n \"aws_access_key_id\",\n \"aws_secret_access_key\",\n \"bucket_name\",\n \"aws_region\",\n \"s3_prefix\",\n ]\n for f_name in aws_fields:\n if f_name in build_config:\n show = f_name != \"aws_format\" or not is_tool_mode\n build_config[f_name][\"show\"] = show\n build_config[f_name][\"advanced\"] = False\n\n elif location == \"Google Drive\":\n gdrive_fields = [\"gdrive_format\", \"service_account_key\", \"folder_id\"]\n for f_name in gdrive_fields:\n if f_name in build_config:\n show = f_name != \"gdrive_format\" or not is_tool_mode\n build_config[f_name][\"show\"] = show\n build_config[f_name][\"advanced\"] = False\n\n return build_config\n\n async def save_to_file(self) -> Message:\n \"\"\"Save the input to a file and upload it, returning a confirmation message.\"\"\"\n # Validate inputs\n if not self.file_name:\n msg = \"File name must be provided.\"\n raise ValueError(msg)\n if not self._get_input_type():\n msg = \"Input type is not set.\"\n raise ValueError(msg)\n\n # Get selected storage location\n storage_location = self._get_selected_storage_location()\n if not storage_location:\n msg = \"Storage location must be selected.\"\n raise ValueError(msg)\n\n # Check if Local storage is disabled in cloud environment\n if storage_location == \"Local\" and is_astra_cloud_environment():\n msg = \"Local storage is not available in cloud environment. Please use AWS or Google Drive.\"\n raise ValueError(msg)\n\n # Route to appropriate save method based on storage location\n if storage_location == \"Local\":\n return await self._save_to_local()\n if storage_location == \"AWS\":\n return await self._save_to_aws()\n if storage_location == \"Google Drive\":\n return await self._save_to_google_drive()\n msg = f\"Unsupported storage location: {storage_location}\"\n raise ValueError(msg)\n\n def _get_input_type(self) -> str:\n \"\"\"Determine the input type based on the provided input.\"\"\"\n # Use exact type checking (type() is) instead of isinstance() to avoid inheritance issues.\n # Since Message inherits from Data, isinstance(message, Data) would return True for Message objects,\n # causing Message inputs to be incorrectly identified as Data type.\n if type(self.input) is DataFrame:\n return \"DataFrame\"\n if type(self.input) is Message:\n return \"Message\"\n if type(self.input) is Data:\n return \"Data\"\n # When invoked by a code agent (e.g. OpenDsStar), the input may be a raw\n # pandas DataFrame rather than Langflow's DataFrame wrapper.\n if isinstance(self.input, pd.DataFrame):\n self.input = DataFrame(self.input)\n return \"DataFrame\"\n # When invoked as a tool, the agent passes a string. Try to parse it as\n # tabular JSON (list of objects) → DataFrame, otherwise wrap as Message.\n if isinstance(self.input, str):\n self.input = self._coerce_string_input(self.input)\n return self._get_input_type()\n msg = f\"Unsupported input type: {type(self.input)}\"\n raise ValueError(msg)\n\n def _coerce_string_input(self, value: str) -> DataFrame | Message:\n \"\"\"Convert a raw string (from agent tool call) into a DataFrame or Message.\n\n Tries to parse as JSON first — a list of objects or a single object becomes\n a DataFrame. Anything else is wrapped in a Message.\n \"\"\"\n try:\n parsed = json.loads(value)\n if isinstance(parsed, list) and parsed and isinstance(parsed[0], dict):\n return DataFrame(pd.DataFrame(parsed))\n if isinstance(parsed, dict):\n return DataFrame(pd.DataFrame([parsed]))\n except (json.JSONDecodeError, ValueError):\n pass\n return Message(text=value)\n\n def _get_default_format(self) -> str:\n \"\"\"Return the default file format based on input type.\"\"\"\n if self._get_input_type() == \"DataFrame\":\n return \"csv\"\n if self._get_input_type() == \"Data\":\n return \"json\"\n if self._get_input_type() == \"Message\":\n return \"json\"\n return \"json\" # Fallback\n\n def _adjust_file_path_with_format(self, path: Path, fmt: str) -> Path:\n \"\"\"Adjust the file path to include the correct extension.\"\"\"\n file_extension = path.suffix.lower().lstrip(\".\")\n if fmt == \"excel\":\n return Path(f\"{path}.xlsx\").expanduser() if file_extension not in [\"xlsx\", \"xls\"] else path\n return Path(f\"{path}.{fmt}\").expanduser() if file_extension != fmt else path\n\n def _is_plain_text_format(self, fmt: str) -> bool:\n \"\"\"Check if a file format is plain text (supports appending).\"\"\"\n plain_text_formats = [\"txt\", \"json\", \"markdown\", \"md\", \"csv\", \"xml\", \"html\", \"yaml\", \"log\", \"tsv\", \"jsonl\"]\n return fmt.lower() in plain_text_formats\n\n async def _upload_file(self, file_path: Path) -> None:\n \"\"\"Upload the saved file using the upload_user_file service.\"\"\"\n from langflow.api.v2.files import upload_user_file\n from langflow.services.database.models.user.crud import get_user_by_id\n\n # Ensure the file exists\n if not file_path.exists():\n msg = f\"File not found: {file_path}\"\n raise FileNotFoundError(msg)\n\n # Upload the file - always use append=False because the local file already contains\n # the correct content (either new or appended locally)\n with file_path.open(\"rb\") as f:\n async with session_scope() as db:\n if not self.user_id:\n msg = \"User ID is required for file saving.\"\n raise ValueError(msg)\n current_user = await get_user_by_id(db, self.user_id)\n\n await upload_user_file(\n file=UploadFile(filename=file_path.name, file=f, size=file_path.stat().st_size),\n session=db,\n current_user=current_user,\n storage_service=get_storage_service(),\n settings_service=get_settings_service(),\n append=False,\n )\n\n def _save_dataframe(self, dataframe: DataFrame, path: Path, fmt: str) -> str:\n \"\"\"Save a DataFrame to the specified file format.\"\"\"\n append_mode = getattr(self, \"append_mode\", False)\n should_append = append_mode and path.exists() and self._is_plain_text_format(fmt)\n\n if fmt == \"csv\":\n dataframe.to_csv(path, index=False, mode=\"a\" if should_append else \"w\", header=not should_append)\n elif fmt == \"excel\":\n dataframe.to_excel(path, index=False, engine=\"openpyxl\")\n elif fmt == \"json\":\n if should_append:\n # Read and parse existing JSON\n existing_data = []\n try:\n existing_content = path.read_text(encoding=\"utf-8\").strip()\n if existing_content:\n parsed = json.loads(existing_content)\n # Handle case where existing content is a single object\n if isinstance(parsed, dict):\n existing_data = [parsed]\n elif isinstance(parsed, list):\n existing_data = parsed\n except (json.JSONDecodeError, FileNotFoundError):\n # Treat parse errors or missing file as empty array\n existing_data = []\n\n # Append new data\n new_records = json.loads(dataframe.to_json(orient=\"records\"))\n existing_data.extend(new_records)\n\n # Write back as a single JSON array\n path.write_text(json.dumps(existing_data, indent=2), encoding=\"utf-8\")\n else:\n dataframe.to_json(path, orient=\"records\", indent=2)\n elif fmt == \"markdown\":\n content = dataframe.to_markdown(index=False)\n if should_append:\n path.write_text(path.read_text(encoding=\"utf-8\") + \"\\n\\n\" + content, encoding=\"utf-8\")\n else:\n path.write_text(content, encoding=\"utf-8\")\n else:\n msg = f\"Unsupported DataFrame format: {fmt}\"\n raise ValueError(msg)\n action = \"appended to\" if should_append else \"saved successfully as\"\n return f\"DataFrame {action} '{path}'\"\n\n def _save_data(self, data: Data, path: Path, fmt: str) -> str:\n \"\"\"Save a Data object to the specified file format.\"\"\"\n append_mode = getattr(self, \"append_mode\", False)\n should_append = append_mode and path.exists() and self._is_plain_text_format(fmt)\n\n if fmt == \"csv\":\n pd.DataFrame(data.data).to_csv(\n path,\n index=False,\n mode=\"a\" if should_append else \"w\",\n header=not should_append,\n )\n elif fmt == \"excel\":\n pd.DataFrame(data.data).to_excel(path, index=False, engine=\"openpyxl\")\n elif fmt == \"json\":\n new_data = jsonable_encoder(data.data)\n if should_append:\n # Read and parse existing JSON\n existing_data = []\n try:\n existing_content = path.read_text(encoding=\"utf-8\").strip()\n if existing_content:\n parsed = json.loads(existing_content)\n # Handle case where existing content is a single object\n if isinstance(parsed, dict):\n existing_data = [parsed]\n elif isinstance(parsed, list):\n existing_data = parsed\n except (json.JSONDecodeError, FileNotFoundError):\n # Treat parse errors or missing file as empty array\n existing_data = []\n\n # Append new data\n if isinstance(new_data, list):\n existing_data.extend(new_data)\n else:\n existing_data.append(new_data)\n\n # Write back as a single JSON array\n path.write_text(json.dumps(existing_data, indent=2), encoding=\"utf-8\")\n else:\n content = orjson.dumps(new_data, option=orjson.OPT_INDENT_2).decode(\"utf-8\")\n path.write_text(content, encoding=\"utf-8\")\n elif fmt == \"markdown\":\n content = pd.DataFrame(data.data).to_markdown(index=False)\n if should_append:\n path.write_text(path.read_text(encoding=\"utf-8\") + \"\\n\\n\" + content, encoding=\"utf-8\")\n else:\n path.write_text(content, encoding=\"utf-8\")\n else:\n msg = f\"Unsupported Data format: {fmt}\"\n raise ValueError(msg)\n action = \"appended to\" if should_append else \"saved successfully as\"\n return f\"Data {action} '{path}'\"\n\n async def _save_message(self, message: Message, path: Path, fmt: str) -> str:\n \"\"\"Save a Message to the specified file format, handling async iterators.\"\"\"\n content = \"\"\n if message.text is None:\n content = \"\"\n elif isinstance(message.text, AsyncIterator):\n async for item in message.text:\n content += str(item) + \" \"\n content = content.strip()\n elif isinstance(message.text, Iterator):\n content = \" \".join(str(item) for item in message.text)\n else:\n content = str(message.text)\n\n append_mode = getattr(self, \"append_mode\", False)\n should_append = append_mode and path.exists() and self._is_plain_text_format(fmt)\n\n if fmt in (\"txt\", \"html\"):\n if should_append:\n path.write_text(path.read_text(encoding=\"utf-8\") + \"\\n\" + content, encoding=\"utf-8\")\n else:\n path.write_text(content, encoding=\"utf-8\")\n elif fmt == \"json\":\n new_message = {\"message\": content}\n if should_append:\n # Read and parse existing JSON\n existing_data = []\n try:\n existing_content = path.read_text(encoding=\"utf-8\").strip()\n if existing_content:\n parsed = json.loads(existing_content)\n # Handle case where existing content is a single object\n if isinstance(parsed, dict):\n existing_data = [parsed]\n elif isinstance(parsed, list):\n existing_data = parsed\n except (json.JSONDecodeError, FileNotFoundError):\n # Treat parse errors or missing file as empty array\n existing_data = []\n\n # Append new message\n existing_data.append(new_message)\n\n # Write back as a single JSON array\n path.write_text(json.dumps(existing_data, indent=2), encoding=\"utf-8\")\n else:\n path.write_text(json.dumps(new_message, indent=2), encoding=\"utf-8\")\n elif fmt == \"markdown\":\n md_content = f\"**Message:**\\n\\n{content}\"\n if should_append:\n path.write_text(path.read_text(encoding=\"utf-8\") + \"\\n\\n\" + md_content, encoding=\"utf-8\")\n else:\n path.write_text(md_content, encoding=\"utf-8\")\n else:\n msg = f\"Unsupported Message format: {fmt}\"\n raise ValueError(msg)\n action = \"appended to\" if should_append else \"saved successfully as\"\n return f\"Message {action} '{path}'\"\n\n def _get_selected_storage_location(self) -> str:\n \"\"\"Get the selected storage location from the SortableListInput.\"\"\"\n if hasattr(self, \"storage_location\") and self.storage_location:\n if isinstance(self.storage_location, list) and len(self.storage_location) > 0:\n return self.storage_location[0].get(\"name\", \"\")\n if isinstance(self.storage_location, dict):\n return self.storage_location.get(\"name\", \"\")\n return \"\"\n\n def _get_file_format_for_location(self, location: str) -> str:\n \"\"\"Get the appropriate file format based on storage location.\n\n If the agent set file_format via tool mode, that takes priority.\n \"\"\"\n agent_format = getattr(self, \"file_format\", None)\n if agent_format:\n return agent_format\n if location == \"Local\":\n return getattr(self, \"local_format\", None) or self._get_default_format()\n if location == \"AWS\":\n return getattr(self, \"aws_format\", \"txt\")\n if location == \"Google Drive\":\n return getattr(self, \"gdrive_format\", \"txt\")\n return self._get_default_format()\n\n async def _save_to_local(self) -> Message:\n \"\"\"Save file to local storage (original functionality).\"\"\"\n file_format = self._get_file_format_for_location(\"Local\")\n\n # Validate file format based on input type\n allowed_formats = (\n self.LOCAL_MESSAGE_FORMAT_CHOICES if self._get_input_type() == \"Message\" else self.LOCAL_DATA_FORMAT_CHOICES\n )\n if file_format not in allowed_formats:\n msg = f\"Invalid file format '{file_format}' for {self._get_input_type()}. Allowed: {allowed_formats}\"\n raise ValueError(msg)\n\n # Prepare file path\n file_path = Path(self.file_name).expanduser()\n if not file_path.parent.exists():\n file_path.parent.mkdir(parents=True, exist_ok=True)\n file_path = self._adjust_file_path_with_format(file_path, file_format)\n\n # Save the input to file based on type\n if self._get_input_type() == \"DataFrame\":\n confirmation = self._save_dataframe(self.input, file_path, file_format)\n elif self._get_input_type() == \"Data\":\n confirmation = self._save_data(self.input, file_path, file_format)\n elif self._get_input_type() == \"Message\":\n confirmation = await self._save_message(self.input, file_path, file_format)\n else:\n msg = f\"Unsupported input type: {self._get_input_type()}\"\n raise ValueError(msg)\n\n # Upload the saved file\n await self._upload_file(file_path)\n\n # Return the final file path and confirmation message\n final_path = Path.cwd() / file_path if not file_path.is_absolute() else file_path\n return Message(text=f\"{confirmation} at {final_path}\")\n\n async def _save_to_aws(self) -> Message:\n \"\"\"Save file to AWS S3 using S3 functionality.\"\"\"\n import os\n\n import boto3\n\n from lfx.base.data.cloud_storage_utils import create_s3_client, validate_aws_credentials\n\n # Get AWS credentials from component inputs or fall back to environment variables\n aws_access_key_id = getattr(self, \"aws_access_key_id\", None)\n if aws_access_key_id and hasattr(aws_access_key_id, \"get_secret_value\"):\n aws_access_key_id = aws_access_key_id.get_secret_value()\n if not aws_access_key_id:\n aws_access_key_id = os.getenv(\"AWS_ACCESS_KEY_ID\")\n\n aws_secret_access_key = getattr(self, \"aws_secret_access_key\", None)\n if aws_secret_access_key and hasattr(aws_secret_access_key, \"get_secret_value\"):\n aws_secret_access_key = aws_secret_access_key.get_secret_value()\n if not aws_secret_access_key:\n aws_secret_access_key = os.getenv(\"AWS_SECRET_ACCESS_KEY\")\n\n bucket_name = getattr(self, \"bucket_name\", None)\n if not bucket_name:\n # Try to get from storage service settings\n settings = get_settings_service().settings\n bucket_name = settings.object_storage_bucket_name\n\n # Validate AWS credentials\n if not aws_access_key_id:\n msg = (\n \"AWS Access Key ID is required for S3 storage. Provide it as a component input \"\n \"or set AWS_ACCESS_KEY_ID environment variable.\"\n )\n raise ValueError(msg)\n if not aws_secret_access_key:\n msg = (\n \"AWS Secret Key is required for S3 storage. Provide it as a component input \"\n \"or set AWS_SECRET_ACCESS_KEY environment variable.\"\n )\n raise ValueError(msg)\n if not bucket_name:\n msg = (\n \"S3 Bucket Name is required for S3 storage. Provide it as a component input \"\n \"or set LANGFLOW_OBJECT_STORAGE_BUCKET_NAME environment variable.\"\n )\n raise ValueError(msg)\n\n # Validate AWS credentials\n validate_aws_credentials(self)\n\n # Create S3 client\n s3_client = create_s3_client(self)\n client_config: dict[str, Any] = {\n \"aws_access_key_id\": str(aws_access_key_id),\n \"aws_secret_access_key\": str(aws_secret_access_key),\n }\n\n # Get region from component input, environment variable, or settings\n aws_region = getattr(self, \"aws_region\", None)\n if not aws_region:\n aws_region = os.getenv(\"AWS_DEFAULT_REGION\") or os.getenv(\"AWS_REGION\")\n if aws_region:\n client_config[\"region_name\"] = str(aws_region)\n\n s3_client = boto3.client(\"s3\", **client_config)\n\n # Extract content\n content = self._extract_content_for_upload()\n file_format = self._get_file_format_for_location(\"AWS\")\n\n # Generate file path\n file_path = f\"{self.file_name}.{file_format}\"\n if hasattr(self, \"s3_prefix\") and self.s3_prefix:\n file_path = f\"{self.s3_prefix.rstrip('/')}/{file_path}\"\n\n # Create temporary file\n import tempfile\n\n with tempfile.NamedTemporaryFile(\n mode=\"w\", encoding=\"utf-8\", suffix=f\".{file_format}\", delete=False\n ) as temp_file:\n temp_file.write(content)\n temp_file_path = temp_file.name\n\n try:\n # Upload to S3\n s3_client.upload_file(temp_file_path, bucket_name, file_path)\n s3_url = f\"s3://{bucket_name}/{file_path}\"\n return Message(text=f\"File successfully uploaded to {s3_url}\")\n finally:\n # Clean up temp file\n if Path(temp_file_path).exists():\n Path(temp_file_path).unlink()\n\n async def _save_to_google_drive(self) -> Message:\n \"\"\"Save file to Google Drive using Google Drive functionality.\"\"\"\n import tempfile\n\n from googleapiclient.http import MediaFileUpload\n\n from lfx.base.data.cloud_storage_utils import create_google_drive_service\n\n # Validate Google Drive credentials\n if not getattr(self, \"service_account_key\", None):\n msg = \"GCP Credentials Secret Key is required for Google Drive storage\"\n raise ValueError(msg)\n if not getattr(self, \"folder_id\", None):\n msg = \"Google Drive Folder ID is required for Google Drive storage\"\n raise ValueError(msg)\n\n # Create Google Drive service with full drive scope (needed for folder operations)\n drive_service, credentials = create_google_drive_service(\n self.service_account_key, scopes=[\"https://www.googleapis.com/auth/drive\"], return_credentials=True\n )\n\n # Extract content and format\n content = self._extract_content_for_upload()\n file_format = self._get_file_format_for_location(\"Google Drive\")\n\n # Handle special Google Drive formats\n if file_format in [\"slides\", \"docs\"]:\n return await self._save_to_google_apps(drive_service, credentials, content, file_format)\n\n # Create temporary file\n file_path = f\"{self.file_name}.{file_format}\"\n with tempfile.NamedTemporaryFile(\n mode=\"w\",\n encoding=\"utf-8\",\n suffix=f\".{file_format}\",\n delete=False,\n ) as temp_file:\n temp_file.write(content)\n temp_file_path = temp_file.name\n\n try:\n # Upload to Google Drive\n # Note: We skip explicit folder verification since it requires broader permissions.\n # If the folder doesn't exist or isn't accessible, the create() call will fail with a clear error.\n file_metadata = {\"name\": file_path, \"parents\": [self.folder_id]}\n media = MediaFileUpload(temp_file_path, resumable=True)\n\n try:\n uploaded_file = (\n drive_service.files().create(body=file_metadata, media_body=media, fields=\"id\").execute()\n )\n except Exception as e:\n msg = (\n f\"Unable to upload file to Google Drive folder '{self.folder_id}'. \"\n f\"Error: {e!s}. \"\n \"Please ensure: 1) The folder ID is correct, 2) The folder exists, \"\n \"3) The service account has been granted access to this folder.\"\n )\n raise ValueError(msg) from e\n\n file_id = uploaded_file.get(\"id\")\n file_url = f\"https://drive.google.com/file/d/{file_id}/view\"\n return Message(text=f\"File successfully uploaded to Google Drive: {file_url}\")\n finally:\n # Clean up temp file\n if Path(temp_file_path).exists():\n Path(temp_file_path).unlink()\n\n async def _save_to_google_apps(self, drive_service, credentials, content: str, app_type: str) -> Message:\n \"\"\"Save content to Google Apps (Slides or Docs).\"\"\"\n import time\n\n if app_type == \"slides\":\n from googleapiclient.discovery import build\n\n slides_service = build(\"slides\", \"v1\", credentials=credentials)\n\n file_metadata = {\n \"name\": self.file_name,\n \"mimeType\": \"application/vnd.google-apps.presentation\",\n \"parents\": [self.folder_id],\n }\n\n created_file = drive_service.files().create(body=file_metadata, fields=\"id\").execute()\n presentation_id = created_file[\"id\"]\n\n time.sleep(2) # Wait for file to be available # noqa: ASYNC251\n\n presentation = slides_service.presentations().get(presentationId=presentation_id).execute()\n slide_id = presentation[\"slides\"][0][\"objectId\"]\n\n # Add content to slide\n requests = [\n {\n \"createShape\": {\n \"objectId\": \"TextBox_01\",\n \"shapeType\": \"TEXT_BOX\",\n \"elementProperties\": {\n \"pageObjectId\": slide_id,\n \"size\": {\n \"height\": {\"magnitude\": 3000000, \"unit\": \"EMU\"},\n \"width\": {\"magnitude\": 6000000, \"unit\": \"EMU\"},\n },\n \"transform\": {\n \"scaleX\": 1,\n \"scaleY\": 1,\n \"translateX\": 1000000,\n \"translateY\": 1000000,\n \"unit\": \"EMU\",\n },\n },\n }\n },\n {\"insertText\": {\"objectId\": \"TextBox_01\", \"insertionIndex\": 0, \"text\": content}},\n ]\n\n slides_service.presentations().batchUpdate(\n presentationId=presentation_id, body={\"requests\": requests}\n ).execute()\n file_url = f\"https://docs.google.com/presentation/d/{presentation_id}/edit\"\n\n elif app_type == \"docs\":\n from googleapiclient.discovery import build\n\n docs_service = build(\"docs\", \"v1\", credentials=credentials)\n\n file_metadata = {\n \"name\": self.file_name,\n \"mimeType\": \"application/vnd.google-apps.document\",\n \"parents\": [self.folder_id],\n }\n\n created_file = drive_service.files().create(body=file_metadata, fields=\"id\").execute()\n document_id = created_file[\"id\"]\n\n time.sleep(2) # Wait for file to be available # noqa: ASYNC251\n\n # Add content to document\n requests = [{\"insertText\": {\"location\": {\"index\": 1}, \"text\": content}}]\n docs_service.documents().batchUpdate(documentId=document_id, body={\"requests\": requests}).execute()\n file_url = f\"https://docs.google.com/document/d/{document_id}/edit\"\n\n return Message(text=f\"File successfully created in Google {app_type.title()}: {file_url}\")\n\n def _extract_content_for_upload(self) -> str:\n \"\"\"Extract content from input for upload to cloud services.\"\"\"\n if self._get_input_type() == \"DataFrame\":\n return self.input.to_csv(index=False)\n if self._get_input_type() == \"Data\":\n if hasattr(self.input, \"data\") and self.input.data:\n if isinstance(self.input.data, dict):\n import json\n\n return json.dumps(self.input.data, indent=2, ensure_ascii=False)\n return str(self.input.data)\n return str(self.input)\n if self._get_input_type() == \"Message\":\n return str(self.input.text) if self.input.text else str(self.input)\n return str(self.input)\n" + "value": "import json\nfrom collections.abc import AsyncIterator, Iterator\nfrom pathlib import Path\nfrom typing import Any\n\nimport orjson\nimport pandas as pd\nfrom fastapi import UploadFile\nfrom fastapi.encoders import jsonable_encoder\n\nfrom lfx.custom import Component\nfrom lfx.inputs import SortableListInput\nfrom lfx.inputs.inputs import DataFrameInput\nfrom lfx.io import BoolInput, DropdownInput, SecretStrInput, StrInput\nfrom lfx.schema import Data, DataFrame, Message\nfrom lfx.services.deps import get_settings_service, get_storage_service, session_scope\nfrom lfx.template.field.base import Output\nfrom lfx.utils.validate_cloud import is_astra_cloud_environment\n\n\ndef _get_storage_location_options():\n \"\"\"Get storage location options, filtering out Local if in Astra cloud environment.\"\"\"\n all_options = [{\"name\": \"AWS\", \"icon\": \"Amazon\"}, {\"name\": \"Google Drive\", \"icon\": \"google\"}]\n if is_astra_cloud_environment():\n return all_options\n return [{\"name\": \"Local\", \"icon\": \"hard-drive\"}, *all_options]\n\n\ndef _get_default_storage_location() -> list[dict[str, str]]:\n \"\"\"Return the default storage selection for the component template.\"\"\"\n return [_get_storage_location_options()[0]]\n\n\ndef _is_default_storage(storage_name: str) -> bool:\n \"\"\"Check whether a storage type is the default selection.\"\"\"\n return _get_default_storage_location()[0][\"name\"] == storage_name\n\n\nclass SaveToFileComponent(Component):\n display_name = \"Write File\"\n description = (\n \"Save data to a file. \"\n \"Arguments: 'input' — the content to save (pass a DataFrame directly, or a JSON string \"\n \"for tabular data, or plain text for messages); \"\n \"'file_name' — the name to save as, without extension (e.g. 'report'); \"\n \"'file_format' — output format: 'csv', 'json', 'txt', 'html', 'excel', 'markdown' (optional). \"\n \"Returns a confirmation with the file path or URL.\"\n )\n documentation: str = \"https://docs.langflow.org/write-file\"\n icon = \"file-text\"\n name = \"SaveToFile\"\n\n # File format options for different storage types\n LOCAL_DATA_FORMAT_CHOICES = [\"csv\", \"excel\", \"json\", \"markdown\"]\n LOCAL_MESSAGE_FORMAT_CHOICES = [\"txt\", \"html\", \"json\", \"markdown\"]\n AWS_FORMAT_CHOICES = [\n \"txt\",\n \"json\",\n \"csv\",\n \"xml\",\n \"html\",\n \"md\",\n \"yaml\",\n \"log\",\n \"tsv\",\n \"jsonl\",\n \"parquet\",\n \"xlsx\",\n \"zip\",\n ]\n GDRIVE_FORMAT_CHOICES = [\"txt\", \"html\", \"json\", \"csv\", \"xlsx\", \"slides\", \"docs\", \"jpg\", \"mp3\"]\n\n inputs = [\n SortableListInput(\n name=\"storage_location\",\n display_name=\"Storage Location\",\n placeholder=\"Select Location\",\n info=\"Choose where to save the file.\",\n options=_get_storage_location_options(),\n real_time_refresh=True,\n limit=1,\n value=_get_default_storage_location(),\n advanced=True,\n ),\n # Common inputs\n DataFrameInput(\n name=\"input\",\n display_name=\"File Content\",\n info=(\n \"The content to save. Accepts a DataFrame, Data, or Message object directly. \"\n 'Can also accept a JSON string (e.g. \\'[{\"col1\": \"val1\"}]\\') which will be '\n \"parsed into a DataFrame, or plain text which will be saved as a Message.\"\n ),\n input_types=[\"Data\", \"JSON\", \"DataFrame\", \"Table\", \"Message\"],\n required=True,\n tool_mode=True,\n ),\n StrInput(\n name=\"file_name\",\n display_name=\"File Name\",\n info=\"File name without extension (e.g. 'report'). Extension is added automatically.\",\n required=True,\n show=True,\n tool_mode=True,\n ),\n StrInput(\n name=\"file_format\",\n display_name=\"File Format (Tool)\",\n info=\"Output format: 'csv', 'json', 'txt', 'html', 'excel', 'markdown'. Overrides pre-configured format.\",\n required=False,\n show=False,\n tool_mode=True,\n ),\n BoolInput(\n name=\"append_mode\",\n display_name=\"Append\",\n info=(\n \"Append to file if it exists (only for Local storage with plain text formats). \"\n \"Not supported for cloud storage (AWS/Google Drive).\"\n ),\n value=False,\n show=_is_default_storage(\"Local\"),\n ),\n # Format inputs (dynamic based on storage location)\n DropdownInput(\n name=\"local_format\",\n display_name=\"File Format\",\n options=list(dict.fromkeys(LOCAL_DATA_FORMAT_CHOICES + LOCAL_MESSAGE_FORMAT_CHOICES)),\n info=\"Select the file format for local storage.\",\n value=\"json\",\n show=_is_default_storage(\"Local\"),\n ),\n DropdownInput(\n name=\"aws_format\",\n display_name=\"File Format\",\n options=AWS_FORMAT_CHOICES,\n info=\"Select the file format for AWS S3 storage.\",\n value=\"txt\",\n show=_is_default_storage(\"AWS\"),\n ),\n DropdownInput(\n name=\"gdrive_format\",\n display_name=\"File Format\",\n options=GDRIVE_FORMAT_CHOICES,\n info=\"Select the file format for Google Drive storage.\",\n value=\"txt\",\n show=_is_default_storage(\"Google Drive\"),\n ),\n # AWS S3 specific inputs\n SecretStrInput(\n name=\"aws_access_key_id\",\n display_name=\"AWS Access Key ID\",\n info=\"AWS Access key ID.\",\n show=_is_default_storage(\"AWS\"),\n advanced=not _is_default_storage(\"AWS\"),\n required=True,\n ),\n SecretStrInput(\n name=\"aws_secret_access_key\",\n display_name=\"AWS Secret Key\",\n info=\"AWS Secret Key.\",\n show=_is_default_storage(\"AWS\"),\n advanced=not _is_default_storage(\"AWS\"),\n required=True,\n ),\n StrInput(\n name=\"bucket_name\",\n display_name=\"S3 Bucket Name\",\n info=\"Enter the name of the S3 bucket.\",\n show=_is_default_storage(\"AWS\"),\n advanced=not _is_default_storage(\"AWS\"),\n required=True,\n ),\n StrInput(\n name=\"aws_region\",\n display_name=\"AWS Region\",\n info=\"AWS region (e.g., us-east-1, eu-west-1).\",\n show=_is_default_storage(\"AWS\"),\n advanced=not _is_default_storage(\"AWS\"),\n ),\n StrInput(\n name=\"s3_prefix\",\n display_name=\"S3 Prefix\",\n info=\"Prefix for all files in S3.\",\n show=_is_default_storage(\"AWS\"),\n advanced=not _is_default_storage(\"AWS\"),\n ),\n # Google Drive specific inputs\n SecretStrInput(\n name=\"service_account_key\",\n display_name=\"GCP Credentials Secret Key\",\n info=\"Your Google Cloud Platform service account JSON key as a secret string (complete JSON content).\",\n show=_is_default_storage(\"Google Drive\"),\n advanced=not _is_default_storage(\"Google Drive\"),\n required=True,\n ),\n StrInput(\n name=\"folder_id\",\n display_name=\"Google Drive Folder ID\",\n info=(\n \"The Google Drive folder ID where the file will be uploaded. \"\n \"The folder must be shared with the service account email.\"\n ),\n required=True,\n show=_is_default_storage(\"Google Drive\"),\n advanced=not _is_default_storage(\"Google Drive\"),\n ),\n ]\n\n outputs = [Output(display_name=\"File Path\", name=\"message\", method=\"save_to_file\")]\n\n def update_build_config(self, build_config, field_value, field_name=None):\n \"\"\"Update build configuration to show/hide fields based on storage location selection.\"\"\"\n # Update options dynamically based on cloud environment\n # This ensures options are refreshed when build_config is updated\n if \"storage_location\" in build_config:\n updated_options = _get_storage_location_options()\n build_config[\"storage_location\"][\"options\"] = updated_options\n\n # When tool_mode is toggled, hide storage-specific format dropdowns\n # (the agent uses the unified file_format input instead)\n if field_name == \"tool_mode\":\n format_fields = [\"local_format\", \"aws_format\", \"gdrive_format\"]\n for f_name in format_fields:\n if f_name in build_config:\n build_config[f_name][\"show\"] = not bool(field_value)\n return build_config\n\n if field_name != \"storage_location\":\n return build_config\n\n # Extract selected storage location\n selected = [location[\"name\"] for location in field_value] if isinstance(field_value, list) else []\n\n # Hide all dynamic fields first\n dynamic_fields = [\n \"file_name\", # Common fields (input is always visible)\n \"append_mode\",\n \"local_format\",\n \"aws_format\",\n \"gdrive_format\",\n \"aws_access_key_id\",\n \"aws_secret_access_key\",\n \"bucket_name\",\n \"aws_region\",\n \"s3_prefix\",\n \"service_account_key\",\n \"folder_id\",\n ]\n\n for f_name in dynamic_fields:\n if f_name in build_config:\n build_config[f_name][\"show\"] = False\n\n # Show fields based on selected storage location\n is_tool_mode = build_config.get(\"tools_metadata\", {}).get(\"show\", False)\n\n if len(selected) == 1:\n location = selected[0]\n\n # Show file_name when any storage location is selected\n if \"file_name\" in build_config:\n build_config[\"file_name\"][\"show\"] = True\n\n # Show append_mode only for Local storage (not supported for cloud storage)\n if \"append_mode\" in build_config:\n build_config[\"append_mode\"][\"show\"] = location == \"Local\"\n\n if location == \"Local\":\n if \"local_format\" in build_config:\n build_config[\"local_format\"][\"show\"] = not is_tool_mode\n\n elif location == \"AWS\":\n aws_fields = [\n \"aws_format\",\n \"aws_access_key_id\",\n \"aws_secret_access_key\",\n \"bucket_name\",\n \"aws_region\",\n \"s3_prefix\",\n ]\n for f_name in aws_fields:\n if f_name in build_config:\n show = f_name != \"aws_format\" or not is_tool_mode\n build_config[f_name][\"show\"] = show\n build_config[f_name][\"advanced\"] = False\n\n elif location == \"Google Drive\":\n gdrive_fields = [\"gdrive_format\", \"service_account_key\", \"folder_id\"]\n for f_name in gdrive_fields:\n if f_name in build_config:\n show = f_name != \"gdrive_format\" or not is_tool_mode\n build_config[f_name][\"show\"] = show\n build_config[f_name][\"advanced\"] = False\n\n return build_config\n\n async def save_to_file(self) -> Message:\n \"\"\"Save the input to a file and upload it, returning a confirmation message.\"\"\"\n # Validate inputs\n if not self.file_name:\n msg = \"File name must be provided.\"\n raise ValueError(msg)\n if not self._get_input_type():\n msg = \"Input type is not set.\"\n raise ValueError(msg)\n\n # Get selected storage location\n storage_location = self._get_selected_storage_location()\n if not storage_location:\n msg = \"Storage location must be selected.\"\n raise ValueError(msg)\n\n # Check if Local storage is disabled in cloud environment\n if storage_location == \"Local\" and is_astra_cloud_environment():\n msg = \"Local storage is not available in cloud environment. Please use AWS or Google Drive.\"\n raise ValueError(msg)\n\n # Route to appropriate save method based on storage location\n if storage_location == \"Local\":\n return await self._save_to_local()\n if storage_location == \"AWS\":\n return await self._save_to_aws()\n if storage_location == \"Google Drive\":\n return await self._save_to_google_drive()\n msg = f\"Unsupported storage location: {storage_location}\"\n raise ValueError(msg)\n\n def _get_input_type(self) -> str:\n \"\"\"Determine the input type based on the provided input.\"\"\"\n # Use exact type checking (type() is) instead of isinstance() to avoid inheritance issues.\n # Since Message inherits from Data, isinstance(message, Data) would return True for Message objects,\n # causing Message inputs to be incorrectly identified as Data type.\n if type(self.input) is DataFrame:\n return \"DataFrame\"\n if type(self.input) is Message:\n return \"Message\"\n if type(self.input) is Data:\n return \"Data\"\n # When invoked by a code agent (e.g. OpenDsStar), the input may be a raw\n # pandas DataFrame rather than Langflow's DataFrame wrapper.\n if isinstance(self.input, pd.DataFrame):\n self.input = DataFrame(self.input)\n return \"DataFrame\"\n # When invoked as a tool, the agent passes a string. Try to parse it as\n # tabular JSON (list of objects) → DataFrame, otherwise wrap as Message.\n if isinstance(self.input, str):\n self.input = self._coerce_string_input(self.input)\n return self._get_input_type()\n msg = f\"Unsupported input type: {type(self.input)}\"\n raise ValueError(msg)\n\n def _coerce_string_input(self, value: str) -> DataFrame | Message:\n \"\"\"Convert a raw string (from agent tool call) into a DataFrame or Message.\n\n Tries to parse as JSON first — a list of objects or a single object becomes\n a DataFrame. Anything else is wrapped in a Message.\n \"\"\"\n try:\n parsed = json.loads(value)\n if isinstance(parsed, list) and parsed and isinstance(parsed[0], dict):\n return DataFrame(pd.DataFrame(parsed))\n if isinstance(parsed, dict):\n return DataFrame(pd.DataFrame([parsed]))\n except (json.JSONDecodeError, ValueError):\n pass\n return Message(text=value)\n\n def _get_default_format(self) -> str:\n \"\"\"Return the default file format based on input type.\"\"\"\n if self._get_input_type() == \"DataFrame\":\n return \"csv\"\n if self._get_input_type() == \"Data\":\n return \"json\"\n if self._get_input_type() == \"Message\":\n return \"json\"\n return \"json\" # Fallback\n\n def _adjust_file_path_with_format(self, path: Path, fmt: str) -> Path:\n \"\"\"Adjust the file path to include the correct extension.\"\"\"\n file_extension = path.suffix.lower().lstrip(\".\")\n if fmt == \"excel\":\n return Path(f\"{path}.xlsx\").expanduser() if file_extension not in [\"xlsx\", \"xls\"] else path\n return Path(f\"{path}.{fmt}\").expanduser() if file_extension != fmt else path\n\n def _is_plain_text_format(self, fmt: str) -> bool:\n \"\"\"Check if a file format is plain text (supports appending).\"\"\"\n plain_text_formats = [\"txt\", \"json\", \"markdown\", \"md\", \"csv\", \"xml\", \"html\", \"yaml\", \"log\", \"tsv\", \"jsonl\"]\n return fmt.lower() in plain_text_formats\n\n async def _upload_file(self, file_path: Path) -> None:\n \"\"\"Upload the saved file using the upload_user_file service.\"\"\"\n from langflow.api.v2.files import upload_user_file\n from langflow.services.database.models.user.crud import get_user_by_id\n\n # Ensure the file exists\n if not file_path.exists():\n msg = f\"File not found: {file_path}\"\n raise FileNotFoundError(msg)\n\n # Upload the file - always use append=False because the local file already contains\n # the correct content (either new or appended locally)\n with file_path.open(\"rb\") as f:\n async with session_scope() as db:\n if not self.user_id:\n msg = \"User ID is required for file saving.\"\n raise ValueError(msg)\n current_user = await get_user_by_id(db, self.user_id)\n\n await upload_user_file(\n file=UploadFile(filename=file_path.name, file=f, size=file_path.stat().st_size),\n session=db,\n current_user=current_user,\n storage_service=get_storage_service(),\n settings_service=get_settings_service(),\n append=False,\n )\n\n def _save_dataframe(self, dataframe: DataFrame, path: Path, fmt: str) -> str:\n \"\"\"Save a DataFrame to the specified file format.\"\"\"\n append_mode = getattr(self, \"append_mode\", False)\n should_append = append_mode and path.exists() and self._is_plain_text_format(fmt)\n\n if fmt == \"csv\":\n dataframe.to_csv(path, index=False, mode=\"a\" if should_append else \"w\", header=not should_append)\n elif fmt == \"excel\":\n dataframe.to_excel(path, index=False, engine=\"openpyxl\")\n elif fmt == \"json\":\n if should_append:\n # Read and parse existing JSON\n existing_data = []\n try:\n existing_content = path.read_text(encoding=\"utf-8\").strip()\n if existing_content:\n parsed = json.loads(existing_content)\n # Handle case where existing content is a single object\n if isinstance(parsed, dict):\n existing_data = [parsed]\n elif isinstance(parsed, list):\n existing_data = parsed\n except (json.JSONDecodeError, FileNotFoundError):\n # Treat parse errors or missing file as empty array\n existing_data = []\n\n # Append new data\n new_records = json.loads(dataframe.to_json(orient=\"records\"))\n existing_data.extend(new_records)\n\n # Write back as a single JSON array\n path.write_text(json.dumps(existing_data, indent=2), encoding=\"utf-8\")\n else:\n dataframe.to_json(path, orient=\"records\", indent=2)\n elif fmt == \"markdown\":\n content = dataframe.to_markdown(index=False)\n if should_append:\n path.write_text(path.read_text(encoding=\"utf-8\") + \"\\n\\n\" + content, encoding=\"utf-8\")\n else:\n path.write_text(content, encoding=\"utf-8\")\n else:\n msg = f\"Unsupported DataFrame format: {fmt}\"\n raise ValueError(msg)\n action = \"appended to\" if should_append else \"saved successfully as\"\n return f\"DataFrame {action} '{path}'\"\n\n def _save_data(self, data: Data, path: Path, fmt: str) -> str:\n \"\"\"Save a Data object to the specified file format.\"\"\"\n append_mode = getattr(self, \"append_mode\", False)\n should_append = append_mode and path.exists() and self._is_plain_text_format(fmt)\n\n if fmt == \"csv\":\n pd.DataFrame(data.data).to_csv(\n path,\n index=False,\n mode=\"a\" if should_append else \"w\",\n header=not should_append,\n )\n elif fmt == \"excel\":\n pd.DataFrame(data.data).to_excel(path, index=False, engine=\"openpyxl\")\n elif fmt == \"json\":\n new_data = jsonable_encoder(data.data)\n if should_append:\n # Read and parse existing JSON\n existing_data = []\n try:\n existing_content = path.read_text(encoding=\"utf-8\").strip()\n if existing_content:\n parsed = json.loads(existing_content)\n # Handle case where existing content is a single object\n if isinstance(parsed, dict):\n existing_data = [parsed]\n elif isinstance(parsed, list):\n existing_data = parsed\n except (json.JSONDecodeError, FileNotFoundError):\n # Treat parse errors or missing file as empty array\n existing_data = []\n\n # Append new data\n if isinstance(new_data, list):\n existing_data.extend(new_data)\n else:\n existing_data.append(new_data)\n\n # Write back as a single JSON array\n path.write_text(json.dumps(existing_data, indent=2), encoding=\"utf-8\")\n else:\n content = orjson.dumps(new_data, option=orjson.OPT_INDENT_2).decode(\"utf-8\")\n path.write_text(content, encoding=\"utf-8\")\n elif fmt == \"markdown\":\n content = pd.DataFrame(data.data).to_markdown(index=False)\n if should_append:\n path.write_text(path.read_text(encoding=\"utf-8\") + \"\\n\\n\" + content, encoding=\"utf-8\")\n else:\n path.write_text(content, encoding=\"utf-8\")\n else:\n msg = f\"Unsupported Data format: {fmt}\"\n raise ValueError(msg)\n action = \"appended to\" if should_append else \"saved successfully as\"\n return f\"Data {action} '{path}'\"\n\n async def _save_message(self, message: Message, path: Path, fmt: str) -> str:\n \"\"\"Save a Message to the specified file format, handling async iterators.\"\"\"\n content = \"\"\n stream = message.text_stream if hasattr(message, \"text_stream\") else None\n if stream is not None and isinstance(stream, AsyncIterator):\n async for item in stream:\n content += str(item) + \" \"\n content = content.strip()\n elif stream is not None and isinstance(stream, Iterator):\n content = \" \".join(str(item) for item in stream)\n else:\n content = str(message.text)\n\n append_mode = getattr(self, \"append_mode\", False)\n should_append = append_mode and path.exists() and self._is_plain_text_format(fmt)\n\n if fmt in (\"txt\", \"html\"):\n if should_append:\n path.write_text(path.read_text(encoding=\"utf-8\") + \"\\n\" + content, encoding=\"utf-8\")\n else:\n path.write_text(content, encoding=\"utf-8\")\n elif fmt == \"json\":\n new_message = {\"message\": content}\n if should_append:\n # Read and parse existing JSON\n existing_data = []\n try:\n existing_content = path.read_text(encoding=\"utf-8\").strip()\n if existing_content:\n parsed = json.loads(existing_content)\n # Handle case where existing content is a single object\n if isinstance(parsed, dict):\n existing_data = [parsed]\n elif isinstance(parsed, list):\n existing_data = parsed\n except (json.JSONDecodeError, FileNotFoundError):\n # Treat parse errors or missing file as empty array\n existing_data = []\n\n # Append new message\n existing_data.append(new_message)\n\n # Write back as a single JSON array\n path.write_text(json.dumps(existing_data, indent=2), encoding=\"utf-8\")\n else:\n path.write_text(json.dumps(new_message, indent=2), encoding=\"utf-8\")\n elif fmt == \"markdown\":\n md_content = f\"**Message:**\\n\\n{content}\"\n if should_append:\n path.write_text(path.read_text(encoding=\"utf-8\") + \"\\n\\n\" + md_content, encoding=\"utf-8\")\n else:\n path.write_text(md_content, encoding=\"utf-8\")\n else:\n msg = f\"Unsupported Message format: {fmt}\"\n raise ValueError(msg)\n action = \"appended to\" if should_append else \"saved successfully as\"\n return f\"Message {action} '{path}'\"\n\n def _get_selected_storage_location(self) -> str:\n \"\"\"Get the selected storage location from the SortableListInput.\"\"\"\n if hasattr(self, \"storage_location\") and self.storage_location:\n if isinstance(self.storage_location, list) and len(self.storage_location) > 0:\n return self.storage_location[0].get(\"name\", \"\")\n if isinstance(self.storage_location, dict):\n return self.storage_location.get(\"name\", \"\")\n return \"\"\n\n def _get_file_format_for_location(self, location: str) -> str:\n \"\"\"Get the appropriate file format based on storage location.\n\n If the agent set file_format via tool mode, that takes priority.\n \"\"\"\n agent_format = getattr(self, \"file_format\", None)\n if agent_format:\n return agent_format\n if location == \"Local\":\n return getattr(self, \"local_format\", None) or self._get_default_format()\n if location == \"AWS\":\n return getattr(self, \"aws_format\", \"txt\")\n if location == \"Google Drive\":\n return getattr(self, \"gdrive_format\", \"txt\")\n return self._get_default_format()\n\n async def _save_to_local(self) -> Message:\n \"\"\"Save file to local storage (original functionality).\"\"\"\n file_format = self._get_file_format_for_location(\"Local\")\n\n # Validate file format based on input type\n allowed_formats = (\n self.LOCAL_MESSAGE_FORMAT_CHOICES if self._get_input_type() == \"Message\" else self.LOCAL_DATA_FORMAT_CHOICES\n )\n if file_format not in allowed_formats:\n msg = f\"Invalid file format '{file_format}' for {self._get_input_type()}. Allowed: {allowed_formats}\"\n raise ValueError(msg)\n\n # Prepare file path\n file_path = Path(self.file_name).expanduser()\n if not file_path.parent.exists():\n file_path.parent.mkdir(parents=True, exist_ok=True)\n file_path = self._adjust_file_path_with_format(file_path, file_format)\n\n # Save the input to file based on type\n if self._get_input_type() == \"DataFrame\":\n confirmation = self._save_dataframe(self.input, file_path, file_format)\n elif self._get_input_type() == \"Data\":\n confirmation = self._save_data(self.input, file_path, file_format)\n elif self._get_input_type() == \"Message\":\n confirmation = await self._save_message(self.input, file_path, file_format)\n else:\n msg = f\"Unsupported input type: {self._get_input_type()}\"\n raise ValueError(msg)\n\n # Upload the saved file\n await self._upload_file(file_path)\n\n # Return the final file path and confirmation message\n final_path = Path.cwd() / file_path if not file_path.is_absolute() else file_path\n return Message(text=f\"{confirmation} at {final_path}\")\n\n async def _save_to_aws(self) -> Message:\n \"\"\"Save file to AWS S3 using S3 functionality.\"\"\"\n import os\n\n import boto3\n\n from lfx.base.data.cloud_storage_utils import create_s3_client, validate_aws_credentials\n\n # Get AWS credentials from component inputs or fall back to environment variables\n aws_access_key_id = getattr(self, \"aws_access_key_id\", None)\n if aws_access_key_id and hasattr(aws_access_key_id, \"get_secret_value\"):\n aws_access_key_id = aws_access_key_id.get_secret_value()\n if not aws_access_key_id:\n aws_access_key_id = os.getenv(\"AWS_ACCESS_KEY_ID\")\n\n aws_secret_access_key = getattr(self, \"aws_secret_access_key\", None)\n if aws_secret_access_key and hasattr(aws_secret_access_key, \"get_secret_value\"):\n aws_secret_access_key = aws_secret_access_key.get_secret_value()\n if not aws_secret_access_key:\n aws_secret_access_key = os.getenv(\"AWS_SECRET_ACCESS_KEY\")\n\n bucket_name = getattr(self, \"bucket_name\", None)\n if not bucket_name:\n # Try to get from storage service settings\n settings = get_settings_service().settings\n bucket_name = settings.object_storage_bucket_name\n\n # Validate AWS credentials\n if not aws_access_key_id:\n msg = (\n \"AWS Access Key ID is required for S3 storage. Provide it as a component input \"\n \"or set AWS_ACCESS_KEY_ID environment variable.\"\n )\n raise ValueError(msg)\n if not aws_secret_access_key:\n msg = (\n \"AWS Secret Key is required for S3 storage. Provide it as a component input \"\n \"or set AWS_SECRET_ACCESS_KEY environment variable.\"\n )\n raise ValueError(msg)\n if not bucket_name:\n msg = (\n \"S3 Bucket Name is required for S3 storage. Provide it as a component input \"\n \"or set LANGFLOW_OBJECT_STORAGE_BUCKET_NAME environment variable.\"\n )\n raise ValueError(msg)\n\n # Validate AWS credentials\n validate_aws_credentials(self)\n\n # Create S3 client\n s3_client = create_s3_client(self)\n client_config: dict[str, Any] = {\n \"aws_access_key_id\": str(aws_access_key_id),\n \"aws_secret_access_key\": str(aws_secret_access_key),\n }\n\n # Get region from component input, environment variable, or settings\n aws_region = getattr(self, \"aws_region\", None)\n if not aws_region:\n aws_region = os.getenv(\"AWS_DEFAULT_REGION\") or os.getenv(\"AWS_REGION\")\n if aws_region:\n client_config[\"region_name\"] = str(aws_region)\n\n s3_client = boto3.client(\"s3\", **client_config)\n\n # Extract content\n content = self._extract_content_for_upload()\n file_format = self._get_file_format_for_location(\"AWS\")\n\n # Generate file path\n file_path = f\"{self.file_name}.{file_format}\"\n if hasattr(self, \"s3_prefix\") and self.s3_prefix:\n file_path = f\"{self.s3_prefix.rstrip('/')}/{file_path}\"\n\n # Create temporary file\n import tempfile\n\n with tempfile.NamedTemporaryFile(\n mode=\"w\", encoding=\"utf-8\", suffix=f\".{file_format}\", delete=False\n ) as temp_file:\n temp_file.write(content)\n temp_file_path = temp_file.name\n\n try:\n # Upload to S3\n s3_client.upload_file(temp_file_path, bucket_name, file_path)\n s3_url = f\"s3://{bucket_name}/{file_path}\"\n return Message(text=f\"File successfully uploaded to {s3_url}\")\n finally:\n # Clean up temp file\n if Path(temp_file_path).exists():\n Path(temp_file_path).unlink()\n\n async def _save_to_google_drive(self) -> Message:\n \"\"\"Save file to Google Drive using Google Drive functionality.\"\"\"\n import tempfile\n\n from googleapiclient.http import MediaFileUpload\n\n from lfx.base.data.cloud_storage_utils import create_google_drive_service\n\n # Validate Google Drive credentials\n if not getattr(self, \"service_account_key\", None):\n msg = \"GCP Credentials Secret Key is required for Google Drive storage\"\n raise ValueError(msg)\n if not getattr(self, \"folder_id\", None):\n msg = \"Google Drive Folder ID is required for Google Drive storage\"\n raise ValueError(msg)\n\n # Create Google Drive service with full drive scope (needed for folder operations)\n drive_service, credentials = create_google_drive_service(\n self.service_account_key, scopes=[\"https://www.googleapis.com/auth/drive\"], return_credentials=True\n )\n\n # Extract content and format\n content = self._extract_content_for_upload()\n file_format = self._get_file_format_for_location(\"Google Drive\")\n\n # Handle special Google Drive formats\n if file_format in [\"slides\", \"docs\"]:\n return await self._save_to_google_apps(drive_service, credentials, content, file_format)\n\n # Create temporary file\n file_path = f\"{self.file_name}.{file_format}\"\n with tempfile.NamedTemporaryFile(\n mode=\"w\",\n encoding=\"utf-8\",\n suffix=f\".{file_format}\",\n delete=False,\n ) as temp_file:\n temp_file.write(content)\n temp_file_path = temp_file.name\n\n try:\n # Upload to Google Drive\n # Note: We skip explicit folder verification since it requires broader permissions.\n # If the folder doesn't exist or isn't accessible, the create() call will fail with a clear error.\n file_metadata = {\"name\": file_path, \"parents\": [self.folder_id]}\n media = MediaFileUpload(temp_file_path, resumable=True)\n\n try:\n uploaded_file = (\n drive_service.files().create(body=file_metadata, media_body=media, fields=\"id\").execute()\n )\n except Exception as e:\n msg = (\n f\"Unable to upload file to Google Drive folder '{self.folder_id}'. \"\n f\"Error: {e!s}. \"\n \"Please ensure: 1) The folder ID is correct, 2) The folder exists, \"\n \"3) The service account has been granted access to this folder.\"\n )\n raise ValueError(msg) from e\n\n file_id = uploaded_file.get(\"id\")\n file_url = f\"https://drive.google.com/file/d/{file_id}/view\"\n return Message(text=f\"File successfully uploaded to Google Drive: {file_url}\")\n finally:\n # Clean up temp file\n if Path(temp_file_path).exists():\n Path(temp_file_path).unlink()\n\n async def _save_to_google_apps(self, drive_service, credentials, content: str, app_type: str) -> Message:\n \"\"\"Save content to Google Apps (Slides or Docs).\"\"\"\n import time\n\n if app_type == \"slides\":\n from googleapiclient.discovery import build\n\n slides_service = build(\"slides\", \"v1\", credentials=credentials)\n\n file_metadata = {\n \"name\": self.file_name,\n \"mimeType\": \"application/vnd.google-apps.presentation\",\n \"parents\": [self.folder_id],\n }\n\n created_file = drive_service.files().create(body=file_metadata, fields=\"id\").execute()\n presentation_id = created_file[\"id\"]\n\n time.sleep(2) # Wait for file to be available # noqa: ASYNC251\n\n presentation = slides_service.presentations().get(presentationId=presentation_id).execute()\n slide_id = presentation[\"slides\"][0][\"objectId\"]\n\n # Add content to slide\n requests = [\n {\n \"createShape\": {\n \"objectId\": \"TextBox_01\",\n \"shapeType\": \"TEXT_BOX\",\n \"elementProperties\": {\n \"pageObjectId\": slide_id,\n \"size\": {\n \"height\": {\"magnitude\": 3000000, \"unit\": \"EMU\"},\n \"width\": {\"magnitude\": 6000000, \"unit\": \"EMU\"},\n },\n \"transform\": {\n \"scaleX\": 1,\n \"scaleY\": 1,\n \"translateX\": 1000000,\n \"translateY\": 1000000,\n \"unit\": \"EMU\",\n },\n },\n }\n },\n {\"insertText\": {\"objectId\": \"TextBox_01\", \"insertionIndex\": 0, \"text\": content}},\n ]\n\n slides_service.presentations().batchUpdate(\n presentationId=presentation_id, body={\"requests\": requests}\n ).execute()\n file_url = f\"https://docs.google.com/presentation/d/{presentation_id}/edit\"\n\n elif app_type == \"docs\":\n from googleapiclient.discovery import build\n\n docs_service = build(\"docs\", \"v1\", credentials=credentials)\n\n file_metadata = {\n \"name\": self.file_name,\n \"mimeType\": \"application/vnd.google-apps.document\",\n \"parents\": [self.folder_id],\n }\n\n created_file = drive_service.files().create(body=file_metadata, fields=\"id\").execute()\n document_id = created_file[\"id\"]\n\n time.sleep(2) # Wait for file to be available # noqa: ASYNC251\n\n # Add content to document\n requests = [{\"insertText\": {\"location\": {\"index\": 1}, \"text\": content}}]\n docs_service.documents().batchUpdate(documentId=document_id, body={\"requests\": requests}).execute()\n file_url = f\"https://docs.google.com/document/d/{document_id}/edit\"\n\n return Message(text=f\"File successfully created in Google {app_type.title()}: {file_url}\")\n\n def _extract_content_for_upload(self) -> str:\n \"\"\"Extract content from input for upload to cloud services.\"\"\"\n if self._get_input_type() == \"DataFrame\":\n return self.input.to_csv(index=False)\n if self._get_input_type() == \"Data\":\n if hasattr(self.input, \"data\") and self.input.data:\n if isinstance(self.input.data, dict):\n import json\n\n return json.dumps(self.input.data, indent=2, ensure_ascii=False)\n return str(self.input.data)\n return str(self.input)\n if self._get_input_type() == \"Message\":\n return str(self.input.text) if self.input.text else str(self.input)\n return str(self.input)\n" }, "file_format": { "_input_type": "StrInput", @@ -118454,6 +118454,6 @@ "num_components": 354, "num_modules": 95 }, - "sha256": "87d65557fa5c4182f09f07270bbe8624d5f9b65a0322c3251f75822b6f1546ac", + "sha256": "57420a9e8822a72ba8e8f8fd0d597f5ff7b9d4f85d965ed40713d108512cc7fe", "version": "1.11.0" } From 54ef780eb62f093cf3b2a9207be2c710cd096609 Mon Sep 17 00:00:00 2001 From: "autofix-ci[bot]" <114827586+autofix-ci[bot]@users.noreply.github.com> Date: Wed, 10 Jun 2026 12:31:41 +0000 Subject: [PATCH 2/3] [autofix.ci] apply automated fixes --- src/lfx/src/lfx/_assets/component_index.json | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/lfx/src/lfx/_assets/component_index.json b/src/lfx/src/lfx/_assets/component_index.json index bce86b23c4..9fc99281be 100644 --- a/src/lfx/src/lfx/_assets/component_index.json +++ b/src/lfx/src/lfx/_assets/component_index.json @@ -70596,7 +70596,7 @@ "icon": "file-text", "legacy": false, "metadata": { - "code_hash": "837120b2e497", + "code_hash": "2897652b1a16", "dependencies": { "dependencies": [ { @@ -70807,7 +70807,7 @@ "show": true, "title_case": false, "type": "code", - "value": "import json\nfrom collections.abc import AsyncIterator, Iterator\nfrom pathlib import Path\nfrom typing import Any\n\nimport orjson\nimport pandas as pd\nfrom fastapi import UploadFile\nfrom fastapi.encoders import jsonable_encoder\n\nfrom lfx.custom import Component\nfrom lfx.inputs import SortableListInput\nfrom lfx.inputs.inputs import DataFrameInput\nfrom lfx.io import BoolInput, DropdownInput, SecretStrInput, StrInput\nfrom lfx.schema import Data, DataFrame, Message\nfrom lfx.services.deps import get_settings_service, get_storage_service, session_scope\nfrom lfx.template.field.base import Output\nfrom lfx.utils.validate_cloud import is_astra_cloud_environment\n\n\ndef _get_storage_location_options():\n \"\"\"Get storage location options, filtering out Local if in Astra cloud environment.\"\"\"\n all_options = [{\"name\": \"AWS\", \"icon\": \"Amazon\"}, {\"name\": \"Google Drive\", \"icon\": \"google\"}]\n if is_astra_cloud_environment():\n return all_options\n return [{\"name\": \"Local\", \"icon\": \"hard-drive\"}, *all_options]\n\n\ndef _get_default_storage_location() -> list[dict[str, str]]:\n \"\"\"Return the default storage selection for the component template.\"\"\"\n return [_get_storage_location_options()[0]]\n\n\ndef _is_default_storage(storage_name: str) -> bool:\n \"\"\"Check whether a storage type is the default selection.\"\"\"\n return _get_default_storage_location()[0][\"name\"] == storage_name\n\n\nclass SaveToFileComponent(Component):\n display_name = \"Write File\"\n description = (\n \"Save data to a file. \"\n \"Arguments: 'input' — the content to save (pass a DataFrame directly, or a JSON string \"\n \"for tabular data, or plain text for messages); \"\n \"'file_name' — the name to save as, without extension (e.g. 'report'); \"\n \"'file_format' — output format: 'csv', 'json', 'txt', 'html', 'excel', 'markdown' (optional). \"\n \"Returns a confirmation with the file path or URL.\"\n )\n documentation: str = \"https://docs.langflow.org/write-file\"\n icon = \"file-text\"\n name = \"SaveToFile\"\n\n # File format options for different storage types\n LOCAL_DATA_FORMAT_CHOICES = [\"csv\", \"excel\", \"json\", \"markdown\"]\n LOCAL_MESSAGE_FORMAT_CHOICES = [\"txt\", \"html\", \"json\", \"markdown\"]\n AWS_FORMAT_CHOICES = [\n \"txt\",\n \"json\",\n \"csv\",\n \"xml\",\n \"html\",\n \"md\",\n \"yaml\",\n \"log\",\n \"tsv\",\n \"jsonl\",\n \"parquet\",\n \"xlsx\",\n \"zip\",\n ]\n GDRIVE_FORMAT_CHOICES = [\"txt\", \"html\", \"json\", \"csv\", \"xlsx\", \"slides\", \"docs\", \"jpg\", \"mp3\"]\n\n inputs = [\n SortableListInput(\n name=\"storage_location\",\n display_name=\"Storage Location\",\n placeholder=\"Select Location\",\n info=\"Choose where to save the file.\",\n options=_get_storage_location_options(),\n real_time_refresh=True,\n limit=1,\n value=_get_default_storage_location(),\n advanced=True,\n ),\n # Common inputs\n DataFrameInput(\n name=\"input\",\n display_name=\"File Content\",\n info=(\n \"The content to save. Accepts a DataFrame, Data, or Message object directly. \"\n 'Can also accept a JSON string (e.g. \\'[{\"col1\": \"val1\"}]\\') which will be '\n \"parsed into a DataFrame, or plain text which will be saved as a Message.\"\n ),\n input_types=[\"Data\", \"JSON\", \"DataFrame\", \"Table\", \"Message\"],\n required=True,\n tool_mode=True,\n ),\n StrInput(\n name=\"file_name\",\n display_name=\"File Name\",\n info=\"File name without extension (e.g. 'report'). Extension is added automatically.\",\n required=True,\n show=True,\n tool_mode=True,\n ),\n StrInput(\n name=\"file_format\",\n display_name=\"File Format (Tool)\",\n info=\"Output format: 'csv', 'json', 'txt', 'html', 'excel', 'markdown'. Overrides pre-configured format.\",\n required=False,\n show=False,\n tool_mode=True,\n ),\n BoolInput(\n name=\"append_mode\",\n display_name=\"Append\",\n info=(\n \"Append to file if it exists (only for Local storage with plain text formats). \"\n \"Not supported for cloud storage (AWS/Google Drive).\"\n ),\n value=False,\n show=_is_default_storage(\"Local\"),\n ),\n # Format inputs (dynamic based on storage location)\n DropdownInput(\n name=\"local_format\",\n display_name=\"File Format\",\n options=list(dict.fromkeys(LOCAL_DATA_FORMAT_CHOICES + LOCAL_MESSAGE_FORMAT_CHOICES)),\n info=\"Select the file format for local storage.\",\n value=\"json\",\n show=_is_default_storage(\"Local\"),\n ),\n DropdownInput(\n name=\"aws_format\",\n display_name=\"File Format\",\n options=AWS_FORMAT_CHOICES,\n info=\"Select the file format for AWS S3 storage.\",\n value=\"txt\",\n show=_is_default_storage(\"AWS\"),\n ),\n DropdownInput(\n name=\"gdrive_format\",\n display_name=\"File Format\",\n options=GDRIVE_FORMAT_CHOICES,\n info=\"Select the file format for Google Drive storage.\",\n value=\"txt\",\n show=_is_default_storage(\"Google Drive\"),\n ),\n # AWS S3 specific inputs\n SecretStrInput(\n name=\"aws_access_key_id\",\n display_name=\"AWS Access Key ID\",\n info=\"AWS Access key ID.\",\n show=_is_default_storage(\"AWS\"),\n advanced=not _is_default_storage(\"AWS\"),\n required=True,\n ),\n SecretStrInput(\n name=\"aws_secret_access_key\",\n display_name=\"AWS Secret Key\",\n info=\"AWS Secret Key.\",\n show=_is_default_storage(\"AWS\"),\n advanced=not _is_default_storage(\"AWS\"),\n required=True,\n ),\n StrInput(\n name=\"bucket_name\",\n display_name=\"S3 Bucket Name\",\n info=\"Enter the name of the S3 bucket.\",\n show=_is_default_storage(\"AWS\"),\n advanced=not _is_default_storage(\"AWS\"),\n required=True,\n ),\n StrInput(\n name=\"aws_region\",\n display_name=\"AWS Region\",\n info=\"AWS region (e.g., us-east-1, eu-west-1).\",\n show=_is_default_storage(\"AWS\"),\n advanced=not _is_default_storage(\"AWS\"),\n ),\n StrInput(\n name=\"s3_prefix\",\n display_name=\"S3 Prefix\",\n info=\"Prefix for all files in S3.\",\n show=_is_default_storage(\"AWS\"),\n advanced=not _is_default_storage(\"AWS\"),\n ),\n # Google Drive specific inputs\n SecretStrInput(\n name=\"service_account_key\",\n display_name=\"GCP Credentials Secret Key\",\n info=\"Your Google Cloud Platform service account JSON key as a secret string (complete JSON content).\",\n show=_is_default_storage(\"Google Drive\"),\n advanced=not _is_default_storage(\"Google Drive\"),\n required=True,\n ),\n StrInput(\n name=\"folder_id\",\n display_name=\"Google Drive Folder ID\",\n info=(\n \"The Google Drive folder ID where the file will be uploaded. \"\n \"The folder must be shared with the service account email.\"\n ),\n required=True,\n show=_is_default_storage(\"Google Drive\"),\n advanced=not _is_default_storage(\"Google Drive\"),\n ),\n ]\n\n outputs = [Output(display_name=\"File Path\", name=\"message\", method=\"save_to_file\")]\n\n def update_build_config(self, build_config, field_value, field_name=None):\n \"\"\"Update build configuration to show/hide fields based on storage location selection.\"\"\"\n # Update options dynamically based on cloud environment\n # This ensures options are refreshed when build_config is updated\n if \"storage_location\" in build_config:\n updated_options = _get_storage_location_options()\n build_config[\"storage_location\"][\"options\"] = updated_options\n\n # When tool_mode is toggled, hide storage-specific format dropdowns\n # (the agent uses the unified file_format input instead)\n if field_name == \"tool_mode\":\n format_fields = [\"local_format\", \"aws_format\", \"gdrive_format\"]\n for f_name in format_fields:\n if f_name in build_config:\n build_config[f_name][\"show\"] = not bool(field_value)\n return build_config\n\n if field_name != \"storage_location\":\n return build_config\n\n # Extract selected storage location\n selected = [location[\"name\"] for location in field_value] if isinstance(field_value, list) else []\n\n # Hide all dynamic fields first\n dynamic_fields = [\n \"file_name\", # Common fields (input is always visible)\n \"append_mode\",\n \"local_format\",\n \"aws_format\",\n \"gdrive_format\",\n \"aws_access_key_id\",\n \"aws_secret_access_key\",\n \"bucket_name\",\n \"aws_region\",\n \"s3_prefix\",\n \"service_account_key\",\n \"folder_id\",\n ]\n\n for f_name in dynamic_fields:\n if f_name in build_config:\n build_config[f_name][\"show\"] = False\n\n # Show fields based on selected storage location\n is_tool_mode = build_config.get(\"tools_metadata\", {}).get(\"show\", False)\n\n if len(selected) == 1:\n location = selected[0]\n\n # Show file_name when any storage location is selected\n if \"file_name\" in build_config:\n build_config[\"file_name\"][\"show\"] = True\n\n # Show append_mode only for Local storage (not supported for cloud storage)\n if \"append_mode\" in build_config:\n build_config[\"append_mode\"][\"show\"] = location == \"Local\"\n\n if location == \"Local\":\n if \"local_format\" in build_config:\n build_config[\"local_format\"][\"show\"] = not is_tool_mode\n\n elif location == \"AWS\":\n aws_fields = [\n \"aws_format\",\n \"aws_access_key_id\",\n \"aws_secret_access_key\",\n \"bucket_name\",\n \"aws_region\",\n \"s3_prefix\",\n ]\n for f_name in aws_fields:\n if f_name in build_config:\n show = f_name != \"aws_format\" or not is_tool_mode\n build_config[f_name][\"show\"] = show\n build_config[f_name][\"advanced\"] = False\n\n elif location == \"Google Drive\":\n gdrive_fields = [\"gdrive_format\", \"service_account_key\", \"folder_id\"]\n for f_name in gdrive_fields:\n if f_name in build_config:\n show = f_name != \"gdrive_format\" or not is_tool_mode\n build_config[f_name][\"show\"] = show\n build_config[f_name][\"advanced\"] = False\n\n return build_config\n\n async def save_to_file(self) -> Message:\n \"\"\"Save the input to a file and upload it, returning a confirmation message.\"\"\"\n # Validate inputs\n if not self.file_name:\n msg = \"File name must be provided.\"\n raise ValueError(msg)\n if not self._get_input_type():\n msg = \"Input type is not set.\"\n raise ValueError(msg)\n\n # Get selected storage location\n storage_location = self._get_selected_storage_location()\n if not storage_location:\n msg = \"Storage location must be selected.\"\n raise ValueError(msg)\n\n # Check if Local storage is disabled in cloud environment\n if storage_location == \"Local\" and is_astra_cloud_environment():\n msg = \"Local storage is not available in cloud environment. Please use AWS or Google Drive.\"\n raise ValueError(msg)\n\n # Route to appropriate save method based on storage location\n if storage_location == \"Local\":\n return await self._save_to_local()\n if storage_location == \"AWS\":\n return await self._save_to_aws()\n if storage_location == \"Google Drive\":\n return await self._save_to_google_drive()\n msg = f\"Unsupported storage location: {storage_location}\"\n raise ValueError(msg)\n\n def _get_input_type(self) -> str:\n \"\"\"Determine the input type based on the provided input.\"\"\"\n # Use exact type checking (type() is) instead of isinstance() to avoid inheritance issues.\n # Since Message inherits from Data, isinstance(message, Data) would return True for Message objects,\n # causing Message inputs to be incorrectly identified as Data type.\n if type(self.input) is DataFrame:\n return \"DataFrame\"\n if type(self.input) is Message:\n return \"Message\"\n if type(self.input) is Data:\n return \"Data\"\n # When invoked by a code agent (e.g. OpenDsStar), the input may be a raw\n # pandas DataFrame rather than Langflow's DataFrame wrapper.\n if isinstance(self.input, pd.DataFrame):\n self.input = DataFrame(self.input)\n return \"DataFrame\"\n # When invoked as a tool, the agent passes a string. Try to parse it as\n # tabular JSON (list of objects) → DataFrame, otherwise wrap as Message.\n if isinstance(self.input, str):\n self.input = self._coerce_string_input(self.input)\n return self._get_input_type()\n msg = f\"Unsupported input type: {type(self.input)}\"\n raise ValueError(msg)\n\n def _coerce_string_input(self, value: str) -> DataFrame | Message:\n \"\"\"Convert a raw string (from agent tool call) into a DataFrame or Message.\n\n Tries to parse as JSON first — a list of objects or a single object becomes\n a DataFrame. Anything else is wrapped in a Message.\n \"\"\"\n try:\n parsed = json.loads(value)\n if isinstance(parsed, list) and parsed and isinstance(parsed[0], dict):\n return DataFrame(pd.DataFrame(parsed))\n if isinstance(parsed, dict):\n return DataFrame(pd.DataFrame([parsed]))\n except (json.JSONDecodeError, ValueError):\n pass\n return Message(text=value)\n\n def _get_default_format(self) -> str:\n \"\"\"Return the default file format based on input type.\"\"\"\n if self._get_input_type() == \"DataFrame\":\n return \"csv\"\n if self._get_input_type() == \"Data\":\n return \"json\"\n if self._get_input_type() == \"Message\":\n return \"json\"\n return \"json\" # Fallback\n\n def _adjust_file_path_with_format(self, path: Path, fmt: str) -> Path:\n \"\"\"Adjust the file path to include the correct extension.\"\"\"\n file_extension = path.suffix.lower().lstrip(\".\")\n if fmt == \"excel\":\n return Path(f\"{path}.xlsx\").expanduser() if file_extension not in [\"xlsx\", \"xls\"] else path\n return Path(f\"{path}.{fmt}\").expanduser() if file_extension != fmt else path\n\n def _is_plain_text_format(self, fmt: str) -> bool:\n \"\"\"Check if a file format is plain text (supports appending).\"\"\"\n plain_text_formats = [\"txt\", \"json\", \"markdown\", \"md\", \"csv\", \"xml\", \"html\", \"yaml\", \"log\", \"tsv\", \"jsonl\"]\n return fmt.lower() in plain_text_formats\n\n async def _upload_file(self, file_path: Path) -> None:\n \"\"\"Upload the saved file using the upload_user_file service.\"\"\"\n from langflow.api.v2.files import upload_user_file\n from langflow.services.database.models.user.crud import get_user_by_id\n\n # Ensure the file exists\n if not file_path.exists():\n msg = f\"File not found: {file_path}\"\n raise FileNotFoundError(msg)\n\n # Upload the file - always use append=False because the local file already contains\n # the correct content (either new or appended locally)\n with file_path.open(\"rb\") as f:\n async with session_scope() as db:\n if not self.user_id:\n msg = \"User ID is required for file saving.\"\n raise ValueError(msg)\n current_user = await get_user_by_id(db, self.user_id)\n\n await upload_user_file(\n file=UploadFile(filename=file_path.name, file=f, size=file_path.stat().st_size),\n session=db,\n current_user=current_user,\n storage_service=get_storage_service(),\n settings_service=get_settings_service(),\n append=False,\n )\n\n def _save_dataframe(self, dataframe: DataFrame, path: Path, fmt: str) -> str:\n \"\"\"Save a DataFrame to the specified file format.\"\"\"\n append_mode = getattr(self, \"append_mode\", False)\n should_append = append_mode and path.exists() and self._is_plain_text_format(fmt)\n\n if fmt == \"csv\":\n dataframe.to_csv(path, index=False, mode=\"a\" if should_append else \"w\", header=not should_append)\n elif fmt == \"excel\":\n dataframe.to_excel(path, index=False, engine=\"openpyxl\")\n elif fmt == \"json\":\n if should_append:\n # Read and parse existing JSON\n existing_data = []\n try:\n existing_content = path.read_text(encoding=\"utf-8\").strip()\n if existing_content:\n parsed = json.loads(existing_content)\n # Handle case where existing content is a single object\n if isinstance(parsed, dict):\n existing_data = [parsed]\n elif isinstance(parsed, list):\n existing_data = parsed\n except (json.JSONDecodeError, FileNotFoundError):\n # Treat parse errors or missing file as empty array\n existing_data = []\n\n # Append new data\n new_records = json.loads(dataframe.to_json(orient=\"records\"))\n existing_data.extend(new_records)\n\n # Write back as a single JSON array\n path.write_text(json.dumps(existing_data, indent=2), encoding=\"utf-8\")\n else:\n dataframe.to_json(path, orient=\"records\", indent=2)\n elif fmt == \"markdown\":\n content = dataframe.to_markdown(index=False)\n if should_append:\n path.write_text(path.read_text(encoding=\"utf-8\") + \"\\n\\n\" + content, encoding=\"utf-8\")\n else:\n path.write_text(content, encoding=\"utf-8\")\n else:\n msg = f\"Unsupported DataFrame format: {fmt}\"\n raise ValueError(msg)\n action = \"appended to\" if should_append else \"saved successfully as\"\n return f\"DataFrame {action} '{path}'\"\n\n def _save_data(self, data: Data, path: Path, fmt: str) -> str:\n \"\"\"Save a Data object to the specified file format.\"\"\"\n append_mode = getattr(self, \"append_mode\", False)\n should_append = append_mode and path.exists() and self._is_plain_text_format(fmt)\n\n if fmt == \"csv\":\n pd.DataFrame(data.data).to_csv(\n path,\n index=False,\n mode=\"a\" if should_append else \"w\",\n header=not should_append,\n )\n elif fmt == \"excel\":\n pd.DataFrame(data.data).to_excel(path, index=False, engine=\"openpyxl\")\n elif fmt == \"json\":\n new_data = jsonable_encoder(data.data)\n if should_append:\n # Read and parse existing JSON\n existing_data = []\n try:\n existing_content = path.read_text(encoding=\"utf-8\").strip()\n if existing_content:\n parsed = json.loads(existing_content)\n # Handle case where existing content is a single object\n if isinstance(parsed, dict):\n existing_data = [parsed]\n elif isinstance(parsed, list):\n existing_data = parsed\n except (json.JSONDecodeError, FileNotFoundError):\n # Treat parse errors or missing file as empty array\n existing_data = []\n\n # Append new data\n if isinstance(new_data, list):\n existing_data.extend(new_data)\n else:\n existing_data.append(new_data)\n\n # Write back as a single JSON array\n path.write_text(json.dumps(existing_data, indent=2), encoding=\"utf-8\")\n else:\n content = orjson.dumps(new_data, option=orjson.OPT_INDENT_2).decode(\"utf-8\")\n path.write_text(content, encoding=\"utf-8\")\n elif fmt == \"markdown\":\n content = pd.DataFrame(data.data).to_markdown(index=False)\n if should_append:\n path.write_text(path.read_text(encoding=\"utf-8\") + \"\\n\\n\" + content, encoding=\"utf-8\")\n else:\n path.write_text(content, encoding=\"utf-8\")\n else:\n msg = f\"Unsupported Data format: {fmt}\"\n raise ValueError(msg)\n action = \"appended to\" if should_append else \"saved successfully as\"\n return f\"Data {action} '{path}'\"\n\n async def _save_message(self, message: Message, path: Path, fmt: str) -> str:\n \"\"\"Save a Message to the specified file format, handling async iterators.\"\"\"\n content = \"\"\n if message.text is None:\n content = \"\"\n elif isinstance(message.text, AsyncIterator):\n async for item in message.text:\n content += str(item) + \" \"\n content = content.strip()\n elif isinstance(message.text, Iterator):\n content = \" \".join(str(item) for item in message.text)\n else:\n content = str(message.text)\n\n append_mode = getattr(self, \"append_mode\", False)\n should_append = append_mode and path.exists() and self._is_plain_text_format(fmt)\n\n if fmt in (\"txt\", \"html\"):\n if should_append:\n path.write_text(path.read_text(encoding=\"utf-8\") + \"\\n\" + content, encoding=\"utf-8\")\n else:\n path.write_text(content, encoding=\"utf-8\")\n elif fmt == \"json\":\n new_message = {\"message\": content}\n if should_append:\n # Read and parse existing JSON\n existing_data = []\n try:\n existing_content = path.read_text(encoding=\"utf-8\").strip()\n if existing_content:\n parsed = json.loads(existing_content)\n # Handle case where existing content is a single object\n if isinstance(parsed, dict):\n existing_data = [parsed]\n elif isinstance(parsed, list):\n existing_data = parsed\n except (json.JSONDecodeError, FileNotFoundError):\n # Treat parse errors or missing file as empty array\n existing_data = []\n\n # Append new message\n existing_data.append(new_message)\n\n # Write back as a single JSON array\n path.write_text(json.dumps(existing_data, indent=2), encoding=\"utf-8\")\n else:\n path.write_text(json.dumps(new_message, indent=2), encoding=\"utf-8\")\n elif fmt == \"markdown\":\n md_content = f\"**Message:**\\n\\n{content}\"\n if should_append:\n path.write_text(path.read_text(encoding=\"utf-8\") + \"\\n\\n\" + md_content, encoding=\"utf-8\")\n else:\n path.write_text(md_content, encoding=\"utf-8\")\n else:\n msg = f\"Unsupported Message format: {fmt}\"\n raise ValueError(msg)\n action = \"appended to\" if should_append else \"saved successfully as\"\n return f\"Message {action} '{path}'\"\n\n def _get_selected_storage_location(self) -> str:\n \"\"\"Get the selected storage location from the SortableListInput.\"\"\"\n if hasattr(self, \"storage_location\") and self.storage_location:\n if isinstance(self.storage_location, list) and len(self.storage_location) > 0:\n return self.storage_location[0].get(\"name\", \"\")\n if isinstance(self.storage_location, dict):\n return self.storage_location.get(\"name\", \"\")\n return \"\"\n\n def _get_file_format_for_location(self, location: str) -> str:\n \"\"\"Get the appropriate file format based on storage location.\n\n If the agent set file_format via tool mode, that takes priority.\n \"\"\"\n agent_format = getattr(self, \"file_format\", None)\n if agent_format:\n return agent_format\n if location == \"Local\":\n return getattr(self, \"local_format\", None) or self._get_default_format()\n if location == \"AWS\":\n return getattr(self, \"aws_format\", \"txt\")\n if location == \"Google Drive\":\n return getattr(self, \"gdrive_format\", \"txt\")\n return self._get_default_format()\n\n async def _save_to_local(self) -> Message:\n \"\"\"Save file to local storage (original functionality).\"\"\"\n file_format = self._get_file_format_for_location(\"Local\")\n\n # Validate file format based on input type\n allowed_formats = (\n self.LOCAL_MESSAGE_FORMAT_CHOICES if self._get_input_type() == \"Message\" else self.LOCAL_DATA_FORMAT_CHOICES\n )\n if file_format not in allowed_formats:\n msg = f\"Invalid file format '{file_format}' for {self._get_input_type()}. Allowed: {allowed_formats}\"\n raise ValueError(msg)\n\n # Prepare file path\n file_path = Path(self.file_name).expanduser()\n if not file_path.parent.exists():\n file_path.parent.mkdir(parents=True, exist_ok=True)\n file_path = self._adjust_file_path_with_format(file_path, file_format)\n\n # Save the input to file based on type\n if self._get_input_type() == \"DataFrame\":\n confirmation = self._save_dataframe(self.input, file_path, file_format)\n elif self._get_input_type() == \"Data\":\n confirmation = self._save_data(self.input, file_path, file_format)\n elif self._get_input_type() == \"Message\":\n confirmation = await self._save_message(self.input, file_path, file_format)\n else:\n msg = f\"Unsupported input type: {self._get_input_type()}\"\n raise ValueError(msg)\n\n # Upload the saved file\n await self._upload_file(file_path)\n\n # Return the final file path and confirmation message\n final_path = Path.cwd() / file_path if not file_path.is_absolute() else file_path\n return Message(text=f\"{confirmation} at {final_path}\")\n\n async def _save_to_aws(self) -> Message:\n \"\"\"Save file to AWS S3 using S3 functionality.\"\"\"\n import os\n\n import boto3\n\n from lfx.base.data.cloud_storage_utils import create_s3_client, validate_aws_credentials\n\n # Get AWS credentials from component inputs or fall back to environment variables\n aws_access_key_id = getattr(self, \"aws_access_key_id\", None)\n if aws_access_key_id and hasattr(aws_access_key_id, \"get_secret_value\"):\n aws_access_key_id = aws_access_key_id.get_secret_value()\n if not aws_access_key_id:\n aws_access_key_id = os.getenv(\"AWS_ACCESS_KEY_ID\")\n\n aws_secret_access_key = getattr(self, \"aws_secret_access_key\", None)\n if aws_secret_access_key and hasattr(aws_secret_access_key, \"get_secret_value\"):\n aws_secret_access_key = aws_secret_access_key.get_secret_value()\n if not aws_secret_access_key:\n aws_secret_access_key = os.getenv(\"AWS_SECRET_ACCESS_KEY\")\n\n bucket_name = getattr(self, \"bucket_name\", None)\n if not bucket_name:\n # Try to get from storage service settings\n settings = get_settings_service().settings\n bucket_name = settings.object_storage_bucket_name\n\n # Validate AWS credentials\n if not aws_access_key_id:\n msg = (\n \"AWS Access Key ID is required for S3 storage. Provide it as a component input \"\n \"or set AWS_ACCESS_KEY_ID environment variable.\"\n )\n raise ValueError(msg)\n if not aws_secret_access_key:\n msg = (\n \"AWS Secret Key is required for S3 storage. Provide it as a component input \"\n \"or set AWS_SECRET_ACCESS_KEY environment variable.\"\n )\n raise ValueError(msg)\n if not bucket_name:\n msg = (\n \"S3 Bucket Name is required for S3 storage. Provide it as a component input \"\n \"or set LANGFLOW_OBJECT_STORAGE_BUCKET_NAME environment variable.\"\n )\n raise ValueError(msg)\n\n # Validate AWS credentials\n validate_aws_credentials(self)\n\n # Create S3 client\n s3_client = create_s3_client(self)\n client_config: dict[str, Any] = {\n \"aws_access_key_id\": str(aws_access_key_id),\n \"aws_secret_access_key\": str(aws_secret_access_key),\n }\n\n # Get region from component input, environment variable, or settings\n aws_region = getattr(self, \"aws_region\", None)\n if not aws_region:\n aws_region = os.getenv(\"AWS_DEFAULT_REGION\") or os.getenv(\"AWS_REGION\")\n if aws_region:\n client_config[\"region_name\"] = str(aws_region)\n\n s3_client = boto3.client(\"s3\", **client_config)\n\n # Extract content\n content = self._extract_content_for_upload()\n file_format = self._get_file_format_for_location(\"AWS\")\n\n # Generate file path\n file_path = f\"{self.file_name}.{file_format}\"\n if hasattr(self, \"s3_prefix\") and self.s3_prefix:\n file_path = f\"{self.s3_prefix.rstrip('/')}/{file_path}\"\n\n # Create temporary file\n import tempfile\n\n with tempfile.NamedTemporaryFile(\n mode=\"w\", encoding=\"utf-8\", suffix=f\".{file_format}\", delete=False\n ) as temp_file:\n temp_file.write(content)\n temp_file_path = temp_file.name\n\n try:\n # Upload to S3\n s3_client.upload_file(temp_file_path, bucket_name, file_path)\n s3_url = f\"s3://{bucket_name}/{file_path}\"\n return Message(text=f\"File successfully uploaded to {s3_url}\")\n finally:\n # Clean up temp file\n if Path(temp_file_path).exists():\n Path(temp_file_path).unlink()\n\n async def _save_to_google_drive(self) -> Message:\n \"\"\"Save file to Google Drive using Google Drive functionality.\"\"\"\n import tempfile\n\n from googleapiclient.http import MediaFileUpload\n\n from lfx.base.data.cloud_storage_utils import create_google_drive_service\n\n # Validate Google Drive credentials\n if not getattr(self, \"service_account_key\", None):\n msg = \"GCP Credentials Secret Key is required for Google Drive storage\"\n raise ValueError(msg)\n if not getattr(self, \"folder_id\", None):\n msg = \"Google Drive Folder ID is required for Google Drive storage\"\n raise ValueError(msg)\n\n # Create Google Drive service with full drive scope (needed for folder operations)\n drive_service, credentials = create_google_drive_service(\n self.service_account_key, scopes=[\"https://www.googleapis.com/auth/drive\"], return_credentials=True\n )\n\n # Extract content and format\n content = self._extract_content_for_upload()\n file_format = self._get_file_format_for_location(\"Google Drive\")\n\n # Handle special Google Drive formats\n if file_format in [\"slides\", \"docs\"]:\n return await self._save_to_google_apps(drive_service, credentials, content, file_format)\n\n # Create temporary file\n file_path = f\"{self.file_name}.{file_format}\"\n with tempfile.NamedTemporaryFile(\n mode=\"w\",\n encoding=\"utf-8\",\n suffix=f\".{file_format}\",\n delete=False,\n ) as temp_file:\n temp_file.write(content)\n temp_file_path = temp_file.name\n\n try:\n # Upload to Google Drive\n # Note: We skip explicit folder verification since it requires broader permissions.\n # If the folder doesn't exist or isn't accessible, the create() call will fail with a clear error.\n file_metadata = {\"name\": file_path, \"parents\": [self.folder_id]}\n media = MediaFileUpload(temp_file_path, resumable=True)\n\n try:\n uploaded_file = (\n drive_service.files().create(body=file_metadata, media_body=media, fields=\"id\").execute()\n )\n except Exception as e:\n msg = (\n f\"Unable to upload file to Google Drive folder '{self.folder_id}'. \"\n f\"Error: {e!s}. \"\n \"Please ensure: 1) The folder ID is correct, 2) The folder exists, \"\n \"3) The service account has been granted access to this folder.\"\n )\n raise ValueError(msg) from e\n\n file_id = uploaded_file.get(\"id\")\n file_url = f\"https://drive.google.com/file/d/{file_id}/view\"\n return Message(text=f\"File successfully uploaded to Google Drive: {file_url}\")\n finally:\n # Clean up temp file\n if Path(temp_file_path).exists():\n Path(temp_file_path).unlink()\n\n async def _save_to_google_apps(self, drive_service, credentials, content: str, app_type: str) -> Message:\n \"\"\"Save content to Google Apps (Slides or Docs).\"\"\"\n import time\n\n if app_type == \"slides\":\n from googleapiclient.discovery import build\n\n slides_service = build(\"slides\", \"v1\", credentials=credentials)\n\n file_metadata = {\n \"name\": self.file_name,\n \"mimeType\": \"application/vnd.google-apps.presentation\",\n \"parents\": [self.folder_id],\n }\n\n created_file = drive_service.files().create(body=file_metadata, fields=\"id\").execute()\n presentation_id = created_file[\"id\"]\n\n time.sleep(2) # Wait for file to be available # noqa: ASYNC251\n\n presentation = slides_service.presentations().get(presentationId=presentation_id).execute()\n slide_id = presentation[\"slides\"][0][\"objectId\"]\n\n # Add content to slide\n requests = [\n {\n \"createShape\": {\n \"objectId\": \"TextBox_01\",\n \"shapeType\": \"TEXT_BOX\",\n \"elementProperties\": {\n \"pageObjectId\": slide_id,\n \"size\": {\n \"height\": {\"magnitude\": 3000000, \"unit\": \"EMU\"},\n \"width\": {\"magnitude\": 6000000, \"unit\": \"EMU\"},\n },\n \"transform\": {\n \"scaleX\": 1,\n \"scaleY\": 1,\n \"translateX\": 1000000,\n \"translateY\": 1000000,\n \"unit\": \"EMU\",\n },\n },\n }\n },\n {\"insertText\": {\"objectId\": \"TextBox_01\", \"insertionIndex\": 0, \"text\": content}},\n ]\n\n slides_service.presentations().batchUpdate(\n presentationId=presentation_id, body={\"requests\": requests}\n ).execute()\n file_url = f\"https://docs.google.com/presentation/d/{presentation_id}/edit\"\n\n elif app_type == \"docs\":\n from googleapiclient.discovery import build\n\n docs_service = build(\"docs\", \"v1\", credentials=credentials)\n\n file_metadata = {\n \"name\": self.file_name,\n \"mimeType\": \"application/vnd.google-apps.document\",\n \"parents\": [self.folder_id],\n }\n\n created_file = drive_service.files().create(body=file_metadata, fields=\"id\").execute()\n document_id = created_file[\"id\"]\n\n time.sleep(2) # Wait for file to be available # noqa: ASYNC251\n\n # Add content to document\n requests = [{\"insertText\": {\"location\": {\"index\": 1}, \"text\": content}}]\n docs_service.documents().batchUpdate(documentId=document_id, body={\"requests\": requests}).execute()\n file_url = f\"https://docs.google.com/document/d/{document_id}/edit\"\n\n return Message(text=f\"File successfully created in Google {app_type.title()}: {file_url}\")\n\n def _extract_content_for_upload(self) -> str:\n \"\"\"Extract content from input for upload to cloud services.\"\"\"\n if self._get_input_type() == \"DataFrame\":\n return self.input.to_csv(index=False)\n if self._get_input_type() == \"Data\":\n if hasattr(self.input, \"data\") and self.input.data:\n if isinstance(self.input.data, dict):\n import json\n\n return json.dumps(self.input.data, indent=2, ensure_ascii=False)\n return str(self.input.data)\n return str(self.input)\n if self._get_input_type() == \"Message\":\n return str(self.input.text) if self.input.text else str(self.input)\n return str(self.input)\n" + "value": "import json\nfrom collections.abc import AsyncIterator, Iterator\nfrom pathlib import Path\nfrom typing import Any\n\nimport orjson\nimport pandas as pd\nfrom fastapi import UploadFile\nfrom fastapi.encoders import jsonable_encoder\n\nfrom lfx.custom import Component\nfrom lfx.inputs import SortableListInput\nfrom lfx.inputs.inputs import DataFrameInput\nfrom lfx.io import BoolInput, DropdownInput, SecretStrInput, StrInput\nfrom lfx.schema import Data, DataFrame, Message\nfrom lfx.services.deps import get_settings_service, get_storage_service, session_scope\nfrom lfx.template.field.base import Output\nfrom lfx.utils.validate_cloud import is_astra_cloud_environment\n\n\ndef _get_storage_location_options():\n \"\"\"Get storage location options, filtering out Local if in Astra cloud environment.\"\"\"\n all_options = [{\"name\": \"AWS\", \"icon\": \"Amazon\"}, {\"name\": \"Google Drive\", \"icon\": \"google\"}]\n if is_astra_cloud_environment():\n return all_options\n return [{\"name\": \"Local\", \"icon\": \"hard-drive\"}, *all_options]\n\n\ndef _get_default_storage_location() -> list[dict[str, str]]:\n \"\"\"Return the default storage selection for the component template.\"\"\"\n return [_get_storage_location_options()[0]]\n\n\ndef _is_default_storage(storage_name: str) -> bool:\n \"\"\"Check whether a storage type is the default selection.\"\"\"\n return _get_default_storage_location()[0][\"name\"] == storage_name\n\n\nclass SaveToFileComponent(Component):\n display_name = \"Write File\"\n description = (\n \"Save data to a file. \"\n \"Arguments: 'input' — the content to save (pass a DataFrame directly, or a JSON string \"\n \"for tabular data, or plain text for messages); \"\n \"'file_name' — the name to save as, without extension (e.g. 'report'); \"\n \"'file_format' — output format: 'csv', 'json', 'txt', 'html', 'excel', 'markdown' (optional). \"\n \"Returns a confirmation with the file path or URL.\"\n )\n documentation: str = \"https://docs.langflow.org/write-file\"\n icon = \"file-text\"\n name = \"SaveToFile\"\n\n # File format options for different storage types\n LOCAL_DATA_FORMAT_CHOICES = [\"csv\", \"excel\", \"json\", \"markdown\"]\n LOCAL_MESSAGE_FORMAT_CHOICES = [\"txt\", \"html\", \"json\", \"markdown\"]\n AWS_FORMAT_CHOICES = [\n \"txt\",\n \"json\",\n \"csv\",\n \"xml\",\n \"html\",\n \"md\",\n \"yaml\",\n \"log\",\n \"tsv\",\n \"jsonl\",\n \"parquet\",\n \"xlsx\",\n \"zip\",\n ]\n GDRIVE_FORMAT_CHOICES = [\"txt\", \"html\", \"json\", \"csv\", \"xlsx\", \"slides\", \"docs\", \"jpg\", \"mp3\"]\n\n inputs = [\n SortableListInput(\n name=\"storage_location\",\n display_name=\"Storage Location\",\n placeholder=\"Select Location\",\n info=\"Choose where to save the file.\",\n options=_get_storage_location_options(),\n real_time_refresh=True,\n limit=1,\n value=_get_default_storage_location(),\n advanced=True,\n ),\n # Common inputs\n DataFrameInput(\n name=\"input\",\n display_name=\"File Content\",\n info=(\n \"The content to save. Accepts a DataFrame, Data, or Message object directly. \"\n 'Can also accept a JSON string (e.g. \\'[{\"col1\": \"val1\"}]\\') which will be '\n \"parsed into a DataFrame, or plain text which will be saved as a Message.\"\n ),\n input_types=[\"Data\", \"JSON\", \"DataFrame\", \"Table\", \"Message\"],\n required=True,\n tool_mode=True,\n ),\n StrInput(\n name=\"file_name\",\n display_name=\"File Name\",\n info=\"File name without extension (e.g. 'report'). Extension is added automatically.\",\n required=True,\n show=True,\n tool_mode=True,\n ),\n StrInput(\n name=\"file_format\",\n display_name=\"File Format (Tool)\",\n info=\"Output format: 'csv', 'json', 'txt', 'html', 'excel', 'markdown'. Overrides pre-configured format.\",\n required=False,\n show=False,\n tool_mode=True,\n ),\n BoolInput(\n name=\"append_mode\",\n display_name=\"Append\",\n info=(\n \"Append to file if it exists (only for Local storage with plain text formats). \"\n \"Not supported for cloud storage (AWS/Google Drive).\"\n ),\n value=False,\n show=_is_default_storage(\"Local\"),\n ),\n # Format inputs (dynamic based on storage location)\n DropdownInput(\n name=\"local_format\",\n display_name=\"File Format\",\n options=list(dict.fromkeys(LOCAL_DATA_FORMAT_CHOICES + LOCAL_MESSAGE_FORMAT_CHOICES)),\n info=\"Select the file format for local storage.\",\n value=\"json\",\n show=_is_default_storage(\"Local\"),\n ),\n DropdownInput(\n name=\"aws_format\",\n display_name=\"File Format\",\n options=AWS_FORMAT_CHOICES,\n info=\"Select the file format for AWS S3 storage.\",\n value=\"txt\",\n show=_is_default_storage(\"AWS\"),\n ),\n DropdownInput(\n name=\"gdrive_format\",\n display_name=\"File Format\",\n options=GDRIVE_FORMAT_CHOICES,\n info=\"Select the file format for Google Drive storage.\",\n value=\"txt\",\n show=_is_default_storage(\"Google Drive\"),\n ),\n # AWS S3 specific inputs\n SecretStrInput(\n name=\"aws_access_key_id\",\n display_name=\"AWS Access Key ID\",\n info=\"AWS Access key ID.\",\n show=_is_default_storage(\"AWS\"),\n advanced=not _is_default_storage(\"AWS\"),\n required=True,\n ),\n SecretStrInput(\n name=\"aws_secret_access_key\",\n display_name=\"AWS Secret Key\",\n info=\"AWS Secret Key.\",\n show=_is_default_storage(\"AWS\"),\n advanced=not _is_default_storage(\"AWS\"),\n required=True,\n ),\n StrInput(\n name=\"bucket_name\",\n display_name=\"S3 Bucket Name\",\n info=\"Enter the name of the S3 bucket.\",\n show=_is_default_storage(\"AWS\"),\n advanced=not _is_default_storage(\"AWS\"),\n required=True,\n ),\n StrInput(\n name=\"aws_region\",\n display_name=\"AWS Region\",\n info=\"AWS region (e.g., us-east-1, eu-west-1).\",\n show=_is_default_storage(\"AWS\"),\n advanced=not _is_default_storage(\"AWS\"),\n ),\n StrInput(\n name=\"s3_prefix\",\n display_name=\"S3 Prefix\",\n info=\"Prefix for all files in S3.\",\n show=_is_default_storage(\"AWS\"),\n advanced=not _is_default_storage(\"AWS\"),\n ),\n # Google Drive specific inputs\n SecretStrInput(\n name=\"service_account_key\",\n display_name=\"GCP Credentials Secret Key\",\n info=\"Your Google Cloud Platform service account JSON key as a secret string (complete JSON content).\",\n show=_is_default_storage(\"Google Drive\"),\n advanced=not _is_default_storage(\"Google Drive\"),\n required=True,\n ),\n StrInput(\n name=\"folder_id\",\n display_name=\"Google Drive Folder ID\",\n info=(\n \"The Google Drive folder ID where the file will be uploaded. \"\n \"The folder must be shared with the service account email.\"\n ),\n required=True,\n show=_is_default_storage(\"Google Drive\"),\n advanced=not _is_default_storage(\"Google Drive\"),\n ),\n ]\n\n outputs = [Output(display_name=\"File Path\", name=\"message\", method=\"save_to_file\")]\n\n def update_build_config(self, build_config, field_value, field_name=None):\n \"\"\"Update build configuration to show/hide fields based on storage location selection.\"\"\"\n # Update options dynamically based on cloud environment\n # This ensures options are refreshed when build_config is updated\n if \"storage_location\" in build_config:\n updated_options = _get_storage_location_options()\n build_config[\"storage_location\"][\"options\"] = updated_options\n\n # When tool_mode is toggled, hide storage-specific format dropdowns\n # (the agent uses the unified file_format input instead)\n if field_name == \"tool_mode\":\n format_fields = [\"local_format\", \"aws_format\", \"gdrive_format\"]\n for f_name in format_fields:\n if f_name in build_config:\n build_config[f_name][\"show\"] = not bool(field_value)\n return build_config\n\n if field_name != \"storage_location\":\n return build_config\n\n # Extract selected storage location\n selected = [location[\"name\"] for location in field_value] if isinstance(field_value, list) else []\n\n # Hide all dynamic fields first\n dynamic_fields = [\n \"file_name\", # Common fields (input is always visible)\n \"append_mode\",\n \"local_format\",\n \"aws_format\",\n \"gdrive_format\",\n \"aws_access_key_id\",\n \"aws_secret_access_key\",\n \"bucket_name\",\n \"aws_region\",\n \"s3_prefix\",\n \"service_account_key\",\n \"folder_id\",\n ]\n\n for f_name in dynamic_fields:\n if f_name in build_config:\n build_config[f_name][\"show\"] = False\n\n # Show fields based on selected storage location\n is_tool_mode = build_config.get(\"tools_metadata\", {}).get(\"show\", False)\n\n if len(selected) == 1:\n location = selected[0]\n\n # Show file_name when any storage location is selected\n if \"file_name\" in build_config:\n build_config[\"file_name\"][\"show\"] = True\n\n # Show append_mode only for Local storage (not supported for cloud storage)\n if \"append_mode\" in build_config:\n build_config[\"append_mode\"][\"show\"] = location == \"Local\"\n\n if location == \"Local\":\n if \"local_format\" in build_config:\n build_config[\"local_format\"][\"show\"] = not is_tool_mode\n\n elif location == \"AWS\":\n aws_fields = [\n \"aws_format\",\n \"aws_access_key_id\",\n \"aws_secret_access_key\",\n \"bucket_name\",\n \"aws_region\",\n \"s3_prefix\",\n ]\n for f_name in aws_fields:\n if f_name in build_config:\n show = f_name != \"aws_format\" or not is_tool_mode\n build_config[f_name][\"show\"] = show\n build_config[f_name][\"advanced\"] = False\n\n elif location == \"Google Drive\":\n gdrive_fields = [\"gdrive_format\", \"service_account_key\", \"folder_id\"]\n for f_name in gdrive_fields:\n if f_name in build_config:\n show = f_name != \"gdrive_format\" or not is_tool_mode\n build_config[f_name][\"show\"] = show\n build_config[f_name][\"advanced\"] = False\n\n return build_config\n\n async def save_to_file(self) -> Message:\n \"\"\"Save the input to a file and upload it, returning a confirmation message.\"\"\"\n # Validate inputs\n if not self.file_name:\n msg = \"File name must be provided.\"\n raise ValueError(msg)\n if not self._get_input_type():\n msg = \"Input type is not set.\"\n raise ValueError(msg)\n\n # Get selected storage location\n storage_location = self._get_selected_storage_location()\n if not storage_location:\n msg = \"Storage location must be selected.\"\n raise ValueError(msg)\n\n # Check if Local storage is disabled in cloud environment\n if storage_location == \"Local\" and is_astra_cloud_environment():\n msg = \"Local storage is not available in cloud environment. Please use AWS or Google Drive.\"\n raise ValueError(msg)\n\n # Route to appropriate save method based on storage location\n if storage_location == \"Local\":\n return await self._save_to_local()\n if storage_location == \"AWS\":\n return await self._save_to_aws()\n if storage_location == \"Google Drive\":\n return await self._save_to_google_drive()\n msg = f\"Unsupported storage location: {storage_location}\"\n raise ValueError(msg)\n\n def _get_input_type(self) -> str:\n \"\"\"Determine the input type based on the provided input.\"\"\"\n # Use exact type checking (type() is) instead of isinstance() to avoid inheritance issues.\n # Since Message inherits from Data, isinstance(message, Data) would return True for Message objects,\n # causing Message inputs to be incorrectly identified as Data type.\n if type(self.input) is DataFrame:\n return \"DataFrame\"\n if type(self.input) is Message:\n return \"Message\"\n if type(self.input) is Data:\n return \"Data\"\n # When invoked by a code agent (e.g. OpenDsStar), the input may be a raw\n # pandas DataFrame rather than Langflow's DataFrame wrapper.\n if isinstance(self.input, pd.DataFrame):\n self.input = DataFrame(self.input)\n return \"DataFrame\"\n # When invoked as a tool, the agent passes a string. Try to parse it as\n # tabular JSON (list of objects) → DataFrame, otherwise wrap as Message.\n if isinstance(self.input, str):\n self.input = self._coerce_string_input(self.input)\n return self._get_input_type()\n msg = f\"Unsupported input type: {type(self.input)}\"\n raise ValueError(msg)\n\n def _coerce_string_input(self, value: str) -> DataFrame | Message:\n \"\"\"Convert a raw string (from agent tool call) into a DataFrame or Message.\n\n Tries to parse as JSON first — a list of objects or a single object becomes\n a DataFrame. Anything else is wrapped in a Message.\n \"\"\"\n try:\n parsed = json.loads(value)\n if isinstance(parsed, list) and parsed and isinstance(parsed[0], dict):\n return DataFrame(pd.DataFrame(parsed))\n if isinstance(parsed, dict):\n return DataFrame(pd.DataFrame([parsed]))\n except (json.JSONDecodeError, ValueError):\n pass\n return Message(text=value)\n\n def _get_default_format(self) -> str:\n \"\"\"Return the default file format based on input type.\"\"\"\n if self._get_input_type() == \"DataFrame\":\n return \"csv\"\n if self._get_input_type() == \"Data\":\n return \"json\"\n if self._get_input_type() == \"Message\":\n return \"json\"\n return \"json\" # Fallback\n\n def _adjust_file_path_with_format(self, path: Path, fmt: str) -> Path:\n \"\"\"Adjust the file path to include the correct extension.\"\"\"\n file_extension = path.suffix.lower().lstrip(\".\")\n if fmt == \"excel\":\n return Path(f\"{path}.xlsx\").expanduser() if file_extension not in [\"xlsx\", \"xls\"] else path\n return Path(f\"{path}.{fmt}\").expanduser() if file_extension != fmt else path\n\n def _is_plain_text_format(self, fmt: str) -> bool:\n \"\"\"Check if a file format is plain text (supports appending).\"\"\"\n plain_text_formats = [\"txt\", \"json\", \"markdown\", \"md\", \"csv\", \"xml\", \"html\", \"yaml\", \"log\", \"tsv\", \"jsonl\"]\n return fmt.lower() in plain_text_formats\n\n async def _upload_file(self, file_path: Path) -> None:\n \"\"\"Upload the saved file using the upload_user_file service.\"\"\"\n from langflow.api.v2.files import upload_user_file\n from langflow.services.database.models.user.crud import get_user_by_id\n\n # Ensure the file exists\n if not file_path.exists():\n msg = f\"File not found: {file_path}\"\n raise FileNotFoundError(msg)\n\n # Upload the file - always use append=False because the local file already contains\n # the correct content (either new or appended locally)\n with file_path.open(\"rb\") as f:\n async with session_scope() as db:\n if not self.user_id:\n msg = \"User ID is required for file saving.\"\n raise ValueError(msg)\n current_user = await get_user_by_id(db, self.user_id)\n\n await upload_user_file(\n file=UploadFile(filename=file_path.name, file=f, size=file_path.stat().st_size),\n session=db,\n current_user=current_user,\n storage_service=get_storage_service(),\n settings_service=get_settings_service(),\n append=False,\n )\n\n def _save_dataframe(self, dataframe: DataFrame, path: Path, fmt: str) -> str:\n \"\"\"Save a DataFrame to the specified file format.\"\"\"\n append_mode = getattr(self, \"append_mode\", False)\n should_append = append_mode and path.exists() and self._is_plain_text_format(fmt)\n\n if fmt == \"csv\":\n dataframe.to_csv(path, index=False, mode=\"a\" if should_append else \"w\", header=not should_append)\n elif fmt == \"excel\":\n dataframe.to_excel(path, index=False, engine=\"openpyxl\")\n elif fmt == \"json\":\n if should_append:\n # Read and parse existing JSON\n existing_data = []\n try:\n existing_content = path.read_text(encoding=\"utf-8\").strip()\n if existing_content:\n parsed = json.loads(existing_content)\n # Handle case where existing content is a single object\n if isinstance(parsed, dict):\n existing_data = [parsed]\n elif isinstance(parsed, list):\n existing_data = parsed\n except (json.JSONDecodeError, FileNotFoundError):\n # Treat parse errors or missing file as empty array\n existing_data = []\n\n # Append new data\n new_records = json.loads(dataframe.to_json(orient=\"records\"))\n existing_data.extend(new_records)\n\n # Write back as a single JSON array\n path.write_text(json.dumps(existing_data, indent=2), encoding=\"utf-8\")\n else:\n dataframe.to_json(path, orient=\"records\", indent=2)\n elif fmt == \"markdown\":\n content = dataframe.to_markdown(index=False)\n if should_append:\n path.write_text(path.read_text(encoding=\"utf-8\") + \"\\n\\n\" + content, encoding=\"utf-8\")\n else:\n path.write_text(content, encoding=\"utf-8\")\n else:\n msg = f\"Unsupported DataFrame format: {fmt}\"\n raise ValueError(msg)\n action = \"appended to\" if should_append else \"saved successfully as\"\n return f\"DataFrame {action} '{path}'\"\n\n def _save_data(self, data: Data, path: Path, fmt: str) -> str:\n \"\"\"Save a Data object to the specified file format.\"\"\"\n append_mode = getattr(self, \"append_mode\", False)\n should_append = append_mode and path.exists() and self._is_plain_text_format(fmt)\n\n if fmt == \"csv\":\n pd.DataFrame(data.data).to_csv(\n path,\n index=False,\n mode=\"a\" if should_append else \"w\",\n header=not should_append,\n )\n elif fmt == \"excel\":\n pd.DataFrame(data.data).to_excel(path, index=False, engine=\"openpyxl\")\n elif fmt == \"json\":\n new_data = jsonable_encoder(data.data)\n if should_append:\n # Read and parse existing JSON\n existing_data = []\n try:\n existing_content = path.read_text(encoding=\"utf-8\").strip()\n if existing_content:\n parsed = json.loads(existing_content)\n # Handle case where existing content is a single object\n if isinstance(parsed, dict):\n existing_data = [parsed]\n elif isinstance(parsed, list):\n existing_data = parsed\n except (json.JSONDecodeError, FileNotFoundError):\n # Treat parse errors or missing file as empty array\n existing_data = []\n\n # Append new data\n if isinstance(new_data, list):\n existing_data.extend(new_data)\n else:\n existing_data.append(new_data)\n\n # Write back as a single JSON array\n path.write_text(json.dumps(existing_data, indent=2), encoding=\"utf-8\")\n else:\n content = orjson.dumps(new_data, option=orjson.OPT_INDENT_2).decode(\"utf-8\")\n path.write_text(content, encoding=\"utf-8\")\n elif fmt == \"markdown\":\n content = pd.DataFrame(data.data).to_markdown(index=False)\n if should_append:\n path.write_text(path.read_text(encoding=\"utf-8\") + \"\\n\\n\" + content, encoding=\"utf-8\")\n else:\n path.write_text(content, encoding=\"utf-8\")\n else:\n msg = f\"Unsupported Data format: {fmt}\"\n raise ValueError(msg)\n action = \"appended to\" if should_append else \"saved successfully as\"\n return f\"Data {action} '{path}'\"\n\n async def _save_message(self, message: Message, path: Path, fmt: str) -> str:\n \"\"\"Save a Message to the specified file format, handling async iterators.\"\"\"\n content = \"\"\n stream = message.text_stream if hasattr(message, \"text_stream\") else None\n if stream is not None and isinstance(stream, AsyncIterator):\n async for item in stream:\n content += str(item) + \" \"\n content = content.strip()\n elif stream is not None and isinstance(stream, Iterator):\n content = \" \".join(str(item) for item in stream)\n else:\n content = str(message.text)\n\n append_mode = getattr(self, \"append_mode\", False)\n should_append = append_mode and path.exists() and self._is_plain_text_format(fmt)\n\n if fmt in (\"txt\", \"html\"):\n if should_append:\n path.write_text(path.read_text(encoding=\"utf-8\") + \"\\n\" + content, encoding=\"utf-8\")\n else:\n path.write_text(content, encoding=\"utf-8\")\n elif fmt == \"json\":\n new_message = {\"message\": content}\n if should_append:\n # Read and parse existing JSON\n existing_data = []\n try:\n existing_content = path.read_text(encoding=\"utf-8\").strip()\n if existing_content:\n parsed = json.loads(existing_content)\n # Handle case where existing content is a single object\n if isinstance(parsed, dict):\n existing_data = [parsed]\n elif isinstance(parsed, list):\n existing_data = parsed\n except (json.JSONDecodeError, FileNotFoundError):\n # Treat parse errors or missing file as empty array\n existing_data = []\n\n # Append new message\n existing_data.append(new_message)\n\n # Write back as a single JSON array\n path.write_text(json.dumps(existing_data, indent=2), encoding=\"utf-8\")\n else:\n path.write_text(json.dumps(new_message, indent=2), encoding=\"utf-8\")\n elif fmt == \"markdown\":\n md_content = f\"**Message:**\\n\\n{content}\"\n if should_append:\n path.write_text(path.read_text(encoding=\"utf-8\") + \"\\n\\n\" + md_content, encoding=\"utf-8\")\n else:\n path.write_text(md_content, encoding=\"utf-8\")\n else:\n msg = f\"Unsupported Message format: {fmt}\"\n raise ValueError(msg)\n action = \"appended to\" if should_append else \"saved successfully as\"\n return f\"Message {action} '{path}'\"\n\n def _get_selected_storage_location(self) -> str:\n \"\"\"Get the selected storage location from the SortableListInput.\"\"\"\n if hasattr(self, \"storage_location\") and self.storage_location:\n if isinstance(self.storage_location, list) and len(self.storage_location) > 0:\n return self.storage_location[0].get(\"name\", \"\")\n if isinstance(self.storage_location, dict):\n return self.storage_location.get(\"name\", \"\")\n return \"\"\n\n def _get_file_format_for_location(self, location: str) -> str:\n \"\"\"Get the appropriate file format based on storage location.\n\n If the agent set file_format via tool mode, that takes priority.\n \"\"\"\n agent_format = getattr(self, \"file_format\", None)\n if agent_format:\n return agent_format\n if location == \"Local\":\n return getattr(self, \"local_format\", None) or self._get_default_format()\n if location == \"AWS\":\n return getattr(self, \"aws_format\", \"txt\")\n if location == \"Google Drive\":\n return getattr(self, \"gdrive_format\", \"txt\")\n return self._get_default_format()\n\n async def _save_to_local(self) -> Message:\n \"\"\"Save file to local storage (original functionality).\"\"\"\n file_format = self._get_file_format_for_location(\"Local\")\n\n # Validate file format based on input type\n allowed_formats = (\n self.LOCAL_MESSAGE_FORMAT_CHOICES if self._get_input_type() == \"Message\" else self.LOCAL_DATA_FORMAT_CHOICES\n )\n if file_format not in allowed_formats:\n msg = f\"Invalid file format '{file_format}' for {self._get_input_type()}. Allowed: {allowed_formats}\"\n raise ValueError(msg)\n\n # Prepare file path\n file_path = Path(self.file_name).expanduser()\n if not file_path.parent.exists():\n file_path.parent.mkdir(parents=True, exist_ok=True)\n file_path = self._adjust_file_path_with_format(file_path, file_format)\n\n # Save the input to file based on type\n if self._get_input_type() == \"DataFrame\":\n confirmation = self._save_dataframe(self.input, file_path, file_format)\n elif self._get_input_type() == \"Data\":\n confirmation = self._save_data(self.input, file_path, file_format)\n elif self._get_input_type() == \"Message\":\n confirmation = await self._save_message(self.input, file_path, file_format)\n else:\n msg = f\"Unsupported input type: {self._get_input_type()}\"\n raise ValueError(msg)\n\n # Upload the saved file\n await self._upload_file(file_path)\n\n # Return the final file path and confirmation message\n final_path = Path.cwd() / file_path if not file_path.is_absolute() else file_path\n return Message(text=f\"{confirmation} at {final_path}\")\n\n async def _save_to_aws(self) -> Message:\n \"\"\"Save file to AWS S3 using S3 functionality.\"\"\"\n import os\n\n import boto3\n\n from lfx.base.data.cloud_storage_utils import create_s3_client, validate_aws_credentials\n\n # Get AWS credentials from component inputs or fall back to environment variables\n aws_access_key_id = getattr(self, \"aws_access_key_id\", None)\n if aws_access_key_id and hasattr(aws_access_key_id, \"get_secret_value\"):\n aws_access_key_id = aws_access_key_id.get_secret_value()\n if not aws_access_key_id:\n aws_access_key_id = os.getenv(\"AWS_ACCESS_KEY_ID\")\n\n aws_secret_access_key = getattr(self, \"aws_secret_access_key\", None)\n if aws_secret_access_key and hasattr(aws_secret_access_key, \"get_secret_value\"):\n aws_secret_access_key = aws_secret_access_key.get_secret_value()\n if not aws_secret_access_key:\n aws_secret_access_key = os.getenv(\"AWS_SECRET_ACCESS_KEY\")\n\n bucket_name = getattr(self, \"bucket_name\", None)\n if not bucket_name:\n # Try to get from storage service settings\n settings = get_settings_service().settings\n bucket_name = settings.object_storage_bucket_name\n\n # Validate AWS credentials\n if not aws_access_key_id:\n msg = (\n \"AWS Access Key ID is required for S3 storage. Provide it as a component input \"\n \"or set AWS_ACCESS_KEY_ID environment variable.\"\n )\n raise ValueError(msg)\n if not aws_secret_access_key:\n msg = (\n \"AWS Secret Key is required for S3 storage. Provide it as a component input \"\n \"or set AWS_SECRET_ACCESS_KEY environment variable.\"\n )\n raise ValueError(msg)\n if not bucket_name:\n msg = (\n \"S3 Bucket Name is required for S3 storage. Provide it as a component input \"\n \"or set LANGFLOW_OBJECT_STORAGE_BUCKET_NAME environment variable.\"\n )\n raise ValueError(msg)\n\n # Validate AWS credentials\n validate_aws_credentials(self)\n\n # Create S3 client\n s3_client = create_s3_client(self)\n client_config: dict[str, Any] = {\n \"aws_access_key_id\": str(aws_access_key_id),\n \"aws_secret_access_key\": str(aws_secret_access_key),\n }\n\n # Get region from component input, environment variable, or settings\n aws_region = getattr(self, \"aws_region\", None)\n if not aws_region:\n aws_region = os.getenv(\"AWS_DEFAULT_REGION\") or os.getenv(\"AWS_REGION\")\n if aws_region:\n client_config[\"region_name\"] = str(aws_region)\n\n s3_client = boto3.client(\"s3\", **client_config)\n\n # Extract content\n content = self._extract_content_for_upload()\n file_format = self._get_file_format_for_location(\"AWS\")\n\n # Generate file path\n file_path = f\"{self.file_name}.{file_format}\"\n if hasattr(self, \"s3_prefix\") and self.s3_prefix:\n file_path = f\"{self.s3_prefix.rstrip('/')}/{file_path}\"\n\n # Create temporary file\n import tempfile\n\n with tempfile.NamedTemporaryFile(\n mode=\"w\", encoding=\"utf-8\", suffix=f\".{file_format}\", delete=False\n ) as temp_file:\n temp_file.write(content)\n temp_file_path = temp_file.name\n\n try:\n # Upload to S3\n s3_client.upload_file(temp_file_path, bucket_name, file_path)\n s3_url = f\"s3://{bucket_name}/{file_path}\"\n return Message(text=f\"File successfully uploaded to {s3_url}\")\n finally:\n # Clean up temp file\n if Path(temp_file_path).exists():\n Path(temp_file_path).unlink()\n\n async def _save_to_google_drive(self) -> Message:\n \"\"\"Save file to Google Drive using Google Drive functionality.\"\"\"\n import tempfile\n\n from googleapiclient.http import MediaFileUpload\n\n from lfx.base.data.cloud_storage_utils import create_google_drive_service\n\n # Validate Google Drive credentials\n if not getattr(self, \"service_account_key\", None):\n msg = \"GCP Credentials Secret Key is required for Google Drive storage\"\n raise ValueError(msg)\n if not getattr(self, \"folder_id\", None):\n msg = \"Google Drive Folder ID is required for Google Drive storage\"\n raise ValueError(msg)\n\n # Create Google Drive service with full drive scope (needed for folder operations)\n drive_service, credentials = create_google_drive_service(\n self.service_account_key, scopes=[\"https://www.googleapis.com/auth/drive\"], return_credentials=True\n )\n\n # Extract content and format\n content = self._extract_content_for_upload()\n file_format = self._get_file_format_for_location(\"Google Drive\")\n\n # Handle special Google Drive formats\n if file_format in [\"slides\", \"docs\"]:\n return await self._save_to_google_apps(drive_service, credentials, content, file_format)\n\n # Create temporary file\n file_path = f\"{self.file_name}.{file_format}\"\n with tempfile.NamedTemporaryFile(\n mode=\"w\",\n encoding=\"utf-8\",\n suffix=f\".{file_format}\",\n delete=False,\n ) as temp_file:\n temp_file.write(content)\n temp_file_path = temp_file.name\n\n try:\n # Upload to Google Drive\n # Note: We skip explicit folder verification since it requires broader permissions.\n # If the folder doesn't exist or isn't accessible, the create() call will fail with a clear error.\n file_metadata = {\"name\": file_path, \"parents\": [self.folder_id]}\n media = MediaFileUpload(temp_file_path, resumable=True)\n\n try:\n uploaded_file = (\n drive_service.files().create(body=file_metadata, media_body=media, fields=\"id\").execute()\n )\n except Exception as e:\n msg = (\n f\"Unable to upload file to Google Drive folder '{self.folder_id}'. \"\n f\"Error: {e!s}. \"\n \"Please ensure: 1) The folder ID is correct, 2) The folder exists, \"\n \"3) The service account has been granted access to this folder.\"\n )\n raise ValueError(msg) from e\n\n file_id = uploaded_file.get(\"id\")\n file_url = f\"https://drive.google.com/file/d/{file_id}/view\"\n return Message(text=f\"File successfully uploaded to Google Drive: {file_url}\")\n finally:\n # Clean up temp file\n if Path(temp_file_path).exists():\n Path(temp_file_path).unlink()\n\n async def _save_to_google_apps(self, drive_service, credentials, content: str, app_type: str) -> Message:\n \"\"\"Save content to Google Apps (Slides or Docs).\"\"\"\n import time\n\n if app_type == \"slides\":\n from googleapiclient.discovery import build\n\n slides_service = build(\"slides\", \"v1\", credentials=credentials)\n\n file_metadata = {\n \"name\": self.file_name,\n \"mimeType\": \"application/vnd.google-apps.presentation\",\n \"parents\": [self.folder_id],\n }\n\n created_file = drive_service.files().create(body=file_metadata, fields=\"id\").execute()\n presentation_id = created_file[\"id\"]\n\n time.sleep(2) # Wait for file to be available # noqa: ASYNC251\n\n presentation = slides_service.presentations().get(presentationId=presentation_id).execute()\n slide_id = presentation[\"slides\"][0][\"objectId\"]\n\n # Add content to slide\n requests = [\n {\n \"createShape\": {\n \"objectId\": \"TextBox_01\",\n \"shapeType\": \"TEXT_BOX\",\n \"elementProperties\": {\n \"pageObjectId\": slide_id,\n \"size\": {\n \"height\": {\"magnitude\": 3000000, \"unit\": \"EMU\"},\n \"width\": {\"magnitude\": 6000000, \"unit\": \"EMU\"},\n },\n \"transform\": {\n \"scaleX\": 1,\n \"scaleY\": 1,\n \"translateX\": 1000000,\n \"translateY\": 1000000,\n \"unit\": \"EMU\",\n },\n },\n }\n },\n {\"insertText\": {\"objectId\": \"TextBox_01\", \"insertionIndex\": 0, \"text\": content}},\n ]\n\n slides_service.presentations().batchUpdate(\n presentationId=presentation_id, body={\"requests\": requests}\n ).execute()\n file_url = f\"https://docs.google.com/presentation/d/{presentation_id}/edit\"\n\n elif app_type == \"docs\":\n from googleapiclient.discovery import build\n\n docs_service = build(\"docs\", \"v1\", credentials=credentials)\n\n file_metadata = {\n \"name\": self.file_name,\n \"mimeType\": \"application/vnd.google-apps.document\",\n \"parents\": [self.folder_id],\n }\n\n created_file = drive_service.files().create(body=file_metadata, fields=\"id\").execute()\n document_id = created_file[\"id\"]\n\n time.sleep(2) # Wait for file to be available # noqa: ASYNC251\n\n # Add content to document\n requests = [{\"insertText\": {\"location\": {\"index\": 1}, \"text\": content}}]\n docs_service.documents().batchUpdate(documentId=document_id, body={\"requests\": requests}).execute()\n file_url = f\"https://docs.google.com/document/d/{document_id}/edit\"\n\n return Message(text=f\"File successfully created in Google {app_type.title()}: {file_url}\")\n\n def _extract_content_for_upload(self) -> str:\n \"\"\"Extract content from input for upload to cloud services.\"\"\"\n if self._get_input_type() == \"DataFrame\":\n return self.input.to_csv(index=False)\n if self._get_input_type() == \"Data\":\n if hasattr(self.input, \"data\") and self.input.data:\n if isinstance(self.input.data, dict):\n import json\n\n return json.dumps(self.input.data, indent=2, ensure_ascii=False)\n return str(self.input.data)\n return str(self.input)\n if self._get_input_type() == \"Message\":\n return str(self.input.text) if self.input.text else str(self.input)\n return str(self.input)\n" }, "file_format": { "_input_type": "StrInput", @@ -118454,6 +118454,6 @@ "num_components": 354, "num_modules": 95 }, - "sha256": "87d65557fa5c4182f09f07270bbe8624d5f9b65a0322c3251f75822b6f1546ac", + "sha256": "57420a9e8822a72ba8e8f8fd0d597f5ff7b9d4f85d965ed40713d108512cc7fe", "version": "1.11.0" } From 6f5de50eb32d6be0a1e6ab098cba9e94907c3354 Mon Sep 17 00:00:00 2001 From: ogabrielluiz Date: Wed, 10 Jun 2026 16:58:51 -0300 Subject: [PATCH 3/3] fix(api/v2): buffer parallel messages in the AG-UI translator instead of dropping them Parallel components stream tokens for different message ids interleaved. The translator tracked a single open message: the first foreign token closed the open message and tombstoned its id, so every later event for it was dropped and its remaining text never reached the client. Tokens for a message that cannot take the wire now buffer until the open message genuinely ends (its add_message finalizer), then flush in arrival order; complete messages landing mid-stream buffer the same way instead of interleaving a second START. end/error drain all buffers before the terminal event. The wire still carries at most one open text message, so the stream stays AG-UI-conformant. --- .../base/langflow/api/v2/agui_translator.py | 149 ++++++++++++++---- .../tests/unit/api/v2/test_agui_translator.py | 127 +++++++++++---- 2 files changed, 217 insertions(+), 59 deletions(-) diff --git a/src/backend/base/langflow/api/v2/agui_translator.py b/src/backend/base/langflow/api/v2/agui_translator.py index 727e71e56e..0fa76a39f2 100644 --- a/src/backend/base/langflow/api/v2/agui_translator.py +++ b/src/backend/base/langflow/api/v2/agui_translator.py @@ -12,6 +12,7 @@ list. The translator is stateful: one instance per run. from __future__ import annotations import json +from dataclasses import dataclass, field from ag_ui.core import ( BaseEvent, @@ -37,6 +38,28 @@ from ag_ui.core import ( _CUSTOM_CONTENT_TYPES = frozenset({"json", "code", "media", "error"}) +@dataclass +class _BufferedMessage: + """A text message waiting for the wire while another message is open. + + AG-UI allows one open text message at a time, but parallel components + stream tokens for different message ids interleaved. Tokens for a message + that cannot open yet accumulate here; ``final_text`` is set when its + ``add_message`` finalizer arrives before it ever reached the wire. + """ + + chunks: list[str] = field(default_factory=list) + final_text: str | None = None + + @property + def complete(self) -> bool: + return self.final_text is not None + + @property + def text(self) -> str: + return self.final_text if self.final_text is not None else "".join(self.chunks) + + class AGUITranslator: """Translates Langflow ``EventManager`` events into AG-UI protocol events. @@ -47,11 +70,16 @@ class AGUITranslator: def __init__(self, run_id: str, thread_id: str) -> None: self.run_id = run_id self.thread_id = thread_id - # Id of the text message currently being streamed by ``token`` events, - # or ``None`` when no message is open. + # Id of the text message currently open on the wire, or ``None``. self._open_message_id: str | None = None - # Message ids already emitted as a complete (non-streamed) text message. + # Message ids whose TEXT_MESSAGE_END has been emitted; the protocol + # considers them closed, so nothing may reopen them. self._emitted_text_message_ids: set[str] = set() + # Messages from parallel components waiting for the wire, in arrival + # order (dict preserves insertion order). Invariant: non-empty only + # while another message holds the wire — closing the open message + # immediately promotes the next buffered one. + self._buffered_messages: dict[str, _BufferedMessage] = {} # Tool-call ids already emitted as TOOL_CALL_START / already resolved # with a TOOL_CALL_RESULT. ``add_message`` can re-fire with the same # (append-only) content_blocks, so emissions must be deduplicated. @@ -95,11 +123,11 @@ class AGUITranslator: # streamed message and must stay transparent, or the message would be # split into multiple START/END pairs reusing an already-ended id. if event_type == "end": - events = self._close_open_message() + events = self._drain_messages() events.append(RunFinishedEvent(run_id=self.run_id, thread_id=self.thread_id)) return events if event_type == "error": - events = self._close_open_message() + events = self._drain_messages() # The ``error`` payload varies by emission path: a full ErrorMessage # dump carries the reason in ``text``; the minimal path sends # ``{"error": str}``. @@ -111,11 +139,13 @@ class AGUITranslator: def _translate_token(self, data: dict) -> list[BaseEvent]: """Map a ``token`` event to text-message events. - The first token of a message opens it with ``TEXT_MESSAGE_START``; a - token for a different message id closes the previous one first. A - token whose id was already ended via ``add_message`` (or a prior - token boundary) is dropped: re-opening it would emit a second - ``TEXT_MESSAGE_START`` for an id the protocol considers closed. + The first token of a message opens it with ``TEXT_MESSAGE_START``. + While a message holds the wire, tokens for other message ids (parallel + components stream interleaved) buffer instead of closing it; the + buffered message flushes when the open one genuinely ends. A token + whose id was already ended via ``add_message`` is dropped: re-opening + it would emit a second ``TEXT_MESSAGE_START`` for an id the protocol + considers closed. """ message_id = str(data.get("id") or "") if not message_id: @@ -123,16 +153,22 @@ class AGUITranslator: # be correlated. Dropping the event is preferable to emitting a # malformed stream with empty message_ids. return [] - if message_id in self._emitted_text_message_ids and self._open_message_id != message_id: + if message_id in self._emitted_text_message_ids: return [] chunk = data.get("chunk", "") - events: list[BaseEvent] = [] - if self._open_message_id != message_id: - events.extend(self._close_open_message()) - events.append(TextMessageStartEvent(message_id=message_id, role="assistant")) + if self._open_message_id == message_id: + return [TextMessageContentEvent(message_id=message_id, delta=chunk)] + if self._open_message_id is None: self._open_message_id = message_id - events.append(TextMessageContentEvent(message_id=message_id, delta=chunk)) - return events + return [ + TextMessageStartEvent(message_id=message_id, role="assistant"), + TextMessageContentEvent(message_id=message_id, delta=chunk), + ] + # Another message holds the wire: buffer this one until it is free. + buffered = self._buffered_messages.setdefault(message_id, _BufferedMessage()) + if not buffered.complete: + buffered.chunks.append(chunk) + return [] def _translate_vertices_sorted(self, data: dict) -> list[BaseEvent]: """Map ``vertices_sorted`` to a ``STATE_SNAPSHOT`` of the node graph. @@ -198,21 +234,32 @@ class AGUITranslator: # Message text. if message_id and message_id == self._open_message_id: - # Finalizer of a token-streamed message: close it. The text was - # already streamed token by token, so it must not be re-emitted now - # or by any later add_message that re-fires for the same id. - self._emitted_text_message_ids.add(message_id) + # Finalizer of the wire-open streamed message: close it. The text + # was already streamed token by token, so it must not be re-emitted + # now or by any later add_message that re-fires for the same id. events.extend(self._close_open_message()) + elif message_id and message_id in self._buffered_messages: + # Finalizer for a message still waiting for the wire: record the + # authoritative full text; the trio is emitted on promotion. + self._buffered_messages[message_id].final_text = ( + data.get("text") or self._buffered_messages[message_id].text + ) else: text = data.get("text") or "" # Skip text-message lifecycle emission without a stable message_id; # tool-call events above are namespaced by block/content index so # they can still ride a missing id, but TEXT_MESSAGE_* cannot. if text and message_id and message_id not in self._emitted_text_message_ids: - self._emitted_text_message_ids.add(message_id) - events.append(TextMessageStartEvent(message_id=message_id, role="assistant")) - events.append(TextMessageContentEvent(message_id=message_id, delta=text)) - events.append(TextMessageEndEvent(message_id=message_id)) + if self._open_message_id is None: + self._emitted_text_message_ids.add(message_id) + events.append(TextMessageStartEvent(message_id=message_id, role="assistant")) + events.append(TextMessageContentEvent(message_id=message_id, delta=text)) + events.append(TextMessageEndEvent(message_id=message_id)) + else: + # A parallel component finished while another message holds + # the wire: buffer the complete message instead of opening + # a second text message mid-stream. + self._buffered_messages[message_id] = _BufferedMessage(final_text=text) return events def _translate_tool_use( @@ -286,17 +333,57 @@ class AGUITranslator: return {"op": "add", "path": f"/nodes/{node_id}", "value": {"status": status, "output": output}} def _close_open_message(self) -> list[BaseEvent]: - """Emit ``TEXT_MESSAGE_END`` for the open message, if any. + """Emit ``TEXT_MESSAGE_END`` for the open message and free the wire. The closed id is recorded in ``_emitted_text_message_ids`` so a later - token (e.g. an interleaved ``A, B, A`` sequence) cannot re-open it - and emit a second ``TEXT_MESSAGE_START`` for an id the protocol - already considers closed. + token cannot re-open it and emit a second ``TEXT_MESSAGE_START`` for + an id the protocol already considers closed. With the wire free, the + next buffered parallel message (if any) is promoted onto it. """ if self._open_message_id is None: return [] closed_id = self._open_message_id - end = TextMessageEndEvent(message_id=closed_id) self._emitted_text_message_ids.add(closed_id) self._open_message_id = None - return [end] + events: list[BaseEvent] = [TextMessageEndEvent(message_id=closed_id)] + events.extend(self._promote_next_buffered()) + return events + + def _promote_next_buffered(self) -> list[BaseEvent]: + """Move buffered parallel messages onto the freed wire, in arrival order. + + Already-complete messages emit their full START/CONTENT/END trio and the + promotion continues; the first still-streaming message replays its + buffered chunks, takes the wire, and stays open for its live tokens. + """ + events: list[BaseEvent] = [] + while self._buffered_messages and self._open_message_id is None: + message_id, buffered = next(iter(self._buffered_messages.items())) + del self._buffered_messages[message_id] + if message_id in self._emitted_text_message_ids: + continue + text = buffered.text + if buffered.complete: + self._emitted_text_message_ids.add(message_id) + if text: + events.append(TextMessageStartEvent(message_id=message_id, role="assistant")) + events.append(TextMessageContentEvent(message_id=message_id, delta=text)) + events.append(TextMessageEndEvent(message_id=message_id)) + continue + self._open_message_id = message_id + events.append(TextMessageStartEvent(message_id=message_id, role="assistant")) + if text: + events.append(TextMessageContentEvent(message_id=message_id, delta=text)) + return events + + def _drain_messages(self) -> list[BaseEvent]: + """Close the open message and flush every buffered one (run boundary). + + At ``end``/``error`` nothing else will free the wire, so buffered + parallel messages flush now — each promoted, emitted, and closed — + rather than being silently lost. + """ + events = self._close_open_message() + while self._open_message_id is not None: + events.extend(self._close_open_message()) + return events diff --git a/src/backend/tests/unit/api/v2/test_agui_translator.py b/src/backend/tests/unit/api/v2/test_agui_translator.py index 0797105bc7..995cea3d86 100644 --- a/src/backend/tests/unit/api/v2/test_agui_translator.py +++ b/src/backend/tests/unit/api/v2/test_agui_translator.py @@ -90,13 +90,31 @@ def test_token_sequence_emits_start_contents_then_end_on_boundary(): assert isinstance(ended[1], RunFinishedEvent) -def test_new_message_id_closes_previous_message_and_opens_new(): +def test_token_for_second_message_is_buffered_until_first_closes(): + """A token for a different message id must not close the streaming one. + + Parallel components stream tokens for different message ids interleaved. + Closing the open message on the first foreign token burned its id, so all + its later tokens were dropped. Instead the foreign message buffers until + the open one genuinely ends (its ``add_message`` finalizer), then flushes. + """ t = AGUITranslator(run_id="r1", thread_id="t1") t.start() t.translate("token", {"chunk": "a", "id": "m1"}) - out = t.translate("token", {"chunk": "b", "id": "m2"}) + buffered = t.translate("token", {"chunk": "b", "id": "m2"}) + # m2 buffers silently; m1 stays open. + assert buffered == [] + + # m1 keeps streaming: its tokens still flow. + more = t.translate("token", {"chunk": "a2", "id": "m1"}) + assert len(more) == 1 + assert isinstance(more[0], TextMessageContentEvent) + assert more[0].message_id == "m1" + + # m1's finalizer closes it and promotes m2 with its buffered content. + out = t.translate("add_message", {"id": "m1", "text": "aa2"}) assert isinstance(out[0], TextMessageEndEvent) assert out[0].message_id == "m1" assert isinstance(out[1], TextMessageStartEvent) @@ -166,32 +184,83 @@ def test_token_for_already_ended_message_id_is_dropped(): ) -def test_token_after_boundary_close_is_dropped(): - """An interleaved token sequence ``A, B, A`` must not re-open id A. +def test_parallel_interleaved_tokens_preserve_all_content(): + """Two components streaming in parallel must not lose either stream. - Switching from id A to id B closes A through ``_close_open_message``. A - later token for A used to slip past the dedup guard because - ``_close_open_message`` was the only finalizer that did not record the - closed id in ``_emitted_text_message_ids``. The translator must treat a - token-boundary close the same way it treats an ``add_message`` close. + Reproduces the reported parallel-components drop: with a single-slot + tracker, START(m2), CONTENT(m2), START(m3) burned m2, so CONTENT(m2) + was dropped and m2's remaining text never reached the client. With + buffering, m2 holds the wire until it genuinely ends; m3's tokens + buffer and flush afterwards. Nothing is dropped, and the wire carries + at most one open text message at a time. + """ + t = AGUITranslator(run_id="r1", thread_id="t1") + sequence = [ + ("token", {"chunk": "Result ", "id": "m2"}), + ("token", {"chunk": "Side ", "id": "m3"}), + ("token", {"chunk": "for X", "id": "m2"}), + ("token", {"chunk": "for Y", "id": "m3"}), + ("add_message", {"id": "m2", "text": "Result for X"}), + ("add_message", {"id": "m3", "text": "Side for Y"}), + ("end", {}), + ] + + out = _run_sequence(t, sequence) + + _assert_well_formed(out) + text_by_message: dict[str, str] = {} + for event in out: + if isinstance(event, TextMessageContentEvent): + text_by_message[event.message_id] = text_by_message.get(event.message_id, "") + event.delta + assert text_by_message["m2"] == "Result for X" + assert text_by_message["m3"] == "Side for Y" + + +def test_add_message_for_other_message_while_streaming_is_buffered(): + """A complete message landing mid-stream must not interleave its trio. + + A parallel non-streaming component can finish (add_message with full + text) while another component holds the wire with an open streamed + message. Emitting START/CONTENT/END for the finished one immediately + would open a second text message mid-stream; instead it buffers and + flushes when the open message closes. """ t = AGUITranslator(run_id="r1", thread_id="t1") t.start() - a1 = t.translate("token", {"chunk": "hi", "id": "m1"}) - assert any(isinstance(e, TextMessageStartEvent) and e.message_id == "m1" for e in a1) + t.translate("token", {"chunk": "streaming...", "id": "m1"}) + parallel_done = t.translate("add_message", {"id": "m9", "text": "finished early"}) + assert all( + not isinstance(e, (TextMessageStartEvent, TextMessageContentEvent, TextMessageEndEvent)) for e in parallel_done + ), f"buffered message leaked text events mid-stream: {parallel_done}" - # Switching to a new id closes m1 via _close_open_message. - b = t.translate("token", {"chunk": "yo", "id": "m2"}) - assert any(isinstance(e, TextMessageEndEvent) and e.message_id == "m1" for e in b) - assert any(isinstance(e, TextMessageStartEvent) and e.message_id == "m2" for e in b) + out = t.translate("add_message", {"id": "m1", "text": "streaming..."}) + assert [type(e) for e in out] == [ + TextMessageEndEvent, + TextMessageStartEvent, + TextMessageContentEvent, + TextMessageEndEvent, + ] + assert out[0].message_id == "m1" + assert out[1].message_id == "m9" + assert out[2].delta == "finished early" - # A late token for m1 must be dropped, not re-open the ended message. - late = t.translate("token", {"chunk": "again", "id": "m1"}) - assert all(not isinstance(e, TextMessageStartEvent) for e in late), ( - f"Token boundary did not mark m1 as ended; emitted: {late}" - ) - assert late == [], f"Expected no events for late token after boundary close; got {late}" + +def test_end_drains_buffered_messages_before_run_finished(): + """A run ending while messages are still buffered must flush them all.""" + t = AGUITranslator(run_id="r1", thread_id="t1") + t.start() + + t.translate("token", {"chunk": "open", "id": "m1"}) + t.translate("token", {"chunk": "waiting", "id": "m2"}) + out = t.translate("end", {}) + + assert isinstance(out[-1], RunFinishedEvent) + m2_content = [e for e in out if isinstance(e, TextMessageContentEvent) and e.message_id == "m2"] + assert len(m2_content) == 1 + assert m2_content[0].delta == "waiting" + ends = [e.message_id for e in out if isinstance(e, TextMessageEndEvent)] + assert ends == ["m1", "m2"] def test_vertices_sorted_emits_state_snapshot_of_all_nodes(): @@ -610,20 +679,22 @@ def _assert_well_formed(events: list) -> None: assert isinstance(events[0], RunStartedEvent) assert isinstance(events[-1], (RunFinishedEvent, RunErrorEvent)) - open_messages: set[str] = set() + open_message: str | None = None seen_messages: set[str] = set() for event in events: if isinstance(event, TextMessageStartEvent): - assert event.message_id not in open_messages, "text message started while already open" + # AG-UI's reference verifier allows at most one open text message; + # interleaved STARTs are rejected by conforming clients. + assert open_message is None, f"text message {event.message_id} started while {open_message} is open" assert event.message_id not in seen_messages, "text message id reused after it ended" - open_messages.add(event.message_id) + open_message = event.message_id seen_messages.add(event.message_id) elif isinstance(event, TextMessageContentEvent): - assert event.message_id in open_messages, "text content for a message that is not open" + assert event.message_id == open_message, "text content for a message that is not open" elif isinstance(event, TextMessageEndEvent): - assert event.message_id in open_messages, "text message ended without being open" - open_messages.discard(event.message_id) - assert not open_messages, "text messages left unclosed" + assert event.message_id == open_message, "text message ended without being open" + open_message = None + assert open_message is None, "text message left unclosed" started_tools: set[str] = set() ended_tools: set[str] = set()