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

edsl/agents/InvigilatorBase.py

@@ -1,21 +1,23 @@
 from abc import ABC, abstractmethod
 import asyncio
-from typing import Coroutine, Dict, Any, Optional
+from typing import Coroutine, Dict, Any, Optional, TYPE_CHECKING
 
-from edsl.prompts.Prompt import Prompt
 from edsl.utilities.decorators import jupyter_nb_handler
 from edsl.data_transfer_models import AgentResponseDict
 
-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
+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_transfer_models import EDSLResultObjectInput
 from edsl.agents.PromptConstructor import PromptConstructor
-
 from edsl.agents.prompt_helpers import PromptPlan
 
 
@@ -29,25 +31,23 @@ 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
 
-        self.raw_model_response = (
-            None  # placeholder for the raw response from the model
-        )
+        # placeholder to store the raw model response
+        self.raw_model_response = None
 
     @property
     def prompt_constructor(self) -> PromptConstructor:
         """Return the prompt constructor."""
         return PromptConstructor(self, prompt_plan=self.prompt_plan)
 
-    def to_dict(self):
+    def to_dict(self, include_cache=False) -> Dict[str, Any]:
         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,43 +104,37 @@ class InvigilatorBase(ABC):
         return {attr: serialize_attribute(attr) for attr in attributes}
 
     @classmethod
-    def from_dict(cls, data):
+    def from_dict(cls, data) -> "InvigilatorBase":
         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])
 
-        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["current_answers"] = data["current_answers"]
+        d["iteration"] = data["iteration"]
+        d["additional_prompt_data"] = data["additional_prompt_data"]
 
-        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,
-        )
+        d = cls(**d)
 
     def __repr__(self) -> str:
         """Return a string representation of the Invigilator.
@@ -149,9 +143,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)}, sidecarmodel={repr(self.sidecar_model)})"
+        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)})"
 
-    def get_failed_task_result(self, failure_reason) -> EDSLResultObjectInput:
+    def get_failed_task_result(self, failure_reason: str) -> EDSLResultObjectInput:
         """Return an AgentResponseDict used in case the question-asking fails.
 
         Possible reasons include:
@@ -172,8 +166,9 @@ 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"),
@@ -205,24 +200,25 @@ 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.enums import InferenceServiceType
-
-        from edsl import Model
+        from edsl.language_models.model import Model
+        from edsl.surveys.Survey import Survey
 
         model = Model("test", canned_response="SPAM!")
 
         if throw_an_exception:
-            model.throw_an_exception = True
+            model.throw_exception = True
         agent = Agent.example()
-        # question = QuestionMultipleChoice.example()
-        from edsl.surveys import Survey
 
         if not survey:
             survey = Survey.example()
@@ -232,14 +228,10 @@ 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(InvigilatorBase):
+        class InvigilatorExample(cls):
             """An example invigilator."""
 
             async def async_answer_question(self):