From b5d5509e24d13c127753f02a42c1503e739eba58 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?R=C3=A9mi=20Louf?= <remi@thetypicalset.com>
Date: Fri, 26 Jul 2024 14:33:06 +0200
Subject: [PATCH 1/5] Update branch name in workflows

---
 .github/workflows/build_docs.yml   | 2 +-
 .github/workflows/publish_docs.yml | 6 +++---
 .github/workflows/tests.yml        | 4 ++--
 3 files changed, 6 insertions(+), 6 deletions(-)

diff --git a/.github/workflows/build_docs.yml b/.github/workflows/build_docs.yml
index 53060fd..db55a28 100644
--- a/.github/workflows/build_docs.yml
+++ b/.github/workflows/build_docs.yml
@@ -2,7 +2,7 @@ name: Build the documentation
 
 on:
   pull_request:
-    branches: [main]
+    branches: [master]
 
 jobs:
   build:
diff --git a/.github/workflows/publish_docs.yml b/.github/workflows/publish_docs.yml
index 4952058..7f9a22e 100644
--- a/.github/workflows/publish_docs.yml
+++ b/.github/workflows/publish_docs.yml
@@ -3,7 +3,7 @@ name: Publish the documentation
 on:
   push:
     branches:
-      - main
+      - master
 
 permissions:
   contents: write
@@ -23,10 +23,10 @@ jobs:
           path: .cache
           restore-keys: |
             mkdocs-material-
-      - run:
+      - run: |
           python -m pip install --upgrade pip
           pip install .[docs]
           pip install -r requirements-doc.txt
       - name: Build documentation
-        run:
+        run: |
           mkdocs gh-deploy --force
diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml
index 1a74c50..8a322a9 100644
--- a/.github/workflows/tests.yml
+++ b/.github/workflows/tests.yml
@@ -2,9 +2,9 @@ name: Tests
 
 on:
   pull_request:
-    branches: [main]
+    branches: [master]
   push:
-    branches: [main]
+    branches: [master]
 
 jobs:
   style:

From a9a5a637ff4b665cd0ef7b27edc44769f030273f Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?R=C3=A9mi=20Louf?= <remi@thetypicalset.com>
Date: Fri, 26 Jul 2024 14:47:46 +0200
Subject: [PATCH 2/5] Update Python version requirement

---
 pyproject.toml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/pyproject.toml b/pyproject.toml
index 4cba33d..b7a049f 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -3,7 +3,7 @@ name = "prompts"
 version = "0.1.0"
 description = "Large Language Models prompting library"
 authors = [{name = "The Outlines developers", email = "contact@dottxt.co"}]
-requires-python = ">= 3.10"
+requires-python = ">= 3.8"
 dependencies = ["jinja2"]
 
 [build-system]

From b7694f1d11d77e795f8c178411aa977a9e2b8a32 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?R=C3=A9mi=20Louf?= <remi@thetypicalset.com>
Date: Fri, 26 Jul 2024 13:17:41 +0200
Subject: [PATCH 3/5] Dispatch template on model name

---
 prompts/templates.py    | 66 +++++++++++++++++++++++++++++++++--------
 tests/test_templates.py | 34 +++++++++++++++++----
 2 files changed, 82 insertions(+), 18 deletions(-)

diff --git a/prompts/templates.py b/prompts/templates.py
index 9e9c12f..82d5cc3 100644
--- a/prompts/templates.py
+++ b/prompts/templates.py
@@ -1,36 +1,43 @@
 import inspect
 import re
-from dataclasses import dataclass
-from typing import Any, Callable, Dict, List, Optional, cast
+from dataclasses import dataclass, field
+from functools import lru_cache
+from typing import Callable, Dict, Hashable, Optional, cast
 
 from jinja2 import Environment, StrictUndefined
 
 
 @dataclass
 class Template:
-    """Represents a prompt function.
+    """Represents a prompt template.
 
-    We return a `Prompt` class instead of a simple function so the
-    template defined in prompt functions can be accessed.
+    A prompt template is a callable that, given a Jinja2 template and a set of values,
+    renders the template using those values. It is recommended to instantiate `Temaplate`
+    using the `template` decorator, which extracts the template from the function's
+    docstring and its variables from the function's signature.
+
+    It is not uncommon that, for the same taks, different models will perform
+    better with different prompt. Here we thus allow to dispatch to associate a
+    prompt with a task and dispatch the prompt based on the model being used; a
+    `Template` instance is thus also a registry that associates model names to
+    other templates.
 
-    TODO: Store variables instead of signature?
 
     Attributes
     ----------
     template
         The template to render.
-    model
-        The model used for inference.
     signature
         The prompt function's signature.
+    registry
+        Registry that maps function names to their respective `Template`
+        instances.
 
     """
 
     template: str
     signature: inspect.Signature
-
-    def __post_init__(self):
-        self.parameters: List[str] = list(self.signature.parameters.keys())
+    registry: Dict[str, Callable] = field(default_factory=dict)
 
     def __call__(self, *args, **kwargs) -> str:
         """Render and return the template.
@@ -47,6 +54,40 @@ def __call__(self, *args, **kwargs) -> str:
     def __str__(self):
         return self.template
 
+    def __getitem__(self, model_name: str):
+        """Get the prompt template corresponding to a model name.
+
+        We return the default template when trying to fetch a
+        template for a model that was not registered.
+
+        Parameters
+        ----------
+        model_name
+            The name of the model whose prompt template we want to retrieve.
+
+        Returns
+        -------
+        The template registered for the model name.
+
+        """
+        try:
+            return self.registry[model_name]
+        except KeyError:
+            return self
+
+    def register(self, model_name: str):
+        """Register the prompt template, as represented by a prompt function,
+        for the model name.
+
+        """
+
+        def wrapper(fn: Callable):
+            tpl = template(fn)
+            self.registry[model_name] = tpl
+            return tpl
+
+        return wrapper
+
 
 def template(fn: Callable) -> Template:
     """Decorate a function that contains a prompt template.
@@ -96,7 +137,8 @@ def template(fn: Callable) -> Template:
     return Template(template, signature)
 
 
-def render(template: str, **values: Optional[Dict[str, Any]]) -> str:
+@lru_cache
+def render(template: str, **values: Optional[Dict[str, Hashable]]) -> str:
     r"""Parse a Jinaj2 template and translate it into an Outlines graph.
 
     This function removes extra whitespaces and linebreaks from templates to
diff --git a/tests/test_templates.py b/tests/test_templates.py
index 6be76af..6c2af9c 100644
--- a/tests/test_templates.py
+++ b/tests/test_templates.py
@@ -67,7 +67,7 @@ def test_render_jinja():
     """
 
     # Notice the newline after the end of the loop
-    examples = ["one", "two"]
+    examples = ("one", "two")
     prompt = render(
         """
         {% for e in examples %}
@@ -78,7 +78,7 @@ def test_render_jinja():
     assert prompt == "Example: one\nExample: two\n"
 
     # We can remove the newline by cloing with -%}
-    examples = ["one", "two"]
+    examples = ("one", "two")
     prompt = render(
         """
         {% for e in examples %}
@@ -102,7 +102,7 @@ def test_render_jinja():
     assert render(tpl, is_true=False) == "final"
 
     # Ignore leading white spaces
-    examples = ["one", "two"]
+    examples = ("one", "two")
     prompt = render(
         """
         {% for e in examples %}
@@ -114,7 +114,7 @@ def test_render_jinja():
     assert prompt == "one\ntwo\n"
 
     # Do not ignore leading white spaces
-    examples = ["one", "two"]
+    examples = ("one", "two")
     prompt = render(
         """
         {% for e in examples %}
@@ -132,7 +132,7 @@ def test_tpl(variable):
         """{{variable}} test"""
 
     assert test_tpl.template == "{{variable}} test"
-    assert test_tpl.parameters == ["variable"]
+    assert list(test_tpl.signature.parameters.keys()) == ["variable"]
 
     with pytest.raises(TypeError):
         test_tpl(v="test")
@@ -157,7 +157,7 @@ def test_kwarg_tpl(var, other_var="other"):
         """{{var}} and {{other_var}}"""
 
     assert test_kwarg_tpl.template == "{{var}} and {{other_var}}"
-    assert test_kwarg_tpl.parameters == ["var", "other_var"]
+    assert list(test_kwarg_tpl.signature.parameters.keys()) == ["var", "other_var"]
 
     p = test_kwarg_tpl("test")
     assert p == "test and other"
@@ -181,3 +181,25 @@ def test_empty(variable):
         @prompts.template
         def test_only_code(variable):
             return variable
+
+
+def test_dispatch():
+
+    @prompts.template
+    def simple_prompt(query: str):
+        """{{ query }}"""
+
+    @simple_prompt.register("provider/name")
+    def simple_prompt_name(query: str):
+        """name: {{ query }}"""
+
+    assert list(simple_prompt.registry.keys()) == ["provider/name"]
+    assert callable(simple_prompt)
+    assert callable(simple_prompt["provider/name"])
+
+    assert simple_prompt("test") == "test"
+    assert simple_prompt_name("test") == "name: test"
+
+    assert simple_prompt("test") == "test"
+    assert simple_prompt["gpt2"]("test") == "test"
+    assert simple_prompt["provider/name"]("test") == "name: test"

From 1f12d668efaefaf132946e0834da96266ff13242 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?R=C3=A9mi=20Louf?= <remi@thetypicalset.com>
Date: Mon, 29 Jul 2024 12:28:37 +0200
Subject: [PATCH 4/5] Update script to build sdist & wheel in CI

---
 .github/scripts/build_sdist_and_wheel.sh | 2 ++
 .github/workflows/tests.yml              | 2 +-
 2 files changed, 3 insertions(+), 1 deletion(-)

diff --git a/.github/scripts/build_sdist_and_wheel.sh b/.github/scripts/build_sdist_and_wheel.sh
index 28b82f4..ab34268 100644
--- a/.github/scripts/build_sdist_and_wheel.sh
+++ b/.github/scripts/build_sdist_and_wheel.sh
@@ -1,3 +1,5 @@
+#!/bin/bash
+
 # Build sdist and wheel
 python -m pip install -U pip
 python -m pip install build
diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml
index 8a322a9..b940c62 100644
--- a/.github/workflows/tests.yml
+++ b/.github/workflows/tests.yml
@@ -43,4 +43,4 @@ jobs:
     steps:
       - uses: actions/checkout@v3
       - name: Build SDist and Wheel
-        run: ./.github/scripts/build_sdist_and_wheel.sh
+        run: bash ./.github/scripts/build_sdist_and_wheel.sh

From be22a05a1b69044ec10f1182c07b391a79752fb2 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?R=C3=A9mi=20Louf?= <remi@thetypicalset.com>
Date: Mon, 29 Jul 2024 12:58:10 +0200
Subject: [PATCH 5/5] Update the documentation

---
 README.md                  |   8 +-
 docs/reference/dispatch.md |   1 +
 docs/reference/index.md    |   1 +
 docs/reference/template.md | 306 +++++++++++++++++++++++++++++++++++++
 mkdocs.yml                 |  25 +++
 5 files changed, 337 insertions(+), 4 deletions(-)
 create mode 100644 docs/reference/dispatch.md
 create mode 100644 docs/reference/index.md
 create mode 100644 docs/reference/template.md

diff --git a/README.md b/README.md
index 25824da..ddd1640 100644
--- a/README.md
+++ b/README.md
@@ -2,12 +2,12 @@
 
 ## Prompt functions
 
-The `prompt` decorator takes a function as an argument whose docstring is a Jinja template, and return a `Prompt` object:
+The `template` decorator takes a function as an argument whose docstring is a Jinja template, and return a `Template` object:
 
 ```python
-from prompts import prompt
+from prompts import template
 
-@prompt
+@template
 def few_shots(instructions, examples, question):
     """{{ instructions }}
 
@@ -21,7 +21,7 @@ def few_shots(instructions, examples, question):
     A: """
 ```
 
-Caling the `Prompt` object renders the Jinja template:
+Caling the `Template` object renders the Jinja template:
 
 ```python
 instructions = "Please answer the following question following the examples" examples = [
diff --git a/docs/reference/dispatch.md b/docs/reference/dispatch.md
new file mode 100644
index 0000000..1728e5b
--- /dev/null
+++ b/docs/reference/dispatch.md
@@ -0,0 +1 @@
+# Model-based prompt dispatching
diff --git a/docs/reference/index.md b/docs/reference/index.md
new file mode 100644
index 0000000..cf5aa07
--- /dev/null
+++ b/docs/reference/index.md
@@ -0,0 +1 @@
+# Reference
diff --git a/docs/reference/template.md b/docs/reference/template.md
new file mode 100644
index 0000000..14baaf1
--- /dev/null
+++ b/docs/reference/template.md
@@ -0,0 +1,306 @@
+# Prompt templating
+
+Prompts provides a powerful domain-specific language to write and manage
+prompts, via what we call *prompt functions*.  Prompt functions are Python
+functions that contain a template for the prompt in their docstring, and their
+arguments correspond to the variables used in the prompt. When called, a prompt
+function returns the template rendered with the values of the arguments.
+
+The aim of prompt functions is to solve several recurrent problems with prompting:
+
+1. **Building complex prompts quickly leads to messy code.** This problem has
+   already been solved in the web development community by using templating, so
+   why not use it here?
+2. **Composing prompts is difficult.** Why not just compose functions?
+3. **Separating prompts from code.** Encapsulation in functions allows a clean
+   separation between prompts and code. Moreover, like any function, prompt
+   functions can be imported from other modules.
+
+Prompts uses the [Jinja templating
+engine](https://jinja.palletsprojects.com/en/3.1.x/) to render prompts, which
+allows to easily compose complex prompts.
+
+!!! warning "Prompt rendering"
+
+    Prompt functions are opinionated when it comes to prompt rendering. These opinions are meant to avoid common prompting errors, but can have unintended consequences if you are doing something unusual. We advise to always print the prompt before using it. You can also [read the
+    reference](#formatting-conventions) section if you want to know more.
+
+## Your first prompt
+
+The following snippet showcases a very simple prompt. The variables between
+curly brackets `{{  }}` are placeholders for the values of the arguments you
+will pass to the prompt function.
+
+=== "Code"
+
+    ```python
+    import prompts
+
+    @prompts.template
+    def greetings(name, question):
+        """Hello, {{ name }}!
+        {{ question }}
+        """
+
+    prompt = greetings("user", "How are you?")
+    print(prompt)
+    ```
+
+=== "Output"
+
+    ```text
+    Hello, user!
+    How are you?
+    ```
+
+If a variable is missing in the function's arguments, Jinja2 will throw an `UndefinedError` exception:
+
+=== "Code"
+
+    ```python
+    import prompts
+
+    @prompts.template
+    def greetings(name):
+        """Hello, {{ surname }}!"""
+
+    prompt = greetings("user")
+    ```
+
+=== "Output"
+
+    ```text
+    Traceback (most recent call last):
+      File "<stdin>", line 9, in <module>
+      File "/home/remi/projects/normal/prompts/prompts.templates.py", line 38, in __call__
+          return render(self.template, **bound_arguments.arguments)
+      File "/home/remi/projects/normal/prompts/prompts.templates.py", line 213, in render
+          return jinja_template.render(**values)
+      File "/home/remi/micromamba/envs/prompts/lib/python3.9/site-packages/jinja2/environment.py", line 1301, in render
+          self.environment.handle_exception()
+      File "/home/remi/micromamba/envs/prompts/lib/python3.9/site-packages/jinja2/environment.py", line 936, in handle_exception
+          raise rewrite_traceback_stack(source=source)
+      File "<template>", line 1, in top-level template code
+      jinja2.exceptions.UndefinedError: 'surname' is undefined
+    ```
+
+## Importing prompt functions
+
+Prompt functions are functions, and thus can be imported from other modules:
+
+=== "prompts.py"
+    ```python
+    import prompts
+
+    @prompts.template
+    def greetings(name, question):
+        """Hello, {{ name }}!
+        {{ question }}
+        """
+    ```
+
+=== "generate.py"
+
+    ```python
+    from .prompts import greetings
+
+    prompt = greetings("John Doe", "How are you today?")
+    ```
+
+=== "Output"
+
+    ```text
+    Hello, John Doe!
+    How are you today?
+    ```
+
+## Few-shot prompting
+
+Few-shot prompting can lead to messy code. Prompt functions allow you to loop
+over lists or dictionaries from the template. In the following example we
+demonstrate how we can generate a prompt by passing a list of dictionaries with
+keys `question` and `answer` to the prompt function:
+
+=== "Code"
+
+    ```python
+    import prompts
+
+    @prompts.template
+    def few_shots(instructions, examples, question):
+        """{{ instructions }}
+
+        Examples
+        --------
+
+        {% for example in examples %}
+        Q: {{ example.question }}
+        A: {{ example.answer }}
+
+        {% endfor %}
+        Question
+        --------
+
+        Q: {{ question }}
+        A:
+        """
+
+    instructions = "Please answer the following question following the examples"
+    examples = [
+        {"question": "2+2=?", "answer":4},
+        {"question": "3+3=?", "answer":6}
+    ]
+    question = "4+4 = ?"
+
+    prompt = few_shots(instructions, examples, question)
+    print(prompt)
+    ```
+
+=== "Output"
+
+    ```text
+    Please answer the following question following the examples
+
+    Examples
+    --------
+
+    Q: 2+2=?
+    A: 4
+
+    Q: 3+3=?
+    A: 6
+
+    Question
+    --------
+
+    Q: 4+4 = ?
+    A:
+    ```
+
+## Conditionals, filters, etc.
+
+Jinja2 has many features beyond looping that are not described here:
+conditionals, filtering, formatting, etc. Please refer to the [Jinja
+documentation](https://jinja.palletsprojects.com/en/3.1.x/>) for more
+information about the syntax of the templating language. The Jinja syntax is
+powerful, and we recommend you take some time to read their documentation if you
+are building complex prompts.
+
+## Formatting conventions
+
+Prompt functions are opinionated when it comes to rendering, and these opinions
+are meant to avoid prompting mistakes and help with formatting.
+
+### Whitespaces
+
+If you have experience working with strings between triple quotes you know that
+indenting has an influence on the string's formatting. Prompt functions adopt
+a few conventions so you don't have to think about indents when writing prompt.
+
+First, whether you start the prompt right after the triple quotes or on the line
+below does not matter for formatting:
+
+=== "Code"
+
+    ```python
+    import prompts
+
+    @prompts.template
+    def prompt1():
+        """My prompt
+        """
+
+    @prompts.template
+    def prompt2():
+        """
+        My prompt
+        """
+
+    print(prompt1())
+    print(prompt2())
+    ```
+
+=== "Output"
+
+    ```text
+    My prompt
+    My prompt
+    ```
+
+Indentation is relative to the second line of the docstring, and leading spaces are removed:
+
+=== "Code"
+
+    ```python
+    import prompts
+
+    @prompts.template
+    def example1():
+        """First line
+        Second line
+        """
+
+    @prompts.template
+    def example2():
+        """
+          Second line
+          Third line
+        """
+
+    @prompts.template
+    def example3():
+        """
+          Second line
+            Third line
+        """
+
+    print(example1())
+    print(example2())
+    print(example3())
+    ```
+
+=== "Output"
+
+    ```text
+    First line
+    Second line
+
+    Second line
+    Third line
+
+    Second line
+      Third line
+    ```
+
+Trailing whitespaces are not removed, unless they follow a linebreak symbol `\` (see [linebreaks](#linebreaks)).
+
+### Linebreaks
+
+You can use the backslash `\` to break a long line of text. It will render as a single line:
+
+=== "Code"
+
+    ```python
+    import prompts
+
+    @prompts.template
+    def example():
+       """
+       Break in \
+       several lines \
+       But respect the indentation
+           on line breaks.
+       And after everything \
+       Goes back to normal
+       """
+
+    print(example())
+    ```
+
+=== "Output"
+
+    ```text
+    Break in several lines But respect the indentation
+        on line breaks.
+    And after everything Goes back to normal
+    ```
diff --git a/mkdocs.yml b/mkdocs.yml
index 296b625..f7dc334 100644
--- a/mkdocs.yml
+++ b/mkdocs.yml
@@ -38,6 +38,27 @@ extra:
       link: https://twitter.com/OutlinesOSS
   generator: false
 
+
+markdown_extensions:
+  - admonition
+  - def_list
+  - attr_list
+  - md_in_html
+  - pymdownx.highlight:
+      anchor_linenums: true
+      line_spans: __span
+      pygments_lang_class: true
+      noclasses: True
+      pygments_style: gruvbox-light
+  - pymdownx.superfences
+  - pymdownx.tabbed:
+      alternate_style: true
+  - pymdownx.inlinehilite
+  - pymdownx.details
+  - pymdownx.emoji:
+      emoji_index: !!python/name:material.extensions.emoji.twemoji
+      emoji_generator: !!python/name:material.extensions.emoji.to_svg
+
 plugins:
   - search
   - section-index
@@ -50,3 +71,7 @@ plugins:
 
 nav:
   - Home: index.md
+  - Docs:
+      - reference/index.md
+      - Prompt template: reference/template.md
+      - Dispatch: reference/dispatch.md