From 6c807a8bef3775e9c1373ce751033bb039a26ca1 Mon Sep 17 00:00:00 2001
From: vinid <chiccobia@gmail.com>
Date: Sun, 1 Sep 2024 15:27:00 -0400
Subject: [PATCH 01/10] temp

---
 textgrad/engine/utils.py | 230 +++++++++++++++++++++++++++++++++++++++
 1 file changed, 230 insertions(+)
 create mode 100644 textgrad/engine/utils.py

diff --git a/textgrad/engine/utils.py b/textgrad/engine/utils.py
new file mode 100644
index 0000000..890d5cd
--- /dev/null
+++ b/textgrad/engine/utils.py
@@ -0,0 +1,230 @@
+import hashlib
+
+try:
+    from openai import OpenAI, AzureOpenAI
+except ImportError:
+    raise ImportError("If you'd like to use OpenAI models, please install the openai package by running `pip install openai`, and add 'OPENAI_API_KEY' to your environment variables.")
+
+import os
+import json
+from abc import ABC, abstractmethod
+import base64
+import platformdirs
+from tenacity import (
+    retry,
+    stop_after_attempt,
+    wait_random_exponential,
+)
+from typing import List, Union
+from functools import wraps
+from .engine_utils import get_image_type_from_bytes
+
+def cached(func):
+    @wraps(func)
+    def wrapper(self, *args, **kwargs):
+        if not self.use_cache:
+            return func(self, *args, **kwargs)
+
+        # get string representation from args and kwargs
+        key = hash(str(args) + str(kwargs))
+        key = hashlib.sha256(f"{key}".encode()).hexdigest()
+
+        if key in self.cache:
+            return self.cache[key]
+
+        result = func(self, *args, **kwargs)
+        self.cache[key] = result
+        return result
+
+    return wrapper
+
+
+class EngineLM(ABC):
+    system_prompt: str = "You are a helpful, creative, and smart assistant."
+    model_string: str
+    is_multimodal: bool
+    use_cache: bool = False
+    cache_path: str
+
+    @abstractmethod
+    def _generate_from_multiple_input(self, prompt, system_prompt=None, **kwargs) -> str:
+        pass
+
+    @abstractmethod
+    def _generate_from_single_prompt(self, prompt, system_prompt=None, **kwargs) -> str:
+        pass
+
+    # TBF this could be simplified to a single generate method
+    def _prepare_generate_from_single_prompt(self, prompt: str, system_prompt: str = None, **kwargs):
+        sys_prompt_arg = system_prompt if system_prompt else self.system_prompt
+        return self._generate_from_single_prompt(prompt, system_prompt=sys_prompt_arg, **kwargs)
+
+    def _prepare_generate_from_multiple_input(self, content: List[Union[str, bytes]], system_prompt=None, **kwargs):
+        sys_prompt_arg = system_prompt if system_prompt else self.system_prompt
+        return self._generate_from_multiple_input(content, system_prompt=sys_prompt_arg, **kwargs)
+
+    def generate(self, content, system_prompt=None, **kwargs):
+        if isinstance(content, str):
+            return self._prepare_generate_from_single_prompt(content, system_prompt=system_prompt, **kwargs)
+
+        elif isinstance(content, list):
+            has_multimodal_input = any(isinstance(item, bytes) for item in content)
+            if (has_multimodal_input) and (not self.is_multimodal):
+                raise NotImplementedError("Multimodal generation is only supported for Claude-3 and beyond.")
+
+            return self._prepare_generate_from_multiple_input(content, system_prompt=system_prompt, **kwargs)
+
+    def __call__(self, *args, **kwargs):
+        pass
+
+class OpenAIEngine(EngineLM):
+    DEFAULT_SYSTEM_PROMPT = "You are a helpful, creative, and smart assistant."
+
+    def __init__(
+            self,
+            model_string: str = "gpt-3.5-turbo-0613",
+            system_prompt: str = DEFAULT_SYSTEM_PROMPT,
+            is_multimodal: bool = False,
+            use_cache: bool = False,
+            base_url: str = None):
+        """
+        :param model_string:
+        :param system_prompt:
+        :param base_url: Used to support Ollama
+        """
+        root = platformdirs.user_cache_dir("textgrad")
+        cache_path = os.path.join(root, f"cache_openai_{model_string}.db")
+
+        super().__init__(system_prompt=system_prompt,
+                         cache_path=cache_path,
+                         use_cache=use_cache,
+                         is_multimodal=is_multimodal,
+                         model_string=model_string)
+
+        self.base_url = base_url
+
+        if not base_url:
+            if os.getenv("OPENAI_API_KEY") is None:
+                raise ValueError(
+                    "Please set the OPENAI_API_KEY environment variable if you'd like to use OpenAI models.")
+
+            self.client = OpenAI(
+                api_key=os.getenv("OPENAI_API_KEY")
+            )
+        else:
+            raise ValueError("Invalid base URL provided. Please use the default OLLAMA base URL or None.")
+
+    @cached
+    def _generate_from_single_prompt(
+            self, prompt: str, system_prompt: str = None, temperature=0, max_tokens=2000, top_p=0.99
+    ):
+
+        response = self.client.chat.completions.create(
+            model=self.model_string,
+            messages=[
+                {"role": "system", "content": system_prompt},
+                {"role": "user", "content": prompt},
+            ],
+            frequency_penalty=0,
+            presence_penalty=0,
+            stop=None,
+            temperature=temperature,
+            max_tokens=max_tokens,
+            top_p=top_p,
+        )
+
+        response = response.choices[0].message.content
+        return response
+
+    @cached
+    def _generate_from_multiple_input(
+            self, content: List[Union[str, bytes]], system_prompt=None, temperature=0, max_tokens=2000, top_p=0.99
+    ):
+        formatted_content = self._format_content(content)
+
+        response = self.client.chat.completions.create(
+            model=self.model_string,
+            messages=[
+                {"role": "system", "content": system_prompt},
+                {"role": "user", "content": formatted_content},
+            ],
+            temperature=temperature,
+            max_tokens=max_tokens,
+            top_p=top_p,
+        )
+
+        response_text = response.choices[0].message.content
+        return response_text
+
+    def __call__(self, prompt, **kwargs):
+        return self.generate(prompt, **kwargs)
+
+    def _format_content(self, content: List[Union[str, bytes]]) -> List[dict]:
+        """Helper function to format a list of strings and bytes into a list of dictionaries to pass as messages to the API.
+        """
+        formatted_content = []
+        for item in content:
+            if isinstance(item, bytes):
+                # For now, bytes are assumed to be images
+                image_type = get_image_type_from_bytes(item)
+                base64_image = base64.b64encode(item).decode('utf-8')
+                formatted_content.append({
+                    "type": "image_url",
+                    "image_url": {
+                        "url": f"data:image/{image_type};base64,{base64_image}"
+                    }
+                })
+            elif isinstance(item, str):
+                formatted_content.append({
+                    "type": "text",
+                    "text": item
+                })
+            else:
+                raise ValueError(f"Unsupported input type: {type(item)}")
+        return formatted_content
+
+
+class OpenAICompatibleEngine(OpenAIEngine):
+    """
+        This is the same as engine.openai.ChatOpenAI, but we pass in an external OpenAI client.
+        """
+
+    DEFAULT_SYSTEM_PROMPT = "You are a helpful, creative, and smart assistant."
+    client = None
+
+    def __init__(
+            self,
+            client: OpenAI,
+            model_string: str = "gpt-3.5-turbo-0613",
+            system_prompt: str = DEFAULT_SYSTEM_PROMPT,
+            is_multimodal: bool = False,
+            use_cache: bool = False,
+            base_url: str = None):
+        """
+        :param client: an OpenAI client object.
+        :param model_string: the model name, used for the cache file name and chat completion requests.
+        :param system_prompt: the system prompt to use in chat completions.
+
+        Example usage with lm-studio local server, but any client that follows the OpenAI API will work.
+
+        ```python
+        from openai import OpenAI
+        from textgrad.engine.local_model_openai_api import ChatExternalClient
+
+        client = OpenAI(base_url="http://localhost:1234/v1", api_key="lm-studio")
+        engine = ChatExternalClient(client=client, model_string="your-model-name")
+        print(engine.generate(max_tokens=40, prompt="What is the meaning of life?"))
+        ```
+
+        """
+
+        if os.getenv("OPENAI_API_KEY") is None:
+            os.environ["OPENAI_API_KEY"] = client.api_key
+
+        self.client = client
+
+        super.__init__(model_string=model_string,
+                       system_prompt=system_prompt,
+                       is_multimodal=is_multimodal,
+                       use_cache=use_cache,
+                       base_url=base_url)
\ No newline at end of file

From 09a71b1e7aabf95ca8170eaf5d1150c10869d54d Mon Sep 17 00:00:00 2001
From: vinid <chiccobia@gmail.com>
Date: Sun, 8 Sep 2024 16:31:54 -0400
Subject: [PATCH 02/10] new engines

---
 textgrad/engine/engine_utils.py              |  16 ---
 textgrad/engine_experimental/__init__.py     |   2 +
 textgrad/engine_experimental/base.py         |  82 ++++++++++++++
 textgrad/engine_experimental/engine_utils.py |  42 ++++++++
 textgrad/engine_experimental/litellm.py      |  61 +++++++++++
 textgrad/engine_experimental/openai.py       | 106 +++++++++++++++++++
 6 files changed, 293 insertions(+), 16 deletions(-)
 delete mode 100644 textgrad/engine/engine_utils.py
 create mode 100644 textgrad/engine_experimental/__init__.py
 create mode 100644 textgrad/engine_experimental/base.py
 create mode 100644 textgrad/engine_experimental/engine_utils.py
 create mode 100644 textgrad/engine_experimental/litellm.py
 create mode 100644 textgrad/engine_experimental/openai.py

diff --git a/textgrad/engine/engine_utils.py b/textgrad/engine/engine_utils.py
deleted file mode 100644
index 51caff9..0000000
--- a/textgrad/engine/engine_utils.py
+++ /dev/null
@@ -1,16 +0,0 @@
-def is_jpeg(data):
-    jpeg_signature = b'\xFF\xD8\xFF'
-    return data.startswith(jpeg_signature)
-
-def is_png(data):
-    png_signature = b'\x89\x50\x4E\x47\x0D\x0A\x1A\x0A'
-    return data.startswith(png_signature)
-
-
-def get_image_type_from_bytes(data):
-    if is_jpeg(data):
-        return "jpeg"
-    elif is_png(data):
-        return "png"
-    else:
-        raise ValueError("Image type not supported, only jpeg and png supported.")
\ No newline at end of file
diff --git a/textgrad/engine_experimental/__init__.py b/textgrad/engine_experimental/__init__.py
new file mode 100644
index 0000000..abe0de9
--- /dev/null
+++ b/textgrad/engine_experimental/__init__.py
@@ -0,0 +1,2 @@
+from textgrad.engine_experimental.openai import OpenAIEngine
+from textgrad.engine_experimental.litellm import LiteLLMEngine
\ No newline at end of file
diff --git a/textgrad/engine_experimental/base.py b/textgrad/engine_experimental/base.py
new file mode 100644
index 0000000..3203a85
--- /dev/null
+++ b/textgrad/engine_experimental/base.py
@@ -0,0 +1,82 @@
+from functools import wraps
+from abc import ABC, abstractmethod
+import hashlib
+from typing import List, Union
+import diskcache as dc
+import platformdirs
+import os
+
+
+
+def cached(func):
+    @wraps(func)
+    def wrapper(self, *args, **kwargs):
+
+        if self.cache is False:
+            return func(self, *args, **kwargs)
+
+        # get string representation from args and kwargs
+        key = hash(str(args) + str(kwargs))
+        key = hashlib.sha256(f"{key}".encode()).hexdigest()
+
+        if key in self.cache:
+            return self.cache[key]
+
+        result = func(self, *args, **kwargs)
+        self.cache[key] = result
+        return result
+
+    return wrapper
+
+
+class EngineLM(ABC):
+    system_prompt: str = "You are a helpful, creative, and smart assistant."
+    model_string: str
+    is_multimodal: bool
+    cache: Union[dc.Cache, bool]
+
+    def __init__(self, model_string: str,
+                 system_prompt: str = "You are a helpful, creative, and smart assistant.",
+                 is_multimodal: bool = False,
+                 cache=Union[dc.Cache, bool]):
+
+        root = platformdirs.user_cache_dir("textgrad")
+        default_cache_path = os.path.join(root, f"cache_model_{model_string}.db")
+
+        self.model_string = model_string
+        self.system_prompt = system_prompt
+        self.is_multimodal = is_multimodal
+
+        if isinstance(cache, dc.Cache):
+            self.cache = cache
+        elif cache is True:
+            self.cache = dc.Cache(default_cache_path)
+        elif cache is False:
+            self.cache = False
+        else:
+            raise ValueError("Cache argument must be a diskcache.Cache object or a boolean.")
+
+    @abstractmethod
+    def _generate_from_multiple_input(self, prompt, system_prompt=None, **kwargs) -> str:
+        pass
+
+    @abstractmethod
+    def _generate_from_single_prompt(self, prompt, system_prompt=None, **kwargs) -> str:
+        pass
+
+    def generate(self, content, system_prompt=Union[str | List[Union[str, bytes]]], **kwargs):
+        sys_prompt_arg = system_prompt if system_prompt else self.system_prompt
+
+        if isinstance(content, str):
+            return self._generate_from_single_prompt(content=content, system_prompt=sys_prompt_arg, **kwargs)
+
+        elif isinstance(content, list):
+            has_multimodal_input = any(isinstance(item, bytes) for item in content)
+            if has_multimodal_input and not self.is_multimodal:
+                raise NotImplementedError("Multimodal generation flag is not set, but multimodal input is provided. "
+                                          "Is this model multimodal?")
+
+            return self._generate_from_multiple_input(content=content, system_prompt=sys_prompt_arg, **kwargs)
+
+    def __call__(self, *args, **kwargs):
+        pass
diff --git a/textgrad/engine_experimental/engine_utils.py b/textgrad/engine_experimental/engine_utils.py
new file mode 100644
index 0000000..9e6b4de
--- /dev/null
+++ b/textgrad/engine_experimental/engine_utils.py
@@ -0,0 +1,42 @@
+from typing import List, Union
+import base64
+
+def is_jpeg(data):
+    jpeg_signature = b'\xFF\xD8\xFF'
+    return data.startswith(jpeg_signature)
+
+def is_png(data):
+    png_signature = b'\x89\x50\x4E\x47\x0D\x0A\x1A\x0A'
+    return data.startswith(png_signature)
+
+def get_image_type_from_bytes(data):
+    if is_jpeg(data):
+        return "jpeg"
+    elif is_png(data):
+        return "png"
+    else:
+        raise ValueError("Image type not supported, only jpeg and png supported.")
+
+def open_ai_like_formatting(content: List[Union[str, bytes]]) -> List[dict]:
+    """Helper function to format a list of strings and bytes into a list of dictionaries to pass as messages to the API.
+    """
+    formatted_content = []
+    for item in content:
+        if isinstance(item, bytes):
+            # For now, bytes are assumed to be images
+            image_type = get_image_type_from_bytes(item)
+            base64_image = base64.b64encode(item).decode('utf-8')
+            formatted_content.append({
+                "type": "image_url",
+                "image_url": {
+                    "url": f"data:image/{image_type};base64,{base64_image}"
+                }
+            })
+        elif isinstance(item, str):
+            formatted_content.append({
+                "type": "text",
+                "text": item
+            })
+        else:
+            raise ValueError(f"Unsupported input type: {type(item)}")
+    return formatted_content
\ No newline at end of file
diff --git a/textgrad/engine_experimental/litellm.py b/textgrad/engine_experimental/litellm.py
new file mode 100644
index 0000000..48b4c10
--- /dev/null
+++ b/textgrad/engine_experimental/litellm.py
@@ -0,0 +1,61 @@
+from litellm import completion
+from textgrad.engine_experimental.base import EngineLM, cached
+import diskcache as dc
+from typing import Union, List
+from .engine_utils import open_ai_like_formatting
+from tenacity import (
+    retry,
+    stop_after_attempt,
+    wait_random_exponential,
+)
+
+class LiteLLMEngine(EngineLM):
+    def lite_llm_generate(self, content, system_prompt=None, **kwargs) -> str:
+        messages = [
+            {"role": "system", "content": system_prompt},
+            {"role": "user", "content": content},
+        ]
+
+        return completion(model=self.model_string,
+                          messages=messages)['choices'][0]['message']['content']
+
+    DEFAULT_SYSTEM_PROMPT = "You are a helpful, creative, and smart assistant."
+
+    def __init__(self,
+                 model_string: str,
+                 system_prompt: str = DEFAULT_SYSTEM_PROMPT,
+                 is_multimodal: bool = True,
+                 cache=Union[dc.Cache, bool]):
+
+        super().__init__(
+            model_string=model_string,
+            system_prompt=system_prompt,
+            is_multimodal=is_multimodal,
+            cache=cache
+        )
+
+    @cached
+    @retry(wait=wait_random_exponential(min=1, max=5), stop=stop_after_attempt(3))
+    def _generate_from_single_prompt(
+            self, content: str, system_prompt: str = None, temperature=0, max_tokens=2000, top_p=0.99
+    ):
+
+        return self.lite_llm_generate(content, system_prompt)
+
+    @cached
+    @retry(wait=wait_random_exponential(min=1, max=5), stop=stop_after_attempt(3))
+    def _generate_from_multiple_input(
+            self, content: List[Union[str, bytes]], system_prompt=None, temperature=0, max_tokens=2000, top_p=0.99
+    ):
+        formatted_content = open_ai_like_formatting(content)
+
+        return self.lite_llm_generate(formatted_content, system_prompt)
+
+    def __call__(self, content, **kwargs):
+        return self.generate(content, **kwargs)
+
+
+
+
+
+
diff --git a/textgrad/engine_experimental/openai.py b/textgrad/engine_experimental/openai.py
new file mode 100644
index 0000000..2fea80f
--- /dev/null
+++ b/textgrad/engine_experimental/openai.py
@@ -0,0 +1,106 @@
+try:
+    from openai import OpenAI
+except ImportError:
+    raise ImportError("If you'd like to use OpenAI models, please install the openai package by running `pip install openai`, and add 'OPENAI_API_KEY' to your environment variables.")
+
+import os
+from typing import List, Union
+from textgrad.engine_experimental.engine_utils import open_ai_like_formatting
+from textgrad.engine_experimental.base import EngineLM, cached
+import diskcache as dc
+from tenacity import (
+    retry,
+    stop_after_attempt,
+    wait_random_exponential,
+)
+
+class OpenAIEngine(EngineLM):
+    DEFAULT_SYSTEM_PROMPT = "You are a helpful, creative, and smart assistant."
+
+    def __init__(self, model_string: str,
+                 system_prompt: str = DEFAULT_SYSTEM_PROMPT,
+                 is_multimodal: bool = False,
+                 cache=Union[dc.Cache, bool]):
+
+        self.validate()
+
+        super().__init__(
+            model_string=model_string,
+            system_prompt=system_prompt,
+            is_multimodal=is_multimodal,
+            cache=cache
+        )
+
+        self.client = OpenAI(
+            api_key=os.getenv("OPENAI_API_KEY")
+        )
+
+    def validate(self) -> None:
+        if os.getenv("OPENAI_API_KEY") is None:
+            raise ValueError(
+                "Please set the OPENAI_API_KEY environment variable if you'd like to use OpenAI models.")
+
+    def openai_call(self, user_content, system_prompt, temperature, max_tokens, top_p):
+        response = self.client.chat.completions.create(
+            model=self.model_string,
+            messages=[
+                {"role": "system", "content": system_prompt},
+                {"role": "user", "content": user_content},
+            ],
+            frequency_penalty=0,
+            presence_penalty=0,
+            stop=None,
+            temperature=temperature,
+            max_tokens=max_tokens,
+            top_p=top_p,
+        )
+
+        return response.choices[0].message.content
+
+    @cached
+    @retry(wait=wait_random_exponential(min=1, max=5), stop=stop_after_attempt(3))
+    def _generate_from_single_prompt(
+            self, content: str, system_prompt: str = None, temperature=0, max_tokens=2000, top_p=0.99
+    ):
+
+        return self.openai_call(content, system_prompt, temperature, max_tokens, top_p)
+
+    @cached
+    @retry(wait=wait_random_exponential(min=1, max=5), stop=stop_after_attempt(3))
+    def _generate_from_multiple_input(
+            self, content: List[Union[str, bytes]], system_prompt=None, temperature=0, max_tokens=2000, top_p=0.99
+    ):
+        formatted_content = open_ai_like_formatting(content)
+
+        return self.openai_call(formatted_content, system_prompt, temperature, max_tokens, top_p)
+
+    def __call__(self, content, **kwargs):
+        return self.generate(content, **kwargs)
+
+
+
+class OpenAICompatibleEngine(OpenAIEngine):
+    """
+        This is the same as engine.openai.ChatOpenAI, but we pass in an external OpenAI client.
+        """
+
+    DEFAULT_SYSTEM_PROMPT = "You are a helpful, creative, and smart assistant."
+    client = None
+
+    def __init__(self,
+                 client,
+                 model_string: str,
+                 system_prompt: str = DEFAULT_SYSTEM_PROMPT,
+                 is_multimodal: bool = False,
+                 cache=Union[dc.Cache, bool]):
+
+            self.client = client
+
+            super().__init__(
+                model_string=model_string,
+                system_prompt=system_prompt,
+                is_multimodal=is_multimodal,
+                cache=cache
+            )
+
+

From d9d9d42d1e23ab686fec716763b1350b455a0164 Mon Sep 17 00:00:00 2001
From: vinid <chiccobia@gmail.com>
Date: Sun, 8 Sep 2024 16:34:37 -0400
Subject: [PATCH 03/10] remove unused

---
 textgrad/engine/engine_utils.py |  16 +++
 textgrad/engine/utils.py        | 230 --------------------------------
 2 files changed, 16 insertions(+), 230 deletions(-)
 create mode 100644 textgrad/engine/engine_utils.py
 delete mode 100644 textgrad/engine/utils.py

diff --git a/textgrad/engine/engine_utils.py b/textgrad/engine/engine_utils.py
new file mode 100644
index 0000000..51caff9
--- /dev/null
+++ b/textgrad/engine/engine_utils.py
@@ -0,0 +1,16 @@
+def is_jpeg(data):
+    jpeg_signature = b'\xFF\xD8\xFF'
+    return data.startswith(jpeg_signature)
+
+def is_png(data):
+    png_signature = b'\x89\x50\x4E\x47\x0D\x0A\x1A\x0A'
+    return data.startswith(png_signature)
+
+
+def get_image_type_from_bytes(data):
+    if is_jpeg(data):
+        return "jpeg"
+    elif is_png(data):
+        return "png"
+    else:
+        raise ValueError("Image type not supported, only jpeg and png supported.")
\ No newline at end of file
diff --git a/textgrad/engine/utils.py b/textgrad/engine/utils.py
deleted file mode 100644
index 890d5cd..0000000
--- a/textgrad/engine/utils.py
+++ /dev/null
@@ -1,230 +0,0 @@
-import hashlib
-
-try:
-    from openai import OpenAI, AzureOpenAI
-except ImportError:
-    raise ImportError("If you'd like to use OpenAI models, please install the openai package by running `pip install openai`, and add 'OPENAI_API_KEY' to your environment variables.")
-
-import os
-import json
-from abc import ABC, abstractmethod
-import base64
-import platformdirs
-from tenacity import (
-    retry,
-    stop_after_attempt,
-    wait_random_exponential,
-)
-from typing import List, Union
-from functools import wraps
-from .engine_utils import get_image_type_from_bytes
-
-def cached(func):
-    @wraps(func)
-    def wrapper(self, *args, **kwargs):
-        if not self.use_cache:
-            return func(self, *args, **kwargs)
-
-        # get string representation from args and kwargs
-        key = hash(str(args) + str(kwargs))
-        key = hashlib.sha256(f"{key}".encode()).hexdigest()
-
-        if key in self.cache:
-            return self.cache[key]
-
-        result = func(self, *args, **kwargs)
-        self.cache[key] = result
-        return result
-
-    return wrapper
-
-
-class EngineLM(ABC):
-    system_prompt: str = "You are a helpful, creative, and smart assistant."
-    model_string: str
-    is_multimodal: bool
-    use_cache: bool = False
-    cache_path: str
-
-    @abstractmethod
-    def _generate_from_multiple_input(self, prompt, system_prompt=None, **kwargs) -> str:
-        pass
-
-    @abstractmethod
-    def _generate_from_single_prompt(self, prompt, system_prompt=None, **kwargs) -> str:
-        pass
-
-    # TBF this could be simplified to a single generate method
-    def _prepare_generate_from_single_prompt(self, prompt: str, system_prompt: str = None, **kwargs):
-        sys_prompt_arg = system_prompt if system_prompt else self.system_prompt
-        return self._generate_from_single_prompt(prompt, system_prompt=sys_prompt_arg, **kwargs)
-
-    def _prepare_generate_from_multiple_input(self, content: List[Union[str, bytes]], system_prompt=None, **kwargs):
-        sys_prompt_arg = system_prompt if system_prompt else self.system_prompt
-        return self._generate_from_multiple_input(content, system_prompt=sys_prompt_arg, **kwargs)
-
-    def generate(self, content, system_prompt=None, **kwargs):
-        if isinstance(content, str):
-            return self._prepare_generate_from_single_prompt(content, system_prompt=system_prompt, **kwargs)
-
-        elif isinstance(content, list):
-            has_multimodal_input = any(isinstance(item, bytes) for item in content)
-            if (has_multimodal_input) and (not self.is_multimodal):
-                raise NotImplementedError("Multimodal generation is only supported for Claude-3 and beyond.")
-
-            return self._prepare_generate_from_multiple_input(content, system_prompt=system_prompt, **kwargs)
-
-    def __call__(self, *args, **kwargs):
-        pass
-
-class OpenAIEngine(EngineLM):
-    DEFAULT_SYSTEM_PROMPT = "You are a helpful, creative, and smart assistant."
-
-    def __init__(
-            self,
-            model_string: str = "gpt-3.5-turbo-0613",
-            system_prompt: str = DEFAULT_SYSTEM_PROMPT,
-            is_multimodal: bool = False,
-            use_cache: bool = False,
-            base_url: str = None):
-        """
-        :param model_string:
-        :param system_prompt:
-        :param base_url: Used to support Ollama
-        """
-        root = platformdirs.user_cache_dir("textgrad")
-        cache_path = os.path.join(root, f"cache_openai_{model_string}.db")
-
-        super().__init__(system_prompt=system_prompt,
-                         cache_path=cache_path,
-                         use_cache=use_cache,
-                         is_multimodal=is_multimodal,
-                         model_string=model_string)
-
-        self.base_url = base_url
-
-        if not base_url:
-            if os.getenv("OPENAI_API_KEY") is None:
-                raise ValueError(
-                    "Please set the OPENAI_API_KEY environment variable if you'd like to use OpenAI models.")
-
-            self.client = OpenAI(
-                api_key=os.getenv("OPENAI_API_KEY")
-            )
-        else:
-            raise ValueError("Invalid base URL provided. Please use the default OLLAMA base URL or None.")
-
-    @cached
-    def _generate_from_single_prompt(
-            self, prompt: str, system_prompt: str = None, temperature=0, max_tokens=2000, top_p=0.99
-    ):
-
-        response = self.client.chat.completions.create(
-            model=self.model_string,
-            messages=[
-                {"role": "system", "content": system_prompt},
-                {"role": "user", "content": prompt},
-            ],
-            frequency_penalty=0,
-            presence_penalty=0,
-            stop=None,
-            temperature=temperature,
-            max_tokens=max_tokens,
-            top_p=top_p,
-        )
-
-        response = response.choices[0].message.content
-        return response
-
-    @cached
-    def _generate_from_multiple_input(
-            self, content: List[Union[str, bytes]], system_prompt=None, temperature=0, max_tokens=2000, top_p=0.99
-    ):
-        formatted_content = self._format_content(content)
-
-        response = self.client.chat.completions.create(
-            model=self.model_string,
-            messages=[
-                {"role": "system", "content": system_prompt},
-                {"role": "user", "content": formatted_content},
-            ],
-            temperature=temperature,
-            max_tokens=max_tokens,
-            top_p=top_p,
-        )
-
-        response_text = response.choices[0].message.content
-        return response_text
-
-    def __call__(self, prompt, **kwargs):
-        return self.generate(prompt, **kwargs)
-
-    def _format_content(self, content: List[Union[str, bytes]]) -> List[dict]:
-        """Helper function to format a list of strings and bytes into a list of dictionaries to pass as messages to the API.
-        """
-        formatted_content = []
-        for item in content:
-            if isinstance(item, bytes):
-                # For now, bytes are assumed to be images
-                image_type = get_image_type_from_bytes(item)
-                base64_image = base64.b64encode(item).decode('utf-8')
-                formatted_content.append({
-                    "type": "image_url",
-                    "image_url": {
-                        "url": f"data:image/{image_type};base64,{base64_image}"
-                    }
-                })
-            elif isinstance(item, str):
-                formatted_content.append({
-                    "type": "text",
-                    "text": item
-                })
-            else:
-                raise ValueError(f"Unsupported input type: {type(item)}")
-        return formatted_content
-
-
-class OpenAICompatibleEngine(OpenAIEngine):
-    """
-        This is the same as engine.openai.ChatOpenAI, but we pass in an external OpenAI client.
-        """
-
-    DEFAULT_SYSTEM_PROMPT = "You are a helpful, creative, and smart assistant."
-    client = None
-
-    def __init__(
-            self,
-            client: OpenAI,
-            model_string: str = "gpt-3.5-turbo-0613",
-            system_prompt: str = DEFAULT_SYSTEM_PROMPT,
-            is_multimodal: bool = False,
-            use_cache: bool = False,
-            base_url: str = None):
-        """
-        :param client: an OpenAI client object.
-        :param model_string: the model name, used for the cache file name and chat completion requests.
-        :param system_prompt: the system prompt to use in chat completions.
-
-        Example usage with lm-studio local server, but any client that follows the OpenAI API will work.
-
-        ```python
-        from openai import OpenAI
-        from textgrad.engine.local_model_openai_api import ChatExternalClient
-
-        client = OpenAI(base_url="http://localhost:1234/v1", api_key="lm-studio")
-        engine = ChatExternalClient(client=client, model_string="your-model-name")
-        print(engine.generate(max_tokens=40, prompt="What is the meaning of life?"))
-        ```
-
-        """
-
-        if os.getenv("OPENAI_API_KEY") is None:
-            os.environ["OPENAI_API_KEY"] = client.api_key
-
-        self.client = client
-
-        super.__init__(model_string=model_string,
-                       system_prompt=system_prompt,
-                       is_multimodal=is_multimodal,
-                       use_cache=use_cache,
-                       base_url=base_url)
\ No newline at end of file

From f81cb8ad0b44695b8069ade8b714ad3fa59433f5 Mon Sep 17 00:00:00 2001
From: vinid <chiccobia@gmail.com>
Date: Sun, 29 Sep 2024 09:44:04 -0400
Subject: [PATCH 04/10] enginesz

---
 README.md                                     |  78 ++++
 .../Tutorial-ExperimentalEngines.ipynb        | 133 ++++++
 ...torial-PrimitivesExperimentalEngines.ipynb | 426 ++++++++++++++++++
 textgrad/config.py                            |   4 +-
 textgrad/engine/__init__.py                   |   5 +
 textgrad/engine_experimental/base.py          |   2 +-
 6 files changed, 645 insertions(+), 3 deletions(-)
 create mode 100644 examples/notebooks/experimental_engines/Tutorial-ExperimentalEngines.ipynb
 create mode 100644 examples/notebooks/experimental_engines/Tutorial-PrimitivesExperimentalEngines.ipynb

diff --git a/README.md b/README.md
index d3dd69e..1a4a263 100644
--- a/README.md
+++ b/README.md
@@ -29,6 +29,47 @@ This API is similar to the Pytorch API, making it simple to adapt to your usecas
 
 ![Analogy with Torch](assets/analogy.png)
 
+
+### Updates:
+
+**29th Sept 2024**:
+
+We are introducing a new engine based on [litellm](https://github.com/BerriAI/litellm). This should allow
+you to use any model you like, as long as it is supported by litellm. This means that now
+**Bedrock, Together, Gemini and even more** are all supported by TextGrad!
+
+In addition to this, with the new engines it should be easy to enable and disable caching.  
+
+We are in the process of testing these new engines and deprecating the old engines. If you have any issues, please let us know!
+
+The new litellm engines can be loaded with the following code: 
+
+An example of loading a litellm engine:
+```python
+engine = get_engine("experimental:gpt-4o", cache=False)
+
+# this also works with
+
+set_backward_engine("experimental:gpt-4o", cache=False)
+```
+
+Be sure to set the relevant environment variables for the new engines!
+
+An example of forward pass:
+```python
+
+import httpx
+from textgrad.engine_experimental.litellm import LiteLLMEngine
+
+LiteLLMEngine("gpt-4o", cache=True).generate(content="hello, what's 3+4", system_prompt="you are an assistant")
+
+image_url = "https://upload.wikimedia.org/wikipedia/commons/a/a7/Camponotus_flavomarginatus_ant.jpg"
+image_data = httpx.get(image_url).content
+```
+
+In the examples folder you will find two new notebooks that show how to use the new engines.
+
+
 ## QuickStart
 If you know PyTorch, you know 80% of TextGrad.
 Let's walk through the key components with a simple example. Say we want to use GPT-4o to solve a simple
@@ -96,7 +137,42 @@ answer
 > :white_check_mark: **answer: It will still take 1 hour to dry 30 shirts under the sun,**
 > **assuming they are all laid out properly to receive equal sunlight.**
 
+### Updates:
+
+**29th Sept 2024**:
+
+We are introducing a new engine based on [litellm](https://github.com/BerriAI/litellm). This should allow
+you to use any model you like, as long as it is supported by litellm. This means that now
+**Bedrock, Together, Gemini and even more** are all supported by TextGrad!
+
+In addition to this, with the new engines it should be easy to enable and disable caching.  
+
+We are in the process of testing these new engines and deprecating the old engines. If you have any issues, please let us know!
 
+The new litellm engines can be loaded with the following code: 
+
+An example of loading a litellm engine:
+```python
+engine = get_engine("experimental:gpt-4o", cache=False)
+
+# this also works with
+
+set_backward_engine("experimental:gpt-4o", cache=False)
+```
+
+An example of forward pass:
+```python
+
+import httpx
+from textgrad.engine_experimental.litellm import LiteLLMEngine
+
+LiteLLMEngine("gpt-4o", cache=True).generate(content="hello, what's 3+4", system_prompt="you are an assistant")
+
+image_url = "https://upload.wikimedia.org/wikipedia/commons/a/a7/Camponotus_flavomarginatus_ant.jpg"
+image_data = httpx.get(image_url).content
+```
+
+In the examples folder you will find two new notebooks that show how to use the new engines.
 
 
 We have many more examples around how TextGrad can optimize all kinds of variables -- code, solutions to problems, molecules, prompts, and all that!
@@ -119,6 +195,8 @@ you need an OpenAI/Anthropic key to run the LLMs).
 
 </div>
 
+
+
 ### Installation
 
 You can install TextGrad using any of the following methods.
diff --git a/examples/notebooks/experimental_engines/Tutorial-ExperimentalEngines.ipynb b/examples/notebooks/experimental_engines/Tutorial-ExperimentalEngines.ipynb
new file mode 100644
index 0000000..7c1e79b
--- /dev/null
+++ b/examples/notebooks/experimental_engines/Tutorial-ExperimentalEngines.ipynb
@@ -0,0 +1,133 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "2661032c-2d9b-43f5-b074-28c119b57a14",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import textgrad\n",
+    "import os\n",
+    "from textgrad.engine_experimental.openai import OpenAIEngine"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "5e74c489-c47e-413c-adae-cd1201f6f94f",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "os.environ[\"OPENAI_API_KEY\"] = \"SOMETHING\"\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "50c02746-dd33-4cb0-896e-4771e4b76ed7",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "OpenAIEngine(\"gpt-4o-mini\", cache=True).generate(content=\"hello, what's 3+4\", system_prompt=\"you are an assistant\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "f0ac186e-5c34-4115-aeda-bbd301be2667",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "b09f7e36-ae4f-4746-8548-c2c189827435",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "2ccce563-3d11-4d20-9c72-05d43cce4f6c",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "3053630a-dbf3-4d3e-b553-5c8ea73e2ccd",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "58db7bb4-7f0f-4517-bba2-60a51b85908b",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "7ee586ac-473e-4807-b66d-8524d08dc236",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import httpx\n",
+    "from textgrad.engine_experimental.litellm import LiteLLMEngine\n",
+    "\n",
+    "LiteLLMEngine(\"gpt-4o\", cache=True).generate(content=\"hello, what's 3+4\", system_prompt=\"you are an assistant\")\n",
+    "\n",
+    "image_url = \"https://upload.wikimedia.org/wikipedia/commons/a/a7/Camponotus_flavomarginatus_ant.jpg\"\n",
+    "image_data = httpx.get(image_url).content"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "e5b1f1d5-8971-4dee-9958-c708ba807921",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "LiteLLMEngine(\"gpt-4o\", cache=True).generate(content=[image_data, \"what is this my boy\"], system_prompt=\"you are an assistant\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "43d9b703-4488-4222-a7fb-773293c13514",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.11.9"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}
diff --git a/examples/notebooks/experimental_engines/Tutorial-PrimitivesExperimentalEngines.ipynb b/examples/notebooks/experimental_engines/Tutorial-PrimitivesExperimentalEngines.ipynb
new file mode 100644
index 0000000..307a479
--- /dev/null
+++ b/examples/notebooks/experimental_engines/Tutorial-PrimitivesExperimentalEngines.ipynb
@@ -0,0 +1,426 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "d66880e8-23d9-4f00-b120-7f289f682511",
+   "metadata": {},
+   "source": [
+    "# TextGrad Tutorials: Primitives\n",
+    "\n",
+    "![TextGrad](https://github.com/vinid/data/blob/master/logo_full.png?raw=true)\n",
+    "\n",
+    "An autograd engine -- for textual gradients!\n",
+    "\n",
+    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/zou-group/TextGrad/blob/main/examples/notebooks/Prompt-Optimization.ipynb)\n",
+    "[![GitHub license](https://img.shields.io/badge/License-MIT-blue.svg)](https://lbesson.mit-license.org/)\n",
+    "[![Arxiv](https://img.shields.io/badge/arXiv-2406.07496-B31B1B.svg)](https://arxiv.org/abs/2406.07496)\n",
+    "[![Documentation Status](https://readthedocs.org/projects/textgrad/badge/?version=latest)](https://textgrad.readthedocs.io/en/latest/?badge=latest)\n",
+    "[![PyPI - Python Version](https://img.shields.io/pypi/pyversions/textgrad)](https://pypi.org/project/textgrad/)\n",
+    "[![PyPI](https://img.shields.io/pypi/v/textgrad)](https://pypi.org/project/textgrad/)\n",
+    "\n",
+    "**Objectives for this tutorial:**\n",
+    "\n",
+    "* Introduce you to the primitives in TextGrad\n",
+    "\n",
+    "**Requirements:**\n",
+    "\n",
+    "* You need to have an OpenAI API key to run this tutorial. This should be set as an environment variable as OPENAI_API_KEY.\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "e8149068-d7ef-4702-b5d1-ed0d7f1271fd",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "!pip install textgrad "
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "36a9c615-17a0-455c-8f9c-f0d25fb8824b",
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2024-06-11T15:43:10.594204491Z",
+     "start_time": "2024-06-11T15:43:10.589328053Z"
+    }
+   },
+   "outputs": [],
+   "source": [
+    "# you might need to restart the notebook after installing textgrad\n",
+    "\n",
+    "from textgrad.engine import get_engine\n",
+    "from textgrad import Variable\n",
+    "from textgrad.optimizer import TextualGradientDescent\n",
+    "from textgrad.loss import TextLoss\n",
+    "from dotenv import load_dotenv\n",
+    "load_dotenv()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "143e3cca-aa39-489d-b2e2-f3e19e4dad7e",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import os\n",
+    "os.environ[\"OPENAI_API_KEY\"] = \"SOMETHING\""
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "8887fbed36c7daf2",
+   "metadata": {
+    "collapsed": false,
+    "jupyter": {
+     "outputs_hidden": false
+    }
+   },
+   "source": [
+    "## Introduction: Variable\n",
+    "\n",
+    "Variables in TextGrad are the metaphorical equivalent of tensors in PyTorch. They are the primary data structure that you will interact with when using TextGrad. \n",
+    "\n",
+    "Variables keep track of gradients and manage the data.\n",
+    "\n",
+    "Variables require two arguments (and there is an optional third one):\n",
+    "\n",
+    "1. `data`: The data that the variable will hold\n",
+    "2. `role_description`: A description of the role of the variable in the computation graph\n",
+    "3. `requires_grad`: (optional) A boolean flag that indicates whether the variable requires gradients"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "c65fb4456d84c8fc",
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2024-06-11T15:43:17.669096228Z",
+     "start_time": "2024-06-11T15:43:17.665325560Z"
+    },
+    "collapsed": false,
+    "jupyter": {
+     "outputs_hidden": false
+    }
+   },
+   "outputs": [],
+   "source": [
+    "x = Variable(\"A sntence with a typo\", role_description=\"The input sentence\", requires_grad=True)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "65857dd50408ebd7",
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2024-06-11T15:43:18.184004948Z",
+     "start_time": "2024-06-11T15:43:18.178187640Z"
+    },
+    "collapsed": false,
+    "jupyter": {
+     "outputs_hidden": false
+    }
+   },
+   "outputs": [],
+   "source": [
+    "x.gradients"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "63f6a6921a1cce6a",
+   "metadata": {
+    "collapsed": false,
+    "jupyter": {
+     "outputs_hidden": false
+    }
+   },
+   "source": [
+    "## Introduction: Engine\n",
+    "\n",
+    "When we talk about the engine in TextGrad, we are referring to an LLM. The engine is an abstraction we use to interact with the model."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "281644022ac1c65d",
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2024-06-11T15:44:32.606319032Z",
+     "start_time": "2024-06-11T15:44:32.561460448Z"
+    },
+    "collapsed": false,
+    "jupyter": {
+     "outputs_hidden": false
+    }
+   },
+   "outputs": [],
+   "source": [
+    "engine = get_engine(\"experimental:gpt-4o\", cache=False)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "0d61c58c-2cc1-4482-8962-e430632cf5f8",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "engine.generate(content=\"hello, what's 3+4\", system_prompt=\"you are an assistant\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "4e41d873-a3b4-4022-b180-bb4bc13b32f7",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import litellm\n",
+    "litellm.set_verbose=True"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "33c7d6eaa115cd6a",
+   "metadata": {
+    "collapsed": false,
+    "jupyter": {
+     "outputs_hidden": false
+    }
+   },
+   "source": [
+    "This object behaves like you would expect an LLM to behave: You can sample generation from the engine using the `generate` method. "
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "37502bf67ef23c53",
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2024-06-11T17:29:41.108552705Z",
+     "start_time": "2024-06-11T17:29:40.294256814Z"
+    },
+    "collapsed": false,
+    "jupyter": {
+     "outputs_hidden": false
+    }
+   },
+   "outputs": [],
+   "source": [
+    "engine.generate(\"Hello how are you?\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "b627edc07c0d3737",
+   "metadata": {
+    "collapsed": false,
+    "jupyter": {
+     "outputs_hidden": false
+    }
+   },
+   "source": [
+    "## Introduction: Loss\n",
+    "\n",
+    "Again, Loss in TextGrad is the metaphorical equivalent of loss in PyTorch. We use Losses in different form in TextGrad but for now we will focus on a simple TextLoss. TextLoss is going to evaluate the loss wrt a string."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "252e0a0152b81f14",
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2024-06-11T15:44:32.894722136Z",
+     "start_time": "2024-06-11T15:44:32.890708561Z"
+    },
+    "collapsed": false,
+    "jupyter": {
+     "outputs_hidden": false
+    }
+   },
+   "outputs": [],
+   "source": [
+    "system_prompt = Variable(\"Evaluate the correctness of this sentence\", role_description=\"The system prompt\")\n",
+    "loss = TextLoss(system_prompt, engine=engine)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "8865cb0d-bab5-4695-9cee-5fc939a8decc",
+   "metadata": {
+    "collapsed": false,
+    "jupyter": {
+     "outputs_hidden": false
+    }
+   },
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "markdown",
+   "id": "6f05ec2bf907b3ba",
+   "metadata": {
+    "collapsed": false,
+    "jupyter": {
+     "outputs_hidden": false
+    }
+   },
+   "source": [
+    "## Introduction: Optimizer\n",
+    "\n",
+    "Keeping on the analogy with PyTorch, the optimizer in TextGrad is the object that will update the parameters of the model. In this case, the parameters are the variables that have `requires_grad` set to `True`.\n",
+    "\n",
+    "**NOTE** This is a text optimizer! It will do all operations with text! "
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "78f93f80b9e3ad36",
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2024-06-11T15:44:33.741130951Z",
+     "start_time": "2024-06-11T15:44:33.734977769Z"
+    },
+    "collapsed": false,
+    "jupyter": {
+     "outputs_hidden": false
+    }
+   },
+   "outputs": [],
+   "source": [
+    "optimizer = TextualGradientDescent(parameters=[x], engine=engine)\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "d26883eb74ce0d01",
+   "metadata": {
+    "collapsed": false,
+    "jupyter": {
+     "outputs_hidden": false
+    }
+   },
+   "source": [
+    "## Putting it all together\n",
+    "\n",
+    "We can now put all the pieces together. We have a variable, an engine, a loss, and an optimizer. We can now run a single optimization step."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "9817e0ae0179376d",
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2024-06-11T15:44:41.730132530Z",
+     "start_time": "2024-06-11T15:44:34.997777872Z"
+    },
+    "collapsed": false,
+    "jupyter": {
+     "outputs_hidden": false
+    }
+   },
+   "outputs": [],
+   "source": [
+    "l = loss(x)\n",
+    "l.backward(engine)\n",
+    "optimizer.step()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "77e3fab0efdd579e",
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2024-06-11T15:44:41.738985151Z",
+     "start_time": "2024-06-11T15:44:41.731989729Z"
+    },
+    "collapsed": false,
+    "jupyter": {
+     "outputs_hidden": false
+    }
+   },
+   "outputs": [],
+   "source": [
+    "x.value"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "6a8aab93b80fb82c",
+   "metadata": {
+    "collapsed": false,
+    "jupyter": {
+     "outputs_hidden": false
+    }
+   },
+   "source": [
+    "While here it is not going to be useful, we can also do multiple optimization steps in a loop! Do not forget to reset the gradients after each step!"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "d6bb8d0dcc2539d1",
+   "metadata": {
+    "ExecuteTime": {
+     "start_time": "2024-06-11T15:44:30.989940227Z"
+    },
+    "collapsed": false,
+    "jupyter": {
+     "outputs_hidden": false
+    }
+   },
+   "outputs": [],
+   "source": [
+    "optimizer.zero_grad()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "a3a84aad4cd58737",
+   "metadata": {
+    "collapsed": false,
+    "jupyter": {
+     "outputs_hidden": false
+    }
+   },
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.11.9"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}
diff --git a/textgrad/config.py b/textgrad/config.py
index 1b01388..0684488 100644
--- a/textgrad/config.py
+++ b/textgrad/config.py
@@ -42,10 +42,10 @@ def get_engine(self):
         """
         return self.engine
 
-def set_backward_engine(engine: Union[EngineLM, str], override: bool = False):
+def set_backward_engine(engine: Union[EngineLM, str], override: bool = False, **kwargs):
     singleton_backward_engine = SingletonBackwardEngine()
     if isinstance(engine, str):
-        engine = get_engine(engine)
+        engine = get_engine(engine, **kwargs)
     singleton_backward_engine.set_engine(engine, override=override)
 
 
diff --git a/textgrad/engine/__init__.py b/textgrad/engine/__init__.py
index 3e0ee24..ea7c6c1 100644
--- a/textgrad/engine/__init__.py
+++ b/textgrad/engine/__init__.py
@@ -1,4 +1,5 @@
 from .base import EngineLM, CachedEngine
+from textgrad.engine_experimental.litellm import LiteLLMEngine
 
 __ENGINE_NAME_SHORTCUTS__ = {
     "opus": "claude-3-opus-20240229",
@@ -34,6 +35,10 @@ def get_engine(engine_name: str, **kwargs) -> EngineLM:
     if "seed" in kwargs and "gpt-4" not in engine_name and "gpt-3.5" not in engine_name and "gpt-35" not in engine_name:
         raise ValueError(f"Seed is currently supported only for OpenAI engines, not {engine_name}")
 
+    # check if engine_name starts with "experimental:"
+    if engine_name.startswith("experimental:"):
+        engine_name = engine_name.split("experimental:")[1]
+        return LiteLLMEngine(model_string=engine_name, **kwargs)
     if engine_name.startswith("azure"):
         from .openai import AzureChatOpenAI
         # remove engine_name "azure-" prefix
diff --git a/textgrad/engine_experimental/base.py b/textgrad/engine_experimental/base.py
index 3203a85..79c9da1 100644
--- a/textgrad/engine_experimental/base.py
+++ b/textgrad/engine_experimental/base.py
@@ -64,7 +64,7 @@ def _generate_from_multiple_input(self, prompt, system_prompt=None, **kwargs) ->
     def _generate_from_single_prompt(self, prompt, system_prompt=None, **kwargs) -> str:
         pass
 
-    def generate(self, content, system_prompt=Union[str | List[Union[str, bytes]]], **kwargs):
+    def generate(self, content, system_prompt: Union[str | List[Union[str, bytes]]] = None, **kwargs):
         sys_prompt_arg = system_prompt if system_prompt else self.system_prompt
 
         if isinstance(content, str):

From 29a195e243ab3ff0d7cfa00709a41d55f5173dba Mon Sep 17 00:00:00 2001
From: vinid <chiccobia@gmail.com>
Date: Sun, 29 Sep 2024 09:47:09 -0400
Subject: [PATCH 05/10] add depend

---
 requirements.txt | 1 +
 1 file changed, 1 insertion(+)

diff --git a/requirements.txt b/requirements.txt
index 10a2483..2199e46 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -7,5 +7,6 @@ datasets>=2.14.6
 diskcache>=5.6.3
 graphviz>=0.20.3
 gdown>=5.2.0
+litellm==1.44.22
 pillow
 httpx
\ No newline at end of file

From 202662c5f781540299679928886ef1cba13a9d19 Mon Sep 17 00:00:00 2001
From: vinid <chiccobia@gmail.com>
Date: Sun, 29 Sep 2024 09:49:51 -0400
Subject: [PATCH 06/10] docstrings

---
 textgrad/engine_experimental/base.py | 14 ++++++++++++++
 1 file changed, 14 insertions(+)

diff --git a/textgrad/engine_experimental/base.py b/textgrad/engine_experimental/base.py
index 79c9da1..e50e613 100644
--- a/textgrad/engine_experimental/base.py
+++ b/textgrad/engine_experimental/base.py
@@ -40,6 +40,19 @@ def __init__(self, model_string: str,
                  is_multimodal: bool = False,
                  cache=Union[dc.Cache, bool]):
 
+        """
+        Base class for the engines.
+
+        :param model_string: The model string to use.
+        :type model_string: str
+        :param system_prompt: The system prompt to use. Defaults to "You are a helpful, creative, and smart assistant."
+        :type system_prompt: str
+        :param is_multimodal: Whether the model is multimodal. Defaults to False.
+        :type is_multimodal: bool
+        :param cache: The cache to use. Defaults to True. Note that cache can also be a diskcache.Cache object.
+        :type cache: Union[diskcache.Cache, bool]
+        """
+
         root = platformdirs.user_cache_dir("textgrad")
         default_cache_path = os.path.join(root, f"cache_model_{model_string}.db")
 
@@ -47,6 +60,7 @@ def __init__(self, model_string: str,
         self.system_prompt = system_prompt
         self.is_multimodal = is_multimodal
 
+        # cache resolution
         if isinstance(cache, dc.Cache):
             self.cache = cache
         elif cache is True:

From 12c62c931d4ac252f57e11efea5e137f90bd0334 Mon Sep 17 00:00:00 2001
From: vinid <chiccobia@gmail.com>
Date: Sun, 29 Sep 2024 10:29:12 -0400
Subject: [PATCH 07/10] readme

---
 README.md | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/README.md b/README.md
index 1a4a263..ed585b1 100644
--- a/README.md
+++ b/README.md
@@ -65,6 +65,8 @@ LiteLLMEngine("gpt-4o", cache=True).generate(content="hello, what's 3+4", system
 
 image_url = "https://upload.wikimedia.org/wikipedia/commons/a/a7/Camponotus_flavomarginatus_ant.jpg"
 image_data = httpx.get(image_url).content
+
+LiteLLMEngine("gpt-4o", cache=True).generate(content=[image_data, "what is this my boy"], system_prompt="you are an assistant")
 ```
 
 In the examples folder you will find two new notebooks that show how to use the new engines.

From 4999f1fff836e842e6fb95d01090f5aef92b2e33 Mon Sep 17 00:00:00 2001
From: vinid <chiccobia@gmail.com>
Date: Sun, 6 Oct 2024 08:03:09 -0400
Subject: [PATCH 08/10] tests and readme

---
 README.md                               | 38 ------------------
 tests/test_basics.py                    | 51 +++++++++++++++++++++++++
 textgrad/engine/__init__.py             |  3 ++
 textgrad/engine_experimental/litellm.py |  2 +-
 4 files changed, 55 insertions(+), 39 deletions(-)

diff --git a/README.md b/README.md
index ed585b1..11e1a57 100644
--- a/README.md
+++ b/README.md
@@ -139,44 +139,6 @@ answer
 > :white_check_mark: **answer: It will still take 1 hour to dry 30 shirts under the sun,**
 > **assuming they are all laid out properly to receive equal sunlight.**
 
-### Updates:
-
-**29th Sept 2024**:
-
-We are introducing a new engine based on [litellm](https://github.com/BerriAI/litellm). This should allow
-you to use any model you like, as long as it is supported by litellm. This means that now
-**Bedrock, Together, Gemini and even more** are all supported by TextGrad!
-
-In addition to this, with the new engines it should be easy to enable and disable caching.  
-
-We are in the process of testing these new engines and deprecating the old engines. If you have any issues, please let us know!
-
-The new litellm engines can be loaded with the following code: 
-
-An example of loading a litellm engine:
-```python
-engine = get_engine("experimental:gpt-4o", cache=False)
-
-# this also works with
-
-set_backward_engine("experimental:gpt-4o", cache=False)
-```
-
-An example of forward pass:
-```python
-
-import httpx
-from textgrad.engine_experimental.litellm import LiteLLMEngine
-
-LiteLLMEngine("gpt-4o", cache=True).generate(content="hello, what's 3+4", system_prompt="you are an assistant")
-
-image_url = "https://upload.wikimedia.org/wikipedia/commons/a/a7/Camponotus_flavomarginatus_ant.jpg"
-image_data = httpx.get(image_url).content
-```
-
-In the examples folder you will find two new notebooks that show how to use the new engines.
-
-
 We have many more examples around how TextGrad can optimize all kinds of variables -- code, solutions to problems, molecules, prompts, and all that!
 
 ### Tutorials
diff --git a/tests/test_basics.py b/tests/test_basics.py
index dba4fa6..0f06b8c 100644
--- a/tests/test_basics.py
+++ b/tests/test_basics.py
@@ -68,6 +68,57 @@ def test_openai_engine():
     os.environ['OPENAI_API_KEY'] = "fake_key"
     engine = ChatOpenAI()
 
+
+def test_set_backward_engine():
+    from textgrad.config import set_backward_engine, SingletonBackwardEngine
+    from textgrad.engine.openai import ChatOpenAI
+    from textgrad.engine_experimental.litellm import LiteLLMEngine
+
+    engine = ChatOpenAI()
+    set_backward_engine(engine, override=False)
+    assert SingletonBackwardEngine().get_engine() == engine
+
+    new_engine = LiteLLMEngine(model_string="gpt-3.5-turbo-0613")
+    set_backward_engine(new_engine, True)
+    assert SingletonBackwardEngine().get_engine() == new_engine
+
+    with pytest.raises(Exception):
+        set_backward_engine(engine, False)
+
+def test_get_engine():
+    from textgrad.engine import get_engine
+    from textgrad.engine.openai import ChatOpenAI
+    from textgrad.engine_experimental.litellm import LiteLLMEngine
+
+    engine = get_engine("gpt-3.5-turbo-0613")
+    assert isinstance(engine, ChatOpenAI)
+
+    engine = get_engine("experimental:claude-3-opus-20240229")
+    assert isinstance(engine, LiteLLMEngine)
+
+    engine = get_engine("experimental:claude-3-opus-20240229", cache=True)
+    assert isinstance(engine, LiteLLMEngine)
+
+    engine = get_engine("experimental:claude-3-opus-20240229", cache=False)
+    assert isinstance(engine, LiteLLMEngine)
+
+    # get local diskcache
+    from diskcache import Cache
+    cache = Cache("./cache")
+
+    engine = get_engine("experimental:claude-3-opus-20240229", cache=cache)
+    assert isinstance(engine, LiteLLMEngine)
+
+    with pytest.raises(ValueError):
+        get_engine("invalid-engine")
+
+    with pytest.raises(ValueError):
+        get_engine("experimental:claude-3-opus-20240229", cache=[1,2,3])
+
+    with pytest.raises(ValueError):
+        get_engine("gpt-4o", cache=True)
+
+
 # Test importing main components from textgrad
 def test_import_main_components():
     from textgrad import Variable, TextualGradientDescent, EngineLM
diff --git a/textgrad/engine/__init__.py b/textgrad/engine/__init__.py
index ea7c6c1..2697faf 100644
--- a/textgrad/engine/__init__.py
+++ b/textgrad/engine/__init__.py
@@ -35,6 +35,9 @@ def get_engine(engine_name: str, **kwargs) -> EngineLM:
     if "seed" in kwargs and "gpt-4" not in engine_name and "gpt-3.5" not in engine_name and "gpt-35" not in engine_name:
         raise ValueError(f"Seed is currently supported only for OpenAI engines, not {engine_name}")
 
+    if "cache" in kwargs and "experimental" not in engine_name:
+        raise ValueError(f"Cache is currently supported only for LiteLLM engines, not {engine_name}")
+
     # check if engine_name starts with "experimental:"
     if engine_name.startswith("experimental:"):
         engine_name = engine_name.split("experimental:")[1]
diff --git a/textgrad/engine_experimental/litellm.py b/textgrad/engine_experimental/litellm.py
index 48b4c10..e28845c 100644
--- a/textgrad/engine_experimental/litellm.py
+++ b/textgrad/engine_experimental/litellm.py
@@ -25,7 +25,7 @@ def __init__(self,
                  model_string: str,
                  system_prompt: str = DEFAULT_SYSTEM_PROMPT,
                  is_multimodal: bool = True,
-                 cache=Union[dc.Cache, bool]):
+                 cache: Union[dc.Cache, bool] = False):
 
         super().__init__(
             model_string=model_string,

From 33d9051b7fbca22be76d8c3cdd5bc06f02de2b1a Mon Sep 17 00:00:00 2001
From: vinid <chiccobia@gmail.com>
Date: Sun, 6 Oct 2024 08:10:54 -0400
Subject: [PATCH 09/10] tests and readme

---
 README.md | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/README.md b/README.md
index 11e1a57..0568efe 100644
--- a/README.md
+++ b/README.md
@@ -38,6 +38,8 @@ We are introducing a new engine based on [litellm](https://github.com/BerriAI/li
 you to use any model you like, as long as it is supported by litellm. This means that now
 **Bedrock, Together, Gemini and even more** are all supported by TextGrad!
 
+This should be seen as experimental but we plan to depreciate the old engines in the future.
+
 In addition to this, with the new engines it should be easy to enable and disable caching.  
 
 We are in the process of testing these new engines and deprecating the old engines. If you have any issues, please let us know!

From fd456ddb53b47b56cc9bf0645d2ec5b5c5120135 Mon Sep 17 00:00:00 2001
From: vinid <chiccobia@gmail.com>
Date: Sun, 6 Oct 2024 08:17:15 -0400
Subject: [PATCH 10/10] update

---
 textgrad/engine_experimental/base.py   | 2 +-
 textgrad/engine_experimental/openai.py | 4 ++--
 2 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/textgrad/engine_experimental/base.py b/textgrad/engine_experimental/base.py
index e50e613..8381446 100644
--- a/textgrad/engine_experimental/base.py
+++ b/textgrad/engine_experimental/base.py
@@ -38,7 +38,7 @@ class EngineLM(ABC):
     def __init__(self, model_string: str,
                  system_prompt: str = "You are a helpful, creative, and smart assistant.",
                  is_multimodal: bool = False,
-                 cache=Union[dc.Cache, bool]):
+                 cache: Union[dc.Cache, bool] = False):
 
         """
         Base class for the engines.
diff --git a/textgrad/engine_experimental/openai.py b/textgrad/engine_experimental/openai.py
index 2fea80f..a73d40a 100644
--- a/textgrad/engine_experimental/openai.py
+++ b/textgrad/engine_experimental/openai.py
@@ -20,7 +20,7 @@ class OpenAIEngine(EngineLM):
     def __init__(self, model_string: str,
                  system_prompt: str = DEFAULT_SYSTEM_PROMPT,
                  is_multimodal: bool = False,
-                 cache=Union[dc.Cache, bool]):
+                 cache: Union[dc.Cache, bool] = False):
 
         self.validate()
 
@@ -92,7 +92,7 @@ def __init__(self,
                  model_string: str,
                  system_prompt: str = DEFAULT_SYSTEM_PROMPT,
                  is_multimodal: bool = False,
-                 cache=Union[dc.Cache, bool]):
+                 cache: Union[dc.Cache, bool] = False):
 
             self.client = client