From fb0712a802711c0f986a9870d1c9aba71ae798f7 Mon Sep 17 00:00:00 2001
From: Joan Antoni RE <joanantoni.re16@gmail.com>
Date: Tue, 23 Jan 2024 19:46:33 +0100
Subject: [PATCH] Add summary kind to summarize endpoint (#1764)

---
 nucliadb/nucliadb/search/search/summarize.py |  1 +
 nucliadb_models/nucliadb_models/search.py    | 12 ++++++++++++
 2 files changed, 13 insertions(+)

diff --git a/nucliadb/nucliadb/search/search/summarize.py b/nucliadb/nucliadb/search/search/summarize.py
index b532e8f1d4..192f6e6336 100644
--- 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
index 5ffff4cd73..900219ee2b 100644
--- a/nucliadb_models/nucliadb_models/search.py
+++ b/nucliadb_models/nucliadb_models/search.py
@@ -782,6 +782,11 @@ 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
@@ -789,6 +794,7 @@ class SummarizeModel(BaseModel):
 
     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