edsl 0.1.60__py3-none-any.whl → 0.1.62__py3-none-any.whl

edsl/scenarios/file_store.py

@@ -512,6 +512,75 @@ class FileStore(Scenario):
         )
         return info
 
+    def offload(self, inplace=False) -> "FileStore":
+        """
+        Offloads base64-encoded content from the FileStore by replacing 'base64_string'
+        with 'offloaded'. This reduces memory usage.
+
+        Args:
+            inplace (bool): If True, modify the current FileStore. If False, return a new one.
+
+        Returns:
+            FileStore: The modified FileStore (either self or a new instance).
+        """
+        if inplace:
+            if hasattr(self, "base64_string"):
+                self.base64_string = "offloaded"
+            return self
+        else:
+            # Create a copy and offload it
+            file_store_dict = self.to_dict()
+            if "base64_string" in file_store_dict:
+                file_store_dict["base64_string"] = "offloaded"
+            return self.__class__.from_dict(file_store_dict)
+
+    def save_to_gcs_bucket(self, signed_url: str) -> dict:
+        """
+        Saves the FileStore's file content to a Google Cloud Storage bucket using a signed URL.
+
+        Args:
+            signed_url (str): The signed URL for uploading to GCS bucket
+
+        Returns:
+            dict: Response from the GCS upload operation
+
+        Raises:
+            ValueError: If base64_string is offloaded or missing
+            requests.RequestException: If the upload fails
+        """
+        import requests
+        import base64
+
+        # Check if content is available
+        if not hasattr(self, "base64_string") or self.base64_string == "offloaded":
+            raise ValueError(
+                "File content is not available (offloaded or missing). Cannot upload to GCS."
+            )
+
+        # Decode base64 content to bytes
+        try:
+            file_content = base64.b64decode(self.base64_string)
+        except Exception as e:
+            raise ValueError(f"Failed to decode base64 content: {e}")
+
+        # Prepare headers with proper content type
+        headers = {
+            "Content-Type": self.mime_type or "application/octet-stream",
+            "Content-Length": str(len(file_content)),
+        }
+
+        # Upload to GCS using the signed URL
+        response = requests.put(signed_url, data=file_content, headers=headers)
+        response.raise_for_status()
+
+        return {
+            "status": "success",
+            "status_code": response.status_code,
+            "file_size": len(file_content),
+            "mime_type": self.mime_type,
+            "file_extension": self.suffix,
+        }
+
     @classmethod
     def pull(cls, url_or_uuid: Union[str, UUID]) -> "FileStore":
         """

edsl/scenarios/scenario.py

@@ -280,6 +280,18 @@ class Scenario(Base, UserDict):
 
         target = self if inplace else Scenario()
 
+        # First check if this Scenario itself has a base64_string (e.g., from FileStore.to_dict())
+        if "base64_string" in self and isinstance(self.get("base64_string"), str):
+            # This is likely a Scenario created from FileStore.to_dict()
+            if inplace:
+                self["base64_string"] = "offloaded"
+            else:
+                # Copy all keys to target
+                for k, v in self.items():
+                    target[k] = v
+                target["base64_string"] = "offloaded"
+            return target
+
         for key, value in self.items():
             if isinstance(value, FileStore):
                 file_store_dict = value.to_dict()
@@ -297,6 +309,227 @@ class Scenario(Base, UserDict):
 
         return target
 
+    def save_to_gcs_bucket(self, signed_url_or_dict) -> dict:
+        """
+        Saves FileStore objects contained within this Scenario to a Google Cloud Storage bucket.
+
+        This method finds all FileStore objects in the Scenario and uploads them to GCS using
+        the provided signed URL(s). If the Scenario itself was created from a FileStore (has
+        base64_string as a top-level key), it uploads that content directly.
+
+        Args:
+            signed_url_or_dict: Either:
+                - str: Single signed URL (for single FileStore or Scenario from FileStore)
+                - dict: Mapping of scenario keys to signed URLs for multiple FileStore objects
+                        e.g., {"video": "signed_url_1", "image": "signed_url_2"}
+
+        Returns:
+            dict: Summary of upload operations performed
+
+        Raises:
+            ValueError: If no uploadable content found or content is offloaded
+            requests.RequestException: If any upload fails
+        """
+        from edsl.scenarios import FileStore
+        import requests
+        import base64
+
+        upload_results = []
+
+        # Case 1: This Scenario was created from a FileStore (has direct base64_string)
+        if "base64_string" in self and isinstance(self.get("base64_string"), str):
+            if self["base64_string"] == "offloaded":
+                raise ValueError("File content is offloaded. Cannot upload to GCS.")
+
+            # For single FileStore scenario, expect string URL
+            if isinstance(signed_url_or_dict, dict):
+                raise ValueError(
+                    "For Scenario created from FileStore, provide a single signed URL string, not a dictionary."
+                )
+
+            signed_url = signed_url_or_dict
+
+            # Get file info from Scenario keys
+            mime_type = self.get("mime_type", "application/octet-stream")
+            suffix = self.get("suffix", "")
+
+            # Decode and upload
+            try:
+                file_content = base64.b64decode(self["base64_string"])
+            except Exception as e:
+                raise ValueError(f"Failed to decode base64 content: {e}")
+
+            headers = {
+                "Content-Type": mime_type,
+                "Content-Length": str(len(file_content)),
+            }
+
+            response = requests.put(signed_url, data=file_content, headers=headers)
+            response.raise_for_status()
+
+            upload_results.append(
+                {
+                    "type": "scenario_filestore_content",
+                    "status": "success",
+                    "status_code": response.status_code,
+                    "file_size": len(file_content),
+                    "mime_type": mime_type,
+                    "file_extension": suffix,
+                }
+            )
+
+        # Case 2: Find FileStore objects in Scenario values
+        else:
+            # Collect all FileStore keys first
+            filestore_keys = [
+                key for key, value in self.items() if isinstance(value, FileStore)
+            ]
+
+            if not filestore_keys:
+                raise ValueError("No FileStore objects found in Scenario to upload.")
+
+            # Handle URL parameter
+            if isinstance(signed_url_or_dict, str):
+                # Single URL provided for multiple FileStore objects - this will cause overwrites
+                if len(filestore_keys) > 1:
+                    raise ValueError(
+                        f"Multiple FileStore objects found ({filestore_keys}) but only one signed URL provided. "
+                        f"Provide a dictionary mapping keys to URLs to avoid overwrites: "
+                        f"{{'{filestore_keys[0]}': 'url1', '{filestore_keys[1]}': 'url2', ...}}"
+                    )
+
+                # Single FileStore object, single URL is fine
+                url_mapping = {filestore_keys[0]: signed_url_or_dict}
+
+            elif isinstance(signed_url_or_dict, dict):
+                # Dictionary of URLs provided
+                missing_keys = set(filestore_keys) - set(signed_url_or_dict.keys())
+                if missing_keys:
+                    raise ValueError(
+                        f"Missing signed URLs for FileStore keys: {list(missing_keys)}"
+                    )
+
+                extra_keys = set(signed_url_or_dict.keys()) - set(filestore_keys)
+                if extra_keys:
+                    raise ValueError(
+                        f"Signed URLs provided for non-FileStore keys: {list(extra_keys)}"
+                    )
+
+                url_mapping = signed_url_or_dict
+
+            else:
+                raise ValueError(
+                    "signed_url_or_dict must be either a string or a dictionary"
+                )
+
+            # Upload each FileStore object
+            for key, value in self.items():
+                if isinstance(value, FileStore):
+                    try:
+                        result = value.save_to_gcs_bucket(url_mapping[key])
+                        result["scenario_key"] = key
+                        result["type"] = "filestore_object"
+                        upload_results.append(result)
+                    except Exception as e:
+                        upload_results.append(
+                            {
+                                "scenario_key": key,
+                                "type": "filestore_object",
+                                "status": "error",
+                                "error": str(e),
+                            }
+                        )
+
+        return {
+            "total_uploads": len(upload_results),
+            "successful_uploads": len(
+                [r for r in upload_results if r.get("status") == "success"]
+            ),
+            "failed_uploads": len(
+                [r for r in upload_results if r.get("status") == "error"]
+            ),
+            "upload_details": upload_results,
+        }
+
+    def get_filestore_info(self) -> dict:
+        """
+        Returns information about FileStore objects present in this Scenario.
+
+        This method is useful for determining how many signed URLs need to be generated
+        and what file extensions/types are present before calling save_to_gcs_bucket().
+
+        Returns:
+            dict: Information about FileStore objects containing:
+                - total_count: Total number of FileStore objects
+                - filestore_keys: List of scenario keys that contain FileStore objects
+                - file_extensions: Dictionary mapping keys to file extensions
+                - file_types: Dictionary mapping keys to MIME types
+                - is_filestore_scenario: Boolean indicating if this Scenario was created from a FileStore
+                - summary: Human-readable summary of files
+
+
+        """
+        from edsl.scenarios import FileStore
+
+        # Check if this Scenario was created from a FileStore
+        is_filestore_scenario = "base64_string" in self and isinstance(
+            self.get("base64_string"), str
+        )
+
+        if is_filestore_scenario:
+            # Single FileStore scenario
+            return {
+                "total_count": 1,
+                "filestore_keys": ["filestore_content"],
+                "file_extensions": {"filestore_content": self.get("suffix", "")},
+                "file_types": {
+                    "filestore_content": self.get(
+                        "mime_type", "application/octet-stream"
+                    )
+                },
+                "is_filestore_scenario": True,
+                "summary": f"Single FileStore content with extension '{self.get('suffix', 'unknown')}'",
+            }
+
+        # Regular Scenario with FileStore objects as values
+        filestore_info = {}
+        file_extensions = {}
+        file_types = {}
+
+        for key, value in self.items():
+            if isinstance(value, FileStore):
+                filestore_info[key] = {
+                    "extension": getattr(value, "suffix", ""),
+                    "mime_type": getattr(
+                        value, "mime_type", "application/octet-stream"
+                    ),
+                    "binary": getattr(value, "binary", True),
+                    "path": getattr(value, "path", "unknown"),
+                }
+                file_extensions[key] = getattr(value, "suffix", "")
+                file_types[key] = getattr(
+                    value, "mime_type", "application/octet-stream"
+                )
+
+        # Generate summary
+        if filestore_info:
+            ext_summary = [f"{key}({ext})" for key, ext in file_extensions.items()]
+            summary = (
+                f"{len(filestore_info)} FileStore objects: {', '.join(ext_summary)}"
+            )
+        else:
+            summary = "No FileStore objects found"
+
+        return {
+            "total_count": len(filestore_info),
+            "filestore_keys": list(filestore_info.keys()),
+            "file_extensions": file_extensions,
+            "file_types": file_types,
+            "is_filestore_scenario": False,
+            "detailed_info": filestore_info,
+            "summary": summary,
+        }
+
     def to_dict(
         self, add_edsl_version: bool = True, offload_base64: bool = False
     ) -> dict:

edsl/scenarios/scenario_list.py

@@ -159,7 +159,15 @@ class ScenarioList(MutableSequence, Base, ScenarioListOperationsMixin):
 
     # Required MutableSequence abstract methods
     def __getitem__(self, index):
-        """Get item at index."""
+        """Get item at index.
+
+        Example:
+            >>> from edsl.scenarios import Scenario, ScenarioList
+            >>> sl = ScenarioList([Scenario({'a': 12})])
+            >>> sl[0]['b'] = 100  # modify in-place
+            >>> sl[0]['b']
+            100
+        """
         if isinstance(index, slice):
             return self.__class__(list(self.data[index]), self.codebook.copy())
         return self.data[index]
@@ -356,7 +364,29 @@ class ScenarioList(MutableSequence, Base, ScenarioListOperationsMixin):
                 new_scenarios.append(Scenario(new_scenario))
 
         return new_scenarios
+
+    @classmethod
+    def from_prompt(self, description: str, name:Optional[str] = "item", target_number:int = 10, verbose = False):
+        from ..questions.question_list import QuestionList
+        q = QuestionList(question_name = name, 
+                         question_text = description + f"\n Please try to return {target_number} examples.")
+        results = q.run(verbose = verbose)
+        return results.select(name).to_scenario_list().expand(name)
     
+
+    def __add__(self, other):
+        if isinstance(other, Scenario):
+            new_list = self.duplicate()
+            new_list.append(other)
+            return new_list
+        elif isinstance(other, ScenarioList):
+            new_list = self.duplicate()
+            for item in other:
+                new_list.append(item)
+        else:
+            raise ScenarioError("Don't know how to combine!")
+        return new_list
+
     @classmethod
     def from_search_terms(cls, search_terms: List[str]) -> ScenarioList:
         """Create a ScenarioList from a list of search terms, using Wikipedia.