Add summary kind to summarize endpoint (#1764)

nuclia · Jan 23, 2024 · fb0712a · fb0712a · github-actions · Jan 23, 2024
1 parent 66baf3f
commit fb0712a
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 0 deletions.
diff --git a/nucliadb/nucliadb/search/search/summarize.py b/nucliadb/nucliadb/search/search/summarize.py
@@ -46,6 +46,7 @@
 async def summarize(kbid: str, request: SummarizeRequest) -> SummarizedResponse:
     predict_request = SummarizeModel()
     predict_request.user_prompt = request.user_prompt
+    predict_request.summary_kind = request.summary_kind
 
     for rid, field_id, extracted_text in await get_extracted_texts(
         kbid, request.resources

diff --git a/nucliadb_models/nucliadb_models/search.py b/nucliadb_models/nucliadb_models/search.py
@@ -782,13 +782,19 @@ class SummarizeResourceModel(BaseModel):
     fields: Dict[str, str] = {}
 
 
+class SummaryKind(str, Enum):
+    SIMPLE = "simple"
+    EXTENDED = "extended"
+
+
 class SummarizeModel(BaseModel):
     """
     Model for the summarize predict api request payload
     """
 
     resources: Dict[str, SummarizeResourceModel] = {}
     user_prompt: Optional[str] = None
+    summary_kind: SummaryKind = SummaryKind.SIMPLE
 
 
 class SummarizeRequest(BaseModel):
@@ -809,6 +815,12 @@ class SummarizeRequest(BaseModel):
         description="Uids of the resources to summarize",
     )
 
+    summary_kind: SummaryKind = Field(
+        default=SummaryKind.SIMPLE,
+        title="Summary kind",
+        description="Option to customize how the summary will be",
+    )
+
 
 class SummarizedResource(BaseModel):
     summary: str