edsl 0.1.39__py3-none-any.whl → 0.1.39.dev1__py3-none-any.whl

edsl/agents/InvigilatorBase.py

@@ -1,23 +1,21 @@
 from abc import ABC, abstractmethod
 import asyncio
-from typing import Coroutine, Dict, Any, Optional, TYPE_CHECKING
+from typing import Coroutine, Dict, Any, Optional
 
+from edsl.prompts.Prompt import Prompt
 from edsl.utilities.decorators import jupyter_nb_handler
 from edsl.data_transfer_models import AgentResponseDict
 
-if TYPE_CHECKING:
-    from edsl.prompts.Prompt import Prompt
-    from edsl.data.Cache import Cache
-    from edsl.questions.QuestionBase import QuestionBase
-    from edsl.scenarios.Scenario import Scenario
-    from edsl.surveys.MemoryPlan import MemoryPlan
-    from edsl.language_models.LanguageModel import LanguageModel
-    from edsl.surveys.Survey import Survey
-    from edsl.agents.Agent import Agent
-    from edsl.language_models.key_management.KeyLookup import KeyLookup
+from edsl.data.Cache import Cache
+
+from edsl.questions.QuestionBase import QuestionBase
+from edsl.scenarios.Scenario import Scenario
+from edsl.surveys.MemoryPlan import MemoryPlan
+from edsl.language_models.LanguageModel import LanguageModel
 
 from edsl.data_transfer_models import EDSLResultObjectInput
 from edsl.agents.PromptConstructor import PromptConstructor
+
 from edsl.agents.prompt_helpers import PromptPlan
 
 
@@ -31,23 +29,25 @@ class InvigilatorBase(ABC):
     'Failed to get response'
 
     This returns an empty prompt because there is no memory the agent needs to have at q0.
+
+
     """
 
     def __init__(
         self,
         agent: "Agent",
-        question: "QuestionBase",
-        scenario: "Scenario",
-        model: "LanguageModel",
-        memory_plan: "MemoryPlan",
+        question: QuestionBase,
+        scenario: Scenario,
+        model: LanguageModel,
+        memory_plan: MemoryPlan,
         current_answers: dict,
         survey: Optional["Survey"],
-        cache: Optional["Cache"] = None,
+        cache: Optional[Cache] = None,
         iteration: Optional[int] = 1,
         additional_prompt_data: Optional[dict] = None,
+        sidecar_model: Optional[LanguageModel] = None,
         raise_validation_errors: Optional[bool] = True,
         prompt_plan: Optional["PromptPlan"] = None,
-        key_lookup: Optional["KeyLookup"] = None,
     ):
         """Initialize a new Invigilator."""
         self.agent = agent
@@ -59,24 +59,24 @@ class InvigilatorBase(ABC):
         self.iteration = iteration
         self.additional_prompt_data = additional_prompt_data
         self.cache = cache
+        self.sidecar_model = sidecar_model
         self.survey = survey
         self.raise_validation_errors = raise_validation_errors
-        self.key_lookup = key_lookup
-
         if prompt_plan is None:
             self.prompt_plan = PromptPlan()
         else:
             self.prompt_plan = prompt_plan
 
-        # placeholder to store the raw model response
-        self.raw_model_response = None
+        self.raw_model_response = (
+            None  # placeholder for the raw response from the model
+        )
 
     @property
     def prompt_constructor(self) -> PromptConstructor:
         """Return the prompt constructor."""
         return PromptConstructor(self, prompt_plan=self.prompt_plan)
 
-    def to_dict(self, include_cache=False) -> Dict[str, Any]:
+    def to_dict(self):
         attributes = [
             "agent",
             "question",
@@ -86,10 +86,10 @@ class InvigilatorBase(ABC):
             "current_answers",
             "iteration",
             "additional_prompt_data",
+            "cache",
+            "sidecar_model",
             "survey",
         ]
-        if include_cache:
-            attributes.append("cache")
 
         def serialize_attribute(attr):
             value = getattr(self, attr)
@@ -104,37 +104,43 @@ class InvigilatorBase(ABC):
         return {attr: serialize_attribute(attr) for attr in attributes}
 
     @classmethod
-    def from_dict(cls, data) -> "InvigilatorBase":
+    def from_dict(cls, data):
         from edsl.agents.Agent import Agent
         from edsl.questions import QuestionBase
         from edsl.scenarios.Scenario import Scenario
         from edsl.surveys.MemoryPlan import MemoryPlan
         from edsl.language_models.LanguageModel import LanguageModel
         from edsl.surveys.Survey import Survey
-        from edsl.data.Cache import Cache
-
-        attributes_to_classes = {
-            "agent": Agent,
-            "question": QuestionBase,
-            "scenario": Scenario,
-            "model": LanguageModel,
-            "memory_plan": MemoryPlan,
-            "survey": Survey,
-            "cache": Cache,
-        }
-        d = {}
-        for attr, cls_ in attributes_to_classes.items():
-            if attr in data and data[attr] is not None:
-                if attr not in data:
-                    d[attr] = {}
-                else:
-                    d[attr] = cls_.from_dict(data[attr])
 
-        d["current_answers"] = data["current_answers"]
-        d["iteration"] = data["iteration"]
-        d["additional_prompt_data"] = data["additional_prompt_data"]
+        agent = Agent.from_dict(data["agent"])
+        question = QuestionBase.from_dict(data["question"])
+        scenario = Scenario.from_dict(data["scenario"])
+        model = LanguageModel.from_dict(data["model"])
+        memory_plan = MemoryPlan.from_dict(data["memory_plan"])
+        survey = Survey.from_dict(data["survey"])
+        current_answers = data["current_answers"]
+        iteration = data["iteration"]
+        additional_prompt_data = data["additional_prompt_data"]
+        cache = Cache.from_dict(data["cache"])
+
+        if data["sidecar_model"] is None:
+            sidecar_model = None
+        else:
+            sidecar_model = LanguageModel.from_dict(data["sidecar_model"])
 
-        d = cls(**d)
+        return cls(
+            agent=agent,
+            question=question,
+            scenario=scenario,
+            model=model,
+            memory_plan=memory_plan,
+            current_answers=current_answers,
+            survey=survey,
+            iteration=iteration,
+            additional_prompt_data=additional_prompt_data,
+            cache=cache,
+            sidecar_model=sidecar_model,
+        )
 
     def __repr__(self) -> str:
         """Return a string representation of the Invigilator.
@@ -143,9 +149,9 @@ class InvigilatorBase(ABC):
         'InvigilatorExample(...)'
 
         """
-        return f"{self.__class__.__name__}(agent={repr(self.agent)}, question={repr(self.question)}, scneario={repr(self.scenario)}, model={repr(self.model)}, memory_plan={repr(self.memory_plan)}, current_answers={repr(self.current_answers)}, iteration{repr(self.iteration)}, additional_prompt_data={repr(self.additional_prompt_data)}, cache={repr(self.cache)})"
+        return f"{self.__class__.__name__}(agent={repr(self.agent)}, question={repr(self.question)}, scneario={repr(self.scenario)}, model={repr(self.model)}, memory_plan={repr(self.memory_plan)}, current_answers={repr(self.current_answers)}, iteration{repr(self.iteration)}, additional_prompt_data={repr(self.additional_prompt_data)}, cache={repr(self.cache)}, sidecarmodel={repr(self.sidecar_model)})"
 
-    def get_failed_task_result(self, failure_reason: str) -> EDSLResultObjectInput:
+    def get_failed_task_result(self, failure_reason) -> EDSLResultObjectInput:
         """Return an AgentResponseDict used in case the question-asking fails.
 
         Possible reasons include:
@@ -166,9 +172,8 @@ class InvigilatorBase(ABC):
         }
         return EDSLResultObjectInput(**data)
 
-    def get_prompts(self) -> Dict[str, "Prompt"]:
+    def get_prompts(self) -> Dict[str, Prompt]:
         """Return the prompt used."""
-        from edsl.prompts.Prompt import Prompt
 
         return {
             "user_prompt": Prompt("NA"),
@@ -200,25 +205,24 @@ class InvigilatorBase(ABC):
         >>> InvigilatorBase.example()
         InvigilatorExample(...)
 
-        >>> InvigilatorBase.example().answer_question()
-        {'message': [{'text': 'SPAM!'}], 'usage': {'prompt_tokens': 1, 'completion_tokens': 1}}
-
-        >>> InvigilatorBase.example(throw_an_exception=True).answer_question()
-        Traceback (most recent call last):
-        ...
-        Exception: This is a test error
         """
         from edsl.agents.Agent import Agent
+        from edsl.questions import QuestionMultipleChoice
         from edsl.scenarios.Scenario import Scenario
+        from edsl.language_models import LanguageModel
         from edsl.surveys.MemoryPlan import MemoryPlan
-        from edsl.language_models.model import Model
-        from edsl.surveys.Survey import Survey
+
+        from edsl.enums import InferenceServiceType
+
+        from edsl import Model
 
         model = Model("test", canned_response="SPAM!")
 
         if throw_an_exception:
-            model.throw_exception = True
+            model.throw_an_exception = True
         agent = Agent.example()
+        # question = QuestionMultipleChoice.example()
+        from edsl.surveys import Survey
 
         if not survey:
             survey = Survey.example()
@@ -228,10 +232,14 @@ class InvigilatorBase(ABC):
 
         question = question or survey.questions[0]
         scenario = scenario or Scenario.example()
+        # memory_plan = None #memory_plan = MemoryPlan()
+        from edsl import Survey
+
         memory_plan = MemoryPlan(survey=survey)
         current_answers = None
+        from edsl.agents.PromptConstructor import PromptConstructor
 
-        class InvigilatorExample(cls):
+        class InvigilatorExample(InvigilatorBase):
             """An example invigilator."""
 
             async def async_answer_question(self):