Update app.py

app.py

@@ -3,20 +3,27 @@ import requests
 import base64
 from pydantic import BaseModel
 from typing import Optional
+import re
 
 app = FastAPI()
 
 # NVIDIA API endpoint and API key
 NVIDIA_API_URL = "https://ai.api.nvidia.com/v1/gr/meta/llama-3.2-90b-vision-instruct/chat/completions"
-API_KEY = "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"  # Replace with your actual API key
+API_KEY = "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"  # Replace securely in production
 
-# Request model for text-based input
+# Request model for single user message
 class TextRequest(BaseModel):
     message: str
     max_tokens: Optional[int] = 512
     temperature: Optional[float] = 1.0
     top_p: Optional[float] = 1.0
 
+# Common pre-prompts
+PRE_PROMPT_MESSAGES = [
+    {"role": "system", "content": "You are a helpful multimodal assistant powered by LLaMA 3.2 Vision-Instruct."},
+    {"role": "assistant", "content": "Hi! You can send text or image-based questions. What would you like to know?"}
+]
+
 # Function to call the NVIDIA API
 def call_nvidia_api(payload: dict):
     headers = {
@@ -28,12 +35,14 @@ def call_nvidia_api(payload: dict):
         raise HTTPException(status_code=response.status_code, detail="NVIDIA API request failed")
     return response.json()
 
-# Endpoint for text-based input
+# /chat/text endpoint: Adds new user message to pre-prompted context
 @app.post("/chat/text")
 async def chat_with_text(request: TextRequest):
+    messages = PRE_PROMPT_MESSAGES + [{"role": "user", "content": request.message}]
+    
     payload = {
         "model": "meta/llama-3.2-90b-vision-instruct",
-        "messages": [{"role": "user", "content": request.message}],
+        "messages": messages,
         "max_tokens": request.max_tokens,
         "temperature": request.temperature,
         "top_p": request.top_p,
@@ -45,12 +54,9 @@ async def chat_with_text(request: TextRequest):
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))
 
-
-
+# /chat/vision endpoint: Handles messages containing image URLs
 @app.post("/chat/vision")
 async def chat_from_text_with_image_url(request: TextRequest):
-    import re
-
     # Detect image URL
     match = re.search(r'https?://\S+\.(jpg|jpeg|png|webp|gif)', request.message)
     if not match:
@@ -65,12 +71,14 @@ async def chat_from_text_with_image_url(request: TextRequest):
     except Exception as e:
         raise HTTPException(status_code=400, detail=f"Failed to fetch image: {e}")
 
-    # Replace image URL with the base64 image
+    # Replace image URL in message
     modified_message = request.message.replace(image_url, img_tag)
 
+    messages = PRE_PROMPT_MESSAGES + [{"role": "user", "content": modified_message}]
+    
     payload = {
         "model": "meta/llama-3.2-90b-vision-instruct",
-        "messages": [{"role": "user", "content": modified_message}],
+        "messages": messages,
         "max_tokens": request.max_tokens,
         "temperature": request.temperature,
         "top_p": request.top_p,
@@ -83,9 +91,14 @@ async def chat_from_text_with_image_url(request: TextRequest):
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))
 
-
-
 # Root endpoint
 @app.get("/")
 async def root():
-    return {"message": "Welcome to the NVIDIA API FastAPI wrapper!"}
+    return {
+        "message": "Welcome to the NVIDIA Vision Chat API!",
+        "endpoints": {
+            "/chat/text": "Send plain text questions (just provide your message).",
+            "/chat/vision": "Send a message with an image URL (e.g. 'What is this? https://example.com/cat.jpg')",
+        },
+        "note": "You do NOT need to include assistant history or system roles — it's pre-injected automatically."
+    }