Repository: deepset-ai/haystack

Issue:
In OpenAIAnswerGenerator avoid tokenizing all documents
We're tokenizing all documents here:
https://github.com/deepset-ai/haystack/blob/c3a38a59c0e7b82afcdfb033e0631868e3c6b896/haystack/nodes/answer_generator/openai.py#L316-L317
and here:
https://github.com/deepset-ai/haystack/blob/c3a38a59c0e7b82afcdfb033e0631868e3c6b896/haystack/nodes/answer_generator/openai.py#L307

Depending on how slow tokenization is, we could tokenize one document at a time and check whether throwing away that document saves enough tokens.  
The call to `count_openai_tokens(text=doc.content, tokenizer=self._tokenizer)` would need to go into the same loop as `n_skipped_tokens += doc_token_len`. What do you think? It's not the focus of this PR so maybe this improvement could become a separate issue?

_Originally posted by @julian-risch in https://github.com/deepset-ai/haystack/pull/4179#discussion_r1115506214_

Yeah, I think this is a good idea. This would avoid a user accidentally sending 100+ docs to the PromptNode or AnswerGenerator and wonder what is taking so long. But this PR is already fairly large so I think opening a separate issue for this would be a good idea.

_Originally posted by @sjrl in https://github.com/deepset-ai/haystack/pull/4179#discussion_r1118950647_


Discussion:
@sjrl any chance you can work on this?
Hey @masci unfortunately, I won't have time to work on this for a few weeks, but I can try and get to it when I do. 

Resolving patch:
diff --git a/haystack/nodes/answer_generator/openai.py b/haystack/nodes/answer_generator/openai.py
--- a/haystack/nodes/answer_generator/openai.py
+++ b/haystack/nodes/answer_generator/openai.py
@@ -273,15 +273,11 @@ def _build_prompt_within_max_length(self, query: str, documents: List[Document])
         skipped_docs = 0
         # If leftover_token_len is negative we have gone past the MAX_TOKENS_LIMIT and the prompt must be trimmed
         if leftover_token_len < 0:
-            n_docs_tokens = [count_openai_tokens(text=doc.content, tokenizer=self._tokenizer) for doc in documents]
-            logger.debug("Number of tokens in documents: %s", n_docs_tokens)
-
-            # Reversing the order of documents b/c we want to throw away less relevant docs first
-            rev_n_docs_tokens = reversed(n_docs_tokens)
             n_skipped_tokens = 0
-            for doc_token_len in rev_n_docs_tokens:
-                n_skipped_tokens += doc_token_len
+            # Reversing the order of documents b/c we want to throw away less relevant docs first
+            for doc in reversed(documents):
                 skipped_docs += 1
+                n_skipped_tokens += count_openai_tokens(text=doc.content, tokenizer=self._tokenizer)
                 # Only skip enough tokens to fit within the MAX_TOKENS_LIMIT
                 if n_skipped_tokens >= abs(leftover_token_len):
                     break


Test patch:
diff --git a/test/nodes/test_generator.py b/test/nodes/test_generator.py
--- a/test/nodes/test_generator.py
+++ b/test/nodes/test_generator.py
@@ -1,6 +1,6 @@
-import os
 import sys
 from typing import List
+from unittest.mock import patch
 
 import pytest
 
@@ -173,3 +173,27 @@ def test_openai_answer_generator_max_token(haystack_openai_config, docs, caplog)
         assert "Skipping all of the provided Documents" in caplog.text
         assert len(prediction["answers"]) == 1
         # Can't easily check content of answer since it is generative and can change between runs
+
+
+# mock tokenizer that splits the string
+class MockTokenizer:
+    def encode(self, *args, **kwargs):
+        return str.split(*args, **kwargs)
+
+    def tokenize(self, *args, **kwargs):
+        return str.split(*args, **kwargs)
+
+
+@pytest.mark.unit
+def test_build_prompt_within_max_length():
+    with patch("haystack.nodes.answer_generator.openai.load_openai_tokenizer") as mock_load_tokenizer:
+        mock_load_tokenizer.return_value = MockTokenizer()
+
+        generator = OpenAIAnswerGenerator(api_key="fake_key", max_tokens=50)
+        generator.MAX_TOKENS_LIMIT = 92
+        query = "query"
+        documents = [Document("most relevant document"), Document("less relevant document")]
+        prompt_str, prompt_docs = generator._build_prompt_within_max_length(query=query, documents=documents)
+
+        assert len(prompt_docs) == 1
+        assert prompt_docs[0] == documents[0]