edsl 0.1.35__py3-none-any.whl → 0.1.36__py3-none-any.whl

edsl/inference_services/models_available_cache.py

@@ -65,7 +65,31 @@ models_available = {
         "meta-llama/Meta-Llama-3-70B-Instruct",
         "openchat/openchat_3.5",
     ],
-    "google": ["gemini-pro"],
+    "google": [
+        "gemini-1.0-pro",
+        "gemini-1.0-pro-001",
+        "gemini-1.0-pro-latest",
+        "gemini-1.0-pro-vision-latest",
+        "gemini-1.5-flash",
+        "gemini-1.5-flash-001",
+        "gemini-1.5-flash-001-tuning",
+        "gemini-1.5-flash-002",
+        "gemini-1.5-flash-8b",
+        "gemini-1.5-flash-8b-001",
+        "gemini-1.5-flash-8b-exp-0827",
+        "gemini-1.5-flash-8b-exp-0924",
+        "gemini-1.5-flash-8b-latest",
+        "gemini-1.5-flash-exp-0827",
+        "gemini-1.5-flash-latest",
+        "gemini-1.5-pro",
+        "gemini-1.5-pro-001",
+        "gemini-1.5-pro-002",
+        "gemini-1.5-pro-exp-0801",
+        "gemini-1.5-pro-exp-0827",
+        "gemini-1.5-pro-latest",
+        "gemini-pro",
+        "gemini-pro-vision",
+    ],
     "bedrock": [
         "amazon.titan-tg1-large",
         "amazon.titan-text-lite-v1",

edsl/inference_services/registry.py

@@ -11,21 +11,29 @@ from edsl.inference_services.AwsBedrock import AwsBedrockService
 from edsl.inference_services.AzureAI import AzureAIService
 from edsl.inference_services.OllamaService import OllamaService
 from edsl.inference_services.TestService import TestService
-from edsl.inference_services.MistralAIService import MistralAIService
 from edsl.inference_services.TogetherAIService import TogetherAIService
 
-default = InferenceServicesCollection(
-    [
-        OpenAIService,
-        AnthropicService,
-        DeepInfraService,
-        GoogleService,
-        GroqService,
-        AwsBedrockService,
-        AzureAIService,
-        OllamaService,
-        TestService,
-        MistralAIService,
-        TogetherAIService,
-    ]
-)
+try:
+    from edsl.inference_services.MistralAIService import MistralAIService
+
+    mistral_available = True
+except Exception as e:
+    mistral_available = False
+
+services = [
+    OpenAIService,
+    AnthropicService,
+    DeepInfraService,
+    GoogleService,
+    GroqService,
+    AwsBedrockService,
+    AzureAIService,
+    OllamaService,
+    TestService,
+    TogetherAIService,
+]
+
+if mistral_available:
+    services.append(MistralAIService)
+
+default = InferenceServicesCollection(services)

edsl/jobs/Jobs.py

@@ -1,8 +1,10 @@
 # """The Jobs class is a collection of agents, scenarios and models and one survey."""
 from __future__ import annotations
 import warnings
+import requests
 from itertools import product
 from typing import Optional, Union, Sequence, Generator
+
 from edsl.Base import Base
 from edsl.exceptions import MissingAPIKeyError
 from edsl.jobs.buckets.BucketCollection import BucketCollection
@@ -10,6 +12,9 @@ from edsl.jobs.interviews.Interview import Interview
 from edsl.jobs.runners.JobsRunnerAsyncio import JobsRunnerAsyncio
 from edsl.utilities.decorators import add_edsl_version, remove_edsl_version
 
+from edsl.data.RemoteCacheSync import RemoteCacheSync
+from edsl.exceptions.coop import CoopServerResponseError
+
 
 class Jobs(Base):
     """
@@ -180,17 +185,15 @@ class Jobs(Base):
                 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)
+
+                prompt_cost = self.estimate_prompt_cost(
+                    system_prompt=system_prompt,
+                    user_prompt=user_prompt,
+                    price_lookup=price_lookup,
+                    inference_service=invigilator.model._inference_service_,
+                    model=invigilator.model.model,
+                )
+                costs.append(prompt_cost["cost"])
 
         d = Dataset(
             [
@@ -205,59 +208,195 @@ class Jobs(Base):
             ]
         )
         return d
-        # if table:
-        #     d.to_scenario_list().print(format="rich")
-        # else:
-        #     return d
 
-    def show_prompts(self) -> None:
+    def show_prompts(self, all=False) -> None:
         """Print the prompts."""
-        self.prompts().to_scenario_list().print(format="rich")
+        if all:
+            self.prompts().to_scenario_list().print(format="rich")
+        else:
+            self.prompts().select(
+                "user_prompt", "system_prompt"
+            ).to_scenario_list().print(format="rich")
+
+    @staticmethod
+    def estimate_prompt_cost(
+        system_prompt: str,
+        user_prompt: str,
+        price_lookup: dict,
+        inference_service: str,
+        model: str,
+    ) -> dict:
+        """Estimates the cost of a prompt. Takes piping into account."""
+
+        def get_piping_multiplier(prompt: str):
+            """Returns 2 if a prompt includes Jinja braces, and 1 otherwise."""
+
+            if "{{" in prompt and "}}" in prompt:
+                return 2
+            return 1
+
+        # Look up prices per token
+        key = (inference_service, model)
+
+        try:
+            relevant_prices = price_lookup[key]
+            output_price_per_token = 1 / float(
+                relevant_prices["output"]["one_usd_buys"]
+            )
+            input_price_per_token = 1 / float(relevant_prices["input"]["one_usd_buys"])
+        except KeyError:
+            # A KeyError is likely to occur if we cannot retrieve prices (the price_lookup dict is empty)
+            # Use a sensible default
+
+            import warnings
+
+            warnings.warn(
+                "Price data could not be retrieved. Using default estimates for input and output token prices. Input: $0.15 / 1M tokens; Output: $0.60 / 1M tokens"
+            )
+
+            output_price_per_token = 0.00000015  # $0.15 / 1M tokens
+            input_price_per_token = 0.00000060  # $0.60 / 1M tokens
+
+        # Compute the number of characters (double if the question involves piping)
+        user_prompt_chars = len(str(user_prompt)) * get_piping_multiplier(
+            str(user_prompt)
+        )
+        system_prompt_chars = len(str(system_prompt)) * get_piping_multiplier(
+            str(system_prompt)
+        )
+
+        # Convert into tokens (1 token approx. equals 4 characters)
+        input_tokens = (user_prompt_chars + system_prompt_chars) // 4
+        output_tokens = input_tokens
+
+        cost = (
+            input_tokens * input_price_per_token
+            + output_tokens * output_price_per_token
+        )
+
+        return {
+            "input_tokens": input_tokens,
+            "output_tokens": output_tokens,
+            "cost": cost,
+        }
+
+    def estimate_job_cost_from_external_prices(self, price_lookup: dict) -> dict:
+        """
+        Estimates the cost of a job according to the following assumptions:
+
+        - 1 token = 4 characters.
+        - Input tokens = output tokens.
+
+        price_lookup is an external pricing dictionary.
+        """
+
+        import pandas as pd
+
+        interviews = self.interviews()
+        data = []
+        for interview in interviews:
+            invigilators = [
+                interview._get_invigilator(question)
+                for question in self.survey.questions
+            ]
+            for invigilator in invigilators:
+                prompts = invigilator.get_prompts()
 
-    def estimate_job_cost(self):
+                # By this point, agent and scenario data has already been added to the prompts
+                user_prompt = prompts["user_prompt"]
+                system_prompt = prompts["system_prompt"]
+                inference_service = invigilator.model._inference_service_
+                model = invigilator.model.model
+
+                prompt_cost = self.estimate_prompt_cost(
+                    system_prompt=system_prompt,
+                    user_prompt=user_prompt,
+                    price_lookup=price_lookup,
+                    inference_service=inference_service,
+                    model=model,
+                )
+
+                data.append(
+                    {
+                        "user_prompt": user_prompt,
+                        "system_prompt": system_prompt,
+                        "estimated_input_tokens": prompt_cost["input_tokens"],
+                        "estimated_output_tokens": prompt_cost["output_tokens"],
+                        "estimated_cost": prompt_cost["cost"],
+                        "inference_service": inference_service,
+                        "model": model,
+                    }
+                )
+
+        df = pd.DataFrame.from_records(data)
+
+        df = (
+            df.groupby(["inference_service", "model"])
+            .agg(
+                {
+                    "estimated_cost": "sum",
+                    "estimated_input_tokens": "sum",
+                    "estimated_output_tokens": "sum",
+                }
+            )
+            .reset_index()
+        )
+
+        estimated_costs_by_model = df.to_dict("records")
+
+        estimated_total_cost = sum(
+            model["estimated_cost"] for model in estimated_costs_by_model
+        )
+        estimated_total_input_tokens = sum(
+            model["estimated_input_tokens"] for model in estimated_costs_by_model
+        )
+        estimated_total_output_tokens = sum(
+            model["estimated_output_tokens"] for model in estimated_costs_by_model
+        )
+
+        output = {
+            "estimated_total_cost": estimated_total_cost,
+            "estimated_total_input_tokens": estimated_total_input_tokens,
+            "estimated_total_output_tokens": estimated_total_output_tokens,
+            "model_costs": estimated_costs_by_model,
+        }
+
+        return output
+
+    def estimate_job_cost(self) -> dict:
+        """
+        Estimates the cost of a job according to the following assumptions:
+
+        - 1 token = 4 characters.
+        - Input tokens = output tokens.
+
+        Fetches prices from Coop.
+        """
         from edsl import Coop
 
         c = Coop()
         price_lookup = c.fetch_prices()
 
-        prompts = self.prompts()
+        return self.estimate_job_cost_from_external_prices(price_lookup=price_lookup)
 
-        text_len = 0
-        for prompt in prompts:
-            text_len += len(str(prompt))
+    @staticmethod
+    def compute_job_cost(job_results: "Results") -> float:
+        """
+        Computes the cost of a completed job in USD.
+        """
+        total_cost = 0
+        for result in job_results:
+            for key in result.raw_model_response:
+                if key.endswith("_cost"):
+                    result_cost = result.raw_model_response[key]
 
-        input_token_aproximations = text_len // 4
+                    question_name = key.removesuffix("_cost")
+                    cache_used = result.cache_used_dict[question_name]
 
-        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,
-        }
+                    if isinstance(result_cost, (int, float)) and not cache_used:
+                        total_cost += result_cost
 
-        return out
+        return total_cost
 
     @staticmethod
     def _get_container_class(object):
@@ -482,7 +621,7 @@ class Jobs(Base):
 
     def _output(self, message) -> None:
         """Check if a Job is verbose. If so, print the message."""
-        if self.verbose:
+        if hasattr(self, "verbose") and self.verbose:
             print(message)
 
     def _check_parameters(self, strict=False, warn=False) -> None:
@@ -559,6 +698,123 @@ class Jobs(Base):
             return False
         return self._raise_validation_errors
 
+    def create_remote_inference_job(
+        self, iterations: int = 1, remote_inference_description: Optional[str] = None
+    ):
+        """ """
+        from edsl.coop.coop import Coop
+
+        coop = Coop()
+        self._output("Remote inference activated. Sending job to server...")
+        remote_job_creation_data = coop.remote_inference_create(
+            self,
+            description=remote_inference_description,
+            status="queued",
+            iterations=iterations,
+        )
+        job_uuid = remote_job_creation_data.get("uuid")
+        print(f"Job sent to server. (Job uuid={job_uuid}).")
+        return remote_job_creation_data
+
+    @staticmethod
+    def check_status(job_uuid):
+        from edsl.coop.coop import Coop
+
+        coop = Coop()
+        return coop.remote_inference_get(job_uuid)
+
+    def poll_remote_inference_job(
+        self, remote_job_creation_data: dict
+    ) -> Union[Results, None]:
+        from edsl.coop.coop import Coop
+        import time
+        from datetime import datetime
+        from edsl.config import CONFIG
+
+        expected_parrot_url = CONFIG.get("EXPECTED_PARROT_URL")
+
+        job_uuid = remote_job_creation_data.get("uuid")
+
+        coop = Coop()
+        job_in_queue = True
+        while job_in_queue:
+            remote_job_data = coop.remote_inference_get(job_uuid)
+            status = remote_job_data.get("status")
+            if status == "cancelled":
+                print("\r" + " " * 80 + "\r", end="")
+                print("Job cancelled by the user.")
+                print(
+                    f"See {expected_parrot_url}/home/remote-inference for more details."
+                )
+                return None
+            elif status == "failed":
+                print("\r" + " " * 80 + "\r", end="")
+                print("Job failed.")
+                print(
+                    f"See {expected_parrot_url}/home/remote-inference for more details."
+                )
+                return None
+            elif status == "completed":
+                results_uuid = remote_job_data.get("results_uuid")
+                results = coop.get(results_uuid, expected_object_type="results")
+                print("\r" + " " * 80 + "\r", end="")
+                url = f"{expected_parrot_url}/content/{results_uuid}"
+                print(f"Job completed and Results stored on Coop: {url}.")
+                return results
+            else:
+                duration = 5
+                time_checked = datetime.now().strftime("%Y-%m-%d %I:%M:%S %p")
+                frames = ["⠋", "⠙", "⠹", "⠸", "⠼", "⠴", "⠦", "⠧", "⠇", "⠏"]
+                start_time = time.time()
+                i = 0
+                while time.time() - start_time < duration:
+                    print(
+                        f"\r{frames[i % len(frames)]} Job status: {status} - last update: {time_checked}",
+                        end="",
+                        flush=True,
+                    )
+                    time.sleep(0.1)
+                    i += 1
+
+    def use_remote_inference(self, disable_remote_inference: bool):
+        if disable_remote_inference:
+            return False
+        if not disable_remote_inference:
+            try:
+                from edsl import Coop
+
+                user_edsl_settings = Coop().edsl_settings
+                return user_edsl_settings.get("remote_inference", False)
+            except requests.ConnectionError:
+                pass
+            except CoopServerResponseError as e:
+                pass
+
+        return False
+
+    def use_remote_cache(self):
+        try:
+            from edsl import Coop
+
+            user_edsl_settings = Coop().edsl_settings
+            return user_edsl_settings.get("remote_caching", False)
+        except requests.ConnectionError:
+            pass
+        except CoopServerResponseError as e:
+            pass
+
+        return False
+
+    def check_api_keys(self):
+        from edsl import Model
+
+        for model in self.models + [Model()]:
+            if not model.has_valid_api_key():
+                raise MissingAPIKeyError(
+                    model_name=str(model.model),
+                    inference_service=model._inference_service_,
+                )
+
     def run(
         self,
         n: int = 1,
@@ -596,91 +852,17 @@ class Jobs(Base):
 
         self.verbose = verbose
 
-        remote_cache = False
-        remote_inference = False
-
-        if not disable_remote_inference:
-            try:
-                coop = Coop()
-                user_edsl_settings = Coop().edsl_settings
-                remote_cache = user_edsl_settings.get("remote_caching", False)
-                remote_inference = user_edsl_settings.get("remote_inference", False)
-            except Exception:
-                pass
-
-        if remote_inference:
-            import time
-            from datetime import datetime
-            from edsl.config import CONFIG
-
-            expected_parrot_url = CONFIG.get("EXPECTED_PARROT_URL")
-
-            self._output("Remote inference activated. Sending job to server...")
-            if remote_cache:
-                self._output(
-                    "Remote caching activated. The remote cache will be used for this job."
-                )
-
-            remote_job_creation_data = coop.remote_inference_create(
-                self,
-                description=remote_inference_description,
-                status="queued",
-                iterations=n,
+        if remote_inference := self.use_remote_inference(disable_remote_inference):
+            remote_job_creation_data = self.create_remote_inference_job(
+                iterations=n, remote_inference_description=remote_inference_description
             )
-            time_queued = datetime.now().strftime("%m/%d/%Y %I:%M:%S %p")
-            job_uuid = remote_job_creation_data.get("uuid")
-            print(f"Remote inference started (Job uuid={job_uuid}).")
-            # print(f"Job queued at {time_queued}.")
-            job_in_queue = True
-            while job_in_queue:
-                remote_job_data = coop.remote_inference_get(job_uuid)
-                status = remote_job_data.get("status")
-                if status == "cancelled":
-                    print("\r" + " " * 80 + "\r", end="")
-                    print("Job cancelled by the user.")
-                    print(
-                        f"See {expected_parrot_url}/home/remote-inference for more details."
-                    )
-                    return None
-                elif status == "failed":
-                    print("\r" + " " * 80 + "\r", end="")
-                    print("Job failed.")
-                    print(
-                        f"See {expected_parrot_url}/home/remote-inference for more details."
-                    )
-                    return None
-                elif status == "completed":
-                    results_uuid = remote_job_data.get("results_uuid")
-                    results = coop.get(results_uuid, expected_object_type="results")
-                    print("\r" + " " * 80 + "\r", end="")
-                    print(
-                        f"Job completed and Results stored on Coop (Results uuid={results_uuid})."
-                    )
-                    return results
-                else:
-                    duration = 5
-                    time_checked = datetime.now().strftime("%Y-%m-%d %I:%M:%S %p")
-                    frames = ["⠋", "⠙", "⠹", "⠸", "⠼", "⠴", "⠦", "⠧", "⠇", "⠏"]
-                    start_time = time.time()
-                    i = 0
-                    while time.time() - start_time < duration:
-                        print(
-                            f"\r{frames[i % len(frames)]} Job status: {status} - last update: {time_checked}",
-                            end="",
-                            flush=True,
-                        )
-                        time.sleep(0.1)
-                        i += 1
-        else:
-            if check_api_keys:
-                from edsl import Model
+            results = self.poll_remote_inference_job(remote_job_creation_data)
+            if results is None:
+                self._output("Job failed.")
+            return results
 
-                for model in self.models + [Model()]:
-                    if not model.has_valid_api_key():
-                        raise MissingAPIKeyError(
-                            model_name=str(model.model),
-                            inference_service=model._inference_service_,
-                        )
+        if check_api_keys:
+            self.check_api_keys()
 
         # handle cache
         if cache is None or cache is True:
@@ -692,51 +874,14 @@ class Jobs(Base):
 
             cache = Cache()
 
-        if not remote_cache:
-            results = self._run_local(
-                n=n,
-                progress_bar=progress_bar,
-                cache=cache,
-                stop_on_exception=stop_on_exception,
-                sidecar_model=sidecar_model,
-                print_exceptions=print_exceptions,
-                raise_validation_errors=raise_validation_errors,
-            )
-
-            results.cache = cache.new_entries_cache()
-
-            self._output(f"There are {len(cache.keys()):,} entries in the local cache.")
-        else:
-            cache_difference = coop.remote_cache_get_diff(cache.keys())
-
-            client_missing_cacheentries = cache_difference.get(
-                "client_missing_cacheentries", []
-            )
-
-            missing_entry_count = len(client_missing_cacheentries)
-            if missing_entry_count > 0:
-                self._output(
-                    f"Updating local cache with {missing_entry_count:,} new "
-                    f"{'entry' if missing_entry_count == 1 else 'entries'} from remote..."
-                )
-                cache.add_from_dict(
-                    {entry.key: entry for entry in client_missing_cacheentries}
-                )
-                self._output("Local cache updated!")
-            else:
-                self._output("No new entries to add to local cache.")
-
-            server_missing_cacheentry_keys = cache_difference.get(
-                "server_missing_cacheentry_keys", []
-            )
-            server_missing_cacheentries = [
-                entry
-                for key in server_missing_cacheentry_keys
-                if (entry := cache.data.get(key)) is not None
-            ]
-            old_entry_keys = [key for key in cache.keys()]
-
-            self._output("Running job...")
+        remote_cache = self.use_remote_cache()
+        with RemoteCacheSync(
+            coop=Coop(),
+            cache=cache,
+            output_func=self._output,
+            remote_cache=remote_cache,
+            remote_cache_description=remote_cache_description,
+        ) as r:
             results = self._run_local(
                 n=n,
                 progress_bar=progress_bar,
@@ -746,32 +891,8 @@ class Jobs(Base):
                 print_exceptions=print_exceptions,
                 raise_validation_errors=raise_validation_errors,
             )
-            self._output("Job completed!")
-
-            new_cache_entries = list(
-                [entry for entry in cache.values() if entry.key not in old_entry_keys]
-            )
-            server_missing_cacheentries.extend(new_cache_entries)
-
-            new_entry_count = len(server_missing_cacheentries)
-            if new_entry_count > 0:
-                self._output(
-                    f"Updating remote cache with {new_entry_count:,} new "
-                    f"{'entry' if new_entry_count == 1 else 'entries'}..."
-                )
-                coop.remote_cache_create_many(
-                    server_missing_cacheentries,
-                    visibility="private",
-                    description=remote_cache_description,
-                )
-                self._output("Remote cache updated!")
-            else:
-                self._output("No new entries to add to remote cache.")
-
-            results.cache = cache.new_entries_cache()
-
-            self._output(f"There are {len(cache.keys()):,} entries in the local cache.")
 
+        results.cache = cache.new_entries_cache()
         return results
 
     def _run_local(self, *args, **kwargs):