edsl 0.1.46__py3-none-any.whl → 0.1.48__py3-none-any.whl

edsl/{agents → invigilators}/question_option_processor.py

@@ -2,6 +2,52 @@ from jinja2 import Environment, meta
 from typing import List, Optional, Union
 
 
+def extract_template_variables(ast) -> List[Union[str, tuple]]:
+    """
+    Extract variable expressions from a Jinja2 AST.
+    
+    Args:
+        ast: Jinja2 AST
+        
+    Returns:
+        List[Union[str, tuple]]: List of variable names or tuples for dotted paths
+    """
+    from jinja2 import nodes
+    from jinja2.visitor import NodeVisitor
+    
+    variables = []
+    
+    class VariableVisitor(NodeVisitor):
+        def visit_Name(self, node):
+            variables.append(node.name)
+        
+        def visit_Getattr(self, node):
+            # For dotted access like scenario.question_options
+            parts = []
+            current = node
+            
+            # Handle the leaf attribute
+            parts.append(node.attr)
+            
+            # Walk up the chain to collect all parts
+            while isinstance(current.node, nodes.Getattr):
+                current = current.node
+                parts.append(current.attr)
+            
+            # Add the root name
+            if isinstance(current.node, nodes.Name):
+                parts.append(current.node.name)
+            
+            # Reverse to get the correct order
+            parts.reverse()
+            variables.append(tuple(parts))
+    
+    visitor = VariableVisitor()
+    visitor.visit(ast)
+    
+    return variables
+
+
 class QuestionOptionProcessor:
     """
     Class that manages the processing of question options.
@@ -16,7 +62,11 @@ class QuestionOptionProcessor:
         return cls(scenario, prior_answers_dict)
 
     def __init__(self, scenario: 'Scenario', prior_answers_dict: dict):
-        self.scenario = scenario 
+        # This handles cases where the question has {{ scenario.key }} - eventually 
+        # we might not allow 'naked' scenario keys w/o the scenario prefix
+        #new_scenario = scenario.copy()
+        #new_scenario.update({'scenario': new_scenario})
+        self.scenario = scenario
         self.prior_answers_dict = prior_answers_dict
 
     @staticmethod
@@ -25,18 +75,23 @@ class QuestionOptionProcessor:
         return [f"<< Option {i} - Placeholder >>" for i in range(1, 4)]
 
     @staticmethod
-    def _parse_template_variable(template_str: str) -> str:
+    def _parse_template_variable(template_str: str) -> Union[str, tuple]:
         """
         Extract the variable name from a template string.
+        If the variable contains dots (e.g., scenario.question_options), 
+        returns a tuple of the path components.
 
         Args:
             template_str (str): Jinja template string
 
         Returns:
-            str: Name of the first undefined variable in the template
+            Union[str, tuple]: Name of the first undefined variable in the template,
+                              or a tuple of path components if the variable contains dots
 
         >>> QuestionOptionProcessor._parse_template_variable("Here are some {{ options }}")
         'options'
+        >>> QuestionOptionProcessor._parse_template_variable("Here are some {{ scenario.question_options }}")
+        ('scenario', 'question_options')
         >>> QuestionOptionProcessor._parse_template_variable("Here are some {{ options }} and {{ other }}")
         Traceback (most recent call last):
         ...
@@ -48,11 +103,13 @@ class QuestionOptionProcessor:
         """
         env = Environment()
         parsed_content = env.parse(template_str)
-        undeclared_variables = list(meta.find_undeclared_variables(parsed_content))
+        undeclared_variables = extract_template_variables(parsed_content)
+        
         if not undeclared_variables:
             raise ValueError("No variables found in template string")
         if len(undeclared_variables) > 1:
             raise ValueError("Multiple variables found in template string")
+        
         return undeclared_variables[0]
 
     @staticmethod
@@ -128,7 +185,7 @@ class QuestionOptionProcessor:
 
         The case where options are provided as a template string:
 
-        >>> question_data = {"question_options": "{{ options }}"}
+        >>> question_data = {"question_options": "{{ scenario.options }}"}
         >>> processor.get_question_options(question_data)
         ['Option 1', 'Option 2']
 
@@ -140,7 +197,7 @@ class QuestionOptionProcessor:
         >>> q0.answer = ["Option 1", "Option 2"]
         >>> mpc.prior_answers_dict = lambda: {'q0': q0}
         >>> processor = QuestionOptionProcessor.from_prompt_constructor(mpc)
-        >>> question_data = {"question_options": "{{ q0 }}"}
+        >>> question_data = {"question_options": "{{ q0.answer }}"}
         >>> processor.get_question_options(question_data)
         ['Option 1', 'Option 2']
 
@@ -156,21 +213,39 @@ class QuestionOptionProcessor:
             return options_entry if options_entry else self._get_default_options()
 
         # Parse template to get variable name
-        option_key = self._parse_template_variable(options_entry)
-
-        # Try getting options from scenario
-        scenario_options = self._get_options_from_scenario(
-            self.scenario, option_key
-        )
-        if scenario_options:
-            return scenario_options
-
-        # Try getting options from prior answers
-        prior_answer_options = self._get_options_from_prior_answers(
-            self.prior_answers_dict, option_key
-        )
-        if prior_answer_options:
-            return prior_answer_options
+        raw_option_key = self._parse_template_variable(options_entry)
+
+        source_type = None
+
+        if isinstance(raw_option_key, tuple):
+            if raw_option_key[0] == 'scenario':
+                source_type = 'scenario'
+                option_key = raw_option_key[-1]
+            else:
+                source_type = 'prior_answers'
+                option_key = raw_option_key[0]
+                #breakpoint()
+        else:
+            option_key = raw_option_key
+
+        #breakpoint()
+
+        if source_type == 'scenario':
+            # Try getting options from scenario
+            scenario_options = self._get_options_from_scenario(
+                self.scenario, option_key
+            )
+            if scenario_options:
+                return scenario_options
+            
+        if source_type == 'prior_answers':
+
+            # Try getting options from prior answers
+            prior_answer_options = self._get_options_from_prior_answers(
+                self.prior_answers_dict, option_key
+            )
+            if prior_answer_options:
+                return prior_answer_options
 
         return self._get_default_options()
 

edsl/{agents/QuestionTemplateReplacementsBuilder.py → invigilators/question_template_replacements_builder.py}

@@ -1,11 +1,12 @@
 from jinja2 import Environment, meta, TemplateSyntaxError
 from typing import Any, Set, TYPE_CHECKING
 
+from ..agents import Agent
+from ..scenarios import Scenario
+
 if TYPE_CHECKING:
-    from edsl.agents.PromptConstructor import PromptConstructor
-    from edsl.scenarios.Scenario import Scenario
-    from edsl.questions.QuestionBase import QuestionBase
-    from edsl.agents.Agent import Agent
+    from .prompt_constructor import PromptConstructor
+    from ..questions import QuestionBase
 
 
 class QuestionTemplateReplacementsBuilder:
@@ -33,17 +34,29 @@ class QuestionTemplateReplacementsBuilder:
 
     def question_file_keys(self):
         """
-        >>> from edsl import QuestionMultipleChoice, Scenario
+        >>> from ..questions import QuestionMultipleChoice
+        >>> from ..scenarios import Scenario
         >>> q = QuestionMultipleChoice(question_text="Do you like school?", question_name = "q0", question_options = ["yes", "no"])
         >>> qtrb = QuestionTemplateReplacementsBuilder(scenario = {"file1": "file1"}, question = q, prior_answers_dict = {'q0': 'q0'}, agent = "agent")
         >>> qtrb.question_file_keys()
         []
-        >>> from edsl import FileStore
+        >>> from ..scenarios import FileStore
         >>> fs = FileStore.example()
+        >>> # Test direct key reference
         >>> q = QuestionMultipleChoice(question_text="What do you think of this file: {{ file1 }}", question_name = "q0", question_options = ["good", "bad"])
         >>> qtrb = QuestionTemplateReplacementsBuilder(scenario = Scenario({"file1": fs}), question = q, prior_answers_dict = {'q0': 'q0'}, agent = "agent")
         >>> qtrb.question_file_keys()
         ['file1']
+        >>> # Test scenario.key reference
+        >>> q = QuestionMultipleChoice(question_text="What do you think of this file: {{ scenario.file2 }}", question_name = "q0", question_options = ["good", "bad"])
+        >>> qtrb = QuestionTemplateReplacementsBuilder(scenario = Scenario({"file2": fs}), question = q, prior_answers_dict = {'q0': 'q0'}, agent = "agent")
+        >>> qtrb.question_file_keys()
+        ['file2']
+        >>> # Test both formats in the same question
+        >>> q = QuestionMultipleChoice(question_text="Compare {{ file1 }} with {{ scenario.file2 }}", question_name = "q0", question_options = ["good", "bad"])
+        >>> qtrb = QuestionTemplateReplacementsBuilder(scenario = Scenario({"file1": fs, "file2": fs}), question = q, prior_answers_dict = {'q0': 'q0'}, agent = "agent")
+        >>> sorted(qtrb.question_file_keys())
+        ['file1', 'file2']
         """
         question_text = self.question.question_text
         file_keys = self._find_file_keys(self.scenario)
@@ -76,8 +89,8 @@ class QuestionTemplateReplacementsBuilder:
         """We need to find all the keys in the scenario that refer to FileStore objects.
         These will be used to append to the prompt a list of files that are part of the scenario.
 
-        >>> from edsl import Scenario
-        >>> from edsl.scenarios.FileStore import FileStore
+        >>> from ..scenarios import Scenario
+        >>> from ..scenarios import FileStore
         >>> import tempfile
         >>> with tempfile.NamedTemporaryFile() as f:
         ...     _ = f.write(b"Hello, world!")
@@ -87,7 +100,7 @@ class QuestionTemplateReplacementsBuilder:
         ...     QuestionTemplateReplacementsBuilder._find_file_keys(scenario)
         ['fs_file']
         """
-        from edsl.scenarios.FileStore import FileStore
+        from ..scenarios import FileStore
 
         file_entries = []
         for key, value in scenario.items():
@@ -100,10 +113,11 @@ class QuestionTemplateReplacementsBuilder:
         question_text: str, scenario_file_keys: list
     ) -> list:
         """
-        Extracts the file keys from a question text.
+        Extracts the file keys from a question text, handling both direct references ({{ file_key }})
+        and scenario-prefixed references ({{ scenario.file_key }}).
 
         >>> from edsl import Scenario
-        >>> from edsl.scenarios.FileStore import FileStore
+        >>> from edsl.scenarios import FileStore
         >>> import tempfile
         >>> with tempfile.NamedTemporaryFile() as f:
         ...     _ = f.write(b"Hello, world!")
@@ -112,15 +126,53 @@ class QuestionTemplateReplacementsBuilder:
         ...     scenario = Scenario({"fs_file": fs, 'a': 1})
         ...     QuestionTemplateReplacementsBuilder._extract_file_keys_from_question_text("{{ fs_file }}", ['fs_file'])
         ['fs_file']
+        >>> with tempfile.NamedTemporaryFile() as f:
+        ...     _ = f.write(b"Hello, world!")
+        ...     _ = f.seek(0)
+        ...     fs = FileStore(f.name)
+        ...     scenario = Scenario({"print": fs, 'a': 1})
+        ...     QuestionTemplateReplacementsBuilder._extract_file_keys_from_question_text("{{ scenario.print }}", ['print'])
+        ['print']
+        >>> with tempfile.NamedTemporaryFile() as f:
+        ...     _ = f.write(b"Hello, world!")
+        ...     _ = f.seek(0)
+        ...     fs = FileStore(f.name)
+        ...     scenario = Scenario({"file1": fs, "file2": fs})
+        ...     sorted(QuestionTemplateReplacementsBuilder._extract_file_keys_from_question_text("Compare {{ file1 }} with {{ scenario.file2 }}", ['file1', 'file2']))
+        ['file1', 'file2']
         """
         variables = QuestionTemplateReplacementsBuilder.get_jinja2_variables(
             question_text
         )
         question_file_keys = []
+        
+        # Direct references: {{ file_key }}
         for var in variables:
             if var in scenario_file_keys:
                 question_file_keys.append(var)
-        return question_file_keys
+                
+        # Scenario-prefixed references: {{ scenario.file_key }}
+        for var in variables:
+            if var == "scenario":
+                # If we find a scenario variable, we need to check for nested references
+                # Create a modified template with just {{ scenario.* }} expressions to isolate them
+                scenario_template = "".join([
+                    "{% for key, value in scenario.items() %}{{ key }}{% endfor %}"
+                ])
+                try:
+                    # This is a check to make sure there's scenario.something syntax in the template
+                    if "scenario." in question_text:
+                        # Extract dot-notation scenario references by parsing the template
+                        import re
+                        scenario_refs = re.findall(r'{{\s*scenario\.(\w+)\s*}}', question_text)
+                        for key in scenario_refs:
+                            if key in scenario_file_keys:
+                                question_file_keys.append(key)
+                except:
+                    # If there's any issue parsing, just continue with what we have
+                    pass
+                
+        return list(set(question_file_keys))  # Remove duplicates
 
     def _scenario_replacements(
         self, replacement_string: str = "<see file {key}>"

edsl/jobs/__init__.py

@@ -1 +1,7 @@
-from edsl.jobs.Jobs import Jobs
+from .jobs import Jobs
+from .jobs import RunConfig, RunParameters, RunEnvironment
+from .remote_inference import JobsRemoteInferenceHandler
+from .jobs_runner_status import JobsRunnerStatusBase
+
+
+__all__ = ["Jobs"]

edsl/jobs/async_interview_runner.py

@@ -1,43 +1,94 @@
+"""
+Asynchronous interview runner module for conducting interviews concurrently.
+
+This module provides functionality to run multiple interviews in parallel
+with controlled concurrency, supporting both error handling and result collection.
+"""
+
 from collections.abc import AsyncGenerator
-from typing import List, TypeVar, Generator, Tuple, TYPE_CHECKING
+from typing import List, Generator, Tuple, TYPE_CHECKING
 from dataclasses import dataclass
 import asyncio
-from contextlib import asynccontextmanager
 from edsl.data_transfer_models import EDSLResultObjectInput
 
-from edsl.results.Result import Result
-from edsl.jobs.interviews.Interview import Interview
-from edsl.config import Config
+from ..results import Result
+from ..interviews import Interview
+from ..config import Config
 config = Config()
 
 if TYPE_CHECKING:
-    from edsl.jobs.Jobs import Jobs
+    from ..jobs import Jobs
 
+from .data_structures import RunConfig
 
 @dataclass
 class InterviewResult:
+    """Container for the result of an interview along with metadata.
+    
+    Attributes:
+        result: The Result object containing the interview answers
+        interview: The Interview object used to conduct the interview
+        order: The original position of this interview in the processing queue
+    """
     result: Result
     interview: Interview
     order: int
 
 
-from edsl.jobs.data_structures import RunConfig
-
-
 class AsyncInterviewRunner:
+    """
+    Runs interviews asynchronously with controlled concurrency.
+    
+    This class manages the parallel execution of multiple interviews while
+    respecting concurrency limits and handling errors appropriately.
+    
+    Examples:
+        >>> from unittest.mock import MagicMock, AsyncMock
+        >>> mock_jobs = MagicMock()
+        >>> mock_run_config = MagicMock()
+        >>> mock_run_config.parameters.n = 1
+        >>> mock_run_config.environment.cache = None
+        >>> runner = AsyncInterviewRunner(mock_jobs, mock_run_config)
+        >>> isinstance(runner._initialized, asyncio.Event)
+        True
+    """
+    
     MAX_CONCURRENT = int(config.EDSL_MAX_CONCURRENT_TASKS)
 
     def __init__(self, jobs: "Jobs", run_config: RunConfig):
+        """
+        Initialize the AsyncInterviewRunner.
+        
+        Args:
+            jobs: The Jobs object that generates interviews
+            run_config: Configuration for running the interviews
+        """
         self.jobs = jobs
         self.run_config = run_config
         self._initialized = asyncio.Event()
 
     def _expand_interviews(self) -> Generator["Interview", None, None]:
-        """Populates self.total_interviews with n copies of each interview.
-
-        It also has to set the cache for each interview.
-
-        :param n: how many times to run each interview.
+        """
+        Create multiple copies of each interview based on the run configuration.
+        
+        This method expands interviews for repeated runs and ensures each has
+        the proper cache configuration.
+        
+        Yields:
+            Interview objects ready to be conducted
+            
+        Examples:
+            >>> from unittest.mock import MagicMock
+            >>> mock_jobs = MagicMock()
+            >>> mock_interview = MagicMock()
+            >>> mock_jobs.generate_interviews.return_value = [mock_interview]
+            >>> mock_run_config = MagicMock()
+            >>> mock_run_config.parameters.n = 2
+            >>> mock_run_config.environment.cache = "mock_cache"
+            >>> runner = AsyncInterviewRunner(mock_jobs, mock_run_config)
+            >>> interviews = list(runner._expand_interviews())
+            >>> len(interviews)
+            2
         """
         for interview in self.jobs.generate_interviews():
             for iteration in range(self.run_config.parameters.n):
@@ -52,21 +103,22 @@ class AsyncInterviewRunner:
     async def _conduct_interview(
         self, interview: "Interview"
     ) -> Tuple["Result", "Interview"]:
-        """Conducts an interview and returns the result object, along with the associated interview.
-
-        We return the interview because it is not populated with exceptions, if any.
-
-        :param interview: the interview to conduct
-        :return: the result of the interview
-
-        'extracted_answers' is a dictionary of the answers to the questions in the interview.
-        This is not the same as the generated_tokens---it can include substantial cleaning and processing / validation.
         """
-        # the model buckets are used to track usage rates
-        # model_buckets = self.bucket_collection[interview.model]
-        # model_buckets = self.run_config.environment.bucket_collection[interview.model]
-
-        # get the results of the interview e.g., {'how_are_you':"Good" 'how_are_you_generated_tokens': "Good"}
+        Asynchronously conduct a single interview.
+        
+        This method performs the interview and creates a Result object with
+        the extracted answers and model responses.
+        
+        Args:
+            interview: The interview to conduct
+            
+        Returns:
+            Tuple containing the Result object and the Interview object
+            
+        Notes:
+            'extracted_answers' contains the processed and validated answers
+            from the interview, which may differ from the raw model output.
+        """
         extracted_answers: dict[str, str]
         model_response_objects: List[EDSLResultObjectInput]
 
@@ -83,10 +135,20 @@ class AsyncInterviewRunner:
     async def run(
         self,
     ) -> AsyncGenerator[tuple[Result, Interview], None]:
-        """Creates and processes tasks asynchronously, yielding results as they complete.
-
-        Uses TaskGroup for structured concurrency and automated cleanup.
-        Results are yielded as they become available while maintaining controlled concurrency.
+        """
+        Run all interviews asynchronously and yield results as they complete.
+        
+        This method processes interviews in chunks based on MAX_CONCURRENT,
+        maintaining controlled concurrency while yielding results as soon as
+        they become available.
+        
+        Yields:
+            Tuples of (Result, Interview) as interviews complete
+            
+        Notes:
+            - Uses structured concurrency patterns for proper resource management
+            - Handles exceptions according to the run configuration
+            - Ensures task cleanup even in case of failures
         """
         interviews = list(self._expand_interviews())
         self._initialized.set()
@@ -102,10 +164,8 @@ class AsyncInterviewRunner:
                 result.order = idx
                 return InterviewResult(result, interview, idx)
             except Exception as e:
-                # breakpoint()
                 if self.run_config.parameters.stop_on_exception:
                     raise
-                # logger.error(f"Task failed with error: {e}")
                 return None
 
         # Process interviews in chunks
@@ -130,7 +190,6 @@ class AsyncInterviewRunner:
             except Exception as e:
                 if self.run_config.parameters.stop_on_exception:
                     raise
-                # logger.error(f"Chunk processing failed with error: {e}")
                 continue
 
             finally:
@@ -138,3 +197,8 @@ class AsyncInterviewRunner:
                 for task in tasks:
                     if not task.done():
                         task.cancel()
+
+
+if __name__ == "__main__":
+    import doctest 
+    doctest.testmod()

edsl/jobs/check_survey_scenario_compatibility.py

@@ -1,5 +1,7 @@
 import warnings
 from typing import TYPE_CHECKING
+from edsl.scenarios import ScenarioList
+from edsl.surveys import Survey
 
 if TYPE_CHECKING:
     from edsl.surveys.Survey import Survey
@@ -15,10 +17,10 @@ class CheckSurveyScenarioCompatibility:
     def check(self, strict: bool = False, warn: bool = False) -> None:
         """Check if the parameters in the survey and scenarios are consistent.
 
-        >>> from edsl.jobs.Jobs import Jobs
-        >>> from edsl.questions.QuestionFreeText import QuestionFreeText
-        >>> from edsl.surveys.Survey import Survey
-        >>> from edsl.scenarios.Scenario import Scenario
+        >>> from edsl.jobs import Jobs
+        >>> from edsl.questions import QuestionFreeText
+        >>> from edsl.surveys import Survey
+        >>> from edsl.scenarios import Scenario
         >>> q = QuestionFreeText(question_text = "{{poo}}", question_name = "ugly_question")
         >>> j = Jobs(survey = Survey(questions=[q]))
         >>> cs = CheckSurveyScenarioCompatibility(j.survey, j.scenarios)
@@ -39,7 +41,7 @@ class CheckSurveyScenarioCompatibility:
 
         >>> q = QuestionFreeText(question_text = "Hello", question_name = "ugly_question")
         >>> s = Scenario({'ugly_question': "B"})
-        >>> from edsl.scenarios.ScenarioList import ScenarioList
+        >>> from edsl.scenarios import ScenarioList
         >>> cs = CheckSurveyScenarioCompatibility(Survey(questions=[q]), ScenarioList([s]))
         >>> cs.check()
         Traceback (most recent call last):