edsl 0.1.38.dev4__py3-none-any.whl → 0.1.39__py3-none-any.whl

edsl/inference_services/TestService.py

@@ -2,7 +2,7 @@ from typing import Any, List, Optional
 import os
 import asyncio
 from edsl.inference_services.InferenceServiceABC import InferenceServiceABC
-from edsl.language_models import LanguageModel
+from edsl.language_models.LanguageModel import LanguageModel
 from edsl.inference_services.rate_limits_cache import rate_limits
 from edsl.utilities.utilities import fix_partial_correct_response
 
@@ -51,6 +51,7 @@ class TestService(InferenceServiceABC):
             @property
             def _canned_response(self):
                 if hasattr(self, "canned_response"):
+
                     return self.canned_response
                 else:
                     return "Hello, world"
@@ -63,15 +64,6 @@ class TestService(InferenceServiceABC):
                 files_list: Optional[List["File"]] = None,
             ) -> dict[str, Any]:
                 await asyncio.sleep(0.1)
-                # return {"message": """{"answer": "Hello, world"}"""}
-
-                if hasattr(self, "func"):
-                    return {
-                        "message": [
-                            {"text": self.func(user_prompt, system_prompt, files_list)}
-                        ],
-                        "usage": {"prompt_tokens": 1, "completion_tokens": 1},
-                    }
 
                 if hasattr(self, "throw_exception") and self.throw_exception:
                     if hasattr(self, "exception_probability"):
@@ -81,6 +73,15 @@ class TestService(InferenceServiceABC):
 
                     if random.random() < p:
                         raise Exception("This is a test error")
+
+                if hasattr(self, "func"):
+                    return {
+                        "message": [
+                            {"text": self.func(user_prompt, system_prompt, files_list)}
+                        ],
+                        "usage": {"prompt_tokens": 1, "completion_tokens": 1},
+                    }
+
                 return {
                     "message": [{"text": f"{self._canned_response}"}],
                     "usage": {"prompt_tokens": 1, "completion_tokens": 1},

edsl/inference_services/TogetherAIService.py

@@ -143,15 +143,17 @@ class TogetherAIService(OpenAIService):
     _async_client_ = openai.AsyncOpenAI
 
     @classmethod
-    def get_model_list(cls):
+    def get_model_list(cls, api_token=None):
         # Togheter.ai has a different response in model list then openai
         # and the OpenAI class returns an error when calling .models.list()
         import requests
         import os
 
         url = "https://api.together.xyz/v1/models?filter=serverless"
-        token = os.getenv(cls._env_key_name_)
-        headers = {"accept": "application/json", "authorization": f"Bearer {token}"}
+        if api_token is None:
+            api_token = os.getenv(cls._env_key_name_)
+
+        headers = {"accept": "application/json", "authorization": f"Bearer {api_token}"}
 
         response = requests.get(url, headers=headers)
         return response.json()

edsl/inference_services/data_structures.py

@@ -0,0 +1,134 @@
+from collections import UserDict, defaultdict, UserList
+from typing import Union, Optional, List
+from edsl.enums import InferenceServiceLiteral
+from dataclasses import dataclass
+
+
+@dataclass
+class LanguageModelInfo:
+    """A dataclass for storing information about a language model.
+
+
+    >>> LanguageModelInfo("gpt-4-1106-preview", "openai")
+    LanguageModelInfo(model_name='gpt-4-1106-preview', service_name='openai')
+
+    >>> model_name, service = LanguageModelInfo.example()
+    >>> model_name
+    'gpt-4-1106-preview'
+
+    >>> LanguageModelInfo.example().service_name
+    'openai'
+
+    """
+
+    model_name: str
+    service_name: str
+
+    def __iter__(self):
+        yield self.model_name
+        yield self.service_name
+
+    def __getitem__(self, key: int) -> str:
+        import warnings
+
+        warnings.warn(
+            "Accessing LanguageModelInfo via index is deprecated. "
+            "Please use .model_name, .service_name, or .index attributes instead.",
+            DeprecationWarning,
+            stacklevel=2,
+        )
+
+        if key == 0:
+            return self.model_name
+        elif key == 1:
+            return self.service_name
+        else:
+            raise IndexError("Index out of range")
+
+    @classmethod
+    def example(cls) -> "LanguageModelInfo":
+        return cls("gpt-4-1106-preview", "openai")
+
+
+class ModelNamesList(UserList):
+    pass
+
+
+class AvailableModels(UserList):
+
+    def __init__(self, data: List[LanguageModelInfo]) -> None:
+        super().__init__(data)
+
+    def __contains__(self, model_name: str) -> bool:
+        for model_entry in self:
+            if model_entry.model_name == model_name:
+                return True
+        return False
+
+    def print(self):
+        return self.to_dataset().print()
+
+    def to_dataset(self):
+        from edsl.scenarios.ScenarioList import ScenarioList
+
+        models, services = zip(
+            *[(model.model_name, model.service_name) for model in self]
+        )
+        return (
+            ScenarioList.from_list("model", models)
+            .add_list("service", services)
+            .to_dataset()
+        )
+
+    def to_model_list(self):
+        from edsl.language_models.ModelList import ModelList
+
+        return ModelList.from_available_models(self)
+
+    def search(
+        self, pattern: str, service_name: Optional[str] = None, regex: bool = False
+    ) -> "AvailableModels":
+        import re
+
+        if not regex:
+            # Escape special regex characters except *
+            pattern = re.escape(pattern).replace(r"\*", ".*")
+
+        try:
+            regex = re.compile(pattern)
+            avm = AvailableModels(
+                [
+                    entry
+                    for entry in self
+                    if regex.search(entry.model_name)
+                    and (service_name is None or entry.service_name == service_name)
+                ]
+            )
+            if len(avm) == 0:
+                raise ValueError(
+                    "No models found matching the search pattern: " + pattern
+                )
+            else:
+                return avm
+        except re.error as e:
+            raise ValueError(f"Invalid regular expression pattern: {e}")
+
+
+class ServiceToModelsMapping(UserDict):
+    def __init__(self, data: dict) -> None:
+        super().__init__(data)
+
+    @property
+    def service_names(self) -> list[str]:
+        return list(self.data.keys())
+
+    def _validate_service_names(self):
+        for service in self.service_names:
+            if service not in InferenceServiceLiteral:
+                raise ValueError(f"Invalid service name: {service}")
+
+    def model_to_services(self) -> dict:
+        self._model_to_service = defaultdict(list)
+        for service, models in self.data.items():
+            for model in models:
+                self._model_to_service[model].append(service)

edsl/jobs/AnswerQuestionFunctionConstructor.py

@@ -0,0 +1,223 @@
+import copy
+import asyncio
+
+from typing import Union, Type, Callable, TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from edsl.questions.QuestionBase import QuestionBase
+    from edsl.jobs.interviews.Interview import Interview
+    from edsl.language_models.key_management.KeyLookup import KeyLookup
+
+from edsl.surveys.base import EndOfSurvey
+from edsl.jobs.tasks.task_status_enum import TaskStatus
+
+from edsl.jobs.FetchInvigilator import FetchInvigilator
+from edsl.exceptions.language_models import LanguageModelNoResponseError
+from edsl.exceptions.questions import QuestionAnswerValidationError
+from edsl.data_transfer_models import AgentResponseDict, EDSLResultObjectInput
+
+from edsl.jobs.Answers import Answers
+
+
+class RetryConfig:
+    from edsl.config import CONFIG
+
+    EDSL_BACKOFF_START_SEC = float(CONFIG.get("EDSL_BACKOFF_START_SEC"))
+    EDSL_BACKOFF_MAX_SEC = float(CONFIG.get("EDSL_BACKOFF_MAX_SEC"))
+    EDSL_MAX_ATTEMPTS = int(CONFIG.get("EDSL_MAX_ATTEMPTS"))
+
+
+class SkipHandler:
+
+    def __init__(self, interview: "Interview"):
+        self.interview = interview
+        self.question_index = self.interview.to_index
+
+        self.skip_function: Callable = (
+            self.interview.survey.rule_collection.skip_question_before_running
+        )
+
+    def should_skip(self, current_question: "QuestionBase") -> bool:
+        """Determine if the current question should be skipped."""
+        current_question_index = self.question_index[current_question.question_name]
+        combined_answers = (
+            self.interview.answers
+            | self.interview.scenario
+            | self.interview.agent["traits"]
+        )
+        return self.skip_function(current_question_index, combined_answers)
+
+    def cancel_skipped_questions(self, current_question: "QuestionBase") -> None:
+        """Cancel the tasks for questions that should be skipped."""
+        current_question_index: int = self.interview.to_index[
+            current_question.question_name
+        ]
+        answers = (
+            self.interview.answers
+            | self.interview.scenario
+            | self.interview.agent["traits"]
+        )
+
+        # Get the index of the next question, which could also be the end of the survey
+        next_question: Union[int, EndOfSurvey] = (
+            self.interview.survey.rule_collection.next_question(
+                q_now=current_question_index,
+                answers=answers,
+            )
+        )
+
+        def cancel_between(start, end):
+            """Cancel the tasks for questions between the start and end indices."""
+            for i in range(start, end):
+                self.interview.tasks[i].cancel()
+
+        if (next_question_index := next_question.next_q) == EndOfSurvey:
+            cancel_between(
+                current_question_index + 1, len(self.interview.survey.questions)
+            )
+            return
+
+        if next_question_index > (current_question_index + 1):
+            cancel_between(current_question_index + 1, next_question_index)
+
+
+class AnswerQuestionFunctionConstructor:
+    """Constructs a function that answers a question and records the answer."""
+
+    def __init__(self, interview: "Interview", key_lookup: "KeyLookup"):
+        self.interview = interview
+        self.key_lookup = key_lookup
+
+        self.had_language_model_no_response_error: bool = False
+        self.question_index = self.interview.to_index
+
+        self.skip_function: Callable = (
+            self.interview.survey.rule_collection.skip_question_before_running
+        )
+
+        self.invigilator_fetcher = FetchInvigilator(
+            self.interview, key_lookup=self.key_lookup
+        )
+        self.skip_handler = SkipHandler(self.interview)
+
+    def _handle_exception(
+        self, e: Exception, invigilator: "InvigilatorBase", task=None
+    ):
+        """Handle an exception that occurred while answering a question."""
+
+        from edsl.jobs.interviews.InterviewExceptionEntry import InterviewExceptionEntry
+
+        answers = copy.copy(
+            self.interview.answers
+        )  # copy to freeze the answers here for logging
+        exception_entry = InterviewExceptionEntry(
+            exception=e,
+            invigilator=invigilator,
+            answers=answers,
+        )
+        if task:
+            task.task_status = TaskStatus.FAILED
+
+        self.interview.exceptions.add(
+            invigilator.question.question_name, exception_entry
+        )
+
+        if self.interview.raise_validation_errors and isinstance(
+            e, QuestionAnswerValidationError
+        ):
+            raise e
+
+        stop_on_exception = getattr(self.interview, "stop_on_exception", False)
+        if stop_on_exception:
+            raise e
+
+    def __call__(self):
+        return self.answer_question_and_record_task
+
+    async def answer_question_and_record_task(
+        self,
+        *,
+        question: "QuestionBase",
+        task=None,
+    ) -> "AgentResponseDict":
+
+        from tenacity import (
+            retry,
+            stop_after_attempt,
+            wait_exponential,
+            retry_if_exception_type,
+            RetryError,
+        )
+
+        @retry(
+            stop=stop_after_attempt(RetryConfig.EDSL_MAX_ATTEMPTS),
+            wait=wait_exponential(
+                multiplier=RetryConfig.EDSL_BACKOFF_START_SEC,
+                max=RetryConfig.EDSL_BACKOFF_MAX_SEC,
+            ),
+            retry=retry_if_exception_type(LanguageModelNoResponseError),
+            reraise=True,
+        )
+        async def attempt_answer():
+            invigilator = self.invigilator_fetcher(question)
+
+            if self.skip_handler.should_skip(question):
+                return invigilator.get_failed_task_result(
+                    failure_reason="Question skipped."
+                )
+
+            try:
+                response: EDSLResultObjectInput = (
+                    await invigilator.async_answer_question()
+                )
+                if response.validated:
+                    self.interview.answers.add_answer(
+                        response=response, question=question
+                    )
+
+                    self.skip_handler.cancel_skipped_questions(question)
+                else:
+                    if (
+                        hasattr(response, "exception_occurred")
+                        and response.exception_occurred
+                    ):
+                        raise response.exception_occurred
+
+            except QuestionAnswerValidationError as e:
+                self._handle_exception(e, invigilator, task)
+                return invigilator.get_failed_task_result(
+                    failure_reason="Question answer validation failed."
+                )
+
+            except asyncio.TimeoutError as e:
+                self._handle_exception(e, invigilator, task)
+                had_language_model_no_response_error = True
+                raise LanguageModelNoResponseError(
+                    f"Language model timed out for question '{question.question_name}.'"
+                )
+
+            except Exception as e:
+                self._handle_exception(e, invigilator, task)
+
+            if "response" not in locals():
+                had_language_model_no_response_error = True
+                raise LanguageModelNoResponseError(
+                    f"Language model did not return a response for question '{question.question_name}.'"
+                )
+
+            if (
+                question.question_name in self.interview.exceptions
+                and had_language_model_no_response_error
+            ):
+                self.interview.exceptions.record_fixed_question(question.question_name)
+
+            return response
+
+        try:
+            return await attempt_answer()
+        except RetryError as retry_error:
+            original_error = retry_error.last_attempt.exception()
+            self._handle_exception(
+                original_error, self.invigilator_fetcher(question), task
+            )
+            raise original_error

edsl/jobs/Answers.py

@@ -1,7 +1,6 @@
 """This module contains the Answers class, which is a helper class to hold the answers to a survey."""
 
 from collections import UserDict
-from rich.table import Table
 from edsl.data_transfer_models import EDSLResultObjectInput
 
 
@@ -22,7 +21,7 @@ class Answers(UserDict):
         if comment:
             self[question.question_name + "_comment"] = comment
 
-    def replace_missing_answers_with_none(self, survey) -> None:
+    def replace_missing_answers_with_none(self, survey: "Survey") -> None:
         """Replace missing answers with None. Answers can be missing if the agent skips a question."""
         for question_name in survey.question_names:
             if question_name not in self:
@@ -37,18 +36,6 @@ class Answers(UserDict):
         """Return an Answers object from a dictionary."""
         return cls(d)
 
-    def rich_print(self):
-        """Display an object as a table."""
-        table = Table(title="Answers")
-        table.add_column("Attribute", style="bold")
-        table.add_column("Value")
-
-        to_display = self
-        for attr_name, attr_value in to_display.items():
-            table.add_row(attr_name, repr(attr_value))
-
-        return table
-
 
 if __name__ == "__main__":
     import doctest

edsl/jobs/FetchInvigilator.py

@@ -0,0 +1,47 @@
+from typing import List, Dict, Any, Optional, TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from edsl.questions.QuestionBase import QuestionBase
+    from edsl.agents.InvigilatorBase import InvigilatorBase
+    from edsl.language_models.key_management.KeyLookup import KeyLookup
+    from edsl.jobs.interviews.Interview import Interview
+
+
+class FetchInvigilator:
+    def __init__(
+        self,
+        interview: "Interview",
+        current_answers: Optional[Dict[str, Any]] = None,
+        key_lookup: Optional["KeyLookup"] = None,
+    ):
+        self.interview = interview
+        if current_answers is None:
+            self.current_answers = self.interview.answers
+        else:
+            self.current_answers = current_answers
+        self.key_lookup = key_lookup
+
+    def get_invigilator(self, question: "QuestionBase") -> "InvigilatorBase":
+        """Return an invigilator for the given question.
+
+        :param question: the question to be answered
+        :param debug: whether to use debug mode, in which case `InvigilatorDebug` is used.
+        """
+
+        invigilator = self.interview.agent.create_invigilator(
+            question=question,
+            scenario=self.interview.scenario,
+            model=self.interview.model,
+            survey=self.interview.survey,
+            memory_plan=self.interview.survey.memory_plan,
+            current_answers=self.current_answers,  # not yet known
+            iteration=self.interview.iteration,
+            cache=self.interview.cache,
+            raise_validation_errors=self.interview.raise_validation_errors,
+            key_lookup=self.key_lookup,
+        )
+        """Return an invigilator for the given question."""
+        return invigilator
+
+    def __call__(self, question):
+        return self.get_invigilator(question)

edsl/jobs/InterviewTaskManager.py

@@ -0,0 +1,98 @@
+from __future__ import annotations
+import asyncio
+from typing import Any, Type, List, Generator, Optional, Union, TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from edsl.questions import QuestionBase
+    from edsl.jobs.tokens.InterviewTokenUsage import InterviewTokenUsage
+    from edsl.jobs.interviews.InterviewStatusDictionary import InterviewStatusDictionary
+    from edsl.jobs.interviews.InterviewStatusLog import InterviewStatusLog
+
+
+class InterviewTaskManager:
+    """Handles creation and management of interview tasks."""
+
+    def __init__(self, survey, iteration=0):
+        from edsl.jobs.tasks.TaskCreators import TaskCreators
+        from edsl.jobs.interviews.InterviewStatusLog import InterviewStatusLog
+
+        self.survey = survey
+        self.iteration = iteration
+        self.task_creators = TaskCreators()
+        self.to_index = {
+            question_name: index
+            for index, question_name in enumerate(self.survey.question_names)
+        }
+        self._task_status_log_dict = InterviewStatusLog()
+
+    def build_question_tasks(
+        self, answer_func, token_estimator, model_buckets
+    ) -> list[asyncio.Task]:
+        """Create tasks for all questions with proper dependencies."""
+        tasks = []
+        for question in self.survey.questions:
+            dependencies = self._get_task_dependencies(tasks, question)
+            task = self._create_single_task(
+                question=question,
+                dependencies=dependencies,
+                answer_func=answer_func,
+                token_estimator=token_estimator,
+                model_buckets=model_buckets,
+            )
+            tasks.append(task)
+        return tuple(tasks)
+
+    def _get_task_dependencies(
+        self, existing_tasks: list[asyncio.Task], question: "QuestionBase"
+    ) -> list[asyncio.Task]:
+        """Get tasks that must be completed before the given question."""
+        dag = self.survey.dag(textify=True)
+        parents = dag.get(question.question_name, [])
+        return [existing_tasks[self.to_index[parent_name]] for parent_name in parents]
+
+    def _create_single_task(
+        self,
+        question: "QuestionBase",
+        dependencies: list[asyncio.Task],
+        answer_func,
+        token_estimator,
+        model_buckets,
+    ) -> asyncio.Task:
+        """Create a single question task with its dependencies."""
+        from edsl.jobs.tasks.QuestionTaskCreator import QuestionTaskCreator
+
+        task_creator = QuestionTaskCreator(
+            question=question,
+            answer_question_func=answer_func,
+            token_estimator=token_estimator,
+            model_buckets=model_buckets,
+            iteration=self.iteration,
+        )
+
+        for dependency in dependencies:
+            task_creator.add_dependency(dependency)
+
+        self.task_creators[question.question_name] = task_creator
+        return task_creator.generate_task()
+
+    @property
+    def task_status_logs(self) -> "InterviewStatusLog":
+        """Return the task status logs for the interview.
+
+        The keys are the question names; the values are the lists of status log changes for each task.
+        """
+        for task_creator in self.task_creators.values():
+            self._task_status_log_dict[
+                task_creator.question.question_name
+            ] = task_creator.status_log
+        return self._task_status_log_dict
+
+    @property
+    def token_usage(self) -> "InterviewTokenUsage":
+        """Determine how many tokens were used for the interview."""
+        return self.task_creators.token_usage
+
+    @property
+    def interview_status(self) -> "InterviewStatusDictionary":
+        """Return a dictionary mapping task status codes to counts."""
+        return self.task_creators.interview_status

edsl/jobs/InterviewsConstructor.py

@@ -0,0 +1,50 @@
+from typing import Generator, TYPE_CHECKING
+from itertools import product
+
+if TYPE_CHECKING:
+    from edsl.jobs.interviews.Interview import Interview
+
+
+class InterviewsConstructor:
+    def __init__(self, jobs: "Jobs", cache: "Cache"):
+        self.jobs = jobs
+        self.cache = cache
+
+    def create_interviews(self) -> Generator["Interview", None, None]:
+        """
+        Generate interviews.
+
+        Note that this sets the agents, model and scenarios if they have not been set. This is a side effect of the method.
+        This is useful because a user can create a job without setting the agents, models, or scenarios, and the job will still run,
+        with us filling in defaults.
+
+        """
+        from edsl.jobs.interviews.Interview import Interview
+
+        agent_index = {
+            hash(agent): index for index, agent in enumerate(self.jobs.agents)
+        }
+        model_index = {
+            hash(model): index for index, model in enumerate(self.jobs.models)
+        }
+        scenario_index = {
+            hash(scenario): index for index, scenario in enumerate(self.jobs.scenarios)
+        }
+
+        for agent, scenario, model in product(
+            self.jobs.agents, self.jobs.scenarios, self.jobs.models
+        ):
+            yield Interview(
+                survey=self.jobs.survey.draw(),
+                agent=agent,
+                scenario=scenario,
+                model=model,
+                cache=self.cache,
+                skip_retry=self.jobs.run_config.parameters.skip_retry,
+                raise_validation_errors=self.jobs.run_config.parameters.raise_validation_errors,
+                indices={
+                    "agent": agent_index[hash(agent)],
+                    "model": model_index[hash(model)],
+                    "scenario": scenario_index[hash(scenario)],
+                },
+            )