docs: display used prompts in swagger docs (#128)

Co-authored-by: Avram Tudor <tudor.avram@8x8.com>
jitsi · Nov 28, 2024 · abf3e7a · abf3e7a
1 parent 2d758b9
commit abf3e7a
Show file tree

Hide file tree

Showing 2 changed files with 27 additions and 4 deletions.
diff --git a/skynet/modules/ttt/summaries/v1/models.py b/skynet/modules/ttt/summaries/v1/models.py
@@ -3,6 +3,9 @@
 
 from pydantic import BaseModel, computed_field, Field
 
+from skynet.modules.ttt.summaries.prompts.action_items import action_items_meeting
+from skynet.modules.ttt.summaries.prompts.summary import summary_meeting
+
 
 class HintType(Enum):
     CONVERSATION = 'conversation'
@@ -23,6 +26,18 @@ class DocumentPayload(BaseModel):
     priority: Priority = Priority.NORMAL
     prompt: str | None = None
 
+
+class ActionItemsDocumentPayload(DocumentPayload):
+    model_config = {
+        'json_schema_extra': {
+            'examples': [
+                {'text': 'Your text here', 'hint': 'text', 'priority': 'normal', 'prompt': action_items_meeting}
+            ]
+        }
+    }
+
+
+class SummaryDocumentPayload(DocumentPayload):
     model_config = {
         'json_schema_extra': {
             'examples': [
@@ -31,7 +46,7 @@ class DocumentPayload(BaseModel):
                     'hint': 'meeting',
                     'max_tokens': 0,
                     'priority': 'normal',
-                    'prompt': 'Summarize the following text',
+                    'prompt': summary_meeting,
                 }
             ]
         }

diff --git a/skynet/modules/ttt/summaries/v1/router.py b/skynet/modules/ttt/summaries/v1/router.py
@@ -6,7 +6,15 @@
 from skynet.utils import get_router
 
 from ..jobs import create_job, get_job as get_job
-from .models import BaseJob, DocumentMetadata, DocumentPayload, JobId, JobType
+from .models import (
+    ActionItemsDocumentPayload,
+    BaseJob,
+    DocumentMetadata,
+    DocumentPayload,
+    JobId,
+    JobType,
+    SummaryDocumentPayload,
+)
 
 router = get_router()
 
@@ -35,7 +43,7 @@ def validate_payload(payload: DocumentPayload) -> None:
 
 @api_version(1)
 @router.post("/action-items", dependencies=[Depends(validate_payload)])
-async def get_action_items(payload: DocumentPayload, request: Request) -> JobId:
+async def get_action_items(payload: ActionItemsDocumentPayload, request: Request) -> JobId:
     """
     Starts a job to extract action items from the given payload.
     """
@@ -45,7 +53,7 @@ async def get_action_items(payload: DocumentPayload, request: Request) -> JobId:
 
 @api_version(1)
 @router.post("/summary", dependencies=[Depends(validate_payload)])
-async def get_summary(payload: DocumentPayload, request: Request) -> JobId:
+async def get_summary(payload: SummaryDocumentPayload, request: Request) -> JobId:
     """
     Starts a job to summarize the given payload.
     """