api: allow prompt override (#95)

jitsi · Aug 26, 2024 · b1c28f8 · b1c28f8
1 parent 2d6ab29
commit b1c28f8
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 1 deletion.
diff --git a/skynet/modules/ttt/summaries/processor.py b/skynet/modules/ttt/summaries/processor.py
@@ -47,7 +47,7 @@ async def process(payload: DocumentPayload, job_type: JobType, model: ChatOpenAI
     if not text:
         return ""
 
-    system_message = hint_type_to_prompt[job_type][payload.hint]
+    system_message = payload.prompt or hint_type_to_prompt[job_type][payload.hint]
     prompt = PromptTemplate(template=system_message, input_variables=['text'])
 
     # this is a rough estimate of the number of tokens in the input text, since llama models will have a different tokenization scheme

diff --git a/skynet/modules/ttt/summaries/v1/models.py b/skynet/modules/ttt/summaries/v1/models.py
@@ -14,6 +14,13 @@ class HintType(Enum):
 class DocumentPayload(BaseModel):
     text: str
     hint: HintType = summary_default_hint_type
+    prompt: str | None = None
+
+    model_config = {
+        'json_schema_extra': {
+            'examples': [{'text': 'Your text here', 'hint': 'text', 'prompt': 'Summarize the following text {text}'}]
+        }
+    }
 
 
 class DocumentMetadata(BaseModel):