From 5f1256ac7e5734c2ea481e72cb7e02c34baf8c43 Mon Sep 17 00:00:00 2001
From: ZanSara <sara.zanzottera@deepset.ai>
Date: Thu, 31 Aug 2023 17:33:12 +0200
Subject: [PATCH] feat: `generators` (2.0) (#5690)

* add generators module

* add tests for module helper

* reno

* add another test

* move into openai

* improve tests
---
 .../preview/components/generators/__init__.py |  0
 .../components/generators/openai/__init__.py  |  0
 .../components/generators/openai/_helpers.py  | 33 +++++++++++++++++++
 .../generators-module-261376beb9c031cc.yaml   |  2 ++
 .../generators/openai/test_openai_helpers.py  | 20 +++++++++++
 test/preview/conftest.py                      | 13 ++++++++
 6 files changed, 68 insertions(+)
 create mode 100644 haystack/preview/components/generators/__init__.py
 create mode 100644 haystack/preview/components/generators/openai/__init__.py
 create mode 100644 haystack/preview/components/generators/openai/_helpers.py
 create mode 100644 releasenotes/notes/generators-module-261376beb9c031cc.yaml
 create mode 100644 test/preview/components/generators/openai/test_openai_helpers.py
 create mode 100644 test/preview/conftest.py

diff --git a/haystack/preview/components/generators/__init__.py b/haystack/preview/components/generators/__init__.py
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/haystack/preview/components/generators/openai/__init__.py b/haystack/preview/components/generators/openai/__init__.py
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/haystack/preview/components/generators/openai/_helpers.py b/haystack/preview/components/generators/openai/_helpers.py
new file mode 100644
index 0000000000..946901b644
--- /dev/null
+++ b/haystack/preview/components/generators/openai/_helpers.py
@@ -0,0 +1,33 @@
+import logging
+
+from haystack.preview.lazy_imports import LazyImport
+
+with LazyImport("Run 'pip install tiktoken'") as tiktoken_import:
+    import tiktoken
+
+
+logger = logging.getLogger(__name__)
+
+
+def enforce_token_limit(prompt: str, tokenizer: "tiktoken.Encoding", max_tokens_limit: int) -> str:
+    """
+    Ensure that the length of the prompt is within the max tokens limit of the model.
+    If needed, truncate the prompt text so that it fits within the limit.
+
+    :param prompt: Prompt text to be sent to the generative model.
+    :param tokenizer: The tokenizer used to encode the prompt.
+    :param max_tokens_limit: The max tokens limit of the model.
+    :return: The prompt text that fits within the max tokens limit of the model.
+    """
+    tiktoken_import.check()
+    tokens = tokenizer.encode(prompt)
+    tokens_count = len(tokens)
+    if tokens_count > max_tokens_limit:
+        logger.warning(
+            "The prompt has been truncated from %s tokens to %s tokens to fit within the max token limit. "
+            "Reduce the length of the prompt to prevent it from being cut off.",
+            tokens_count,
+            max_tokens_limit,
+        )
+        prompt = tokenizer.decode(tokens[:max_tokens_limit])
+    return prompt
diff --git a/releasenotes/notes/generators-module-261376beb9c031cc.yaml b/releasenotes/notes/generators-module-261376beb9c031cc.yaml
new file mode 100644
index 0000000000..0c57e6bc7b
--- /dev/null
+++ b/releasenotes/notes/generators-module-261376beb9c031cc.yaml
@@ -0,0 +1,2 @@
+preview:
+  - Add generators module for LLM generator components.
diff --git a/test/preview/components/generators/openai/test_openai_helpers.py b/test/preview/components/generators/openai/test_openai_helpers.py
new file mode 100644
index 0000000000..23a66117d1
--- /dev/null
+++ b/test/preview/components/generators/openai/test_openai_helpers.py
@@ -0,0 +1,20 @@
+import pytest
+
+from haystack.preview.components.generators.openai._helpers import enforce_token_limit
+
+
+@pytest.mark.unit
+def test_enforce_token_limit_above_limit(caplog, mock_tokenizer):
+    prompt = enforce_token_limit("This is a test prompt.", tokenizer=mock_tokenizer, max_tokens_limit=3)
+    assert prompt == "This is a"
+    assert caplog.records[0].message == (
+        "The prompt has been truncated from 5 tokens to 3 tokens to fit within the max token "
+        "limit. Reduce the length of the prompt to prevent it from being cut off."
+    )
+
+
+@pytest.mark.unit
+def test_enforce_token_limit_below_limit(caplog, mock_tokenizer):
+    prompt = enforce_token_limit("This is a test prompt.", tokenizer=mock_tokenizer, max_tokens_limit=100)
+    assert prompt == "This is a test prompt."
+    assert not caplog.records
diff --git a/test/preview/conftest.py b/test/preview/conftest.py
new file mode 100644
index 0000000000..b8abfa41a6
--- /dev/null
+++ b/test/preview/conftest.py
@@ -0,0 +1,13 @@
+from unittest.mock import Mock
+import pytest
+
+
+@pytest.fixture()
+def mock_tokenizer():
+    """
+    Tokenizes the string by splitting on spaces.
+    """
+    tokenizer = Mock()
+    tokenizer.encode = lambda text: text.split()
+    tokenizer.decode = lambda tokens: " ".join(tokens)
+    return tokenizer