Update main.py

main.py

@@ -532,41 +532,21 @@ def create_response(content: str, model: str, finish_reason: Optional[str] = Non
         "usage": None,  # To be filled in non-streaming responses
     }
 
-def extract_image_from_content(content: str) -> Optional[Tuple[str, str]]:
+def extract_all_images_from_content(content_items: List[ContentItem]) -> List[Tuple[str, str]]:
     """
-    Extracts the first image from the content string.
-    Returns a tuple of (alt_text, image_data_uri) if found, else None.
-    """
-    # Regex to match markdown image syntax: ![Alt Text](image_url)
-    match = re.search(r'!\[([^\]]*)\]\((data:image/\w+;base64,[^\)]+)\)', content)
-    if match:
-        alt_text = match.group(1)
-        image_data_uri = match.group(2)
-        return alt_text, image_data_uri
-    return None
-
-def extract_all_images_from_content(content: str) -> List[Tuple[str, str]]:
-    """
-    Extracts all images from the content string.
+    Extracts all images from the content list.
     Returns a list of tuples containing (alt_text, image_data_uri).
     """
-    # Regex to match markdown image syntax: ![Alt Text](image_url)
-    matches = re.findall(r'!\[([^\]]*)\]\((data:image/\w+;base64,[^\)]+)\)', content)
-    return matches if matches else []
-
-async def analyze_image(image_data_uri: str) -> str:
-    """
-    Placeholder function to analyze the image.
-    Replace this with actual image analysis logic or API calls.
-    """
-    # Extract base64 data
-    image_data = image_data_uri.split(",")[1]
-    # Decode and process the image as needed
-    # For example, send it to an external API
-    # Here, we'll return a dummy response
-    await asyncio.sleep(1)  # Simulate processing delay
-    return "Image analysis result: The image depicts a beautiful sunset over the mountains."
-
+    images = []
+    for item in content_items:
+        if isinstance(item, ImageContent):
+            alt_text = item.image_url.get('alt', '')  # Optional alt text
+            image_data_uri = item.image_url.get('url', '')
+            if image_data_uri:
+                images.append((alt_text, image_data_uri))
+    return images
+
+# 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)):
     client_ip = req.client.host
@@ -588,7 +568,7 @@ async def chat_completions(request: ChatRequest, req: Request, api_key: str = De
         for msg in request.messages:
             if msg.role == "user":
                 # Extract all images from the message content
-                images = extract_all_images_from_content(" ".join([item.text if item.type == "text" else item.image_url['url'] for item in msg.content]))
+                images = extract_all_images_from_content(msg.content)
                 for alt_text, image_data_uri in images:
                     # Analyze the image
                     analysis_result = await analyze_image(image_data_uri)
@@ -598,35 +578,94 @@ async def chat_completions(request: ChatRequest, req: Request, api_key: str = De
         assistant_content += "Based on the image you provided, here are the insights..."
 
         # Calculate token usage (simple approximation)
-        prompt_tokens = sum(len(" ".join([item.text if item.type == "text" else item.image_url['url'] for item in msg.content]).split()) for msg in request.messages)
+        prompt_tokens = sum(len(" ".join([item.text if isinstance(item, TextContent) else item.image_url['url'] for item in msg.content]).split()) for msg in request.messages)
         completion_tokens = len(assistant_content.split())
         total_tokens = prompt_tokens + completion_tokens
         estimated_cost = calculate_estimated_cost(prompt_tokens, completion_tokens)
 
         logger.info(f"Completed response generation for API key: {api_key} | IP: {client_ip}")
 
-        return {
-            "id": f"chatcmpl-{uuid.uuid4()}",
-            "object": "chat.completion",
-            "created": int(datetime.now().timestamp()),
-            "model": request.model,
-            "choices": [
-                {
-                    "message": {
-                        "role": "assistant",
-                        "content": assistant_content.strip()
-                    },
-                    "finish_reason": "stop",
-                    "index": 0
-                }
-            ],
-            "usage": {
-                "prompt_tokens": prompt_tokens,
-                "completion_tokens": completion_tokens,
-                "total_tokens": total_tokens,
-                "estimated_cost": estimated_cost
-            },
-        }
+        if request.stream:
+            async def generate():
+                try:
+                    for msg in request.messages:
+                        if msg.role == "user":
+                            images = extract_all_images_from_content(msg.content)
+                            for alt_text, image_data_uri in images:
+                                analysis_result = await analyze_image(image_data_uri)
+                                response_chunk = {
+                                    "id": f"chatcmpl-{uuid.uuid4()}",
+                                    "object": "chat.completion.chunk",
+                                    "created": int(datetime.now().timestamp()),
+                                    "model": request.model,
+                                    "choices": [
+                                        {
+                                            "index": 0,
+                                            "delta": {"content": analysis_result + "\n", "role": "assistant"},
+                                            "finish_reason": None,
+                                        }
+                                    ],
+                                    "usage": None,
+                                }
+                                yield f"data: {json.dumps(response_chunk)}\n\n"
+
+                    # Final message
+                    final_response = {
+                        "id": f"chatcmpl-{uuid.uuid4()}",
+                        "object": "chat.completion",
+                        "created": int(datetime.now().timestamp()),
+                        "model": request.model,
+                        "choices": [
+                            {
+                                "message": {
+                                    "role": "assistant",
+                                    "content": assistant_content.strip()
+                                },
+                                "finish_reason": "stop",
+                                "index": 0
+                            }
+                        ],
+                        "usage": {
+                            "prompt_tokens": prompt_tokens,
+                            "completion_tokens": completion_tokens,
+                            "total_tokens": total_tokens,
+                            "estimated_cost": estimated_cost
+                        },
+                    }
+                    yield f"data: {json.dumps(final_response)}\n\n"
+                    yield "data: [DONE]\n\n"
+                except HTTPException as he:
+                    error_response = {"error": he.detail}
+                    yield f"data: {json.dumps(error_response)}\n\n"
+                except Exception as e:
+                    logger.exception(f"Error during streaming response generation from IP: {client_ip}.")
+                    error_response = {"error": str(e)}
+                    yield f"data: {json.dumps(error_response)}\n\n"
+
+            return StreamingResponse(generate(), media_type="text/event-stream")
+        else:
+            return {
+                "id": f"chatcmpl-{uuid.uuid4()}",
+                "object": "chat.completion",
+                "created": int(datetime.now().timestamp()),
+                "model": request.model,
+                "choices": [
+                    {
+                        "message": {
+                            "role": "assistant",
+                            "content": assistant_content.strip()
+                        },
+                        "finish_reason": "stop",
+                        "index": 0
+                    }
+                ],
+                "usage": {
+                    "prompt_tokens": prompt_tokens,
+                    "completion_tokens": completion_tokens,
+                    "total_tokens": total_tokens,
+                    "estimated_cost": estimated_cost
+                },
+            }
     except ModelNotWorkingException as e:
         logger.warning(f"Model not working: {e} | IP: {client_ip}")
         raise HTTPException(status_code=503, detail=str(e))