edsl 0.1.14__py3-none-any.whl → 0.1.40__py3-none-any.whl

edsl/agents/Invigilator.py

@@ -1,409 +1,285 @@
-from abc import ABC, abstractmethod
-import asyncio
-import json
-from typing import Coroutine, Dict, Any
-from collections import UserDict
-
-from edsl.exceptions import AgentRespondedWithBadJSONError
-from edsl.prompts.Prompt import Prompt
-from edsl.utilities.decorators import sync_wrapper, jupyter_nb_handler
-from edsl.prompts.registry import get_classes
-from edsl.exceptions import QuestionScenarioRenderError
-
-from edsl.data_transfer_models import AgentResponseDict
-from edsl.exceptions.agents import FailedTaskException
-
-
-class InvigilatorBase(ABC):
-    """An invigiator (someone who administers an exam) is a class that is responsible for administering a question to an agent."""
-
-    def __init__(
-        self, agent, question, scenario, model, memory_plan, current_answers: dict
-    ):
-        self.agent = agent
-        self.question = question
-        self.scenario = scenario
-        self.model = model
-        self.memory_plan = memory_plan
-        self.current_answers = current_answers
-
-    def get_failed_task_result(self):
-        return AgentResponseDict(
-            answer=None,
-            comment="Failed to get response",
-            question_name=self.question.question_name,
-            prompts=self.get_prompts(),
-        )
-
-    def get_prompts(self) -> Dict[str, Prompt]:
-        return {
-            "user_prompt": Prompt("NA").text,
-            "system_prompt": Prompt("NA").text,
-        }
+"""Module for creating Invigilators, which are objects to administer a question to an Agent."""
 
-    @classmethod
-    def example(cls):
-        """Returns an example invigilator."""
-        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.enums import LanguageModelType, InferenceServiceType
-
-        class TestLanguageModelGood(LanguageModel):
-            _model_ = LanguageModelType.TEST.value
-            _parameters_ = {"temperature": 0.5}
-            _inference_service_ = InferenceServiceType.TEST.value
-
-            async def async_execute_model_call(
-                self, user_prompt: str, system_prompt: str
-            ) -> dict[str, Any]:
-                await asyncio.sleep(0.1)
-                return {"message": """{"answer": "SPAM!"}"""}
-
-            def parse_response(self, raw_response: dict[str, Any]) -> str:
-                return raw_response["message"]
-
-        model = TestLanguageModelGood()
-        agent = Agent.example()
-        question = QuestionMultipleChoice.example()
-        scenario = Scenario.example()
-        #        model = LanguageModel.example()
-        memory_plan = None
-        current_answers = None
-        return cls(
-            agent=agent,
-            question=question,
-            scenario=scenario,
-            model=model,
-            memory_plan=memory_plan,
-            current_answers=current_answers,
-        )
+from typing import Dict, Any, Optional, TYPE_CHECKING
 
-    @abstractmethod
-    async def async_answer_question(self):
-        "This is the async method that actually answers the question."
-        pass
+from edsl.utilities.decorators import sync_wrapper
+from edsl.exceptions.questions import QuestionAnswerValidationError
+from edsl.agents.InvigilatorBase import InvigilatorBase
+from edsl.data_transfer_models import AgentResponseDict, EDSLResultObjectInput
 
-    @jupyter_nb_handler
-    def answer_question(self) -> Coroutine:
-        async def main():
-            results = await asyncio.gather(self.async_answer_question())
-            return results[0]  # Since there's only one task, return its result
+if TYPE_CHECKING:
+    from edsl.prompts.Prompt import Prompt
+    from edsl.scenarios.Scenario import Scenario
+    from edsl.surveys.Survey import Survey
 
-        return main()
 
-    def create_memory_prompt(self, question_name):
-        """Creates a memory for the agent."""
-        return self.memory_plan.get_memory_prompt_fragment(
-            question_name, self.current_answers
-        )
+NA = "Not Applicable"
 
 
 class InvigilatorAI(InvigilatorBase):
     """An invigilator that uses an AI model to answer questions."""
 
-    async def async_answer_question(self, failed=False) -> AgentResponseDict:
-        data = {
-            "agent": self.agent,
-            "question": self.question,
-            "scenario": self.scenario,
-        }
-        # This calls the self.async_get_response method w/ the prompts
-        # The raw response is a dictionary.
-        raw_response = await self.async_get_response(**self.get_prompts())
-        assert "raw_model_response" in raw_response
-        response = self._format_raw_response(
-            **(
-                data
-                | {
-                    "raw_response": raw_response,
-                    "raw_model_response": raw_response["raw_model_response"],
-                }
-            )
-        )
-        return response
+    def get_prompts(self) -> Dict[str, "Prompt"]:
+        """Return the prompts used."""
+        return self.prompt_constructor.get_prompts()
 
-    async def async_get_response(self, user_prompt: Prompt, system_prompt: Prompt):
-        """Calls the LLM and gets a response. Used in the `answer_question` method."""
-        try:
-            response = await self.model.async_get_response(
-                user_prompt.text, system_prompt.text
-            )
-        except json.JSONDecodeError as e:
-            raise AgentRespondedWithBadJSONError(
-                f"Returned bad JSON: {e}"
-                f"Prompt: {user_prompt}"
-                f"System Prompt: {system_prompt}"
-            )
-
-        return response
-
-    def _format_raw_response(
-        self, agent, question, scenario, raw_response, raw_model_response
-    ) -> AgentResponseDict:
-        response = question.validate_answer(raw_response)
-        comment = response.get("comment", "")
-        answer_code = response["answer"]
-        answer = question.translate_answer_code_to_answer(answer_code, scenario)
-        raw_model_response = raw_model_response
-        data = {
-            "answer": answer,
-            "comment": comment,
-            "question_name": question.question_name,
-            "prompts": {k: v.to_dict() for k, v in self.get_prompts().items()},
-            "cached_response": raw_response["cached_response"],
-            "usage": raw_response.get("usage", {}),
-            "raw_model_response": raw_model_response,
+    async def async_get_agent_response(self) -> AgentResponseDict:
+        prompts = self.get_prompts()
+        params = {
+            "user_prompt": prompts["user_prompt"].text,
+            "system_prompt": prompts["system_prompt"].text,
         }
-        return AgentResponseDict(**data)
-
-    get_response = sync_wrapper(async_get_response)
-
-    def construct_system_prompt(self) -> Prompt:
-        """Constructs the system prompt for the LLM call."""
-        applicable_prompts = get_classes(
-            component_type="agent_instructions",
-            model=self.model.model,
-        )
-        if len(applicable_prompts) == 0:
-            raise Exception("No applicable prompts found")
-
-        agent_instructions = applicable_prompts[0](text=self.agent.instruction)
-
-        if not hasattr(self.agent, "agent_persona"):
-            applicable_prompts = get_classes(
-                component_type="agent_persona",
-                model=self.model.model,
-            )
-            persona_prompt_template = applicable_prompts[0]()
-        else:
-            persona_prompt_template = self.agent.agent_persona
-
-        if undefined := persona_prompt_template.undefined_template_variables(
-            self.agent.traits
-            | {"traits": self.agent.traits}
-            | {"codebook": self.agent.codebook}
-            | {"traits": self.agent.traits}
-        ):
-            raise QuestionScenarioRenderError(
-                f"Agent persona still has variables that were not rendered: {undefined}"
-            )
-
-        persona_prompt = persona_prompt_template.render(
-            self.agent.traits | {"traits": self.agent.traits},
-            codebook=self.agent.codebook,
-            traits=self.agent.traits,
-        )
-
-        if persona_prompt.has_variables:
-            raise QuestionScenarioRenderError(
-                "Agent persona still has variables that were not rendered."
-            )
-
-        return (
-            agent_instructions
-            + " " * int(len(persona_prompt.text) > 0)
-            + persona_prompt
+        if "encoded_image" in prompts:
+            params["encoded_image"] = prompts["encoded_image"]
+            raise NotImplementedError("encoded_image not implemented")
+
+        if "files_list" in prompts:
+            params["files_list"] = prompts["files_list"]
+
+        params.update({"iteration": self.iteration, "cache": self.cache})
+        params.update({"invigilator": self})
+
+        if self.key_lookup:
+            self.model.set_key_lookup(self.key_lookup)
+
+        return await self.model.async_get_response(**params)
+
+    def store_response(self, agent_response_dict: AgentResponseDict) -> None:
+        """Store the response in the invigilator, in case it is needed later because of validation failure."""
+        self.raw_model_response = agent_response_dict.model_outputs.response
+        self.generated_tokens = agent_response_dict.edsl_dict.generated_tokens
+        self.cache_key = agent_response_dict.model_outputs.cache_key
+
+    async def async_answer_question(self) -> EDSLResultObjectInput:
+        """Answer a question using the AI model.
+
+        >>> i = InvigilatorAI.example()
+        """
+        agent_response_dict: AgentResponseDict = await self.async_get_agent_response()
+        self.store_response(agent_response_dict)
+        return self._extract_edsl_result_entry_and_validate(agent_response_dict)
+
+    def _remove_from_cache(self, cache_key) -> None:
+        """Remove an entry from the cache."""
+        if cache_key:
+            del self.cache.data[cache_key]
+
+    def _determine_answer(self, raw_answer: str) -> Any:
+        """Determine the answer from the raw answer.
+
+        >>> i = InvigilatorAI.example()
+        >>> i._determine_answer("SPAM!")
+        'SPAM!'
+
+        >>> from edsl.questions import QuestionMultipleChoice
+        >>> q = QuestionMultipleChoice(question_text = "How are you?", question_name = "how_are_you", question_options = ["Good", "Bad"], use_code = True)
+        >>> i = InvigilatorAI.example(question = q)
+        >>> i._determine_answer("1")
+        'Bad'
+        >>> i._determine_answer("0")
+        'Good'
+
+        This shows how the answer can depend on scenario details
+
+        >>> from edsl import Scenario
+        >>> s = Scenario({'feeling_options':['Good', 'Bad']})
+        >>> q = QuestionMultipleChoice(question_text = "How are you?", question_name = "how_are_you", question_options = "{{ feeling_options }}", use_code = True)
+        >>> i = InvigilatorAI.example(question = q, scenario = s)
+        >>> i._determine_answer("1")
+        'Bad'
+
+        >>> from edsl import QuestionList, QuestionMultipleChoice, Survey
+        >>> q1 = QuestionList(question_name = "favs", question_text = "What are your top 3 colors?")
+        >>> q2 = QuestionMultipleChoice(question_text = "What is your favorite color?", question_name = "best", question_options = "{{ favs.answer }}", use_code = True)
+        >>> survey = Survey([q1, q2])
+        >>> i = InvigilatorAI.example(question = q2, scenario = s, survey = survey)
+        >>> i.current_answers = {"favs": ["Green", "Blue", "Red"]}
+        >>> i._determine_answer("2")
+        'Red'
+        """
+        substitution_dict = self._prepare_substitution_dict(
+            self.survey, self.current_answers, self.scenario
         )
-
-    def get_question_instructions(self) -> Prompt:
-        """Gets the instructions for the question."""
-        applicable_prompts = get_classes(
-            component_type="question_instructions",
-            question_type=self.question.question_type,
-            model=self.model.model,
+        return self.question._translate_answer_code_to_answer(
+            raw_answer, substitution_dict
         )
-        ## Get the question instructions and renders with the scenario & question.data
-        question_prompt = applicable_prompts[0]()
 
-        undefined_template_variables = question_prompt.undefined_template_variables(
-            self.question.data | self.scenario
-        )
-        if undefined_template_variables:
-            print(undefined_template_variables)
-            raise QuestionScenarioRenderError(
-                "Question instructions still has variables"
-            )
-
-        return question_prompt.render(self.question.data | self.scenario)
-
-    def construct_user_prompt(self) -> Prompt:
-        """Gets the user prompt for the LLM call."""
-        user_prompt = self.get_question_instructions()
-        if self.memory_plan is not None:
-            user_prompt += self.create_memory_prompt(self.question.question_name)
-        return user_prompt
-
-    def get_prompts(self) -> Dict[str, Prompt]:
-        """Gets the prompts for the LLM call."""
-        system_prompt = self.construct_system_prompt()
-        user_prompt = self.construct_user_prompt()
-        return {
-            "user_prompt": user_prompt,
-            "system_prompt": system_prompt,
-        }
+    @staticmethod
+    def _prepare_substitution_dict(
+        survey: "Survey", current_answers: dict, scenario: "Scenario"
+    ) -> Dict[str, Any]:
+        """Prepares a substitution dictionary for the question based on the survey, current answers, and scenario.
+
+        This is necessary beause sometimes the model's answer to a question could depend on details in
+        the prompt that were provided by the answer to a previous question or a scenario detail.
+
+        Note that the question object is getting the answer & a the comment appended to it, as the
+        jinja2 template might be referencing these values with a dot notation.
+
+        """
+        question_dict = survey.duplicate().question_names_to_questions()
+
+        # iterates through the current answers and updates the question_dict (which is all questions)
+        for other_question, answer in current_answers.items():
+            if other_question in question_dict:
+                question_dict[other_question].answer = answer
+            else:
+                # it might be a comment
+                if (
+                    new_question := other_question.split("_comment")[0]
+                ) in question_dict:
+                    question_dict[new_question].comment = answer
+
+        return {**question_dict, **scenario}
+
+    def _extract_edsl_result_entry_and_validate(
+        self, agent_response_dict: AgentResponseDict
+    ) -> EDSLResultObjectInput:
+        """Extract the EDSL result entry and validate it."""
+        edsl_dict = agent_response_dict.edsl_dict._asdict()
+        exception_occurred = None
+        validated = False
+        try:
+            # if the question has jinja parameters, it is easier to make a new question with the parameters
+            if self.question.parameters:
+                prior_answers_dict = self.prompt_constructor.prior_answers_dict()
+
+                # question options have be treated differently because of dynamic question
+                # this logic is all in the prompt constructor
+                if "question_options" in self.question.data:
+                    new_question_options = self.prompt_constructor.get_question_options(
+                        self.question.data
+                    )
+                    if new_question_options != self.question.data["question_options"]:
+                        # I don't love this direct writing but it seems to work
+                        self.question.question_options = new_question_options
+
+                question_with_validators = self.question.render(
+                    self.scenario | prior_answers_dict
+                )
+                question_with_validators.use_code = self.question.use_code
+            else:
+                question_with_validators = self.question
+
+            validated_edsl_dict = question_with_validators._validate_answer(edsl_dict)
+            answer = self._determine_answer(validated_edsl_dict["answer"])
+            comment = validated_edsl_dict.get("comment", "")
+            validated = True
+        except QuestionAnswerValidationError as e:
+            answer = None
+            comment = "The response was not valid."
+            # if self.raise_validation_errors:
+            exception_occurred = e
+        except Exception as non_validation_error:
+            answer = None
+            comment = "Some other error occurred."
+            exception_occurred = non_validation_error
+        finally:
+            # even if validation failes, we still return the result
+            data = {
+                "answer": answer,
+                "comment": comment,
+                "generated_tokens": agent_response_dict.edsl_dict.generated_tokens,
+                "question_name": self.question.question_name,
+                "prompts": self.get_prompts(),
+                "cached_response": agent_response_dict.model_outputs.cached_response,
+                "raw_model_response": agent_response_dict.model_outputs.response,
+                "cache_used": agent_response_dict.model_outputs.cache_used,
+                "cache_key": agent_response_dict.model_outputs.cache_key,
+                "validated": validated,
+                "exception_occurred": exception_occurred,
+                "cost": agent_response_dict.model_outputs.cost,
+            }
+            result = EDSLResultObjectInput(**data)
+            return result
 
     answer_question = sync_wrapper(async_answer_question)
 
 
-class InvigilatorDebug(InvigilatorBase):
-    async def async_answer_question(self) -> AgentResponseDict:
-        results = self.question.simulate_answer(human_readable=True)
-        results["prompts"] = self.get_prompts()
-        results["question_name"] = self.question.question_name
-        results["comment"] = "Debug comment"
-        return AgentResponseDict(**results)
+class InvigilatorHuman(InvigilatorBase):
+    """An invigilator for when a human is answering the question."""
 
-    def get_prompts(self) -> Dict[str, Prompt]:
-        return {
-            "user_prompt": Prompt("NA").text,
-            "system_prompt": Prompt("NA").text,
-        }
+    validate_response: bool = False
+    translate_response: bool = False
 
+    async def async_answer_question(self, iteration: int = 0) -> AgentResponseDict:
+        """Return the answer to the question."""
+        comment = "This is a real survey response from a human."
 
-class InvigilatorHuman(InvigilatorBase):
-    async def async_answer_question(self) -> AgentResponseDict:
-        data = {
-            "comment": "This is a real survey response from a human.",
-            "answer": None,
-            "prompts": self.get_prompts(),
-            "question_name": self.question.question_name,
-        }
+        def __repr__(self):
+            return f"{self.literal}"
+
+        exception_occurred = None
+        validated = False
         try:
             answer = self.agent.answer_question_directly(self.question, self.scenario)
-            return AgentResponseDict(**(data | {"answer": answer}))
+            self.raw_model_response = answer
+
+            if self.validate_response:
+                _ = self.question._validate_answer({"answer": answer})
+            if self.translate_response:
+                answer = self.question._translate_answer_code_to_answer(
+                    answer, self.scenario
+                )
+            validated = True
+        except QuestionAnswerValidationError as e:
+            answer = None
+            if self.raise_validation_errors:
+                exception_occurred = e
         except Exception as e:
-            agent_response_dict = AgentResponseDict(
-                **(data | {"answer": None, "comment": str(e)})
-            )
-            raise FailedTaskException(
-                f"Failed to get response. The exception is {str(e)}",
-                agent_response_dict,
-            ) from e
+            answer = None
+            if self.raise_validation_errors:
+                exception_occurred = e
+        finally:
+            data = {
+                "generated_tokens": NA,  # NotApplicable(),
+                "question_name": self.question.question_name,
+                "prompts": self.get_prompts(),
+                "cached_response": NA,
+                "raw_model_response": NA,
+                "cache_used": NA,
+                "cache_key": NA,
+                "answer": answer,
+                "comment": comment,
+                "validated": validated,
+                "exception_occurred": exception_occurred,
+            }
+            return EDSLResultObjectInput(**data)
 
 
 class InvigilatorFunctional(InvigilatorBase):
-    async def async_answer_question(self) -> AgentResponseDict:
+    """A Invigilator for when the question has a answer_question_directly function."""
+
+    async def async_answer_question(self, iteration: int = 0) -> AgentResponseDict:
+        """Return the answer to the question."""
         func = self.question.answer_question_directly
-        data = {
-            "comment": "Functional.",
-            "prompts": self.get_prompts(),
-            "question_name": self.question.question_name,
-        }
-        try:
-            answer = func(scenario=self.scenario, agent_traits=self.agent.traits)
-            return AgentResponseDict(**(data | {"answer": answer}))
-        except Exception as e:
-            agent_response_dict = AgentResponseDict(
-                **(data | {"answer": None, "comment": str(e)})
-            )
-            raise FailedTaskException(
-                f"Failed to get response. The exception is {str(e)}",
-                agent_response_dict,
-            ) from e
-
-    def get_prompts(self) -> Dict[str, Prompt]:
+        answer = func(scenario=self.scenario, agent_traits=self.agent.traits)
+
+        return EDSLResultObjectInput(
+            generated_tokens=str(answer),
+            question_name=self.question.question_name,
+            prompts=self.get_prompts(),
+            cached_response=NA,
+            raw_model_response=NA,
+            cache_used=NA,
+            cache_key=NA,
+            answer=answer["answer"],
+            comment="This is the result of a functional question.",
+            validated=True,
+            exception_occurred=None,
+        )
+
+    def get_prompts(self) -> Dict[str, "Prompt"]:
+        from edsl.prompts.Prompt import Prompt
+
+        """Return the prompts used."""
         return {
-            "user_prompt": Prompt("NA").text,
-            "system_prompt": Prompt("NA").text,
+            "user_prompt": Prompt("NA"),
+            "system_prompt": Prompt("NA"),
         }
 
 
 if __name__ == "__main__":
-    from edsl.enums import LanguageModelType
-
-    from edsl.agents.Agent import Agent
-
-    a = Agent(
-        instruction="You are a happy-go lucky agent.",
-        traits={"feeling": "happy", "age": "Young at heart"},
-        codebook={"feeling": "Feelings right now", "age": "Age in years"},
-        trait_presentation_template="",
-    )
-
-    class MockModel:
-        model = LanguageModelType.GPT_4.value
-
-    class MockQuestion:
-        question_type = "free_text"
-        question_text = "How are you feeling?"
-        question_name = "feelings_question"
-        data = {
-            "question_name": "feelings",
-            "question_text": "How are you feeling?",
-            "question_type": "feelings_question",
-        }
+    import doctest
 
-    i = InvigilatorAI(
-        agent=a,
-        question=MockQuestion(),
-        scenario={},
-        model=MockModel(),
-        memory_plan=None,
-        current_answers=None,
-    )
-    print(i.get_prompts()["system_prompt"])
-    assert i.get_prompts()["system_prompt"].text == "You are a happy-go lucky agent."
-
-    ###############
-    ## Render one
-    ###############
-
-    a = Agent(
-        instruction="You are a happy-go lucky agent.",
-        traits={"feeling": "happy", "age": "Young at heart"},
-        codebook={"feeling": "Feelings right now", "age": "Age in years"},
-        trait_presentation_template="You are feeling {{ feeling }}.",
-    )
-
-    i = InvigilatorAI(
-        agent=a,
-        question=MockQuestion(),
-        scenario={},
-        model=MockModel(),
-        memory_plan=None,
-        current_answers=None,
-    )
-    print(i.get_prompts()["system_prompt"])
-
-    assert (
-        i.get_prompts()["system_prompt"].text
-        == "You are a happy-go lucky agent. You are feeling happy."
-    )
-    try:
-        assert i.get_prompts()["system_prompt"].unused_traits(a.traits) == ["age"]
-    except AssertionError:
-        unused_traits = i.get_prompts()["system_prompt"].unused_traits(a.traits)
-        print(f"System prompt: {i.get_prompts()['system_prompt']}")
-        print(f"Agent traits: {a.traits}")
-        print(f"Unused_traits: {unused_traits}")
-        # breakpoint()
-
-    ###############
-    ## Render one
-    ###############
-
-    a = Agent(
-        instruction="You are a happy-go lucky agent.",
-        traits={"feeling": "happy", "age": "Young at heart"},
-        codebook={"feeling": "Feelings right now", "age": "Age in years"},
-        trait_presentation_template="You are feeling {{ feeling }}. You eat lots of {{ food }}.",
-    )
-
-    i = InvigilatorAI(
-        agent=a,
-        question=MockQuestion(),
-        scenario={},
-        model=MockModel(),
-        memory_plan=None,
-        current_answers=None,
-    )
-    print(i.get_prompts()["system_prompt"])
-
-    ## Should raise a QuestionScenarioRenderError
-    assert (
-        i.get_prompts()["system_prompt"].text
-        == "You are a happy-go lucky agent. You are feeling happy."
-    )
+    doctest.testmod(optionflags=doctest.ELLIPSIS)