edsl 0.1.47__py3-none-any.whl → 0.1.49__py3-none-any.whl

edsl/base/data_transfer_models.py

@@ -0,0 +1,114 @@
+from collections import UserDict
+from typing import NamedTuple, Dict, List, Optional, Any
+from dataclasses import dataclass, fields
+
+
+class ModelInputs(NamedTuple):
+    "This is what was send by the agent to the model"
+    user_prompt: str
+    system_prompt: str
+    encoded_image: Optional[str] = None
+
+
+class EDSLOutput(NamedTuple):
+    "This is the edsl dictionary that is returned by the model"
+    answer: Any
+    generated_tokens: str
+    comment: Optional[str] = None
+
+
+class ModelResponse(NamedTuple):
+    "This is the metadata that is returned by the model and includes info about the cache"
+    response: dict
+    cache_used: bool
+    cache_key: str
+    cached_response: Optional[Dict[str, Any]] = None
+    cost: Optional[float] = None
+
+
+class AgentResponseDict(NamedTuple):
+    edsl_dict: EDSLOutput
+    model_inputs: ModelInputs
+    model_outputs: ModelResponse
+
+
+class EDSLResultObjectInput(NamedTuple):
+    generated_tokens: str
+    question_name: str
+    prompts: dict
+    cached_response: str
+    raw_model_response: str
+    cache_used: bool
+    cache_key: str
+    answer: Any
+    comment: str
+    validated: bool = False
+    exception_occurred: Exception = None
+    cost: Optional[float] = None
+
+
+@dataclass
+class ImageInfo:
+    file_path: str
+    file_name: str
+    image_format: str
+    file_size: int
+    encoded_image: str
+
+    def __repr__(self):
+        import reprlib
+
+        reprlib_instance = reprlib.Repr()
+        reprlib_instance.maxstring = 30  # Limit the string length for the encoded image
+
+        # Get all fields except encoded_image
+        field_reprs = [
+            f"{f.name}={getattr(self, f.name)!r}"
+            for f in fields(self)
+            if f.name != "encoded_image"
+        ]
+
+        # Add the reprlib-restricted encoded_image field
+        field_reprs.append(f"encoded_image={reprlib_instance.repr(self.encoded_image)}")
+
+        # Join everything to create the repr
+        return f"{self.__class__.__name__}({', '.join(field_reprs)})"
+
+
+class Answers(UserDict):
+    """Helper class to hold the answers to a survey."""
+
+    def add_answer(
+        self, response: EDSLResultObjectInput, question: "QuestionBase"
+    ) -> None:
+        """Add a response to the answers dictionary."""
+        answer = response.answer
+        comment = response.comment
+        generated_tokens = response.generated_tokens
+        # record the answer
+        if generated_tokens:
+            self[question.question_name + "_generated_tokens"] = generated_tokens
+        self[question.question_name] = answer
+        if comment:
+            self[question.question_name + "_comment"] = comment
+
+    def replace_missing_answers_with_none(self, survey: "Survey") -> None:
+        """Replace missing answers with None. Answers can be missing if the agent skips a question."""
+        for question_name in survey.question_names:
+            if question_name not in self:
+                self[question_name] = None
+
+    def to_dict(self):
+        """Return a dictionary of the answers."""
+        return self.data
+
+    @classmethod
+    def from_dict(cls, d):
+        """Return an Answers object from a dictionary."""
+        return cls(d)
+
+
+if __name__ == "__main__":
+    import doctest
+
+    doctest.testmod()

edsl/base/enums.py

@@ -0,0 +1,215 @@
+"""Enums for the different types of questions, language models, and inference services."""
+
+from enum import Enum
+from typing import Literal
+
+
+class EnumWithChecks(Enum):
+    """Base class for all enums with checks."""
+
+    @classmethod
+    def is_value_valid(cls, value):
+        """Check if the value is valid."""
+        return any(value == item.value for item in cls)
+
+
+class QuestionType(EnumWithChecks):
+    """Enum for the question types."""
+
+    MULTIPLE_CHOICE = "multiple_choice"
+    YES_NO = "yes_no"
+    FREE_TEXT = "free_text"
+    RANK = "rank"
+    BUDGET = "budget"
+    CHECKBOX = "checkbox"
+    EXTRACT = "extract"
+    FUNCTIONAL = "functional"
+    LIST = "list"
+    NUMERICAL = "numerical"
+    TOP_K = "top_k"
+    LIKERT_FIVE = "likert_five"
+    LINEAR_SCALE = "linear_scale"
+
+
+# https://huggingface.co/meta-llama/Llama-2-70b-chat-hf
+
+
+# class LanguageModelType(EnumWithChecks):
+#     """Enum for the language model types."""
+
+#     GPT_4 = "gpt-4-1106-preview"
+#     GPT_3_5_Turbo = "gpt-3.5-turbo"
+#     LLAMA_2_70B_CHAT_HF = "llama-2-70b-chat-hf"
+#     LLAMA_2_13B_CHAT_HF = "llama-2-13b-chat-hf"
+#     GEMINI_PRO = "gemini_pro"
+#     MIXTRAL_8x7B_INSTRUCT = "mixtral-8x7B-instruct-v0.1"
+#     TEST = "test"
+#     ANTHROPIC_3_OPUS = "claude-3-opus-20240229"
+#     ANTHROPIC_3_SONNET = "claude-3-sonnet-20240229"
+#     ANTHROPIC_3_HAIKU = "claude-3-haiku-20240307"
+#     DBRX_INSTRUCT = "dbrx-instruct"
+
+
+class InferenceServiceType(EnumWithChecks):
+    """Enum for the inference service types."""
+
+    BEDROCK = "bedrock"
+    DEEP_INFRA = "deep_infra"
+    REPLICATE = "replicate"
+    OPENAI = "openai"
+    GOOGLE = "google"
+    TEST = "test"
+    ANTHROPIC = "anthropic"
+    GROQ = "groq"
+    AZURE = "azure"
+    OLLAMA = "ollama"
+    MISTRAL = "mistral"
+    TOGETHER = "together"
+    PERPLEXITY = "perplexity"
+    DEEPSEEK = "deepseek"
+    XAI = "xai"
+
+
+# unavoidable violation of the DRY principle but it is necessary
+# checked w/ a unit test to make sure consistent with services in enums.py
+InferenceServiceLiteral = Literal[
+    "bedrock",
+    "deep_infra",
+    "replicate",
+    "openai",
+    "google",
+    "test",
+    "anthropic",
+    "groq",
+    "azure",
+    "ollama",
+    "mistral",
+    "together",
+    "perplexity",
+    "deepseek",
+    "xai",
+]
+
+available_models_urls = {
+    "anthropic": "https://docs.anthropic.com/en/docs/about-claude/models",
+    "openai": "https://platform.openai.com/docs/models/gp",
+    "groq": "https://console.groq.com/docs/models",
+    "google": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models",
+}
+
+
+service_to_api_keyname = {
+    InferenceServiceType.DEEP_INFRA.value: "DEEP_INFRA_API_KEY",
+    InferenceServiceType.REPLICATE.value: "TBD",
+    InferenceServiceType.OPENAI.value: "OPENAI_API_KEY",
+    InferenceServiceType.GOOGLE.value: "GOOGLE_API_KEY",
+    InferenceServiceType.TEST.value: "TBD",
+    InferenceServiceType.ANTHROPIC.value: "ANTHROPIC_API_KEY",
+    InferenceServiceType.GROQ.value: "GROQ_API_KEY",
+    InferenceServiceType.BEDROCK.value: ["AWS_ACCESS_KEY_ID", "AWS_SECRET_ACCESS_KEY"],
+    InferenceServiceType.MISTRAL.value: "MISTRAL_API_KEY",
+    InferenceServiceType.TOGETHER.value: "TOGETHER_API_KEY",
+    InferenceServiceType.PERPLEXITY.value: "PERPLEXITY_API_KEY",
+    InferenceServiceType.DEEPSEEK.value: "DEEPSEEK_API_KEY",
+    InferenceServiceType.XAI.value: "XAI_API_KEY",
+}
+
+
+class TokenPricing:
+    def __init__(
+        self,
+        *,
+        model_name,
+        prompt_token_price_per_k: float,
+        completion_token_price_per_k: float,
+    ):
+        self.model_name = model_name
+        self.prompt_token_price = prompt_token_price_per_k / 1_000.0
+        self.completion_token_price = completion_token_price_per_k / 1_000.0
+
+    def __eq__(self, other):
+        if not isinstance(other, TokenPricing):
+            return False
+        return (
+            self.model_name == other.model_name
+            and self.prompt_token_price == other.prompt_token_price
+            and self.completion_token_price == other.completion_token_price
+        )
+    
+    @classmethod
+    def example(cls) -> "TokenPricing":
+        """Return an example TokenPricing object."""
+        return cls(
+            model_name="fake_model",
+            prompt_token_price_per_k=0.01,
+            completion_token_price_per_k=0.03,
+        )
+
+pricing = {
+    "dbrx-instruct": TokenPricing(
+        model_name="dbrx-instruct",
+        prompt_token_price_per_k=0.0,
+        completion_token_price_per_k=0.0,
+    ),
+    "claude-3-opus-20240229": TokenPricing(
+        model_name="claude-3-opus-20240229",
+        prompt_token_price_per_k=0.0,
+        completion_token_price_per_k=0.0,
+    ),
+    "claude-3-haiku-20240307": TokenPricing(
+        model_name="claude-3-haiku-20240307",
+        prompt_token_price_per_k=0.0,
+        completion_token_price_per_k=0.0,
+    ),
+    "claude-3-sonnet-20240229": TokenPricing(
+        model_name="claude-3-sonnet-20240229",
+        prompt_token_price_per_k=0.0,
+        completion_token_price_per_k=0.0,
+    ),
+    "gpt-3.5-turbo": TokenPricing(
+        model_name="gpt-3.5-turbo",
+        prompt_token_price_per_k=0.0005,
+        completion_token_price_per_k=0.0015,
+    ),
+    "gpt-4-1106-preview": TokenPricing(
+        model_name="gpt-4",
+        prompt_token_price_per_k=0.01,
+        completion_token_price_per_k=0.03,
+    ),
+    "test": TokenPricing(
+        model_name="test",
+        prompt_token_price_per_k=0.0,
+        completion_token_price_per_k=0.0,
+    ),
+    "gemini_pro": TokenPricing(
+        model_name="gemini_pro",
+        prompt_token_price_per_k=0.0,
+        completion_token_price_per_k=0.0,
+    ),
+    "llama-2-13b-chat-hf": TokenPricing(
+        model_name="llama-2-13b-chat-hf",
+        prompt_token_price_per_k=0.0,
+        completion_token_price_per_k=0.0,
+    ),
+    "llama-2-70b-chat-hf": TokenPricing(
+        model_name="llama-2-70b-chat-hf",
+        prompt_token_price_per_k=0.0,
+        completion_token_price_per_k=0.0,
+    ),
+    "mixtral-8x7B-instruct-v0.1": TokenPricing(
+        model_name="mixtral-8x7B-instruct-v0.1",
+        prompt_token_price_per_k=0.0,
+        completion_token_price_per_k=0.0,
+    ),
+}
+
+
+def get_token_pricing(model_name):
+    if model_name in pricing:
+        return pricing[model_name]
+    else:
+        return TokenPricing(
+            model_name=model_name,
+            prompt_token_price_per_k=0.0,
+            completion_token_price_per_k=0.0,
+        )

edsl/base.py

@@ -0,0 +1,8 @@
+"""Base module wrapper for backward compatibility.
+
+This module re-exports everything from the edsl.base module.
+"""
+
+# Re-export everything from edsl.base module
+from edsl.base import *
+from edsl.base.base_class import *

edsl/buckets/__init__.py

@@ -0,0 +1,25 @@
+"""
+Buckets module for managing rate limits of language model API requests.
+
+This module provides a robust rate-limiting system for language model API calls,
+implementing the token bucket algorithm to manage both requests-per-minute and
+tokens-per-minute limits. It supports both local (in-process) and remote
+(distributed) rate limiting through a client-server architecture.
+
+Key components:
+- TokenBucket: Core rate-limiting class implementing the token bucket algorithm
+- ModelBuckets: Manages rate limits for a specific language model, containing
+  separate buckets for requests and tokens
+- BucketCollection: Manages multiple ModelBuckets instances across different
+  language model services
+
+The module also includes a FastAPI server implementation (token_bucket_api) and
+client (token_bucket_client) for distributed rate limiting scenarios where
+multiple processes or machines need to share rate limits.
+"""
+
+from .bucket_collection import BucketCollection
+from .model_buckets import ModelBuckets
+from .token_bucket import TokenBucket
+
+__all__ = ["BucketCollection", "ModelBuckets", "TokenBucket"]

edsl/buckets/bucket_collection.py

@@ -0,0 +1,324 @@
+"""
+BucketCollection module for managing rate limits across multiple language models.
+
+This module provides the BucketCollection class, which manages rate limits for
+multiple language models, organizing them by service provider. It ensures that
+API rate limits are respected while allowing models from the same service to
+share the same rate limit buckets.
+"""
+
+from typing import Optional, TYPE_CHECKING, Dict, List, Any, Tuple
+from collections import UserDict
+from threading import RLock
+import matplotlib.pyplot as plt
+from matplotlib.figure import Figure
+
+from .token_bucket import TokenBucket
+from .model_buckets import ModelBuckets
+from ..jobs.decorators import synchronized_class
+
+if TYPE_CHECKING:
+    from ..language_models import LanguageModel
+    from ..key_management import KeyLookup
+    
+@synchronized_class
+class BucketCollection(UserDict):
+    """
+    Collection of ModelBuckets for managing rate limits across multiple language models.
+    
+    BucketCollection is a thread-safe dictionary-like container that maps language models
+    to their corresponding ModelBuckets objects. It helps manage rate limits for multiple
+    models, organizing them by service provider to ensure that API rate limits are
+    respected across all models using the same service.
+    
+    The class maps models to services, and services to buckets, allowing models from
+    the same service to share rate limit buckets. This approach ensures accurate
+    rate limiting when multiple models use the same underlying service.
+    
+    Attributes:
+        infinity_buckets (bool): If True, all buckets have infinite capacity and refill rate
+        models_to_services (dict): Maps model names to their service provider names
+        services_to_buckets (dict): Maps service names to their ModelBuckets instances
+        remote_url (str, optional): URL for remote token bucket server if using distributed mode
+        
+    Example:
+        >>> from edsl import Model
+        >>> bucket_collection = BucketCollection()
+        >>> model = Model('gpt-4')
+        >>> bucket_collection.add_model(model)
+        >>> # Now rate limits for the model are being tracked
+    """
+
+    def __init__(self, infinity_buckets: bool = False):
+        """
+        Initialize a new BucketCollection.
+        
+        Creates a new BucketCollection to manage rate limits across multiple language
+        models. If infinity_buckets is True, all buckets will have unlimited capacity
+        and refill rate, effectively bypassing rate limits.
+        
+        Args:
+            infinity_buckets: If True, creates buckets with unlimited capacity
+                             and refill rate (default: False)
+                             
+        Example:
+            >>> # Create a standard bucket collection with rate limiting
+            >>> bucket_collection = BucketCollection()
+            >>> # Create a bucket collection with unlimited capacity (for testing)
+            >>> unlimited_collection = BucketCollection(infinity_buckets=True)
+        """
+        super().__init__()
+        self.infinity_buckets = infinity_buckets
+        self.models_to_services = {}  # Maps model names to service names
+        self.services_to_buckets = {} # Maps service names to ModelBuckets
+        self._lock = RLock()
+
+        # Check for remote token bucket server URL in environment
+        from edsl.config import CONFIG
+        import os
+
+        url = os.environ.get("EDSL_REMOTE_TOKEN_BUCKET_URL", None)
+
+        if url == "None" or url is None:
+            self.remote_url = None
+        else:
+            self.remote_url = url
+
+    @classmethod
+    def from_models(
+        cls, models_list: List["LanguageModel"], infinity_buckets: bool = False
+    ) -> "BucketCollection":
+        """
+        Create a BucketCollection pre-populated with a list of models.
+        
+        This factory method creates a new BucketCollection and adds multiple
+        models to it at initialization time.
+        
+        Args:
+            models_list: List of LanguageModel instances to add to the collection
+            infinity_buckets: If True, creates buckets with unlimited capacity
+                             and refill rate (default: False)
+                             
+        Returns:
+            A new BucketCollection containing the specified models
+            
+        Example:
+            >>> from edsl import Model
+            >>> models = [Model('gpt-4'), Model('gpt-3.5-turbo')]
+            >>> collection = BucketCollection.from_models(models)
+        """
+        bucket_collection = cls(infinity_buckets=infinity_buckets)
+        for model in models_list:
+            bucket_collection.add_model(model)
+        return bucket_collection
+
+    def get_tokens(
+        self, model: 'LanguageModel', bucket_type: str, num_tokens: int
+    ) -> int:
+        """
+        [DEPRECATED] Get the number of tokens remaining in the bucket.
+        
+        This method is deprecated and will raise an exception if called.
+        It is kept for reference purposes only.
+        
+        Args:
+            model: The language model to get tokens for
+            bucket_type: The type of bucket ('requests' or 'tokens')
+            num_tokens: The number of tokens to retrieve
+            
+        Raises:
+            Exception: This method is deprecated
+            
+        Example:
+            >>> bucket_collection = BucketCollection()
+            >>> from edsl import Model 
+            >>> m = Model('test')
+            >>> bucket_collection.add_model(m)
+            >>> # The following would raise an exception:
+            >>> # bucket_collection.get_tokens(m, 'tokens', 10)
+        """
+        raise Exception("This method is deprecated and should not be used")
+        # The following code is kept for reference only
+        # relevant_bucket = getattr(self[model], bucket_type)
+        # return relevant_bucket.get_tokens(num_tokens)
+
+    def __repr__(self) -> str:
+        """
+        Generate a string representation of the BucketCollection.
+        
+        Returns:
+            String representation showing the collection's contents
+        """
+        return f"BucketCollection({self.data})"
+
+    def add_model(self, model: "LanguageModel") -> None:
+        """
+        Add a language model to the bucket collection.
+        
+        This method adds a language model to the BucketCollection, creating the
+        necessary token buckets for its service provider if they don't already exist.
+        Models from the same service share the same buckets.
+        
+        Args:
+            model: The LanguageModel instance to add to the collection
+            
+        Example:
+            >>> from edsl import Model 
+            >>> model = Model('gpt-4')
+            >>> bucket_collection = BucketCollection()
+            >>> bucket_collection.add_model(model)
+        """
+        # Calculate tokens-per-second (TPS) and requests-per-second (RPS) rates
+        if not self.infinity_buckets:
+            TPS = model.tpm / 60.0  # Convert tokens-per-minute to tokens-per-second
+            RPS = model.rpm / 60.0  # Convert requests-per-minute to requests-per-second
+        else:
+            TPS = float("inf")  # Infinite tokens per second
+            RPS = float("inf")  # Infinite requests per second
+
+        # If this is a new model we haven't seen before
+        if model.model not in self.models_to_services:
+            service = model._inference_service_
+            
+            # If this is a new service we haven't created buckets for yet
+            if service not in self.services_to_buckets:
+                # Create request rate limiting bucket
+                requests_bucket = TokenBucket(
+                    bucket_name=service,
+                    bucket_type="requests",
+                    capacity=RPS,
+                    refill_rate=RPS,
+                    remote_url=self.remote_url,
+                )
+                
+                # Create token rate limiting bucket
+                tokens_bucket = TokenBucket(
+                    bucket_name=service,
+                    bucket_type="tokens",
+                    capacity=TPS,
+                    refill_rate=TPS,
+                    remote_url=self.remote_url,
+                )
+                
+                # Store the buckets for this service
+                self.services_to_buckets[service] = ModelBuckets(
+                    requests_bucket, tokens_bucket
+                )
+                
+            # Map this model to its service and buckets
+            self.models_to_services[model.model] = service
+            self[model] = self.services_to_buckets[service]
+        else:
+            # Model already exists, just retrieve its existing buckets
+            self[model] = self.services_to_buckets[self.models_to_services[model.model]]
+
+    def update_from_key_lookup(self, key_lookup: "KeyLookup") -> None:
+        """
+        Update bucket rate limits based on information from KeyLookup.
+        
+        This method updates the capacity and refill rates of all buckets based on
+        the RPM (requests per minute) and TPM (tokens per minute) limits specified
+        in the provided KeyLookup. This is useful when API keys are rotated or
+        rate limits change.
+        
+        Args:
+            key_lookup: KeyLookup object containing service rate limit information
+            
+        Example:
+            >>> from edsl.key_management import KeyLookup
+            >>> key_lookup = KeyLookup()  # Assume this has rate limit info
+            >>> bucket_collection = BucketCollection()
+            >>> # Add some models to the collection
+            >>> bucket_collection.update_from_key_lookup(key_lookup)
+            >>> # Now rate limits are updated based on key_lookup
+        """
+        # Skip updates if we're using infinite buckets
+        if self.infinity_buckets:
+            return
+            
+        # Update each service with new rate limits
+        for model_name, service in self.models_to_services.items():
+            if service in key_lookup:
+                # Update request rate limits if available
+                if key_lookup[service].rpm is not None:
+                    new_rps = key_lookup[service].rpm / 60.0  # Convert to per-second
+                    new_requests_bucket = TokenBucket(
+                        bucket_name=service,
+                        bucket_type="requests",
+                        capacity=new_rps,
+                        refill_rate=new_rps,
+                        remote_url=self.remote_url,
+                    )
+                    self.services_to_buckets[service].requests_bucket = new_requests_bucket
+
+                # Update token rate limits if available
+                if key_lookup[service].tpm is not None:
+                    new_tps = key_lookup[service].tpm / 60.0  # Convert to per-second
+                    new_tokens_bucket = TokenBucket(
+                        bucket_name=service,
+                        bucket_type="tokens",
+                        capacity=new_tps,
+                        refill_rate=new_tps,
+                        remote_url=self.remote_url,
+                    )
+                    self.services_to_buckets[service].tokens_bucket = new_tokens_bucket
+
+    def visualize(self) -> Dict["LanguageModel", Tuple[Figure, Figure]]:
+        """
+        Visualize the token and request buckets for all models.
+        
+        This method generates visualization plots for each model's token and
+        request buckets, which can be useful for monitoring rate limit usage
+        and debugging rate limiting issues.
+        
+        Returns:
+            Dictionary mapping language models to tuples of (request_plot, token_plot)
+            
+        Example:
+            >>> bucket_collection = BucketCollection()
+            >>> # Add some models
+            >>> plots = bucket_collection.visualize()
+            >>> # Now you can display or save these plots
+        """
+        plots = {}
+        for model in self:
+            plots[model] = self[model].visualize()
+        return plots
+
+
+# Examples and doctests
+if __name__ == "__main__":
+    import doctest
+    
+    # Example showing how to use BucketCollection
+    def example_usage():
+        """
+        Example demonstrating how to use BucketCollection:
+        
+        >>> from edsl import Model
+        >>> # Create models
+        >>> gpt4 = Model('gpt-4')
+        >>> gpt35 = Model('gpt-3.5-turbo')
+        >>> claude = Model('claude-3-opus-20240229')
+        >>> 
+        >>> # Create bucket collection
+        >>> collection = BucketCollection()
+        >>> 
+        >>> # Add models to the collection
+        >>> collection.add_model(gpt4)
+        >>> collection.add_model(gpt35)
+        >>> collection.add_model(claude)
+        >>> 
+        >>> # Models from the same service share rate limits
+        >>> print(collection[gpt4] is collection[gpt35])  # Both OpenAI
+        True
+        >>> print(collection[gpt4] is collection[claude])  # Different services
+        False
+        >>> 
+        >>> # Visualize rate limits
+        >>> # plots = collection.visualize()
+        """
+        pass
+    
+    # Run doctests
+    doctest.testmod(optionflags=doctest.ELLIPSIS)