made sure the selected_documents match the final output

botisan-ai · Dec 30, 2021 · 07b7cce · 07b7cce
1 parent 56e1e2e
commit 07b7cce
Show file tree

Hide file tree

Showing 4 changed files with 85 additions and 4 deletions.
diff --git a/Pipfile b/Pipfile
@@ -9,6 +9,7 @@ python-dotenv = "*"
 boto3 = "*"
 jwcrypto = "*"
 pydantic = "*"
+regex = "*"
 
 [dev-packages]
 autopep8 = "*"

diff --git a/Pipfile.lock b/Pipfile.lock
diff --git a/functions/api/better_answers/__init__.py b/functions/api/better_answers/__init__.py
@@ -70,7 +70,7 @@ def handler(event: Dict[str, Any], context) -> Dict[str, Any]:
                 },
             }
 
-        prompt = better_answers_prompt(
+        prompt, idx = better_answers_prompt(
             input.model,
             input.max_tokens,
             input.question,
@@ -98,7 +98,7 @@ def handler(event: Dict[str, Any], context) -> Dict[str, Any]:
             model=completion_response.get('model'),
             prompt=prompt if input.return_prompt else None,
             search_model=input.search_model,
-            selected_documents=[SearchDocument(**d.dict(exclude={'metadata'})) if not input.return_metadata else d for d in selected_documents],
+            selected_documents=[SearchDocument(**d.dict(exclude={'metadata'})) if not input.return_metadata else d for d in selected_documents[idx:]],
         )
 
         return {

diff --git a/functions/api/better_answers/prompt.py b/functions/api/better_answers/prompt.py
@@ -47,4 +47,4 @@ def better_answers_prompt(
         i += 1
         prompt = combined_prompt(example_context_prompt, question, selected_documents[i:])
 
-    return prompt
+    return prompt, i