edsl 0.1.33.dev3__py3-none-any.whl → 0.1.34.dev1__py3-none-any.whl

edsl/inference_services/GoogleService.py

@@ -1,25 +1,54 @@
 import os
-import aiohttp
-import json
-from typing import Any
+from typing import Any, Dict, List, Optional
+import google
+import google.generativeai as genai
+from google.generativeai.types import GenerationConfig
+from google.api_core.exceptions import InvalidArgument
+
 from edsl.exceptions import MissingAPIKeyError
 from edsl.language_models.LanguageModel import LanguageModel
-
 from edsl.inference_services.InferenceServiceABC import InferenceServiceABC
 
+safety_settings = [
+    {
+        "category": "HARM_CATEGORY_HARASSMENT",
+        "threshold": "BLOCK_NONE",
+    },
+    {
+        "category": "HARM_CATEGORY_HATE_SPEECH",
+        "threshold": "BLOCK_NONE",
+    },
+    {
+        "category": "HARM_CATEGORY_SEXUALLY_EXPLICIT",
+        "threshold": "BLOCK_NONE",
+    },
+    {
+        "category": "HARM_CATEGORY_DANGEROUS_CONTENT",
+        "threshold": "BLOCK_NONE",
+    },
+]
+
 
 class GoogleService(InferenceServiceABC):
     _inference_service_ = "google"
     key_sequence = ["candidates", 0, "content", "parts", 0, "text"]
-    usage_sequence = ["usageMetadata"]
-    input_token_name = "promptTokenCount"
-    output_token_name = "candidatesTokenCount"
+    usage_sequence = ["usage_metadata"]
+    input_token_name = "prompt_token_count"
+    output_token_name = "candidates_token_count"
 
     model_exclude_list = []
 
+    # @classmethod
+    # def available(cls) -> List[str]:
+    #     return ["gemini-pro", "gemini-1.5-pro", "gemini-1.5-flash", "gemini-1.0-pro"]
+
     @classmethod
-    def available(cls):
-        return ["gemini-pro", "gemini-1.5-pro", "gemini-1.5-flash", "gemini-1.0-pro"]
+    def available(cls) -> List[str]:
+        model_list = []
+        for m in genai.list_models():
+            if "generateContent" in m.supported_generation_methods:
+                model_list.append(m.name.split("/")[-1])
+        return model_list
 
     @classmethod
     def create_model(
@@ -47,33 +76,79 @@ class GoogleService(InferenceServiceABC):
                 "stopSequences": [],
             }
 
+            api_token = None
+            model = None
+
+            @classmethod
+            def initialize(cls):
+                if cls.api_token is None:
+                    cls.api_token = os.getenv("GOOGLE_API_KEY")
+                    if not cls.api_token:
+                        raise MissingAPIKeyError(
+                            "GOOGLE_API_KEY environment variable is not set"
+                        )
+                    genai.configure(api_key=cls.api_token)
+                    cls.generative_model = genai.GenerativeModel(
+                        cls._model_, safety_settings=safety_settings
+                    )
+
+            def __init__(self, *args, **kwargs):
+                super().__init__(*args, **kwargs)
+                self.initialize()
+
+            def get_generation_config(self) -> GenerationConfig:
+                return GenerationConfig(
+                    temperature=self.temperature,
+                    top_p=self.topP,
+                    top_k=self.topK,
+                    max_output_tokens=self.maxOutputTokens,
+                    stop_sequences=self.stopSequences,
+                )
+
             async def async_execute_model_call(
-                self, user_prompt: str, system_prompt: str = ""
-            ) -> dict[str, Any]:
-                # self.api_token = os.getenv("GOOGLE_API_KEY")
-                combined_prompt = user_prompt + system_prompt
-                url = f"https://generativelanguage.googleapis.com/v1beta/models/{model_name}:generateContent?key={self.api_token}"
-                headers = {"Content-Type": "application/json"}
-                data = {
-                    "contents": [{"parts": [{"text": combined_prompt}]}],
-                    "generationConfig": {
-                        "temperature": self.temperature,
-                        "topK": self.topK,
-                        "topP": self.topP,
-                        "maxOutputTokens": self.maxOutputTokens,
-                        "stopSequences": self.stopSequences,
-                    },
-                }
-                # print(combined_prompt)
-                async with aiohttp.ClientSession() as session:
-                    async with session.post(
-                        url, headers=headers, data=json.dumps(data)
-                    ) as response:
-                        raw_response_text = await response.text()
-                        return json.loads(raw_response_text)
+                self,
+                user_prompt: str,
+                system_prompt: str = "",
+                files_list: Optional["Files"] = None,
+            ) -> Dict[str, Any]:
+                generation_config = self.get_generation_config()
 
-        LLM.__name__ = model_name
+                if files_list is None:
+                    files_list = []
+
+                if (
+                    system_prompt is not None
+                    and system_prompt != ""
+                    and self._model_ != "gemini-pro"
+                ):
+                    try:
+                        self.generative_model = genai.GenerativeModel(
+                            self._model_,
+                            safety_settings=safety_settings,
+                            system_instruction=system_prompt,
+                        )
+                    except InvalidArgument as e:
+                        print(
+                            f"This model, {self._model_}, does not support system_instruction"
+                        )
+                        print("Will add system_prompt to user_prompt")
+                        user_prompt = f"{system_prompt}\n{user_prompt}"
 
+                combined_prompt = [user_prompt]
+                for file in files_list:
+                    if "google" not in file.external_locations:
+                        _ = file.upload_google()
+                    gen_ai_file = google.generativeai.types.file_types.File(
+                        file.external_locations["google"]
+                    )
+                    combined_prompt.append(gen_ai_file)
+
+                response = await self.generative_model.generate_content_async(
+                    combined_prompt, generation_config=generation_config
+                )
+                return response.to_dict()
+
+        LLM.__name__ = model_name
         return LLM
 
 

edsl/inference_services/MistralAIService.py

@@ -1,5 +1,5 @@
 import os
-from typing import Any, List
+from typing import Any, List, Optional
 from edsl.inference_services.InferenceServiceABC import InferenceServiceABC
 from edsl.language_models.LanguageModel import LanguageModel
 import asyncio
@@ -95,7 +95,10 @@ class MistralAIService(InferenceServiceABC):
                 return cls.async_client()
 
             async def async_execute_model_call(
-                self, user_prompt: str, system_prompt: str = ""
+                self,
+                user_prompt: str,
+                system_prompt: str = "",
+                files_list: Optional[List["FileStore"]] = None,
             ) -> dict[str, Any]:
                 """Calls the Mistral API and returns the API response."""
                 s = self.async_client()

edsl/inference_services/OpenAIService.py

@@ -168,13 +168,14 @@ class OpenAIService(InferenceServiceABC):
                 self,
                 user_prompt: str,
                 system_prompt: str = "",
-                encoded_image=None,
+                files_list: Optional[List["Files"]] = None,
                 invigilator: Optional[
                     "InvigilatorAI"
                 ] = None,  # TBD - can eventually be used for function-calling
             ) -> dict[str, Any]:
                 """Calls the OpenAI API and returns the API response."""
-                if encoded_image:
+                if files_list:
+                    encoded_image = files_list[0].base64_string
                     content = [{"type": "text", "text": user_prompt}]
                     content.append(
                         {

edsl/inference_services/TestService.py

@@ -1,4 +1,4 @@
-from typing import Any, List
+from typing import Any, List, Optional
 import os
 import asyncio
 from edsl.inference_services.InferenceServiceABC import InferenceServiceABC
@@ -59,11 +59,20 @@ class TestService(InferenceServiceABC):
                 self,
                 user_prompt: str,
                 system_prompt: str,
-                encoded_image=None,
+                # func: Optional[callable] = None,
+                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"):
                         p = self.exception_probability

edsl/inference_services/TogetherAIService.py

@@ -1,7 +1,7 @@
 import aiohttp
 import json
 import requests
-from typing import Any, List
+from typing import Any, List, Optional
 
 # from edsl.inference_services.InferenceServiceABC import InferenceServiceABC
 from edsl.language_models import LanguageModel

edsl/jobs/interviews/Interview.py

@@ -105,9 +105,9 @@ class Interview(InterviewStatusMixin):
         self.debug = debug
         self.iteration = iteration
         self.cache = cache
-        self.answers: dict[
-            str, str
-        ] = Answers()  # will get filled in as interview progresses
+        self.answers: dict[str, str] = (
+            Answers()
+        )  # will get filled in as interview progresses
         self.sidecar_model = sidecar_model
 
         # self.stop_on_exception = False
@@ -248,17 +248,24 @@ class Interview(InterviewStatusMixin):
 
     def _get_estimated_request_tokens(self, question) -> float:
         """Estimate the number of tokens that will be required to run the focal task."""
+        from edsl.scenarios.FileStore import FileStore
+
         invigilator = self._get_invigilator(question=question)
         # TODO: There should be a way to get a more accurate estimate.
         combined_text = ""
+        file_tokens = 0
         for prompt in invigilator.get_prompts().values():
             if hasattr(prompt, "text"):
                 combined_text += prompt.text
             elif isinstance(prompt, str):
                 combined_text += prompt
+            elif isinstance(prompt, list):
+                for file in prompt:
+                    if isinstance(file, FileStore):
+                        file_tokens += file.size * 0.25
             else:
                 raise ValueError(f"Prompt is of type {type(prompt)}")
-        return len(combined_text) / 4.0
+        return len(combined_text) / 4.0 + file_tokens
 
     async def _answer_question_and_record_task(
         self,
@@ -296,6 +303,9 @@ class Interview(InterviewStatusMixin):
                     self.answers.add_answer(response=response, question=question)
                     self._cancel_skipped_questions(question)
                 else:
+                    # When a question is not validated, it is not added to the answers.
+                    # this should also cancel and dependent children questions.
+                    # Is that happening now?
                     if (
                         hasattr(response, "exception_occurred")
                         and response.exception_occurred
@@ -418,11 +428,11 @@ class Interview(InterviewStatusMixin):
         """
         current_question_index: int = self.to_index[current_question.question_name]
 
-        next_question: Union[
-            int, EndOfSurvey
-        ] = self.survey.rule_collection.next_question(
-            q_now=current_question_index,
-            answers=self.answers | self.scenario | self.agent["traits"],
+        next_question: Union[int, EndOfSurvey] = (
+            self.survey.rule_collection.next_question(
+                q_now=current_question_index,
+                answers=self.answers | self.scenario | self.agent["traits"],
+            )
         )
 
         next_question_index = next_question.next_q

edsl/jobs/runners/JobsRunnerAsyncio.py

@@ -8,10 +8,10 @@ from typing import Coroutine, List, AsyncGenerator, Optional, Union, Generator
 from contextlib import contextmanager
 from collections import UserList
 
-from edsl.results.Results import Results
 from rich.live import Live
 from rich.console import Console
 
+from edsl.results.Results import Results
 from edsl import shared_globals
 from edsl.jobs.interviews.Interview import Interview
 from edsl.jobs.runners.JobsRunnerStatus import JobsRunnerStatus
@@ -22,6 +22,8 @@ from edsl.utilities.decorators import jupyter_nb_handler
 from edsl.data.Cache import Cache
 from edsl.results.Result import Result
 from edsl.results.Results import Results
+from edsl.language_models.LanguageModel import LanguageModel
+from edsl.data.Cache import Cache
 
 
 class StatusTracker(UserList):
@@ -50,10 +52,10 @@ class JobsRunnerAsyncio:
 
     async def run_async_generator(
         self,
-        cache: "Cache",
+        cache: Cache,
         n: int = 1,
         stop_on_exception: bool = False,
-        sidecar_model: Optional["LanguageModel"] = None,
+        sidecar_model: Optional[LanguageModel] = None,
         total_interviews: Optional[List["Interview"]] = None,
         raise_validation_errors: bool = False,
     ) -> AsyncGenerator["Result", None]:
@@ -104,7 +106,7 @@ class JobsRunnerAsyncio:
                     interview.cache = self.cache
                     yield interview
 
-    async def run_async(self, cache: Optional["Cache"] = None, n: int = 1) -> Results:
+    async def run_async(self, cache: Optional[Cache] = None, n: int = 1) -> Results:
         """Used for some other modules that have a non-standard way of running interviews."""
         self.jobs_runner_status = JobsRunnerStatus(self, n=n)
         self.cache = Cache() if cache is None else cache
@@ -291,6 +293,8 @@ class JobsRunnerAsyncio:
 
         self.jobs_runner_status = JobsRunnerStatus(self, n=n)
 
+        stop_event = threading.Event()
+
         async def process_results(cache):
             """Processes results from interviews."""
             async for result in self.run_async_generator(
@@ -303,20 +307,37 @@ class JobsRunnerAsyncio:
                 self.results.append(result)
             self.completed = True
 
-        def run_progress_bar():
+        def run_progress_bar(stop_event):
             """Runs the progress bar in a separate thread."""
-            self.jobs_runner_status.update_progress()
+            self.jobs_runner_status.update_progress(stop_event)
 
         if progress_bar:
-            progress_thread = threading.Thread(target=run_progress_bar)
+            progress_thread = threading.Thread(
+                target=run_progress_bar, args=(stop_event,)
+            )
             progress_thread.start()
 
-        with cache as c:
-            await process_results(cache=c)
-
-        if progress_bar:
-            progress_thread.join()
-
-        return self.process_results(
-            raw_results=self.results, cache=cache, print_exceptions=print_exceptions
-        )
+        exception_to_raise = None
+        try:
+            with cache as c:
+                await process_results(cache=c)
+        except KeyboardInterrupt:
+            print("Keyboard interrupt received. Stopping gracefully...")
+            stop_event.set()
+        except Exception as e:
+            if stop_on_exception:
+                exception_to_raise = e
+            stop_event.set()
+        finally:
+            stop_event.set()
+            if progress_bar:
+                # self.jobs_runner_status.stop_event.set()
+                if progress_thread:
+                    progress_thread.join()
+
+            if exception_to_raise:
+                raise exception_to_raise
+
+            return self.process_results(
+                raw_results=self.results, cache=cache, print_exceptions=print_exceptions
+            )

edsl/jobs/runners/JobsRunnerStatus.py

@@ -265,14 +265,15 @@ class JobsRunnerStatus:
             table.add_row(pretty_name, value)
         return table
 
-    def update_progress(self):
+    def update_progress(self, stop_event):
         layout, progress, task_ids = self.generate_layout()
 
         with Live(
             layout, refresh_per_second=int(1 / self.refresh_rate), transient=True
         ) as live:
-            while len(self.completed_interviews) < len(
-                self.jobs_runner.total_interviews
+            while (
+                len(self.completed_interviews) < len(self.jobs_runner.total_interviews)
+                and not stop_event.is_set()
             ):
                 completed_tasks = len(self.completed_interviews)
                 total_tasks = len(self.jobs_runner.total_interviews)

edsl/jobs/tasks/QuestionTaskCreator.py

@@ -156,19 +156,6 @@ class QuestionTaskCreator(UserList):
             self.tokens_bucket.turbo_mode_off()
             self.requests_bucket.turbo_mode_off()
 
-        # breakpoint()
-        # _ = results.pop("cached_response", None)
-
-        # tracker = self.cached_token_usage if self.from_cache else self.new_token_usage
-
-        # TODO: This is hacky. The 'func' call should return an object that definitely has a 'usage' key.
-        # usage = results.get("usage", {"prompt_tokens": 0, "completion_tokens": 0})
-        # prompt_tokens = usage.get("prompt_tokens", 0)
-        # completion_tokens = usage.get("completion_tokens", 0)
-        # tracker.add_tokens(
-        #    prompt_tokens=prompt_tokens, completion_tokens=completion_tokens
-        # )
-
         return results
 
     @classmethod
@@ -249,6 +236,7 @@ class QuestionTaskCreator(UserList):
                 f"Required tasks failed for {self.question.question_name}"
             ) from e
 
+        # this only runs if all the dependencies are successful
         return await self._run_focal_task()
 
 

edsl/language_models/LanguageModel.py

@@ -440,7 +440,7 @@ class LanguageModel(
         system_prompt: str,
         cache: "Cache",
         iteration: int = 0,
-        encoded_image=None,
+        files_list=None,
     ) -> ModelResponse:
         """Handle caching of responses.
 
@@ -462,16 +462,18 @@ class LanguageModel(
         >>> m._get_intended_model_call_outcome(user_prompt = "Hello", system_prompt = "hello", cache = Cache())
         ModelResponse(...)"""
 
-        if encoded_image:
-            # the image has is appended to the user_prompt for hash-lookup purposes
-            image_hash = hashlib.md5(encoded_image.encode()).hexdigest()
-            user_prompt += f" {image_hash}"
+        if files_list:
+            files_hash = "+".join([str(hash(file)) for file in files_list])
+            # print(f"Files hash: {files_hash}")
+            user_prompt_with_hashes = user_prompt + f" {files_hash}"
+        else:
+            user_prompt_with_hashes = user_prompt
 
         cache_call_params = {
             "model": str(self.model),
             "parameters": self.parameters,
             "system_prompt": system_prompt,
-            "user_prompt": user_prompt,
+            "user_prompt": user_prompt_with_hashes,
             "iteration": iteration,
         }
         cached_response, cache_key = cache.fetch(**cache_call_params)
@@ -487,7 +489,8 @@ class LanguageModel(
             params = {
                 "user_prompt": user_prompt,
                 "system_prompt": system_prompt,
-                **({"encoded_image": encoded_image} if encoded_image else {}),
+                "files_list": files_list
+                #**({"encoded_image": encoded_image} if encoded_image else {}),
             }
             # response = await f(**params)
             response = await asyncio.wait_for(f(**params), timeout=TIMEOUT)
@@ -531,7 +534,7 @@ class LanguageModel(
         system_prompt: str,
         cache: "Cache",
         iteration: int = 1,
-        encoded_image=None,
+        files_list: Optional[List['File']] = None,
     ) -> dict:
         """Get response, parse, and return as string.
 
@@ -547,7 +550,7 @@ class LanguageModel(
             "system_prompt": system_prompt,
             "iteration": iteration,
             "cache": cache,
-            **({"encoded_image": encoded_image} if encoded_image else {}),
+            "files_list": files_list,
         }
         model_inputs = ModelInputs(user_prompt=user_prompt, system_prompt=system_prompt)
         model_outputs = await self._async_get_intended_model_call_outcome(**params)

edsl/language_models/utilities.py

@@ -1,5 +1,5 @@
 import asyncio
-from typing import Any
+from typing import Any, Optional, List
 from edsl import Survey
 from edsl.config import CONFIG
 from edsl.enums import InferenceServiceType
@@ -40,7 +40,8 @@ def create_language_model(
         _tpm = 1000000000000
 
         async def async_execute_model_call(
-            self, user_prompt: str, system_prompt: str
+            self, user_prompt: str, system_prompt: str, 
+            files_list: Optional[List[Any]] = None
         ) -> dict[str, Any]:
             question_number = int(
                 user_prompt.split("XX")[1]

edsl/questions/QuestionBase.py

@@ -44,6 +44,13 @@ class QuestionBase(
     _answering_instructions = None
     _question_presentation = None
 
+    @property
+    def response_model(self) -> type["BaseModel"]:
+        if self._response_model is not None:
+            return self._response_model
+        else:
+            return self.create_response_model()
+
     # region: Validation and simulation methods
     @property
     def response_validator(self) -> "ResponseValidatorBase":
@@ -98,7 +105,9 @@ class QuestionBase(
         comment: Optional[str]
         generated_tokens: Optional[str]
 
-    def _validate_answer(self, answer: dict) -> ValidatedAnswer:
+    def _validate_answer(
+        self, answer: dict, replacement_dict: dict = None
+    ) -> ValidatedAnswer:
         """Validate the answer.
         >>> from edsl.exceptions import QuestionAnswerValidationError
         >>> from edsl import QuestionFreeText as Q
@@ -106,7 +115,7 @@ class QuestionBase(
         {'answer': 'Hello', 'generated_tokens': 'Hello'}
         """
 
-        return self.response_validator.validate(answer)
+        return self.response_validator.validate(answer, replacement_dict)
 
     # endregion
 

edsl/questions/QuestionBaseGenMixin.py

@@ -95,6 +95,34 @@ class QuestionBaseGenMixin:
             questions.append(QuestionBase.from_dict(new_data))
         return questions
 
+    def render(self, replacement_dict: dict) -> "QuestionBase":
+        """Render the question components as jinja2 templates with the replacement dictionary."""
+        from jinja2 import Environment
+        from edsl import Scenario
+
+        strings_only_replacement_dict = {
+            k: v for k, v in replacement_dict.items() if not isinstance(v, Scenario)
+        }
+
+        def render_string(value: str) -> str:
+            if value is None or not isinstance(value, str):
+                return value
+            else:
+                try:
+                    return (
+                        Environment()
+                        .from_string(value)
+                        .render(strings_only_replacement_dict)
+                    )
+                except Exception as e:
+                    import warnings
+
+                    warnings.warn("Failed to render string: " + value)
+                    # breakpoint()
+                    return value
+
+        return self.apply_function(render_string)
+
     def apply_function(self, func: Callable, exclude_components=None) -> QuestionBase:
         """Apply a function to the question parts
 

edsl/questions/QuestionCheckBox.py

@@ -245,7 +245,7 @@ class QuestionCheckBox(QuestionBase):
 
         scenario = scenario or Scenario()
         translated_options = [
-            Template(option).render(scenario) for option in self.question_options
+            Template(str(option)).render(scenario) for option in self.question_options
         ]
         translated_codes = []
         for answer_code in answer_codes:

edsl/questions/QuestionMultipleChoice.py

@@ -163,7 +163,11 @@ class QuestionMultipleChoice(QuestionBase):
     # Answer methods
     ################
 
-    def create_response_model(self):
+    def create_response_model(self, replacement_dict: dict = None):
+        if replacement_dict is None:
+            replacement_dict = {}
+            # The replacement dict that could be from scenario, current answers, etc. to populate the response model
+
         if self.use_code:
             return create_response_model(
                 list(range(len(self.question_options))), self.permissive

edsl/questions/ResponseValidatorABC.py

@@ -92,7 +92,11 @@ class ResponseValidatorABC(ABC):
         generated_tokens: Optional[str]
 
     def validate(
-        self, raw_edsl_answer_dict: RawEdslAnswerDict, fix=False, verbose=False
+        self,
+        raw_edsl_answer_dict: RawEdslAnswerDict,
+        fix=False,
+        verbose=False,
+        replacement_dict: dict = None,
     ) -> EdslAnswerDict:
         """This is the main validation function.