Add model instructions and output format

main.py

@@ -17,6 +17,8 @@ app = FastAPI()
 
 class Prompt(BaseModel):
     text: str
+    instructions: str | None = None   # optional system instructions
+    output_format: str | None = None
 
 def query(payload):
     response = requests.post(API_URL, headers=headers, json=payload)
@@ -28,13 +30,31 @@ def read_root():
 
 @app.post("/generate")
 async def generate_text(prompt: Prompt):
+    messages = []
+
+    # Add system instructions if provided
+    if prompt.instructions:
+        messages.append({
+            "role": "system",
+            "content": prompt.instructions
+        })
+
+    # Add user message
+    user_content = prompt.text
+
+    # If an output format is specified, append formatting instructions
+    if prompt.output_format:
+        user_content += f"\n\nPlease always format the response as: {prompt.output_format}"
+
+    messages.append({
+        "role": "user",
+        "content": user_content
+    })
+
+    # Call HF API
     response = query({
-        "messages": [
-            {
-                "role": "user",
-                "content": prompt.text
-            }
-        ],
+        "messages": messages,
         "model": "meta-llama/Llama-3.2-3B-Instruct"
     })
+
     return {"response": response["choices"][0]["message"]}