lfx-nightly 0.2.0.dev25__py3-none-any.whl

lfx/components/twelvelabs/convert_astra_results.py

@@ -0,0 +1,84 @@
+from typing import Any
+
+from lfx.custom import Component
+from lfx.io import HandleInput, Output
+from lfx.schema import Data
+from lfx.schema.message import Message
+
+
+class ConvertAstraToTwelveLabs(Component):
+    """Convert Astra DB search results to TwelveLabs Pegasus inputs."""
+
+    display_name = "Convert Astra DB to Pegasus Input"
+    description = "Converts Astra DB search results to inputs compatible with TwelveLabs Pegasus."
+    icon = "TwelveLabs"
+    name = "ConvertAstraToTwelveLabs"
+    documentation = "https://github.com/twelvelabs-io/twelvelabs-developer-experience/blob/main/integrations/Langflow/TWELVE_LABS_COMPONENTS_README.md"
+
+    inputs = [
+        HandleInput(
+            name="astra_results",
+            display_name="Astra DB Results",
+            input_types=["Data"],
+            info="Search results from Astra DB component",
+            required=True,
+            is_list=True,
+        )
+    ]
+
+    outputs = [
+        Output(
+            name="index_id",
+            display_name="Index ID",
+            type_=Message,
+            method="get_index_id",
+        ),
+        Output(
+            name="video_id",
+            display_name="Video ID",
+            type_=Message,
+            method="get_video_id",
+        ),
+    ]
+
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+        self._video_id = None
+        self._index_id = None
+
+    def build(self, **kwargs: Any) -> None:  # noqa: ARG002 - Required for parent class compatibility
+        """Process the Astra DB results and extract TwelveLabs index information."""
+        if not self.astra_results:
+            return
+
+        # Convert to list if single item
+        results = self.astra_results if isinstance(self.astra_results, list) else [self.astra_results]
+
+        # Try to extract index information from metadata
+        for doc in results:
+            if not isinstance(doc, Data):
+                continue
+
+            # Get the metadata, handling the nested structure
+            metadata = {}
+            if hasattr(doc, "metadata") and isinstance(doc.metadata, dict):
+                # Handle nested metadata using .get() method
+                metadata = doc.metadata.get("metadata", doc.metadata)
+
+            # Extract index_id and video_id
+            self._index_id = metadata.get("index_id")
+            self._video_id = metadata.get("video_id")
+
+            # If we found both, we can stop searching
+            if self._index_id and self._video_id:
+                break
+
+    def get_video_id(self) -> Message:
+        """Return the extracted video ID as a Message."""
+        self.build()
+        return Message(text=self._video_id if self._video_id else "")
+
+    def get_index_id(self) -> Message:
+        """Return the extracted index ID as a Message."""
+        self.build()
+        return Message(text=self._index_id if self._index_id else "")

lfx/components/twelvelabs/pegasus_index.py

@@ -0,0 +1,311 @@
+import time
+from concurrent.futures import ThreadPoolExecutor
+from pathlib import Path
+from typing import Any
+
+from tenacity import retry, stop_after_attempt, wait_exponential
+from twelvelabs import TwelveLabs
+
+from lfx.custom import Component
+from lfx.inputs import DataInput, DropdownInput, SecretStrInput, StrInput
+from lfx.io import Output
+from lfx.schema import Data
+
+
+class TwelveLabsError(Exception):
+    """Base exception for TwelveLabs errors."""
+
+
+class IndexCreationError(TwelveLabsError):
+    """Error raised when there's an issue with an index."""
+
+
+class TaskError(TwelveLabsError):
+    """Error raised when a task fails."""
+
+
+class TaskTimeoutError(TwelveLabsError):
+    """Error raised when a task times out."""
+
+
+class PegasusIndexVideo(Component):
+    """Indexes videos using TwelveLabs Pegasus API and adds the video ID to metadata."""
+
+    display_name = "TwelveLabs Pegasus Index Video"
+    description = "Index videos using TwelveLabs and add the video_id to metadata."
+    icon = "TwelveLabs"
+    name = "TwelveLabsPegasusIndexVideo"
+    documentation = "https://github.com/twelvelabs-io/twelvelabs-developer-experience/blob/main/integrations/Langflow/TWELVE_LABS_COMPONENTS_README.md"
+
+    inputs = [
+        DataInput(
+            name="videodata",
+            display_name="Video Data",
+            info="Video Data objects (from VideoFile or SplitVideo)",
+            is_list=True,
+            required=True,
+        ),
+        SecretStrInput(
+            name="api_key", display_name="TwelveLabs API Key", info="Enter your TwelveLabs API Key.", required=True
+        ),
+        DropdownInput(
+            name="model_name",
+            display_name="Model",
+            info="Pegasus model to use for indexing",
+            options=["pegasus1.2"],
+            value="pegasus1.2",
+            advanced=False,
+        ),
+        StrInput(
+            name="index_name",
+            display_name="Index Name",
+            info="Name of the index to use. If the index doesn't exist, it will be created.",
+            required=False,
+        ),
+        StrInput(
+            name="index_id",
+            display_name="Index ID",
+            info="ID of an existing index to use. If provided, index_name will be ignored.",
+            required=False,
+        ),
+    ]
+
+    outputs = [
+        Output(
+            display_name="Indexed Data", name="indexed_data", method="index_videos", output_types=["Data"], is_list=True
+        ),
+    ]
+
+    def _get_or_create_index(self, client: TwelveLabs) -> tuple[str, str]:
+        """Get existing index or create new one.
+
+        Returns (index_id, index_name).
+        """
+        # First check if index_id is provided and valid
+        if hasattr(self, "index_id") and self.index_id:
+            try:
+                index = client.index.retrieve(id=self.index_id)
+            except (ValueError, KeyError) as e:
+                if not hasattr(self, "index_name") or not self.index_name:
+                    error_msg = "Invalid index ID provided and no index name specified for fallback"
+                    raise IndexCreationError(error_msg) from e
+            else:
+                return self.index_id, index.name
+
+        # If index_name is provided, try to find it
+        if hasattr(self, "index_name") and self.index_name:
+            try:
+                # List all indexes and find by name
+                indexes = client.index.list()
+                for idx in indexes:
+                    if idx.name == self.index_name:
+                        return idx.id, idx.name
+
+                # If we get here, index wasn't found - create it
+                index = client.index.create(
+                    name=self.index_name,
+                    models=[
+                        {
+                            "name": self.model_name if hasattr(self, "model_name") else "pegasus1.2",
+                            "options": ["visual", "audio"],
+                        }
+                    ],
+                )
+            except (ValueError, KeyError) as e:
+                error_msg = f"Error with index name {self.index_name}"
+                raise IndexCreationError(error_msg) from e
+            else:
+                return index.id, index.name
+
+        # If we get here, neither index_id nor index_name was provided
+        error_msg = "Either index_name or index_id must be provided"
+        raise IndexCreationError(error_msg)
+
+    def on_task_update(self, task: Any, video_path: str) -> None:
+        """Callback for task status updates.
+
+        Updates the component status with the current task status.
+        """
+        video_name = Path(video_path).name
+        status_msg = f"Indexing {video_name}... Status: {task.status}"
+        self.status = status_msg
+
+    @retry(stop=stop_after_attempt(5), wait=wait_exponential(multiplier=1, min=5, max=60), reraise=True)
+    def _check_task_status(
+        self,
+        client: TwelveLabs,
+        task_id: str,
+        video_path: str,
+    ) -> Any:
+        """Check task status once.
+
+        Makes a single API call to check the status of a task.
+        """
+        task = client.task.retrieve(id=task_id)
+        self.on_task_update(task, video_path)
+        return task
+
+    def _wait_for_task_completion(
+        self, client: TwelveLabs, task_id: str, video_path: str, max_retries: int = 120, sleep_time: int = 10
+    ) -> Any:
+        """Wait for task completion with timeout and improved error handling.
+
+        Polls the task status until completion or timeout.
+        """
+        retries = 0
+        consecutive_errors = 0
+        max_consecutive_errors = 5
+        video_name = Path(video_path).name
+
+        while retries < max_retries:
+            try:
+                self.status = f"Checking task status for {video_name} (attempt {retries + 1})"
+                task = self._check_task_status(client, task_id, video_path)
+
+                if task.status == "ready":
+                    self.status = f"Indexing for {video_name} completed successfully!"
+                    return task
+                if task.status == "failed":
+                    error_msg = f"Task failed for {video_name}: {getattr(task, 'error', 'Unknown error')}"
+                    self.status = error_msg
+                    raise TaskError(error_msg)
+                if task.status == "error":
+                    error_msg = f"Task encountered an error for {video_name}: {getattr(task, 'error', 'Unknown error')}"
+                    self.status = error_msg
+                    raise TaskError(error_msg)
+
+                time.sleep(sleep_time)
+                retries += 1
+                elapsed_time = retries * sleep_time
+                self.status = f"Indexing {video_name}... {elapsed_time}s elapsed"
+
+            except (ValueError, KeyError) as e:
+                consecutive_errors += 1
+                error_msg = f"Error checking task status for {video_name}: {e!s}"
+                self.status = error_msg
+
+                if consecutive_errors >= max_consecutive_errors:
+                    too_many_errors = f"Too many consecutive errors checking task status for {video_name}"
+                    raise TaskError(too_many_errors) from e
+
+                time.sleep(sleep_time * (2**consecutive_errors))
+                continue
+
+        timeout_msg = f"Timeout waiting for indexing of {video_name} after {max_retries * sleep_time} seconds"
+        self.status = timeout_msg
+        raise TaskTimeoutError(timeout_msg)
+
+    def _upload_video(self, client: TwelveLabs, video_path: str, index_id: str) -> str:
+        """Upload a single video and return its task ID.
+
+        Uploads a video file to the specified index and returns the task ID.
+        """
+        video_name = Path(video_path).name
+        with Path(video_path).open("rb") as video_file:
+            self.status = f"Uploading {video_name} to index {index_id}..."
+            task = client.task.create(index_id=index_id, file=video_file)
+            task_id = task.id
+            self.status = f"Upload complete for {video_name}. Task ID: {task_id}"
+            return task_id
+
+    def index_videos(self) -> list[Data]:
+        """Indexes each video and adds the video_id to its metadata."""
+        if not self.videodata:
+            self.status = "No video data provided."
+            return []
+
+        if not self.api_key:
+            error_msg = "TwelveLabs API Key is required"
+            raise IndexCreationError(error_msg)
+
+        if not (hasattr(self, "index_name") and self.index_name) and not (hasattr(self, "index_id") and self.index_id):
+            error_msg = "Either index_name or index_id must be provided"
+            raise IndexCreationError(error_msg)
+
+        client = TwelveLabs(api_key=self.api_key)
+        indexed_data_list: list[Data] = []
+
+        # Get or create the index
+        try:
+            index_id, index_name = self._get_or_create_index(client)
+            self.status = f"Using index: {index_name} (ID: {index_id})"
+        except IndexCreationError as e:
+            self.status = f"Failed to get/create TwelveLabs index: {e!s}"
+            raise
+
+        # First, validate all videos and create a list of valid ones
+        valid_videos: list[tuple[Data, str]] = []
+        for video_data_item in self.videodata:
+            if not isinstance(video_data_item, Data):
+                self.status = f"Skipping invalid data item: {video_data_item}"
+                continue
+
+            video_info = video_data_item.data
+            if not isinstance(video_info, dict):
+                self.status = f"Skipping item with invalid data structure: {video_info}"
+                continue
+
+            video_path = video_info.get("text")
+            if not video_path or not isinstance(video_path, str):
+                self.status = f"Skipping item with missing or invalid video path: {video_info}"
+                continue
+
+            if not Path(video_path).exists():
+                self.status = f"Video file not found, skipping: {video_path}"
+                continue
+
+            valid_videos.append((video_data_item, video_path))
+
+        if not valid_videos:
+            self.status = "No valid videos to process."
+            return []
+
+        # Upload all videos first and collect their task IDs
+        upload_tasks: list[tuple[Data, str, str]] = []  # (data_item, video_path, task_id)
+        for data_item, video_path in valid_videos:
+            try:
+                task_id = self._upload_video(client, video_path, index_id)
+                upload_tasks.append((data_item, video_path, task_id))
+            except (ValueError, KeyError) as e:
+                self.status = f"Failed to upload {video_path}: {e!s}"
+                continue
+
+        # Now check all tasks in parallel using a thread pool
+        with ThreadPoolExecutor(max_workers=min(10, len(upload_tasks))) as executor:
+            futures = []
+            for data_item, video_path, task_id in upload_tasks:
+                future = executor.submit(self._wait_for_task_completion, client, task_id, video_path)
+                futures.append((data_item, video_path, future))
+
+            # Process results as they complete
+            for data_item, video_path, future in futures:
+                try:
+                    completed_task = future.result()
+                    if completed_task.status == "ready":
+                        video_id = completed_task.video_id
+                        video_name = Path(video_path).name
+                        self.status = f"Video {video_name} indexed successfully. Video ID: {video_id}"
+
+                        # Add video_id to the metadata
+                        video_info = data_item.data
+                        if "metadata" not in video_info:
+                            video_info["metadata"] = {}
+                        elif not isinstance(video_info["metadata"], dict):
+                            self.status = f"Warning: Overwriting non-dict metadata for {video_path}"
+                            video_info["metadata"] = {}
+
+                        video_info["metadata"].update(
+                            {"video_id": video_id, "index_id": index_id, "index_name": index_name}
+                        )
+
+                        updated_data_item = Data(data=video_info)
+                        indexed_data_list.append(updated_data_item)
+                except (TaskError, TaskTimeoutError) as e:
+                    self.status = f"Failed to process {video_path}: {e!s}"
+
+        if not indexed_data_list:
+            self.status = "No videos were successfully indexed."
+        else:
+            self.status = f"Finished indexing {len(indexed_data_list)}/{len(self.videodata)} videos."
+
+        return indexed_data_list