Update app.py

app.py

@@ -94,15 +94,14 @@
 #         content={"error": "Method not allowed. Please check the API documentation."}
 #     )
 
-
 # app.py
 import os
 import logging
+from typing import Optional
 from fastapi import FastAPI, HTTPException
 from fastapi.responses import JSONResponse
 from pydantic import BaseModel
 from huggingface_hub import InferenceClient
-from typing import Optional
 
 # Set up logging
 logging.basicConfig(level=logging.INFO)
@@ -111,8 +110,8 @@ logger = logging.getLogger(__name__)
 # Initialize FastAPI app
 app = FastAPI(
     title="LLM Chat API",
-    description="API for getting chat responses from Llama model with image support",
-    version="1.1.0"
+    description="API for getting chat responses from Llama model (supports text and image input)",
+    version="1.0.0"
 )
 
 class ChatRequest(BaseModel):
@@ -130,64 +129,39 @@ def llm_chat_response(text: str, image_url: Optional[str] = None) -> str:
         if not HF_TOKEN:
             logger.error("HF_TOKEN not found in environment variables")
             raise HTTPException(status_code=500, detail="HF_TOKEN not configured")
-
+        
         logger.info("Initializing InferenceClient...")
         client = InferenceClient(
             provider="sambanova",
             api_key=HF_TOKEN
         )
-
-        # Prepare content list for the message
-        content = [
-            {
-                "type": "text",
-                "text": text + " describe in one line only"
-            }
-        ]
-
-        # Add image to content if provided
+        
+        # Build the messages payload dynamically based on whether an image URL is provided.
+        # If only text is provided, add an instruction for a one-line description.
+        message_content = [{
+            "type": "text",
+            "text": text + ("" if image_url else " describe in one line only")
+        }]
+        
         if image_url:
-            logger.info(f"Adding image URL to request: {image_url}")
-            content.append({
+            message_content.append({
                 "type": "image_url",
-                "image_url": {
-                    "url": image_url
-                }
+                "image_url": {"url": image_url}
             })
-
-        messages = [
-            {
-                "role": "user",
-                "content": content
-            }
-        ]
-
-        logger.info("Sending request to model...")
-        logger.info(f"Request payload: {messages}")
         
+        messages = [{
+            "role": "user",
+            "content": message_content
+        }]
+        
+        logger.info("Sending request to model...")
         completion = client.chat.completions.create(
             model="meta-llama/Llama-3.2-11B-Vision-Instruct",
             messages=messages,
             max_tokens=500
         )
-        
-        logger.info(f"Response received: {completion}")
-        
-        # Check the structure of the response and extract content
-        if hasattr(completion, 'choices') and len(completion.choices) > 0:
-            message = completion.choices[0].message
-            # Handle different response formats
-            if isinstance(message, dict) and 'content' in message:
-                return message['content']
-            elif hasattr(message, 'content'):
-                return message.content
-            else:
-                logger.error(f"Unexpected message format: {message}")
-                return str(message)
-        else:
-            logger.error(f"Unexpected completion format: {completion}")
-            return str(completion)
-            
+        return completion.choices[0].message['content']
+
     except Exception as e:
         logger.error(f"Error in llm_chat_response: {str(e)}")
         raise HTTPException(status_code=500, detail=str(e))
@@ -197,8 +171,7 @@ async def chat(request: ChatRequest):
     try:
         logger.info(f"Received chat request with text: {request.text}")
         if request.image_url:
-            logger.info(f"Image URL included: {request.image_url}")
-        
+            logger.info(f"Image URL provided: {request.image_url}")
         response = llm_chat_response(request.text, request.image_url)
         return ChatResponse(response=response, status="success")
     except HTTPException as he:
@@ -210,7 +183,7 @@ async def chat(request: ChatRequest):
 
 @app.get("/")
 async def root():
-    return {"message": "Welcome to the LLM Chat API with image support. Use POST /chat endpoint to get responses."}
+    return {"message": "Welcome to the LLM Chat API. Use POST /chat endpoint with 'text' and optionally 'image_url' for queries."}
 
 @app.exception_handler(404)
 async def not_found_handler(request, exc):
@@ -224,4 +197,4 @@ async def method_not_allowed_handler(request, exc):
     return JSONResponse(
         status_code=405,
         content={"error": "Method not allowed. Please check the API documentation."}
-    )
+    )