Update main.py

main.py

@@ -582,6 +582,24 @@ async def analyze_image(image_data_uri: str) -> str:
         logger.error(f"Failed to analyze image: {e}")
         raise HTTPException(status_code=400, detail="Failed to process the provided image.")
 
+# Helper Function for Token Counting
+def count_prompt_tokens(request: ChatRequest) -> int:
+    """
+    Counts the number of tokens in the prompt (input messages).
+    Handles both string and list types for the 'content' field.
+    """
+    total = 0
+    for msg in request.messages:
+        if isinstance(msg.content, str):
+            total += count_tokens(msg.content)
+        elif isinstance(msg.content, list):
+            for item in msg.content:
+                if isinstance(item, TextContent):
+                    total += count_tokens(item.text)
+                elif isinstance(item, ImageContent):
+                    total += count_tokens(item.image_url['url'])
+    return total
+
 # Endpoint: POST /v1/chat/completions
 @app.post("/v1/chat/completions", dependencies=[Depends(rate_limiter_per_ip)])
 async def chat_completions(request: ChatRequest, req: Request, api_key: str = Depends(get_api_key)):
@@ -613,8 +631,8 @@ async def chat_completions(request: ChatRequest, req: Request, api_key: str = De
         # Example response content
         assistant_content += "Based on the image you provided, here are the insights..."
 
-        # Calculate token usage (simple approximation)
-        prompt_tokens = sum(count_tokens(" ".join([item.text if isinstance(item, TextContent) else item.image_url['url'] for item in msg.content]) ) for msg in request.messages)
+        # Calculate token usage using the helper function
+        prompt_tokens = count_prompt_tokens(request)
         completion_tokens = count_tokens(assistant_content)
         total_tokens = prompt_tokens + completion_tokens
         estimated_cost = calculate_estimated_cost(prompt_tokens, completion_tokens)