Fix prepdocs compatibility with openai key and add test (#605)

* Remove defaults for getenv

* Remove print

* missing output

* Add tests and fix prepdocs issue

* rm uneeded print

Author: pamelafox

scripts/prepdocs.py

@@ -30,7 +30,14 @@
 )
 from azure.storage.blob import BlobServiceClient
 from pypdf import PdfReader, PdfWriter
-from tenacity import retry, stop_after_attempt, wait_random_exponential
+from tenacity import (
+    retry,
+    retry_if_exception_type,
+    stop_after_attempt,
+    wait_random_exponential,
+)
+
+args = argparse.Namespace(verbose=False)
 
 MAX_SECTION_LENGTH = 1000
 SENTENCE_SEARCH_LIMIT = 100
@@ -225,7 +232,7 @@ def filename_to_id(filename):
     filename_hash = base64.b16encode(filename.encode('utf-8')).decode('ascii')
     return f"file-{filename_ascii}-{filename_hash}"
 
-def create_sections(filename, page_map, use_vectors):
+def create_sections(filename, page_map, use_vectors, embedding_deployment: str = None):
     file_id = filename_to_id(filename)
     for i, (content, pagenum) in enumerate(split_text(page_map, filename)):
         section = {
@@ -236,16 +243,16 @@ def create_sections(filename, page_map, use_vectors):
             "sourcefile": filename
         }
         if use_vectors:
-            section["embedding"] = compute_embedding(content)
+            section["embedding"] = compute_embedding(content, embedding_deployment)
         yield section
 
 def before_retry_sleep(retry_state):
     if args.verbose: print("Rate limited on the OpenAI embeddings API, sleeping before retrying...")
 
-@retry(wait=wait_random_exponential(min=15, max=60), stop=stop_after_attempt(15), before_sleep=before_retry_sleep)
-def compute_embedding(text):
+@retry(retry=retry_if_exception_type(openai.error.RateLimitError), wait=wait_random_exponential(min=15, max=60), stop=stop_after_attempt(15), before_sleep=before_retry_sleep)
+def compute_embedding(text, embedding_deployment):
     refresh_openai_token()
-    return openai.Embedding.create(engine=args.openaideployment, input=text)["data"][0]["embedding"]
+    return openai.Embedding.create(engine=embedding_deployment, input=text)["data"][0]["embedding"]
 
 @retry(wait=wait_random_exponential(min=15, max=60), stop=stop_after_attempt(15), before_sleep=before_retry_sleep)
 def compute_embedding_in_batch(texts):
@@ -314,7 +321,7 @@ def update_embeddings_in_batch(sections):
         if args.verbose: print(f"Batch Completed. Batch size  {len(batch_queue)} Token count {token_count}")
         for emb, item in zip(emb_responses, batch_queue):
             batch_response[item["id"]] = emb
-    
+
     for s in copy_s:
         s["embedding"] = batch_response[s["id"]]
         yield s
@@ -355,14 +362,18 @@ def remove_from_index(filename):
         # It can take a few seconds for search results to reflect changes, so wait a bit
         time.sleep(2)
 
-# refresh open ai token every 5 minutes
+
 def refresh_openai_token():
-    if open_ai_token_cache[CACHE_KEY_TOKEN_TYPE] == 'azure_ad' and open_ai_token_cache[CACHE_KEY_CREATED_TIME] + 300 < time.time():
+    """
+    Refresh OpenAI token every 5 minutes
+    """
+    if CACHE_KEY_TOKEN_TYPE in open_ai_token_cache and open_ai_token_cache[CACHE_KEY_TOKEN_TYPE] == 'azure_ad' and open_ai_token_cache[CACHE_KEY_CREATED_TIME] + 300 < time.time():
         token_cred = open_ai_token_cache[CACHE_KEY_TOKEN_CRED]
         openai.api_key = token_cred.get_token("https://cognitiveservices.azure.com/.default").token
         open_ai_token_cache[CACHE_KEY_CREATED_TIME] = time.time()
 
-def read_files(path_pattern: str, use_vectors: bool, vectors_batch_support: bool):
+
+def read_files(path_pattern: str, use_vectors: bool, vectors_batch_support: bool, embedding_deployment: str = None):
     """
     Recursively read directory structure under `path_pattern`
     and execute indexing for the individual files
@@ -380,8 +391,7 @@ def read_files(path_pattern: str, use_vectors: bool, vectors_batch_support: bool
                 if not args.skipblobs:
                     upload_blobs(filename)
                 page_map = get_document_text(filename)
-                sections = create_sections(os.path.basename(filename), page_map, use_vectors and not vectors_batch_support)
-                print (use_vectors and vectors_batch_support)
+                sections = create_sections(os.path.basename(filename), page_map, use_vectors and not vectors_batch_support, embedding_deployment)
                 if use_vectors and vectors_batch_support:
                     sections = update_embeddings_in_batch(sections)
                 index_sections(os.path.basename(filename), sections)
@@ -456,4 +466,4 @@ def read_files(path_pattern: str, use_vectors: bool, vectors_batch_support: bool
             create_search_index()
 
         print("Processing files...")
-        read_files(args.files, use_vectors, compute_vectors_in_batch)
+        read_files(args.files, use_vectors, compute_vectors_in_batch, args.openaideployment)

tests/test_prepdocs.py

@@ -1,4 +1,7 @@
-from scripts.prepdocs import filename_to_id
+import openai
+import pytest
+import tenacity
+from scripts.prepdocs import args, compute_embedding, filename_to_id
 
 
 def test_filename_to_id():
@@ -8,3 +11,57 @@ def test_filename_to_id():
     assert filename_to_id("foo\u00A9.txt") == "file-foo__txt-666F6FC2A92E747874"
     # test filenaming starting with unicode
     assert filename_to_id("ファイル名.pdf") == "file-______pdf-E38395E382A1E382A4E383ABE5908D2E706466"
+
+
+def test_compute_embedding_success(monkeypatch, capsys):
+    monkeypatch.setattr(args, "verbose", True)
+    def mock_create(*args, **kwargs):
+        # From https://platform.openai.com/docs/api-reference/embeddings/create
+        return {
+            "object": "list",
+            "data": [
+                {
+                "object": "embedding",
+                "embedding": [
+                    0.0023064255,
+                    -0.009327292,
+                    -0.0028842222,
+                ],
+                "index": 0
+                }
+            ],
+            "model": "text-embedding-ada-002",
+            "usage": {
+                "prompt_tokens": 8,
+                "total_tokens": 8
+            }
+        }
+
+    monkeypatch.setattr(openai.Embedding, "create", mock_create)
+    assert compute_embedding("foo", "ada") == [
+                    0.0023064255,
+                    -0.009327292,
+                    -0.0028842222,
+                ]
+
+
+def test_compute_embedding_ratelimiterror(monkeypatch, capsys):
+    monkeypatch.setattr(args, "verbose", True)
+    def mock_create(*args, **kwargs):
+        raise openai.error.RateLimitError
+    monkeypatch.setattr(openai.Embedding, "create", mock_create)
+    monkeypatch.setattr(tenacity.nap.time, "sleep", lambda x: None)
+    with pytest.raises(tenacity.RetryError):
+        compute_embedding("foo", "ada")
+    captured = capsys.readouterr()
+    assert captured.out.count("Rate limited on the OpenAI embeddings API") == 14
+
+
+def test_compute_embedding_autherror(monkeypatch, capsys):
+    monkeypatch.setattr(args, "verbose", True)
+    def mock_create(*args, **kwargs):
+        raise openai.error.AuthenticationError
+    monkeypatch.setattr(openai.Embedding, "create", mock_create)
+    monkeypatch.setattr(tenacity.nap.time, "sleep", lambda x: None)
+    with pytest.raises(openai.error.AuthenticationError):
+        compute_embedding("foo", "ada")