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

edsl/jobs/buckets/TokenBucketAPI.py

@@ -0,0 +1,211 @@
+from fastapi import FastAPI, HTTPException
+from pydantic import BaseModel
+from typing import Union, Dict
+from typing import Union, List, Any, Optional
+from threading import RLock
+from edsl.jobs.buckets.TokenBucket import TokenBucket  # Original implementation
+
+
+def safe_float_for_json(value: float) -> Union[float, str]:
+    """Convert float('inf') to 'infinity' for JSON serialization.
+
+    Args:
+        value: The float value to convert
+
+    Returns:
+        Either the original float or the string 'infinity' if the value is infinite
+    """
+    if value == float("inf"):
+        return "infinity"
+    return value
+
+
+app = FastAPI()
+
+# In-memory storage for TokenBucket instances
+buckets: Dict[str, TokenBucket] = {}
+
+
+class TokenBucketCreate(BaseModel):
+    bucket_name: str
+    bucket_type: str
+    capacity: Union[int, float]
+    refill_rate: Union[int, float]
+
+
+@app.get("/buckets")
+async def list_buckets(
+    bucket_type: Optional[str] = None,
+    bucket_name: Optional[str] = None,
+    include_logs: bool = False,
+):
+    """List all buckets and their current status.
+
+    Args:
+        bucket_type: Optional filter by bucket type
+        bucket_name: Optional filter by bucket name
+        include_logs: Whether to include the full logs in the response
+    """
+    result = {}
+
+    for bucket_id, bucket in buckets.items():
+        # Apply filters if specified
+        if bucket_type and bucket.bucket_type != bucket_type:
+            continue
+        if bucket_name and bucket.bucket_name != bucket_name:
+            continue
+
+        # Get basic bucket info
+        bucket_info = {
+            "bucket_name": bucket.bucket_name,
+            "bucket_type": bucket.bucket_type,
+            "tokens": bucket.tokens,
+            "capacity": bucket.capacity,
+            "refill_rate": bucket.refill_rate,
+            "turbo_mode": bucket.turbo_mode,
+            "num_requests": bucket.num_requests,
+            "num_released": bucket.num_released,
+            "tokens_returned": bucket.tokens_returned,
+        }
+        for k, v in bucket_info.items():
+            if isinstance(v, float):
+                bucket_info[k] = safe_float_for_json(v)
+
+        # Only include logs if requested
+        if include_logs:
+            bucket_info["log"] = bucket.log
+
+        result[bucket_id] = bucket_info
+
+    return result
+
+
+@app.post("/bucket/{bucket_id}/add_tokens")
+async def add_tokens(bucket_id: str, amount: float):
+    """Add tokens to an existing bucket."""
+    if bucket_id not in buckets:
+        raise HTTPException(status_code=404, detail="Bucket not found")
+
+    if not isinstance(amount, (int, float)) or amount != amount:  # Check for NaN
+        raise HTTPException(status_code=400, detail="Invalid amount specified")
+
+    if amount == float("inf") or amount == float("-inf"):
+        raise HTTPException(status_code=400, detail="Amount cannot be infinite")
+
+    bucket = buckets[bucket_id]
+    bucket.add_tokens(amount)
+
+    # Ensure we return a JSON-serializable float
+    current_tokens = float(bucket.tokens)
+    if not -1e308 <= current_tokens <= 1e308:  # Check if within JSON float bounds
+        current_tokens = 0.0  # or some other reasonable default
+
+    return {"status": "success", "current_tokens": safe_float_for_json(current_tokens)}
+
+
+# @app.post("/bucket")
+# async def create_bucket(bucket: TokenBucketCreate):
+#     bucket_id = f"{bucket.bucket_name}_{bucket.bucket_type}"
+#     if bucket_id in buckets:
+#         raise HTTPException(status_code=400, detail="Bucket already exists")
+
+#     # Create an actual TokenBucket instance
+#     buckets[bucket_id] = TokenBucket(
+#         bucket_name=bucket.bucket_name,
+#         bucket_type=bucket.bucket_type,
+#         capacity=bucket.capacity,
+#         refill_rate=bucket.refill_rate,
+#     )
+#     return {"status": "created"}
+
+
+@app.post("/bucket")
+async def create_bucket(bucket: TokenBucketCreate):
+    if (
+        not isinstance(bucket.capacity, (int, float))
+        or bucket.capacity != bucket.capacity
+    ):  # Check for NaN
+        raise HTTPException(status_code=400, detail="Invalid capacity value")
+    if (
+        not isinstance(bucket.refill_rate, (int, float))
+        or bucket.refill_rate != bucket.refill_rate
+    ):  # Check for NaN
+        raise HTTPException(status_code=400, detail="Invalid refill rate value")
+    if bucket.capacity == float("inf") or bucket.refill_rate == float("inf"):
+        raise HTTPException(status_code=400, detail="Values cannot be infinite")
+    bucket_id = f"{bucket.bucket_name}_{bucket.bucket_type}"
+    if bucket_id in buckets:
+        # Instead of error, return success with "existing" status
+        return {
+            "status": "existing",
+            "bucket": {
+                "capacity": safe_float_for_json(buckets[bucket_id].capacity),
+                "refill_rate": safe_float_for_json(buckets[bucket_id].refill_rate),
+            },
+        }
+
+    # Create a new bucket
+    buckets[bucket_id] = TokenBucket(
+        bucket_name=bucket.bucket_name,
+        bucket_type=bucket.bucket_type,
+        capacity=bucket.capacity,
+        refill_rate=bucket.refill_rate,
+    )
+    return {"status": "created"}
+
+
+@app.post("/bucket/{bucket_id}/get_tokens")
+async def get_tokens(bucket_id: str, amount: float, cheat_bucket_capacity: bool = True):
+    if bucket_id not in buckets:
+        raise HTTPException(status_code=404, detail="Bucket not found")
+
+    bucket = buckets[bucket_id]
+    await bucket.get_tokens(amount, cheat_bucket_capacity)
+    return {"status": "success"}
+
+
+@app.post("/bucket/{bucket_id}/turbo_mode/{state}")
+async def set_turbo_mode(bucket_id: str, state: bool):
+    if bucket_id not in buckets:
+        raise HTTPException(status_code=404, detail="Bucket not found")
+
+    bucket = buckets[bucket_id]
+    if state:
+        bucket.turbo_mode_on()
+    else:
+        bucket.turbo_mode_off()
+    return {"status": "success"}
+
+
+@app.get("/bucket/{bucket_id}/status")
+async def get_bucket_status(bucket_id: str):
+    if bucket_id not in buckets:
+        raise HTTPException(status_code=404, detail="Bucket not found")
+
+    bucket = buckets[bucket_id]
+    status = {
+        "tokens": bucket.tokens,
+        "capacity": bucket.capacity,
+        "refill_rate": bucket.refill_rate,
+        "turbo_mode": bucket.turbo_mode,
+        "num_requests": bucket.num_requests,
+        "num_released": bucket.num_released,
+        "tokens_returned": bucket.tokens_returned,
+        "log": bucket.log,
+    }
+    for k, v in status.items():
+        if isinstance(v, float):
+            status[k] = safe_float_for_json(v)
+
+    for index, entry in enumerate(status["log"]):
+        ts, value = entry
+        status["log"][index] = (ts, safe_float_for_json(value))
+
+    # print(status)
+    return status
+
+
+if __name__ == "__main__":
+    import uvicorn
+
+    uvicorn.run(app, host="0.0.0.0", port=8001)

edsl/jobs/buckets/TokenBucketClient.py

@@ -0,0 +1,191 @@
+from typing import Union, Optional
+import asyncio
+import time
+import aiohttp
+
+
+class TokenBucketClient:
+    """REST API client version of TokenBucket that maintains the same interface
+    by delegating to a server running the original TokenBucket implementation."""
+
+    def __init__(
+        self,
+        *,
+        bucket_name: str,
+        bucket_type: str,
+        capacity: Union[int, float],
+        refill_rate: Union[int, float],
+        api_base_url: str = "http://localhost:8000",
+    ):
+        self.bucket_name = bucket_name
+        self.bucket_type = bucket_type
+        self.capacity = capacity
+        self.refill_rate = refill_rate
+        self.api_base_url = api_base_url
+        self.bucket_id = f"{bucket_name}_{bucket_type}"
+
+        # Initialize the bucket on the server
+        asyncio.run(self._create_bucket())
+
+        # Cache some values locally
+        self.creation_time = time.monotonic()
+        self.turbo_mode = False
+
+    async def _create_bucket(self):
+        async with aiohttp.ClientSession() as session:
+            payload = {
+                "bucket_name": self.bucket_name,
+                "bucket_type": self.bucket_type,
+                "capacity": self.capacity,
+                "refill_rate": self.refill_rate,
+            }
+            async with session.post(
+                f"{self.api_base_url}/bucket",
+                json=payload,
+            ) as response:
+                if response.status != 200:
+                    raise ValueError(f"Unexpected error: {await response.text()}")
+
+                result = await response.json()
+                if result["status"] == "existing":
+                    # Update our local values to match the existing bucket
+                    self.capacity = float(result["bucket"]["capacity"])
+                    self.refill_rate = float(result["bucket"]["refill_rate"])
+
+    def turbo_mode_on(self):
+        """Set the refill rate to infinity."""
+        asyncio.run(self._set_turbo_mode(True))
+        self.turbo_mode = True
+
+    def turbo_mode_off(self):
+        """Restore the refill rate to its original value."""
+        asyncio.run(self._set_turbo_mode(False))
+        self.turbo_mode = False
+
+    async def add_tokens(self, amount: Union[int, float]):
+        """Add tokens to the bucket."""
+        async with aiohttp.ClientSession() as session:
+            async with session.post(
+                f"{self.api_base_url}/bucket/{self.bucket_id}/add_tokens",
+                params={"amount": amount},
+            ) as response:
+                if response.status != 200:
+                    raise ValueError(f"Failed to add tokens: {await response.text()}")
+
+    async def _set_turbo_mode(self, state: bool):
+        async with aiohttp.ClientSession() as session:
+            async with session.post(
+                f"{self.api_base_url}/bucket/{self.bucket_id}/turbo_mode/{str(state).lower()}"
+            ) as response:
+                if response.status != 200:
+                    raise ValueError(
+                        f"Failed to set turbo mode: {await response.text()}"
+                    )
+
+    async def get_tokens(
+        self, amount: Union[int, float] = 1, cheat_bucket_capacity=True
+    ) -> None:
+        async with aiohttp.ClientSession() as session:
+            async with session.post(
+                f"{self.api_base_url}/bucket/{self.bucket_id}/get_tokens",
+                params={
+                    "amount": amount,
+                    "cheat_bucket_capacity": int(cheat_bucket_capacity),
+                },
+            ) as response:
+                if response.status != 200:
+                    raise ValueError(f"Failed to get tokens: {await response.text()}")
+
+    def get_throughput(self, time_window: Optional[float] = None) -> float:
+        status = asyncio.run(self._get_status())
+        now = time.monotonic()
+
+        if time_window is None:
+            start_time = self.creation_time
+        else:
+            start_time = now - time_window
+
+        if start_time < self.creation_time:
+            start_time = self.creation_time
+
+        elapsed_time = now - start_time
+
+        if elapsed_time == 0:
+            return status["num_released"] / 0.001
+
+        return (status["num_released"] / elapsed_time) * 60
+
+    async def _get_status(self) -> dict:
+        async with aiohttp.ClientSession() as session:
+            async with session.get(
+                f"{self.api_base_url}/bucket/{self.bucket_id}/status"
+            ) as response:
+                if response.status != 200:
+                    raise ValueError(
+                        f"Failed to get bucket status: {await response.text()}"
+                    )
+                return await response.json()
+
+    def __add__(self, other) -> "TokenBucketClient":
+        """Combine two token buckets."""
+        return TokenBucketClient(
+            bucket_name=self.bucket_name,
+            bucket_type=self.bucket_type,
+            capacity=min(self.capacity, other.capacity),
+            refill_rate=min(self.refill_rate, other.refill_rate),
+            api_base_url=self.api_base_url,
+        )
+
+    @property
+    def tokens(self) -> float:
+        """Get the number of tokens remaining in the bucket."""
+        status = asyncio.run(self._get_status())
+        return float(status["tokens"])
+
+    def wait_time(self, requested_tokens: Union[float, int]) -> float:
+        """Calculate the time to wait for the requested number of tokens."""
+        # self.refill()  # Update the current token count
+        if self.tokens >= float(requested_tokens):
+            return 0.0
+        try:
+            return (requested_tokens - self.tokens) / self.refill_rate
+        except Exception as e:
+            raise ValueError(f"Error calculating wait time: {e}")
+
+    # def wait_time(self, num_tokens: Union[int, float]) -> float:
+    #     return 0  # TODO - Need to implement this on the server side
+
+    def visualize(self):
+        """Visualize the token bucket over time."""
+        status = asyncio.run(self._get_status())
+        times, tokens = zip(*status["log"])
+        start_time = times[0]
+        times = [t - start_time for t in times]
+
+        from matplotlib import pyplot as plt
+
+        plt.figure(figsize=(10, 6))
+        plt.plot(times, tokens, label="Tokens Available")
+        plt.xlabel("Time (seconds)", fontsize=12)
+        plt.ylabel("Number of Tokens", fontsize=12)
+        details = f"{self.bucket_name} ({self.bucket_type}) Bucket Usage Over Time\nCapacity: {self.capacity:.1f}, Refill Rate: {self.refill_rate:.1f}/second"
+        plt.title(details, fontsize=14)
+        plt.legend()
+        plt.grid(True)
+        plt.tight_layout()
+        plt.show()
+
+
+if __name__ == "__main__":
+    import doctest
+
+    doctest.testmod()
+    # bucket = TokenBucketClient(
+    #     bucket_name="test", bucket_type="test", capacity=100, refill_rate=10
+    # )
+    # asyncio.run(bucket.get_tokens(50))
+    # time.sleep(1)  # Wait for 1 second
+    # asyncio.run(bucket.get_tokens(30))
+    # throughput = bucket.get_throughput(1)
+    # print(throughput)
+    # bucket.visualize()

edsl/jobs/check_survey_scenario_compatibility.py

@@ -0,0 +1,85 @@
+import warnings
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from edsl.surveys.Survey import Survey
+    from edsl.scenarios.ScenarioList import ScenarioList
+
+
+class CheckSurveyScenarioCompatibility:
+
+    def __init__(self, survey: "Survey", scenarios: "ScenarioList"):
+        self.survey = survey
+        self.scenarios = scenarios
+
+    def check(self, strict: bool = False, warn: bool = False) -> None:
+        """Check if the parameters in the survey and scenarios are consistent.
+
+        >>> from edsl.jobs.Jobs import Jobs
+        >>> from edsl.questions.QuestionFreeText import QuestionFreeText
+        >>> from edsl.surveys.Survey import Survey
+        >>> from edsl.scenarios.Scenario import Scenario
+        >>> q = QuestionFreeText(question_text = "{{poo}}", question_name = "ugly_question")
+        >>> j = Jobs(survey = Survey(questions=[q]))
+        >>> cs = CheckSurveyScenarioCompatibility(j.survey, j.scenarios)
+        >>> with warnings.catch_warnings(record=True) as w:
+        ...     cs.check(warn = True)
+        ...     assert len(w) == 1
+        ...     assert issubclass(w[-1].category, UserWarning)
+        ...     assert "The following parameters are in the survey but not in the scenarios" in str(w[-1].message)
+
+        >>> q = QuestionFreeText(question_text = "{{poo}}", question_name = "ugly_question")
+        >>> s = Scenario({'plop': "A", 'poo': "B"})
+        >>> j = Jobs(survey = Survey(questions=[q])).by(s)
+        >>> cs = CheckSurveyScenarioCompatibility(j.survey, j.scenarios)
+        >>> cs.check(strict = True)
+        Traceback (most recent call last):
+        ...
+        ValueError: The following parameters are in the scenarios but not in the survey: {'plop'}
+
+        >>> q = QuestionFreeText(question_text = "Hello", question_name = "ugly_question")
+        >>> s = Scenario({'ugly_question': "B"})
+        >>> from edsl.scenarios.ScenarioList import ScenarioList
+        >>> cs = CheckSurveyScenarioCompatibility(Survey(questions=[q]), ScenarioList([s]))
+        >>> cs.check()
+        Traceback (most recent call last):
+        ...
+        ValueError: The following names are in both the survey question_names and the scenario keys: {'ugly_question'}. This will create issues.
+        """
+        survey_parameters: set = self.survey.parameters
+        scenario_parameters: set = self.scenarios.parameters
+
+        msg0, msg1, msg2 = None, None, None
+
+        # look for key issues
+        if intersection := set(self.scenarios.parameters) & set(
+            self.survey.question_names
+        ):
+            msg0 = f"The following names are in both the survey question_names and the scenario keys: {intersection}. This will create issues."
+
+            raise ValueError(msg0)
+
+        if in_survey_but_not_in_scenarios := survey_parameters - scenario_parameters:
+            msg1 = f"The following parameters are in the survey but not in the scenarios: {in_survey_but_not_in_scenarios}"
+        if in_scenarios_but_not_in_survey := scenario_parameters - survey_parameters:
+            msg2 = f"The following parameters are in the scenarios but not in the survey: {in_scenarios_but_not_in_survey}"
+
+        if msg1 or msg2:
+            message = "\n".join(filter(None, [msg1, msg2]))
+            if strict:
+                raise ValueError(message)
+            else:
+                if warn:
+                    warnings.warn(message)
+
+        if self.scenarios.has_jinja_braces:
+            warnings.warn(
+                "The scenarios have Jinja braces ({{ and }}). Converting to '<<' and '>>'. If you want a different conversion, use the convert_jinja_braces method first to modify the scenario."
+            )
+            self.scenarios = self.scenarios._convert_jinja_braces()
+
+
+if __name__ == "__main__":
+    import doctest
+
+    doctest.testmod()

edsl/jobs/data_structures.py

@@ -0,0 +1,120 @@
+from typing import Optional, Literal
+from dataclasses import dataclass, asdict
+
+# from edsl.data_transfer_models import VisibilityType
+from edsl.data.Cache import Cache
+from edsl.jobs.buckets.BucketCollection import BucketCollection
+from edsl.language_models.key_management.KeyLookup import KeyLookup
+from edsl.jobs.runners.JobsRunnerStatus import JobsRunnerStatus
+
+VisibilityType = Literal["private", "public", "unlisted"]
+from edsl.Base import Base
+
+
+@dataclass
+class RunEnvironment:
+    cache: Optional[Cache] = None
+    bucket_collection: Optional[BucketCollection] = None
+    key_lookup: Optional[KeyLookup] = None
+    jobs_runner_status: Optional["JobsRunnerStatus"] = None
+
+
+@dataclass
+class RunParameters(Base):
+    n: int = 1
+    progress_bar: bool = False
+    stop_on_exception: bool = False
+    check_api_keys: bool = False
+    verbose: bool = True
+    print_exceptions: bool = True
+    remote_cache_description: Optional[str] = None
+    remote_inference_description: Optional[str] = None
+    remote_inference_results_visibility: Optional[VisibilityType] = "unlisted"
+    skip_retry: bool = False
+    raise_validation_errors: bool = False
+    disable_remote_cache: bool = False
+    disable_remote_inference: bool = False
+    job_uuid: Optional[str] = None
+
+    def to_dict(self, add_edsl_version=False) -> dict:
+        d = asdict(self)
+        if add_edsl_version:
+            from edsl import __version__
+
+            d["edsl_version"] = __version__
+            d["edsl_class_name"] = "RunConfig"
+        return d
+
+    @classmethod
+    def from_dict(cls, data: dict) -> "RunConfig":
+        return cls(**data)
+
+    def code(self):
+        return f"RunConfig(**{self.to_dict()})"
+
+    @classmethod
+    def example(cls) -> "RunConfig":
+        return cls()
+
+
+@dataclass
+class RunConfig:
+    environment: RunEnvironment
+    parameters: RunParameters
+
+    def add_environment(self, environment: RunEnvironment):
+        self.environment = environment
+
+    def add_bucket_collection(self, bucket_collection: BucketCollection):
+        self.environment.bucket_collection = bucket_collection
+
+    def add_cache(self, cache: Cache):
+        self.environment.cache = cache
+
+    def add_key_lookup(self, key_lookup: KeyLookup):
+        self.environment.key_lookup = key_lookup
+
+
+"""This module contains the Answers class, which is a helper class to hold the answers to a survey."""
+
+from collections import UserDict
+from edsl.data_transfer_models import EDSLResultObjectInput
+
+
+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/jobs/decorators.py

@@ -0,0 +1,35 @@
+from functools import wraps
+from threading import RLock
+import inspect
+
+
+def synchronized_class(wrapped_class):
+    """Class decorator that makes all methods thread-safe."""
+
+    # Add a lock to the class
+    setattr(wrapped_class, "_lock", RLock())
+
+    # Get all methods from the class
+    for name, method in inspect.getmembers(wrapped_class, inspect.isfunction):
+        # Skip magic methods except __getitem__, __setitem__, __delitem__
+        if name.startswith("__") and name not in [
+            "__getitem__",
+            "__setitem__",
+            "__delitem__",
+        ]:
+            continue
+
+        # Create synchronized version of the method
+        def create_synchronized_method(method):
+            @wraps(method)
+            def synchronized_method(*args, **kwargs):
+                instance = args[0]  # first arg is self
+                with instance._lock:
+                    return method(*args, **kwargs)
+
+            return synchronized_method
+
+        # Replace the original method with synchronized version
+        setattr(wrapped_class, name, create_synchronized_method(method))
+
+    return wrapped_class