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

edsl/inference_services/AzureAI.py

@@ -1,5 +1,5 @@
 import os
-from typing import Any
+from typing import Any, Optional, List
 import re
 from openai import AsyncAzureOpenAI
 from edsl.inference_services.InferenceServiceABC import InferenceServiceABC
@@ -122,7 +122,10 @@ class AzureAIService(InferenceServiceABC):
             _tpm = cls.get_tpm(cls)
 
             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 Azure OpenAI API and returns the API response."""
 

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/Jobs.py

@@ -145,14 +145,21 @@ class Jobs(Base):
         >>> Jobs.example().prompts()
         Dataset(...)
         """
+        from edsl import Coop
+
+        c = Coop()
+        price_lookup = c.fetch_prices()
 
         interviews = self.interviews()
         # data = []
         interview_indices = []
-        question_indices = []
+        question_names = []
         user_prompts = []
         system_prompts = []
         scenario_indices = []
+        agent_indices = []
+        models = []
+        costs = []
         from edsl.results.Dataset import Dataset
 
         for interview_index, interview in enumerate(interviews):
@@ -160,23 +167,97 @@ class Jobs(Base):
                 interview._get_invigilator(question)
                 for question in self.survey.questions
             ]
-            # list(interview._build_invigilators(debug=False))
             for _, invigilator in enumerate(invigilators):
                 prompts = invigilator.get_prompts()
-                user_prompts.append(prompts["user_prompt"])
-                system_prompts.append(prompts["system_prompt"])
+                user_prompt = prompts["user_prompt"]
+                system_prompt = prompts["system_prompt"]
+                user_prompts.append(user_prompt)
+                system_prompts.append(system_prompt)
+                agent_index = self.agents.index(invigilator.agent)
+                agent_indices.append(agent_index)
                 interview_indices.append(interview_index)
-                scenario_indices.append(invigilator.scenario)
-                question_indices.append(invigilator.question.question_name)
-        return Dataset(
+                scenario_index = self.scenarios.index(invigilator.scenario)
+                scenario_indices.append(scenario_index)
+                models.append(invigilator.model.model)
+                question_names.append(invigilator.question.question_name)
+                # cost calculation
+                key = (invigilator.model._inference_service_, invigilator.model.model)
+                relevant_prices = price_lookup[key]
+                inverse_output_price = relevant_prices["output"]["one_usd_buys"]
+                inverse_input_price = relevant_prices["input"]["one_usd_buys"]
+                input_tokens = len(str(user_prompt) + str(system_prompt)) // 4
+                output_tokens = len(str(user_prompt) + str(system_prompt)) // 4
+                cost = input_tokens / float(
+                    inverse_input_price
+                ) + output_tokens / float(inverse_output_price)
+                costs.append(cost)
+
+        d = Dataset(
             [
-                {"interview_index": interview_indices},
-                {"question_index": question_indices},
                 {"user_prompt": user_prompts},
-                {"scenario_index": scenario_indices},
                 {"system_prompt": system_prompts},
+                {"interview_index": interview_indices},
+                {"question_name": question_names},
+                {"scenario_index": scenario_indices},
+                {"agent_index": agent_indices},
+                {"model": models},
+                {"estimated_cost": costs},
             ]
         )
+        return d
+        # if table:
+        #     d.to_scenario_list().print(format="rich")
+        # else:
+        #     return d
+
+    def show_prompts(self) -> None:
+        """Print the prompts."""
+        self.prompts().to_scenario_list().print(format="rich")
+
+    def estimate_job_cost(self):
+        from edsl import Coop
+
+        c = Coop()
+        price_lookup = c.fetch_prices()
+
+        prompts = self.prompts()
+
+        text_len = 0
+        for prompt in prompts:
+            text_len += len(str(prompt))
+
+        input_token_aproximations = text_len // 4
+
+        aproximation_cost = {}
+        total_cost = 0
+        for model in self.models:
+            key = (model._inference_service_, model.model)
+            relevant_prices = price_lookup[key]
+            inverse_output_price = relevant_prices["output"]["one_usd_buys"]
+            inverse_input_price = relevant_prices["input"]["one_usd_buys"]
+
+            aproximation_cost[key] = {
+                "input": input_token_aproximations / float(inverse_input_price),
+                "output": input_token_aproximations / float(inverse_output_price),
+            }
+            ##TODO curenlty we approximate the number of output tokens with the number
+            # of input tokens. A better solution will be to compute the quesiton answer options length and sum them
+            # to compute the output tokens
+
+            total_cost += input_token_aproximations / float(inverse_input_price)
+            total_cost += input_token_aproximations / float(inverse_output_price)
+
+        # multiply_factor = len(self.agents or [1]) * len(self.scenarios or [1])
+        multiply_factor = 1
+        out = {
+            "input_token_aproximations": input_token_aproximations,
+            "models_costs": aproximation_cost,
+            "estimated_total_cost": total_cost * multiply_factor,
+            "multiply_factor": multiply_factor,
+            "single_config_cost": total_cost,
+        }
+
+        return out
 
     @staticmethod
     def _get_container_class(object):

edsl/jobs/interviews/Interview.py

@@ -3,6 +3,7 @@
 from __future__ import annotations
 import asyncio
 from typing import Any, Type, List, Generator, Optional, Union
+import copy
 
 from tenacity import (
     retry,
@@ -99,7 +100,9 @@ class Interview(InterviewStatusMixin):
 
         """
         self.agent = agent
-        self.survey = survey
+        # what I would like to do
+        self.survey = copy.deepcopy(survey)  # survey  copy.deepcopy(survey)
+        # self.survey = survey
         self.scenario = scenario
         self.model = model
         self.debug = debug
@@ -248,17 +251,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 +306,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

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
@@ -171,19 +173,19 @@ class JobsRunnerAsyncio:
 
         prompt_dictionary = {}
         for answer_key_name in answer_key_names:
-            prompt_dictionary[answer_key_name + "_user_prompt"] = (
-                question_name_to_prompts[answer_key_name]["user_prompt"]
-            )
-            prompt_dictionary[answer_key_name + "_system_prompt"] = (
-                question_name_to_prompts[answer_key_name]["system_prompt"]
-            )
+            prompt_dictionary[
+                answer_key_name + "_user_prompt"
+            ] = question_name_to_prompts[answer_key_name]["user_prompt"]
+            prompt_dictionary[
+                answer_key_name + "_system_prompt"
+            ] = question_name_to_prompts[answer_key_name]["system_prompt"]
 
         raw_model_results_dictionary = {}
         for result in valid_results:
             question_name = result.question_name
-            raw_model_results_dictionary[question_name + "_raw_model_response"] = (
-                result.raw_model_response
-            )
+            raw_model_results_dictionary[
+                question_name + "_raw_model_response"
+            ] = result.raw_model_response
             raw_model_results_dictionary[question_name + "_cost"] = result.cost
             one_use_buys = (
                 "NA"
@@ -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()