[FIX] Supporting multi chunk for large documents (#335)

* Simple and subquestion retriever * Migration file * Removing Query modification for Simple Retriever Co-authored-by: Hari John Kuriakose <[email protected]> Signed-off-by: harini-venkataraman <[email protected]> * Removing query changes ti simple retriever * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fixing pre-commit issues * Pre-commit issue fixes --------- Signed-off-by: harini-venkataraman <[email protected]> Co-authored-by: Hari John Kuriakose <[email protected]> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Zipstack · May 23, 2024 · afd2238 · afd2238
1 parent 307bc4f
commit afd2238
Show file tree

Hide file tree

Showing 5 changed files with 61 additions and 22 deletions.
diff --git a/..._studio/prompt_profile_manager/migrations/0011_alter_profilemanager_retrieval_strategy.py b/..._studio/prompt_profile_manager/migrations/0011_alter_profilemanager_retrieval_strategy.py
@@ -0,0 +1,28 @@
+# Generated by Django 4.2.1 on 2024-05-14 09:58
+
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        (
+            "prompt_profile_manager",
+            "0010_alter_profilemanager_retrieval_strategy_and_more",
+        ),
+    ]
+
+    operations = [
+        migrations.AlterField(
+            model_name="profilemanager",
+            name="retrieval_strategy",
+            field=models.TextField(
+                blank=True,
+                choices=[
+                    ("simple", "Simple retrieval"),
+                    ("subquestion", "Subquestion retrieval"),
+                ],
+                db_comment="Field to store the retrieval strategy for prompts",
+            ),
+        ),
+    ]
diff --git a/backend/prompt_studio/prompt_profile_manager/models.py b/backend/prompt_studio/prompt_profile_manager/models.py
@@ -13,6 +13,7 @@ class ProfileManager(BaseModel):
 
     class RetrievalStrategy(models.TextChoices):
         SIMPLE = "simple", "Simple retrieval"
+        SUBQUESTION = "subquestion", "Subquestion retrieval"
 
     profile_id = models.UUIDField(primary_key=True, default=uuid.uuid4, editable=False)
     profile_name = models.TextField(blank=False)

diff --git a/backend/prompt_studio/prompt_studio_core/static/select_choices.json b/backend/prompt_studio/prompt_studio_core/static/select_choices.json
@@ -23,7 +23,8 @@
         "openai_embedding":"openai_embedding"
     },
     "retrieval_strategy":{
-        "simple":"simple"
+        "simple":"simple",
+        "subquestion":"subquestion"
     },
     "vector_store":{
         "Postgres pg_vector":"Postgres pg_vector",

diff --git a/prompt-service/src/unstract/prompt_service/constants.py b/prompt-service/src/unstract/prompt_service/constants.py
@@ -23,6 +23,7 @@ class PromptServiceContants:
     ASSERTION_FAILURE_PROMPT = "assertion_failure_prompt"
     RETRIEVAL_STRATEGY = "retrieval-strategy"
     SIMPLE = "simple"
+    SUBQUESTION = "subquestion"
     TYPE = "type"
     NUMBER = "number"
     EMAIL = "email"

diff --git a/prompt-service/src/unstract/prompt_service/main.py b/prompt-service/src/unstract/prompt_service/main.py
@@ -361,11 +361,12 @@ def prompt_processor() -> Any:
             )
 
             if output[PSKeys.RETRIEVAL_STRATEGY] == PSKeys.SIMPLE:
-                answer, context = simple_retriver(
-                    output,
-                    doc_id,
-                    llm,
-                    vector_index,
+                answer, context = run_retrieval(
+                    output, doc_id, llm, vector_index, PSKeys.SIMPLE
+                )
+            elif output[PSKeys.RETRIEVAL_STRATEGY] == PSKeys.SUBQUESTION:
+                answer, context = run_retrieval(
+                    output, doc_id, llm, vector_index, PSKeys.SUBQUESTION
                 )
             else:
                 app.logger.info(
@@ -656,27 +657,41 @@ def prompt_processor() -> Any:
     return response
 
 
-def simple_retriver(  # type:ignore
+def run_retrieval(  # type:ignore
     output: dict[str, Any],
     doc_id: str,
     llm: LLM,
     vector_index,
+    retrieval_type: str,
 ) -> tuple[str, str]:
     prompt = construct_prompt_for_engine(
         preamble=output["preamble"],
         prompt=output["promptx"],
         postamble=output["postamble"],
         grammar_list=output["grammar"],
     )
-    subq_prompt = (
-        f"Generate a sub-question from the following verbose prompt that will"
-        f" help extract relevant documents from a vector store:\n\n{prompt}"
-    )
-    answer = run_completion(
-        llm=llm,
-        prompt=subq_prompt,
+    if retrieval_type is PSKeys.SUBQUESTION:
+        subq_prompt = (
+            f"Generate a sub-question from the following verbose prompt that will"
+            f" help extract relevant documents from a vector store:\n\n{prompt}"
+        )
+        prompt = run_completion(
+            llm=llm,
+            prompt=subq_prompt,
+        )
+    context = _retrieve_context(output, doc_id, vector_index, prompt)
+
+    answer = construct_and_run_prompt(  # type:ignore
+        output,
+        llm,
+        context,
+        "promptx",
     )
 
+    return (answer, context)
+
+
+def _retrieve_context(output, doc_id, vector_index, answer) -> str:
     retriever = vector_index.as_retriever(
         similarity_top_k=output[PSKeys.SIMILARITY_TOP_K],
         filters=MetadataFilters(
@@ -695,14 +710,7 @@ def simple_retriver(  # type:ignore
             text += node.get_content() + "\n"
         else:
             app.logger.info("Node score is less than 0.6. " f"Ignored: {node.score}")
-
-    answer = construct_and_run_prompt(  # type:ignore
-        output,
-        llm,
-        text,
-        "promptx",
-    )
-    return (answer, text)
+    return text
 
 
 def construct_and_run_prompt(