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

edsl/data/SQLiteDict.py

@@ -0,0 +1,292 @@
+from __future__ import annotations
+import json
+from typing import Any, Generator, Optional, Union
+
+from edsl.config import CONFIG
+from edsl.data.CacheEntry import CacheEntry
+from edsl.data.orm import Base, Data
+
+
+class SQLiteDict:
+    """
+    A dictionary-like object that is an interface for an local database.
+    - You can use SQLiteDict as a regular dictionary.
+    - Supports only SQLite for now.
+    """
+
+    def __init__(self, db_path: Optional[str] = None):
+        """
+
+        >>> temp_db_path = SQLiteDict._get_temp_path()
+        >>> SQLiteDict(f"sqlite:///{temp_db_path}")  # Use the temp file for SQLite
+        SQLiteDict(db_path='...')
+        >>> import os; os.unlink(temp_db_path)  # Clean up the temp file after the test
+
+        """
+        from sqlalchemy.exc import SQLAlchemyError
+        from sqlalchemy.orm import sessionmaker
+        from sqlalchemy import create_engine
+
+        self.db_path = db_path or CONFIG.get("EDSL_DATABASE_PATH")
+        if not self.db_path.startswith("sqlite:///"):
+            self.db_path = f"sqlite:///{self.db_path}"
+        try:
+            from edsl.data.orm import Base, Data
+
+            self.engine = create_engine(self.db_path, echo=False, future=True)
+            Base.metadata.create_all(self.engine)
+            self.Session = sessionmaker(bind=self.engine)
+        except SQLAlchemyError as e:
+            raise Exception(
+                f"""Database initialization error: {e}. The attempted DB path was {db_path}"""
+            ) from e
+
+    @classmethod
+    def _get_temp_path(self):
+        import tempfile
+        import os
+
+        _, temp_db_path = tempfile.mkstemp(suffix=".db")
+        return temp_db_path
+
+    def __setitem__(self, key: str, value: CacheEntry) -> None:
+        """
+        Stores a key-value pair.
+
+        >>> d = SQLiteDict.example()
+        >>> d["foo"] = CacheEntry.example()
+        """
+        if not isinstance(value, CacheEntry):
+            raise ValueError(f"Value must be a CacheEntry object (got {type(value)}).")
+        with self.Session() as db:
+            from edsl.data.orm import Base, Data
+
+            db.merge(Data(key=key, value=json.dumps(value.to_dict())))
+            db.commit()
+
+    def __getitem__(self, key: str) -> CacheEntry:
+        """
+        Gets a value for a given key.
+        - Raises a KeyError if the key is not found.
+
+        >>> d = SQLiteDict.example()
+        >>> d["foo"] = CacheEntry.example()
+        >>> d["foo"] == CacheEntry.example()
+        True
+        """
+        with self.Session() as db:
+            from edsl.data.orm import Base, Data
+
+            value = db.query(Data).filter_by(key=key).first()
+            if not value:
+                raise KeyError(f"Key '{key}' not found.")
+            return CacheEntry.from_dict(json.loads(value.value))
+
+    def get(self, key: str, default: Optional[Any] = None) -> Union[CacheEntry, Any]:
+        """
+        Gets the value for a given key
+        - Returns the `default` value if the key is not found.
+
+        >>> d = SQLiteDict.example()
+        >>> d.get("foo", "bar")
+        'bar'
+        """
+        try:
+            return self[key]
+        except KeyError:
+            return default
+
+    def __bool__(self) -> bool:
+        """This is so likes like
+        self.data = data or {} 'work' as expected
+        """
+        return True
+
+    def update(
+        self,
+        new_d: Union[dict, SQLiteDict],
+        overwrite: Optional[bool] = False,
+        max_batch_size: Optional[int] = 100,
+    ) -> None:
+        """
+        Update the dictionary with the values from another dictionary.
+
+        :param new_d: The dictionary to update the current dictionary with.
+        :param overwrite: If `overwrite` is False, existing values will not be overwritten.
+        :param max_batch_size: The maximum number of items to update in a single transaction.
+
+        - If `overwrite` is True, existing values will be overwritten.
+
+        >>> d = SQLiteDict.example()
+        >>> d.update({"foo": CacheEntry.example()})
+        >>> d["foo"] == CacheEntry.example()
+        True
+        """
+        if not (isinstance(new_d, dict) or isinstance(new_d, SQLiteDict)):
+            raise ValueError(
+                f"new_d must be a dict or SQLiteDict object (got {type(new_d)})"
+            )
+        current_batch = 0
+        with self.Session() as db:
+            for key, value in new_d.items():
+                if current_batch == max_batch_size:
+                    db.commit()
+                    current_batch = 0
+                current_batch += 1
+                if key in self:
+                    if overwrite:
+                        db.merge(Data(key=key, value=json.dumps(value.to_dict())))
+                else:
+                    db.merge(Data(key=key, value=json.dumps(value.to_dict())))
+            db.commit()
+
+    def values(self) -> Generator[CacheEntry, None, None]:
+        """
+        Returns a generator that yields the values in the cache.
+
+        >>> d = SQLiteDict.example()
+        >>> d["foo"] = CacheEntry.example()
+        >>> list(d.values()) == [CacheEntry.example()]
+        True
+        """
+        with self.Session() as db:
+            for instance in db.query(Data).all():
+                yield CacheEntry.from_dict(json.loads(instance.value))
+
+    def items(self) -> Generator[tuple[str, CacheEntry], None, None]:
+        """
+        Returns a generator that yields the items in the cache.
+
+        >>> d = SQLiteDict.example()
+        >>> d["foo"] = CacheEntry.example()
+        >>> list(d.items()) == [("foo", CacheEntry.example())]
+        True
+        """
+        with self.Session() as db:
+            for instance in db.query(Data).all():
+                yield (instance.key, CacheEntry.from_dict(json.loads(instance.value)))
+
+    def to_dict(self):
+        """
+        Returns the cache as a dictionary.
+
+        >>> d = SQLiteDict.example()
+        >>> d["foo"] = CacheEntry.example()
+        >>> d.to_dict() == {"foo": CacheEntry.example()}
+        True
+        """
+        return dict(self.items())
+
+    def __delitem__(self, key: str) -> None:
+        """
+        Deletes the value for a given key.
+
+        >>> d = SQLiteDict.example()
+        >>> d["foo"] = CacheEntry.example()
+        >>> del d["foo"]
+        >>> d.get("foo", "missing")
+        'missing'
+        """
+        with self.Session() as db:
+            instance = db.query(Data).filter_by(key=key).one_or_none()
+            if instance:
+                db.delete(instance)
+                db.commit()
+            else:
+                raise KeyError(f"Key '{key}' not found.")
+
+    def __contains__(self, key: str) -> bool:
+        """
+        Checks if the dict contains the given key.
+
+        >>> d = SQLiteDict.example()
+        >>> d["foo"] = CacheEntry.example()
+        >>> "foo" in d
+        True
+        >>> "bar" in d
+        False
+        """
+        with self.Session() as db:
+            return db.query(Data).filter_by(key=key).first() is not None
+
+    def __iter__(self) -> Generator[str, None, None]:
+        """
+        Returns a generator that yields the keys in the dict.
+
+        >>> d = SQLiteDict.example()
+        >>> d["foo"] = CacheEntry.example()
+        >>> list(iter(d)) == ["foo"]
+        True
+        """
+        with self.Session() as db:
+            for instance in db.query(Data).all():
+                yield instance.key
+
+    def __len__(self) -> int:
+        """
+        Returns the number of items in the cache.
+
+        >>> d = SQLiteDict.example()
+        >>> len(d)
+        0
+        >>> d["foo"] = CacheEntry.example()
+        >>> len(d)
+        1
+        """
+        with self.Session() as db:
+            return db.query(Data).count()
+
+    def keys(self) -> Generator[str, None, None]:
+        """
+        Returns a generator that yields the keys in the cache.
+
+        >>> d = SQLiteDict.example()
+        >>> d["foo"] = CacheEntry.example()
+        >>> list(d.keys()) == ["foo"]
+        True
+        """
+        return self.__iter__()
+
+    def __repr__(self) -> str:
+        return f"{self.__class__.__name__}(db_path={self.db_path!r})"
+
+    @classmethod
+    def example(cls) -> SQLiteDict:
+        """
+        Returns an example SQLiteDict object.
+        - The example SQLiteDict is empty and stored in memory.
+
+        >>> SQLiteDict.example()
+        SQLiteDict(db_path='sqlite:///:memory:')
+        """
+        return cls(db_path="sqlite:///:memory:")
+
+
+def main():
+    from edsl.data.CacheEntry import CacheEntry
+    from edsl.data.SQLiteDict import SQLiteDict
+
+    d = SQLiteDict.example()
+    d["foo"] = CacheEntry.example()
+    d["foo"]
+    d.get("foo")
+    d.get("poo")
+    d.get("poo", "not found")
+    d.update({"poo": CacheEntry.example()})
+    d["poo"]
+    len(d)
+    list(d.keys())
+    list(d.values())
+    list(d.items())
+    assert "poo" in d
+    assert "loo" not in d
+    del d["poo"]
+    assert len(d) == 1
+    repr(d)
+    d
+
+
+if __name__ == "__main__":
+    import doctest
+
+    doctest.testmod()

edsl/data/__init__.py

@@ -1,3 +1,5 @@
-from edsl.data.orm import LLMOutputDataDB, Base
-from edsl.data.Database import Database, database
-from edsl.data.crud import CRUD, CRUDOperations
+# from edsl.data.CacheEntry import CacheEntry
+# from edsl.data.SQLiteDict import SQLiteDict
+from edsl.data.Cache import Cache
+
+# from edsl.data.CacheHandler import CacheHandler

edsl/data/orm.py

@@ -1,37 +1,10 @@
-from sqlalchemy import Column, String, Integer, Text, Index, Sequence
-from sqlalchemy.orm import declarative_base
+from sqlalchemy import Column, String
+from sqlalchemy.ext.declarative import declarative_base
 
 Base = declarative_base()
 
 
-class LLMOutputDataDB(Base):
-    __tablename__ = "responses"
-
-    # Primary key
-    id = Column(Integer, Sequence("response_id_seq"), primary_key=True)
-
-    # Fields
-    model = Column(String(100), nullable=False)
-    parameters = Column(Text, nullable=False)
-    system_prompt = Column(Text, nullable=False)
-    prompt = Column(Text, nullable=False)
-    output = Column(Text, nullable=False)
-
-    # Index for faster queries
-    idx_responses_fields = Index(
-        "idx_responses_fields", "prompt", "system_prompt", "model", "parameters"
-    )
-
-
-class ResultDB(Base):
-    __tablename__ = "result"
-
-    id = Column(Integer, primary_key=True)
-    job_uuid = Column(String, nullable=False)
-    result_uuid = Column(String, nullable=False)
-    agent = Column(Text, nullable=False)
-    scenario = Column(Text, nullable=False)
-    model = Column(Text, nullable=False)
-    answer = Column(Text, nullable=False)
-
-    idx_job_uuid = Index("idx_job_uuid", "job_uuid")
+class Data(Base):
+    __tablename__ = "data"
+    key = Column(String, primary_key=True)
+    value = Column(String)

edsl/data_transfer_models.py

@@ -1,27 +1,74 @@
-from collections import UserDict
-
-
-class AgentResponseDict(UserDict):
-    def __init__(
-        self,
-        *,
-        question_name,
-        answer,
-        comment,
-        prompts,
-        usage=None,
-        cached_response=None,
-        raw_model_response=None
-    ):
-        usage = usage or {"prompt_tokens": 0, "completion_tokens": 0}
-        super().__init__(
-            {
-                "answer": answer,
-                "comment": comment,
-                "question_name": question_name,
-                "prompts": prompts,
-                "usage": usage,
-                "cached_response": cached_response,
-                "raw_model_response": raw_model_response,
-            }
-        )
+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)})"

edsl/enums.py

@@ -1,13 +1,21 @@
+"""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"
@@ -26,20 +34,169 @@ class QuestionType(EnumWithChecks):
 # https://huggingface.co/meta-llama/Llama-2-70b-chat-hf
 
 
-class LanguageModelType(EnumWithChecks):
-    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"
+# 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"
+
+
+# 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",
+]
+
+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.BEDROCK.value: "TBD",
+    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",
+}
+
+
+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
+        )
+
+
+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/exceptions/BaseException.py

@@ -0,0 +1,21 @@
+class BaseException(Exception):
+    relevant_doc = "https://docs.expectedparrot.com/"
+
+    def __init__(self, message, *, show_docs=True):
+        # Format main error message
+        formatted_message = [message.strip()]
+
+        # Add documentation links if requested
+        if show_docs:
+            if hasattr(self, "relevant_doc"):
+                formatted_message.append(
+                    f"\nFor more information, see:\n{self.relevant_doc}"
+                )
+            if hasattr(self, "relevant_notebook"):
+                formatted_message.append(
+                    f"\nFor a usage example, see:\n{self.relevant_notebook}"
+                )
+
+        # Join with double newlines for clear separation
+        final_message = "\n\n".join(formatted_message)
+        super().__init__(final_message)