Update app.py

app.py

@@ -1,8 +1,9 @@
-from fastapi import FastAPI, File, UploadFile, HTTPException
+from fastapi import FastAPI, HTTPException
 import requests
 import base64
 from pydantic import BaseModel
-from typing import Optional
+from typing import Optional, List
+
 import re
 
 app = FastAPI()
@@ -11,20 +12,20 @@ app = FastAPI()
 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 securely in production
 
-# Request model for single user message
+class ChatMessage(BaseModel):
+    role: str  # "user" or "assistant" or "system"
+    content: str
+
 class TextRequest(BaseModel):
-    message: str
+    messages: List[ChatMessage]
     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 = {
         "Authorization": f"Bearer {API_KEY}",
@@ -35,11 +36,11 @@ def call_nvidia_api(payload: dict):
         raise HTTPException(status_code=response.status_code, detail="NVIDIA API request failed")
     return response.json()
 
-# /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}]
-    
+    # Combine pre-prompt with user-provided history
+    messages = PRE_PROMPT_MESSAGES + [msg.dict() for msg in request.messages]
+
     payload = {
         "model": "meta/llama-3.2-90b-vision-instruct",
         "messages": messages,
@@ -54,28 +55,31 @@ 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):
-    # Detect image URL
-    match = re.search(r'https?://\S+\.(jpg|jpeg|png|webp|gif)', request.message)
-    if not match:
-        raise HTTPException(status_code=400, detail="No image URL found")
+    # Find image URLs in the last user message(s)
+    new_messages = []
+    for msg in request.messages:
+        if msg.role == "user":
+            # Replace all image URLs in user messages with base64 img tags
+            def replace_img_url(match):
+                url = match.group(0)
+                try:
+                    img_resp = requests.get(url)
+                    img_resp.raise_for_status()
+                    b64 = base64.b64encode(img_resp.content).decode("utf-8")
+                    return f'<img src="data:image/png;base64,{b64}" />'
+                except Exception:
+                    return url  # fallback to original URL if fetch fails
 
-    image_url = match.group(0)
-    try:
-        img_response = requests.get(image_url)
-        img_response.raise_for_status()
-        base64_image = base64.b64encode(img_response.content).decode("utf-8")
-        img_tag = f'<img src="data:image/png;base64,{base64_image}" />'
-    except Exception as e:
-        raise HTTPException(status_code=400, detail=f"Failed to fetch image: {e}")
+            content_with_img = re.sub(r'https?://\S+\.(jpg|jpeg|png|webp|gif)', replace_img_url, msg.content)
+            new_messages.append({"role": "user", "content": content_with_img})
+        else:
+            new_messages.append(msg.dict())
 
-    # Replace image URL in message
-    modified_message = request.message.replace(image_url, img_tag)
+    messages = PRE_PROMPT_MESSAGES + new_messages
 
-    messages = PRE_PROMPT_MESSAGES + [{"role": "user", "content": modified_message}]
-    
     payload = {
         "model": "meta/llama-3.2-90b-vision-instruct",
         "messages": messages,
@@ -90,15 +94,3 @@ async def chat_from_text_with_image_url(request: TextRequest):
         return {"response": response["choices"][0]["message"]["content"]}
     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 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."
-    }