edsl 0.1.41__py3-none-any.whl → 0.1.43__py3-none-any.whl

edsl/results/DatasetExportMixin.py

@@ -7,7 +7,6 @@ from typing import Optional, Tuple, Union, List
 
 from edsl.results.file_exports import CSVExport, ExcelExport, JSONLExport, SQLiteExport
 
-
 class DatasetExportMixin:
     """Mixin class for exporting Dataset objects."""
 
@@ -220,23 +219,45 @@ class DatasetExportMixin:
         )
         return exporter.export()
 
-    def _db(self, remove_prefix: bool = True):
+    def _db(self, remove_prefix: bool = True, shape: str = "wide") -> "sqlalchemy.engine.Engine":
         """Create a SQLite database in memory and return the connection.
 
         Args:
-            shape: The shape of the data in the database (wide or long)
             remove_prefix: Whether to remove the prefix from the column names
+            shape: The shape of the data in the database ("wide" or "long")
 
         Returns:
             A database connection
+        >>> from sqlalchemy import text
+        >>> from edsl import Results 
+        >>> engine = Results.example()._db()
+        >>> len(engine.execute(text("SELECT * FROM self")).fetchall())
+        4
+        >>> engine = Results.example()._db(shape = "long")
+        >>> len(engine.execute(text("SELECT * FROM self")).fetchall())
+        172
         """
-        from sqlalchemy import create_engine
+        from sqlalchemy import create_engine, text
 
         engine = create_engine("sqlite:///:memory:")
-        if remove_prefix:
+        if remove_prefix and shape == "wide":
             df = self.remove_prefix().to_pandas(lists_as_strings=True)
         else:
             df = self.to_pandas(lists_as_strings=True)
+
+        if shape == "long":
+            # Melt the dataframe to convert it to long format
+            df = df.melt(
+                var_name='key', 
+                value_name='value'
+            )
+            # Add a row number column for reference
+            df.insert(0, 'row_number', range(1, len(df) + 1))
+            
+            # Split the key into data_type and key
+            df['data_type'] = df['key'].apply(lambda x: x.split('.')[0] if '.' in x else None)
+            df['key'] = df['key'].apply(lambda x: '.'.join(x.split('.')[1:]) if '.' in x else x)
+
         df.to_sql(
             "self",
             engine,
@@ -251,6 +272,7 @@ class DatasetExportMixin:
         transpose: bool = None,
         transpose_by: str = None,
         remove_prefix: bool = True,
+        shape: str = "wide",
     ) -> Union["pd.DataFrame", str]:
         """Execute a SQL query and return the results as a DataFrame.
 
@@ -268,10 +290,17 @@ class DatasetExportMixin:
         Returns:
             DataFrame, CSV string, list, or LaTeX string depending on parameters
 
+       Examples:
+           >>> from edsl import Results
+           >>> r = Results.example(); 
+           >>> len(r.sql("SELECT * FROM self", shape = "wide"))
+           4
+           >>> len(r.sql("SELECT * FROM self", shape = "long"))
+           172
         """
         import pandas as pd
 
-        conn = self._db(remove_prefix=remove_prefix)
+        conn = self._db(remove_prefix=remove_prefix, shape=shape)
         df = pd.read_sql_query(query, conn)
 
         # Transpose the DataFrame if transpose is True

edsl/results/Result.py

@@ -78,7 +78,6 @@ class Result(Base, UserDict):
         self.question_to_attributes = (
             question_to_attributes or self._create_question_to_attributes(survey)
         )
-
         data = {
             "agent": agent,
             "scenario": scenario,
@@ -87,7 +86,7 @@ class Result(Base, UserDict):
             "answer": answer,
             "prompt": prompt or {},
             "raw_model_response": raw_model_response or {},
-            "question_to_attributes": question_to_attributes,
+            "question_to_attributes": self.question_to_attributes,
             "generated_tokens": generated_tokens or {},
             "comments_dict": comments_dict or {},
             "cache_used_dict": cache_used_dict or {},
@@ -154,7 +153,9 @@ class Result(Base, UserDict):
     @staticmethod
     def _create_model_sub_dict(model) -> dict:
         return {
-            "model": model.parameters | {"model": model.model},
+            "model": model.parameters
+            | {"model": model.model}
+            | {"inference_service": model._inference_service_},
         }
 
     @staticmethod
@@ -365,6 +366,10 @@ class Result(Base, UserDict):
                         else prompt_obj.to_dict()
                     )
                 d[key] = new_prompt_dict
+            
+        if self.indices is not None:
+            d["indices"] = self.indices
+
         if add_edsl_version:
             from edsl import __version__
 
@@ -414,6 +419,7 @@ class Result(Base, UserDict):
             comments_dict=json_dict.get("comments_dict", {}),
             cache_used_dict=json_dict.get("cache_used_dict", {}),
             cache_keys=json_dict.get("cache_keys", {}),
+            indices = json_dict.get("indices", None)
         )
         return result
 

edsl/results/Results.py

@@ -38,6 +38,64 @@ from edsl.results.ResultsGGMixin import ResultsGGMixin
 from edsl.results.results_fetch_mixin import ResultsFetchMixin
 from edsl.utilities.remove_edsl_version import remove_edsl_version
 
+def ensure_fetched(method):
+    """A decorator that checks if remote data is loaded, and if not, attempts to fetch it."""
+    def wrapper(self, *args, **kwargs):
+        if not self._fetched:
+            # If not fetched, try fetching now.
+            # (If you know you have job info stored in self.job_info)
+            self.fetch_remote(self.job_info)
+        return method(self, *args, **kwargs)
+    return wrapper
+
+def ensure_ready(method):
+    """
+    Decorator for Results methods.
+    
+    If the Results object is not ready, for most methods we return a NotReadyObject.
+    However, for __repr__ (and other methods that need to return a string), we return
+    the string representation of NotReadyObject.
+    """
+    from functools import wraps
+
+    @wraps(method)
+    def wrapper(self, *args, **kwargs):
+        if self.completed:
+            return method(self, *args, **kwargs)
+        # Attempt to fetch remote data
+        try:
+            if hasattr(self, "job_info"):
+                self.fetch_remote(self.job_info)
+        except Exception as e:
+            print(f"Error during fetch_remote in {method.__name__}: {e}")
+        if not self.completed:
+            not_ready = NotReadyObject(name = method.__name__, job_info = self.job_info)
+            # For __repr__, ensure we return a string
+            if method.__name__ == "__repr__" or method.__name__ == "__str__":
+                return not_ready.__repr__()
+            return not_ready
+        return method(self, *args, **kwargs)
+
+    return wrapper
+
+class NotReadyObject:
+    """A placeholder object that prints a message when any attribute is accessed."""
+    def __init__(self, name: str, job_info: RemoteJobInfo):
+        self.name = name
+        self.job_info = job_info
+        #print(f"Not ready to call {name}")
+
+    def __repr__(self):
+        message = f"""Results not ready - job still running on server."""
+        for key, value in self.job_info.creation_data.items():
+            message += f"\n{key}: {value}"
+        return message
+
+    def __getattr__(self, _):
+        return self
+    
+    def __call__(self, *args, **kwargs):
+        return self
 
 class Mixins(
     ResultsExportMixin,
@@ -93,6 +151,16 @@ class Results(UserList, Mixins, Base):
         "cache_keys",
     ]
 
+    @classmethod
+    def from_job_info(cls, job_info: dict) -> Results:
+        """
+        Instantiate a `Results` object from a job info dictionary.
+        """
+        results = cls()
+        results.completed = False
+        results.job_info = job_info
+        return results
+
     def __init__(
         self,
         survey: Optional[Survey] = None,
@@ -112,6 +180,8 @@ class Results(UserList, Mixins, Base):
         :param total_results: An integer representing the total number of results.
         :cache: A Cache object.
         """
+        self.completed = True
+        self._fetching = False
         super().__init__(data)
         from edsl.data.Cache import Cache
         from edsl.jobs.tasks.TaskHistory import TaskHistory
@@ -315,7 +385,22 @@ class Results(UserList, Mixins, Base):
             data=self.data + other.data,
             created_columns=self.created_columns,
         )
-
+    
+    def _repr_html_(self):
+        if not self.completed:
+            if hasattr(self, "job_info"):
+                self.fetch_remote(self.job_info)
+            
+            if not self.completed:
+                return f"Results not ready to call"
+        
+        return super()._repr_html_()
+
+    # @ensure_ready
+    # def __str__(self):
+    #     super().__str__()    
+    
+    @ensure_ready
     def __repr__(self) -> str:
         return f"Results(data = {self.data}, survey = {repr(self.survey)}, created_columns = {self.created_columns})"
 
@@ -647,7 +732,7 @@ class Results(UserList, Mixins, Base):
 
         >>> r = Results.example()
         >>> r.model_keys
-        ['frequency_penalty', 'logprobs', 'max_tokens', 'model', 'model_index', 'presence_penalty', 'temperature', 'top_logprobs', 'top_p']
+        ['frequency_penalty', 'inference_service', 'logprobs', 'max_tokens', 'model', 'model_index', 'presence_penalty', 'temperature', 'top_logprobs', 'top_p']
         """
         return sorted(self._data_type_to_keys["model"])
 
@@ -732,6 +817,7 @@ class Results(UserList, Mixins, Base):
 
         return self.recode(column, recode_function=f, new_var_name=new_var_name)
 
+    @ensure_ready
     def recode(
         self, column: str, recode_function: Optional[Callable], new_var_name=None
     ) -> Results:
@@ -760,6 +846,7 @@ class Results(UserList, Mixins, Base):
             created_columns=self.created_columns + [new_var_name],
         )
 
+    @ensure_ready
     def add_column(self, column_name: str, values: list) -> Results:
         """Adds columns to Results
 
@@ -780,6 +867,7 @@ class Results(UserList, Mixins, Base):
             created_columns=self.created_columns + [column_name],
         )
 
+    @ensure_ready
     def add_columns_from_dict(self, columns: List[dict]) -> Results:
         """Adds columns to Results from a list of dictionaries.
 
@@ -829,6 +917,7 @@ class Results(UserList, Mixins, Base):
         evaluator.functions.update(int=int, float=float)
         return evaluator
 
+    @ensure_ready
     def mutate(
         self, new_var_string: str, functions_dict: Optional[dict] = None
     ) -> Results:
@@ -879,6 +968,7 @@ class Results(UserList, Mixins, Base):
             created_columns=self.created_columns + [var_name],
         )
 
+    @ensure_ready
     def add_column(self, column_name: str, values: list) -> Results:
         """Adds columns to Results
 
@@ -899,6 +989,7 @@ class Results(UserList, Mixins, Base):
             created_columns=self.created_columns + [column_name],
         )
 
+    @ensure_ready
     def rename(self, old_name: str, new_name: str) -> Results:
         """Rename an answer column in a Results object.
 
@@ -916,6 +1007,7 @@ class Results(UserList, Mixins, Base):
 
         return self
 
+    @ensure_ready
     def shuffle(self, seed: Optional[str] = "edsl") -> Results:
         """Shuffle the results.
 
@@ -932,6 +1024,7 @@ class Results(UserList, Mixins, Base):
         random.shuffle(new_data)
         return Results(survey=self.survey, data=new_data, created_columns=None)
 
+    @ensure_ready
     def sample(
         self,
         n: Optional[int] = None,
@@ -971,6 +1064,7 @@ class Results(UserList, Mixins, Base):
 
         return Results(survey=self.survey, data=new_data, created_columns=None)
 
+    @ensure_ready
     def select(self, *columns: Union[str, list[str]]) -> Results:
         """
         Select data from the results and format it.
@@ -1004,6 +1098,7 @@ class Results(UserList, Mixins, Base):
         )
         return selector.select(*columns)
 
+    @ensure_ready
     def sort_by(self, *columns: str, reverse: bool = False) -> Results:
         """Sort the results by one or more columns."""
         import warnings
@@ -1019,6 +1114,7 @@ class Results(UserList, Mixins, Base):
             return column.split(".")
         return self._key_to_data_type[column], column
 
+    @ensure_ready
     def order_by(self, *columns: str, reverse: bool = False) -> Results:
         """Sort the results by one or more columns.
 
@@ -1055,6 +1151,7 @@ class Results(UserList, Mixins, Base):
         new_data = sorted(self.data, key=sort_key, reverse=reverse)
         return Results(survey=self.survey, data=new_data, created_columns=None)
 
+    @ensure_ready
     def filter(self, expression: str) -> Results:
         """
         Filter based on the given expression and returns the filtered `Results`.
@@ -1156,6 +1253,7 @@ class Results(UserList, Mixins, Base):
         """Display an object as a table."""
         pass
 
+    @ensure_ready
     def __str__(self):
         data = self.to_dict()["data"]
         return json.dumps(data, indent=4)
@@ -1178,6 +1276,86 @@ class Results(UserList, Mixins, Base):
         [1, 1, 0, 0]
         """
         return [r.score(f) for r in self.data]
+    
+
+    def fetch_remote(self, job_info: "RemoteJobInfo") -> None:
+        """
+        Fetches the remote Results object using the provided RemoteJobInfo and updates this instance with the remote data.
+        
+        This is useful when you have a Results object that was created locally but want to sync it with 
+        the latest data from the remote server.
+        
+        Args:
+            job_info: RemoteJobInfo object containing the job_uuid and other remote job details
+        
+        """
+        #print("Calling fetch_remote")
+        try:
+            from edsl.coop.coop import Coop
+            from edsl.jobs.JobsRemoteInferenceHandler import JobsRemoteInferenceHandler
+            
+            # Get the remote job data
+            remote_job_data = JobsRemoteInferenceHandler.check_status(job_info.job_uuid)
+            
+            if remote_job_data.get("status") not in ["completed", "failed"]:
+                return False
+                #            
+            results_uuid = remote_job_data.get("results_uuid")
+            if not results_uuid:
+                raise ResultsError("No results_uuid found in remote job data")
+            
+            # Fetch the remote Results object
+            coop = Coop()
+            remote_results = coop.get(results_uuid, expected_object_type="results")
+            
+            # Update this instance with remote data
+            self.data = remote_results.data
+            self.survey = remote_results.survey
+            self.created_columns = remote_results.created_columns
+            self.cache = remote_results.cache
+            self.task_history = remote_results.task_history
+            self.completed = True
+            
+            # Set job_uuid and results_uuid from remote data
+            self.job_uuid = job_info.job_uuid
+            if hasattr(remote_results, 'results_uuid'):
+                self.results_uuid = remote_results.results_uuid
+
+            return True
+
+        except Exception as e:
+            raise ResultsError(f"Failed to fetch remote results: {str(e)}")
+
+    def fetch(self, polling_interval: float = 1.0) -> Results:
+        """
+        Polls the server for job completion and updates this Results instance with the completed data.
+        
+        Args:
+            polling_interval: Number of seconds to wait between polling attempts (default: 1.0)
+        
+        Returns:
+            self: The updated Results instance
+        """
+        if not hasattr(self, "job_info"):
+            raise ResultsError("No job info available - this Results object wasn't created from a remote job")
+        
+        from edsl.jobs.JobsRemoteInferenceHandler import JobsRemoteInferenceHandler
+        
+        try:
+            # Get the remote job data
+            remote_job_data = JobsRemoteInferenceHandler.check_status(self.job_info.job_uuid)
+            
+            while remote_job_data.get("status") not in ["completed", "failed"]:
+                import time
+                time.sleep(polling_interval)
+                remote_job_data = JobsRemoteInferenceHandler.check_status(self.job_info.job_uuid)
+                
+            # Once complete, fetch the full results
+            self.fetch_remote(self.job_info)
+            return self
+            
+        except Exception as e:
+            raise ResultsError(f"Failed to fetch remote results: {str(e)}")
 
 
 def main():  # pragma: no cover

edsl/results/ResultsGGMixin.py

@@ -5,46 +5,113 @@ import tempfile
 from typing import Optional
 
 
+class GGPlot:
+    """A class to handle ggplot2 plot display and saving."""
+    
+    def __init__(self, r_code: str, width: float = 6, height: float = 4):
+        """Initialize with R code and dimensions."""
+        self.r_code = r_code
+        self.width = width
+        self.height = height
+        self._svg_data = None
+        self._saved = False  # Track if the plot was saved
+        
+    def _execute_r_code(self, save_command: str = ""):
+        """Execute R code with optional save command."""
+        full_r_code = self.r_code + save_command
+        
+        result = subprocess.run(
+            ["Rscript", "-"],
+            input=full_r_code,
+            text=True,
+            stdout=subprocess.PIPE,
+            stderr=subprocess.PIPE,
+        )
+
+        if result.returncode != 0:
+            if result.returncode == 127:
+                raise RuntimeError(
+                    "Rscript is probably not installed. Please install R from https://cran.r-project.org/"
+                )
+            else:
+                raise RuntimeError(
+                    f"An error occurred while running Rscript: {result.stderr}"
+                )
+
+        if result.stderr:
+            print("Error in R script:", result.stderr)
+            
+        return result
+        
+    def save(self, filename: str):
+        """Save the plot to a file."""
+        format = filename.split('.')[-1].lower()
+        if format not in ['svg', 'png']:
+            raise ValueError("Only 'svg' and 'png' formats are supported")
+            
+        save_command = f'\nggsave("{filename}", plot = last_plot(), width = {self.width}, height = {self.height}, device = "{format}")'
+        self._execute_r_code(save_command)
+        
+        self._saved = True
+        print(f"File saved to: {filename}")
+        return None  # Return None instead of self
+        
+    def _repr_html_(self):
+        """Display the plot in a Jupyter notebook."""
+        # Don't display if the plot was saved
+        if self._saved:
+            return None
+            
+        import tempfile
+        
+        # Generate SVG if we haven't already
+        if self._svg_data is None:
+            # Create temporary SVG file
+            with tempfile.NamedTemporaryFile(suffix='.svg') as tmp:
+                save_command = f'\nggsave("{tmp.name}", plot = last_plot(), width = {self.width}, height = {self.height}, device = "svg")'
+                self._execute_r_code(save_command)
+                with open(tmp.name, 'r') as f:
+                    self._svg_data = f.read()
+                    
+        return self._svg_data
+
 class ResultsGGMixin:
     """Mixin class for ggplot2 plotting."""
 
     def ggplot2(
         self,
         ggplot_code: str,
-        filename: str = None,
         shape="wide",
         sql: str = None,
         remove_prefix: bool = True,
         debug: bool = False,
         height=4,
         width=6,
-        format="svg",
         factor_orders: Optional[dict] = None,
     ):
         """Create a ggplot2 plot from a DataFrame.
 
+        Returns a GGPlot object that can be displayed in a notebook or saved to a file.
+
         :param ggplot_code: The ggplot2 code to execute.
-        :param filename: The filename to save the plot to.
         :param shape: The shape of the data in the DataFrame (wide or long).
         :param sql: The SQL query to execute beforehand to manipulate the data.
         :param remove_prefix: Whether to remove the prefix from the column names.
         :param debug: Whether to print the R code instead of executing it.
         :param height: The height of the plot in inches.
         :param width: The width of the plot in inches.
-        :param format: The format to save the plot in (png or svg).
         :param factor_orders: A dictionary of factor columns and their order.
         """
-
         if sql == None:
             sql = "select * from self"
 
         if shape == "long":
             df = self.sql(sql, shape="long")
         elif shape == "wide":
-            df = self.sql(sql, shape="wide", remove_prefix=remove_prefix)
+            df = self.sql(sql, remove_prefix=remove_prefix)
 
         # Convert DataFrame to CSV format
-        csv_data = df.to_csv(index=False)
+        csv_data = df.to_csv().text
 
         # Embed the CSV data within the R script
         csv_data_escaped = csv_data.replace("\n", "\\n").replace("'", "\\'")
@@ -52,70 +119,60 @@ class ResultsGGMixin:
 
         if factor_orders is not None:
             for factor, order in factor_orders.items():
-                # read_csv_code += f"""self${{{factor}}} <- factor(self${{{factor}}}, levels=c({','.join(['"{}"'.format(x) for x in order])}))"""
-
                 level_string = ", ".join([f'"{x}"' for x in order])
                 read_csv_code += (
                     f"self${factor} <- factor(self${factor}, levels=c({level_string}))"
                 )
                 read_csv_code += "\n"
 
-        # Load ggplot2 library
-        load_ggplot2 = "library(ggplot2)\n"
-
-        # Check if a filename is provided for the plot, if not create a temporary one
-        if not filename:
-            filename = tempfile.mktemp(suffix=f".{format}")
-
-        # Combine all R script parts
-        full_r_code = load_ggplot2 + read_csv_code + ggplot_code
-
-        # Add command to save the plot to a file
-        full_r_code += f'\nggsave("{filename}", plot = last_plot(), width = {width}, height = {height}, device = "{format}")'
+        # Load ggplot2 library and combine all R script parts
+        full_r_code = "library(ggplot2)\n" + read_csv_code + ggplot_code
 
         if debug:
             print(full_r_code)
             return
 
-        result = subprocess.run(
-            ["Rscript", "-"],
-            input=full_r_code,
-            text=True,
-            stdout=subprocess.PIPE,
-            stderr=subprocess.PIPE,
-        )
-
-        if result.returncode != 0:
-            if result.returncode == 127:  # 'command not found'
-                raise RuntimeError(
-                    "Rscript is probably not installed. Please install R from https://cran.r-project.org/"
-                )
-            else:
-                raise RuntimeError(
-                    f"An error occurred while running Rscript: {result.stderr}"
-                )
-
-        if result.stderr:
-            print("Error in R script:", result.stderr)
-        else:
-            self._display_plot(filename, width, height)
+        return GGPlot(full_r_code, width=width, height=height)
 
     def _display_plot(self, filename: str, width: float, height: float):
-        """Display the plot in the notebook."""
-        import matplotlib.pyplot as plt
-        import matplotlib.image as mpimg
-
-        if filename.endswith(".png"):
-            img = mpimg.imread(filename)
-            plt.figure(
-                figsize=(width, height)
-            )  # Set the figure size (width, height) in inches
-            plt.imshow(img)
-            plt.axis("off")
-            plt.show()
-        elif filename.endswith(".svg"):
-            from IPython.display import SVG, display
-
-            display(SVG(filename=filename))
+        """Display the plot in the notebook or open in system viewer if running from terminal."""
+        try:
+            # Try to import IPython-related modules
+            import matplotlib.pyplot as plt
+            import matplotlib.image as mpimg
+            from IPython import get_ipython
+
+            # Check if we're in a notebook environment
+            if get_ipython() is not None:
+                if filename.endswith(".png"):
+                    img = mpimg.imread(filename)
+                    plt.figure(figsize=(width, height))
+                    plt.imshow(img)
+                    plt.axis("off")
+                    plt.show()
+                elif filename.endswith(".svg"):
+                    from IPython.display import SVG, display
+                    display(SVG(filename=filename))
+                else:
+                    print("Unsupported file format. Please provide a PNG or SVG file.")
+                return
+
+        except ImportError:
+            pass
+
+        # If we're not in a notebook or imports failed, open with system viewer
+        import platform
+        import os
+        
+        system = platform.system()
+        if system == 'Darwin':       # macOS
+            if filename.endswith('.svg'):
+                subprocess.run(['open', '-a', 'Preview', filename])
+            else:
+                subprocess.run(['open', filename])
+        elif system == 'Linux':
+            subprocess.run(['xdg-open', filename])
+        elif system == 'Windows':
+            os.startfile(filename)
         else:
-            print("Unsupported file format. Please provide a PNG or SVG file.")
+            print(f"File saved to: {filename}")

edsl/scenarios/PdfExtractor.py

@@ -2,14 +2,11 @@ import os
 
 
 class PdfExtractor:
-    def __init__(self, pdf_path: str, parent_object: object):
+    def __init__(self, pdf_path: str):
         self.pdf_path = pdf_path
-        self.constructor = parent_object.__class__
+        #self.constructor = parent_object.__class__
 
-    def get_object(self) -> object:
-        return self.constructor(self._get_pdf_dict())
-
-    def _get_pdf_dict(self) -> dict:
+    def get_pdf_dict(self) -> dict:
         # Ensure the file exists
         import fitz