Feat: Build a student homework assistant API

.gitignore

@@ -0,0 +1,11 @@
+# Python-generated files
+__pycache__/
+*.py[oc]
+build/
+dist/
+wheels/
+*.egg-info
+
+# Virtual environments
+.venv
+.env

.python-version

@@ -0,0 +1 @@
+3.11

api/chat.py

@@ -0,0 +1,59 @@
+from fastapi import APIRouter, HTTPException, Depends
+from fastapi.responses import StreamingResponse
+import json
+from typing import AsyncGenerator
+
+from models.chat import ChatRequest, ChatResponse, StreamResponse
+from services.chat_service import ChatService
+
+chat_router = APIRouter(prefix="/chat", tags=["chat"])
+
+def get_chat_service():
+    return ChatService()
+
+@chat_router.post("/", response_model=ChatResponse)
+async def chat_completion(
+    request: ChatRequest,
+    chat_service: ChatService = Depends(get_chat_service)
+):
+    """
+    Process chat completion request
+    """
+    try:
+        response = await chat_service.process_chat(request)
+        return response
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
+@chat_router.post("/stream")
+async def chat_stream(
+    request: ChatRequest,
+    chat_service: ChatService = Depends(get_chat_service)
+):
+    """
+    Stream chat completion response
+    """
+    try:
+        async def generate_stream() -> AsyncGenerator[str, None]:
+            async for chunk in chat_service.stream_chat(request):
+                yield f"data: {json.dumps(chunk.model_dump())}\n\n"
+            
+            # Send final message
+            final_chunk = StreamResponse(finished=True)
+            yield f"data: {json.dumps(final_chunk.model_dump())}\n\n"
+
+        return StreamingResponse(
+            generate_stream(),
+            media_type="text/plain",
+            headers={"Cache-Control": "no-cache", "Connection": "keep-alive"}
+        )
+    except Exception as e:
+        error_chunk = StreamResponse(error=str(e), finished=True)
+        async def error_stream():
+            yield f"data: {json.dumps(error_chunk.model_dump())}\n\n"
+        
+        return StreamingResponse(
+            error_stream(),
+            media_type="text/plain",
+            status_code=500
+        )

main.py

@@ -0,0 +1,38 @@
+from fastapi import FastAPI, HTTPException
+from fastapi.staticfiles import StaticFiles
+from fastapi.responses import FileResponse
+from fastapi.middleware.cors import CORSMiddleware
+from dotenv import load_dotenv
+
+from api.chat import chat_router
+
+load_dotenv()
+
+app = FastAPI(
+    title="Chat Assistant API",
+    description="FastAPI backend for chat assistant with LiteLLM and Gemini",
+    version="1.0.0"
+)
+
+# CORS middleware for frontend integration
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],  # Configure this properly for production
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+# Include routers
+app.include_router(chat_router, prefix="/api")
+
+# Mount static files for the frontend
+app.mount("/static", StaticFiles(directory="static"), name="static")
+
+@app.get("/")
+async def root():
+    return FileResponse("static/index.html")
+
+@app.get("/health")
+async def health_check():
+    return {"status": "healthy"}

models/chat.py

@@ -0,0 +1,32 @@
+from pydantic import BaseModel, Field
+from typing import List, Dict, Literal, Optional, Union
+
+class MessageContent(BaseModel):
+    type: str = Field(..., description="Content type: 'text' or 'image_url'")
+    text: Optional[str] = Field(None, description="Text content")
+    image_url: Optional[Dict[str, str]] = Field(None, description="Image URL object with 'url' key")
+
+class ChatMessage(BaseModel):
+    role: str = Field(..., description="Role of the message sender (system, user, assistant)")
+    content: Union[str, List[MessageContent]] = Field(..., description="Message content - string for simple text or list for mixed content")
+
+class ChatSettings(BaseModel):
+    temperature: float = Field(1.0, ge=0.0, le=2.0)
+    top_p: float = Field(0.8, ge=0.0, le=1.0)
+    reasoning_effort: Literal["low", "medium", "high"] = "low"
+
+class ChatRequest(BaseModel):
+    messages: List[ChatMessage] = Field(..., description="List of messages in the conversation")
+    settings: Optional[ChatSettings] = Field(default_factory=ChatSettings)
+    command: Optional[Literal["search", "url_context"]] = None
+
+class ChatResponse(BaseModel):
+    content: str
+    reasoning_content: Optional[str] = None
+    error: Optional[str] = None
+
+class StreamResponse(BaseModel):
+    content: Optional[str] = None
+    reasoning_content: Optional[str] = None
+    finished: bool = False
+    error: Optional[str] = None

prompts.py

@@ -0,0 +1,8 @@
+# prompts.py
+# Replace this with your actual system prompt
+
+SYSTEM_PROMPT = """You are a helpful AI assistant. You can analyze text and images, answer questions, and help with various tasks.
+
+When provided with images, analyze them carefully and describe what you see. When asked to search for information, use the available tools to find current information.
+
+Be helpful, accurate, and concise in your responses."""

pyproject.toml

@@ -0,0 +1,14 @@
+[project]
+name = "student-homework-assistant"
+version = "0.1.0"
+description = "Add your description here"
+readme = "README.md"
+requires-python = ">=3.11"
+dependencies = [
+    "chainlit>=2.7.2",
+    "fastapi>=0.116.1",
+    "litellm>=1.76.2",
+    "pillow>=11.3.0",
+    "pydantic-settings>=2.10.1",
+    "uvicorn>=0.35.0",
+]

services/chat_service.py

@@ -0,0 +1,137 @@
+import os
+from typing import List, Dict, Any, Optional, AsyncGenerator
+from litellm import acompletion
+
+from models.chat import ChatRequest, ChatResponse, StreamResponse
+from prompts import SYSTEM_PROMPT
+
+class ChatService:
+    def __init__(self):
+        self.model = f"gemini/{os.getenv('MODEL', 'gemini-2.5-flash')}"
+        self.api_key = os.getenv("GOOGLE_API_KEY")
+        
+        if not self.api_key:
+            raise ValueError("GOOGLE_API_KEY environment variable is required")
+
+    async def process_chat(self, request: ChatRequest) -> ChatResponse:
+        """Process a complete chat request and return the response"""
+        try:
+            # Prepare messages with system prompt
+            messages = self._prepare_messages(request)
+            
+            # Prepare tools based on command
+            tools = self._prepare_tools(request.command)
+            
+            # Make the API call
+            response = await acompletion(
+                model=self.model,
+                messages=messages,
+                temperature=request.settings.temperature,
+                top_p=request.settings.top_p,
+                reasoning_effort=request.settings.reasoning_effort,
+                api_key=self.api_key,
+                tools=tools,
+                stream=False
+            )
+            
+            content = response.choices[0].message.content or ""
+            reasoning_content = getattr(response.choices[0].message, "reasoning_content", None)
+            
+            return ChatResponse(
+                content=content,
+                reasoning_content=reasoning_content
+            )
+            
+        except Exception as e:
+            return ChatResponse(
+                content="",
+                error=str(e)
+            )
+
+    async def stream_chat(self, request: ChatRequest) -> AsyncGenerator[StreamResponse, None]:
+        """Stream chat response"""
+        try:
+            # Prepare messages with system prompt
+            messages = self._prepare_messages(request)
+            
+            # Prepare tools based on command
+            tools = self._prepare_tools(request.command)
+            
+            # Make the streaming API call
+            response = await acompletion(
+                model=self.model,
+                messages=messages,
+                temperature=request.settings.temperature,
+                top_p=request.settings.top_p,
+                reasoning_effort=request.settings.reasoning_effort,
+                api_key=self.api_key,
+                tools=tools,
+                stream=True
+            )
+            
+            async for part in response:
+                chunk_response = StreamResponse()
+                
+                if hasattr(part.choices[0].delta, "reasoning_content"):
+                    reasoning = part.choices[0].delta.reasoning_content or ""
+                    if reasoning:
+                        chunk_response.reasoning_content = reasoning
+                
+                content = part.choices[0].delta.content or ""
+                if content:
+                    chunk_response.content = content
+                
+                if chunk_response.content or chunk_response.reasoning_content:
+                    yield chunk_response
+                    
+        except Exception as e:
+            yield StreamResponse(error=str(e), finished=True)
+
+    def _prepare_messages(self, request: ChatRequest) -> List[Dict[str, Any]]:
+        """Prepare messages with system prompt"""
+        messages = []
+        
+        # Add system message first
+        messages.append({
+            "role": "system", 
+            "content": SYSTEM_PROMPT
+        })
+        
+        # Add all messages from the request
+        for message in request.messages:
+            if message.role == "system":
+                continue  # Skip additional system messages since we already added one
+                
+            # Convert message to dict format expected by litellm
+            if isinstance(message.content, str):
+                # Simple text message
+                messages.append({
+                    "role": message.role,
+                    "content": message.content
+                })
+            else:
+                # Complex content with text and/or images
+                content_parts = []
+                for part in message.content:
+                    if hasattr(part, 'dict'):
+                        content_parts.append(part.dict())
+                    elif isinstance(part, dict):
+                        content_parts.append(part)
+                
+                messages.append({
+                    "role": message.role,
+                    "content": content_parts
+                })
+        
+        return messages
+
+    def _prepare_tools(self, command: str) -> List[Dict[str, Any]]:
+        """Prepare tools based on the command"""
+        tools = []
+        
+        if command == "search":
+            tools = [{"googleSearch": {}}]
+        elif command == "url_context":
+            tools = [{"urlContext": {}}]
+            
+        return tools

static/app.js

@@ -0,0 +1,239 @@
+const messagesEl = document.getElementById('messages');
+const formEl = document.getElementById('chat-form');
+const inputEl = document.getElementById('user-input');
+const sendBtn = document.getElementById('send-btn');
+const statusEl = document.getElementById('status');
+const commandEl = document.getElementById('command');
+const imageInput = document.getElementById('image-input');
+const addImageBtn = document.getElementById('add-image-btn');
+const clearImagesBtn = document.getElementById('clear-images-btn');
+const imagePreviews = document.getElementById('image-previews');
+
+// Conversation state
+let conversation = [];
+let pendingImages = []; // data URLs
+
+function renderMessage(message, isStreaming = false) {
+  const wrapper = document.createElement('div');
+  wrapper.className = `message role-${message.role}`;
+
+  const avatar = document.createElement('div');
+  avatar.className = 'avatar';
+  avatar.textContent = message.role === 'user' ? 'U' : 'A';
+
+  const bubble = document.createElement('div');
+  bubble.className = 'bubble';
+  if (isStreaming) bubble.classList.add('assistant-streaming');
+  // Text
+  let text = '';
+  if (typeof message.content === 'string') {
+    text = message.content;
+  } else if (Array.isArray(message.content)) {
+    const textPart = message.content.find(p => p.type === 'text');
+    text = textPart?.text || '';
+  }
+  const textNode = document.createElement('div');
+  textNode.innerText = text;
+  bubble.appendChild(textNode);
+
+  // Images
+  if (Array.isArray(message.content)) {
+    const images = message.content.filter(p => p.type === 'image_url');
+    if (images.length) {
+      const imgWrap = document.createElement('div');
+      imgWrap.className = 'bubble-images';
+      for (const img of images) {
+        const el = document.createElement('img');
+        el.src = img.image_url?.url || '';
+        imgWrap.appendChild(el);
+      }
+      bubble.appendChild(imgWrap);
+    }
+  }
+
+  wrapper.appendChild(avatar);
+  wrapper.appendChild(bubble);
+  messagesEl.appendChild(wrapper);
+  messagesEl.scrollTop = messagesEl.scrollHeight;
+  return bubble; // return bubble for incremental updates
+}
+
+function setBusy(busy) {
+  sendBtn.disabled = busy;
+  statusEl.textContent = busy ? 'Thinking…' : 'Ready';
+}
+
+function buildRequestPayload(command) {
+  return {
+    messages: conversation,
+    settings: {
+      temperature: 1.0,
+      top_p: 0.8,
+      reasoning_effort: 'low'
+    },
+    command: command || null
+  };
+}
+
+async function sendChat(command) {
+  setBusy(true);
+  const userMsg = conversation[conversation.length - 1];
+  renderMessage(userMsg);
+  const assistantMsg = { role: 'assistant', content: '' };
+  const assistantBubble = renderMessage(assistantMsg, true);
+
+  try {
+    const res = await fetch('/api/chat/stream', {
+      method: 'POST',
+      headers: { 'Content-Type': 'application/json' },
+      body: JSON.stringify(buildRequestPayload(command))
+    });
+
+    if (!res.ok || !res.body) {
+      throw new Error(`Request failed: ${res.status}`);
+    }
+
+    const reader = res.body.getReader();
+    const decoder = new TextDecoder();
+    let assistantText = '';
+    while (true) {
+      const { value, done } = await reader.read();
+      if (done) break;
+      const chunk = decoder.decode(value, { stream: true });
+
+      // SSE-style lines: data: {json}\n\n
+      const lines = chunk.split('\n');
+      for (const line of lines) {
+        if (!line.startsWith('data:')) continue;
+        const jsonStr = line.replace('data: ', '');
+        if (!jsonStr) continue;
+        try {
+          const data = JSON.parse(jsonStr);
+          if (data.error) {
+            throw new Error(data.error);
+          }
+          if (data.content) {
+            assistantText += data.content;
+            assistantBubble.firstChild.innerText = assistantText;
+          }
+          if (data.reasoning_content) {
+            // Optionally, we could display reasoning in a muted style
+          }
+          if (data.finished) {
+            assistantBubble.classList.remove('assistant-streaming');
+            // Push final assistant message to history
+            conversation.push({ role: 'assistant', content: assistantText });
+          }
+        } catch (e) {
+          // non-JSON line, ignore
+        }
+      }
+    }
+  } catch (err) {
+    assistantBubble.classList.remove('assistant-streaming');
+    assistantBubble.firstChild.innerText = `Error: ${err.message}`;
+    assistantBubble.classList.add('error');
+  } finally {
+    setBusy(false);
+    try { localStorage.setItem('sha_conversation', JSON.stringify(conversation)); } catch {}
+  }
+}
+
+formEl.addEventListener('submit', (e) => {
+  e.preventDefault();
+  const text = inputEl.value.trim();
+  const command = commandEl.value;
+  if (!text && pendingImages.length === 0) return;
+
+  const contentParts = [];
+  if (text) contentParts.push({ type: 'text', text });
+  for (const url of pendingImages) {
+    contentParts.push({ type: 'image_url', image_url: { url } });
+  }
+
+  const userMsg = contentParts.length === 1 && text
+    ? { role: 'user', content: text }
+    : { role: 'user', content: contentParts };
+
+  conversation.push(userMsg);
+
+  inputEl.value = '';
+  clearPendingImages();
+  sendChat(command);
+});
+
+// Auto-resize textarea
+inputEl.addEventListener('input', () => {
+  inputEl.style.height = 'auto';
+  inputEl.style.height = Math.min(inputEl.scrollHeight, 120) + 'px';
+});
+
+// Focus input on load and restore session
+window.addEventListener('load', () => {
+  inputEl.focus();
+  try {
+    const saved = localStorage.getItem('sha_conversation');
+    if (saved) {
+      conversation = JSON.parse(saved);
+      messagesEl.innerHTML = '';
+      for (const msg of conversation) {
+        renderMessage(msg);
+      }
+    }
+  } catch {}
+});
+
+// Image handling
+addImageBtn.addEventListener('click', () => imageInput.click());
+imageInput.addEventListener('change', async (e) => {
+  const files = Array.from(e.target.files || []);
+  for (const file of files) {
+    const dataUrl = await fileToDataUrl(file);
+    pendingImages.push(dataUrl);
+  }
+  renderPendingImages();
+  imageInput.value = '';
+});
+
+clearImagesBtn.addEventListener('click', () => {
+  pendingImages = [];
+  renderPendingImages();
+});
+
+function fileToDataUrl(file) {
+  return new Promise((resolve, reject) => {
+    const reader = new FileReader();
+    reader.onload = () => resolve(reader.result);
+    reader.onerror = reject;
+    reader.readAsDataURL(file);
+  });
+}
+
+function renderPendingImages() {
+  imagePreviews.innerHTML = '';
+  pendingImages.forEach((url, idx) => {
+    const wrap = document.createElement('div');
+    wrap.className = 'image-preview';
+    const img = document.createElement('img');
+    img.src = url;
+    const btn = document.createElement('button');
+    btn.className = 'remove';
+    btn.type = 'button';
+    btn.textContent = '×';
+    btn.addEventListener('click', () => {
+      pendingImages.splice(idx, 1);
+      renderPendingImages();
+    });
+    wrap.appendChild(img);
+    wrap.appendChild(btn);
+    imagePreviews.appendChild(wrap);
+  });
+}
+
+function clearPendingImages() {
+  pendingImages = [];
+  renderPendingImages();
+}
+
+
+

static/index.html

@@ -0,0 +1,47 @@
+<!DOCTYPE html>
+<html lang="en">
+<head>
+    <meta charset="UTF-8">
+    <meta name="viewport" content="width=device-width, initial-scale=1.0">
+    <title>Student Homework Assistant</title>
+    <link rel="stylesheet" href="/static/styles.css">
+    <link rel="preconnect" href="https://fonts.googleapis.com">
+    <link rel="preconnect" href="https://fonts.gstatic.com" crossorigin>
+    <link href="https://fonts.googleapis.com/css2?family=Inter:wght@400;500;600;700&display=swap" rel="stylesheet">
+</head>
+<body>
+    <div class="app">
+        <header class="app-header">
+            <div class="brand">Student Homework Assistant</div>
+            <div class="status" id="status">Ready</div>
+        </header>
+
+        <main class="chat-container">
+            <div id="messages" class="messages" aria-live="polite"></div>
+
+            <form id="chat-form" class="chat-input" autocomplete="off">
+                <select id="command" class="command-select" title="Command">
+                    <option value="">Chat</option>
+                    <option value="search">Search</option>
+                    <option value="url_context">URL Context</option>
+                </select>
+                <div class="input-area">
+                    <div id="image-previews" class="image-previews" aria-live="polite"></div>
+                    <textarea id="user-input" placeholder="Ask a question..." rows="1"></textarea>
+                </div>
+                <div class="actions">
+                    <input id="image-input" type="file" accept="image/*" hidden>
+                    <button id="add-image-btn" type="button" class="icon-btn" title="Add image" aria-label="Add image">📷</button>
+                    <button id="clear-images-btn" type="button" class="icon-btn" title="Clear images" aria-label="Clear images">✖️</button>
+                    <button id="send-btn" type="submit" class="send-btn">Send</button>
+                </div>
+            </form>
+        </main>
+    </div>
+
+    <script src="/static/app.js"></script>
+</body>
+<!-- More content could be added here (settings, history) -->
+</html>
+
+

static/styles.css

@@ -0,0 +1,66 @@
+:root {
+  --bg: #0b1220;
+  --panel: #121a2b;
+  --muted: #93a1ba;
+  --text: #e6ebf5;
+  --acc: #5b8cff;
+  --acc-2: #21d4a3;
+  --danger: #ff6b6b;
+  --border: #223154;
+}
+
+* { box-sizing: border-box; }
+html, body { height: 100%; }
+body {
+  margin: 0;
+  font-family: Inter, system-ui, -apple-system, Segoe UI, Roboto, Arial, sans-serif;
+  color: var(--text);
+  background: radial-gradient(1200px 800px at 80% -10%, #1a2a4a55, transparent 70%),
+             radial-gradient(1200px 800px at -10% 120%, #163a3a55, transparent 70%),
+             var(--bg);
+}
+
+.app { max-width: 920px; margin: 0 auto; padding: 16px; height: 100%; display: flex; flex-direction: column; }
+.app-header { display: flex; justify-content: space-between; align-items: center; padding: 12px 16px; border: 1px solid var(--border); background: linear-gradient(180deg, #17233b 0%, #101827 100%); border-radius: 14px; box-shadow: 0 10px 30px rgba(0,0,0,0.25) }
+.brand { font-weight: 700; letter-spacing: 0.2px; }
+.status { color: var(--muted); font-size: 13px; }
+
+.chat-container { flex: 1; display: flex; flex-direction: column; gap: 12px; margin-top: 14px; }
+.messages { flex: 1; overflow-y: auto; padding: 12px; border: 1px solid var(--border); background: #0e1627aa; border-radius: 14px; scroll-behavior: smooth; }
+
+.message { display: grid; grid-template-columns: 36px 1fr; gap: 10px; padding: 10px 8px; border-bottom: 1px dashed #23324f; }
+.message:last-child { border-bottom: none; }
+.avatar { width: 36px; height: 36px; border-radius: 10px; align-self: start; background: #152341; display:flex; align-items:center; justify-content:center; color:#9db4ff; font-weight:700; }
+.role-user .avatar{ background:#182b4d; color:#bfe1ff }
+.role-assistant .avatar{ background:#193b2f; color:#7ff0c4 }
+.bubble { padding: 10px 12px; border-radius: 12px; background: #13203b; border: 1px solid #23324f; }
+.role-user .bubble { background: #162343; }
+.role-assistant .bubble { background: #132f2a; border-color: #1f4b40; }
+
+.chat-input { display: grid; grid-template-columns: 140px 1fr auto; gap: 10px; padding: 12px; border: 1px solid var(--border); background: #0e1627aa; border-radius: 14px; align-items: end; }
+.command-select { height: 40px; border-radius: 10px; border: 1px solid var(--border); background: #101a30; color: var(--text); padding: 0 10px; }
+textarea { width: 100%; resize: none; max-height: 120px; min-height: 40px; padding: 10px 12px; border-radius: 10px; border: 1px solid var(--border); background: #101a30; color: var(--text); outline: none; }
+.send-btn { height: 40px; padding: 0 16px; border-radius: 10px; border: 1px solid #3a57a8; background: linear-gradient(180deg, #496eea 0%, #3754bd 100%); color: white; font-weight: 600; }
+.send-btn:disabled{ opacity: .7; cursor: not-allowed; }
+
+.assistant-streaming { background: repeating-linear-gradient(90deg, #132f2a, #132f2a 10px, #13392f 10px, #13392f 20px); animation: barber 1.2s linear infinite; }
+@keyframes barber { from { background-position: 0 0; } to { background-position: 40px 0; } }
+
+.error { color: var(--danger); }
+
+/* Input area */
+.input-area { display: flex; flex-direction: column; gap: 8px; }
+.actions { display: flex; align-items: center; gap: 8px; }
+.icon-btn { height: 40px; padding: 0 12px; border-radius: 10px; border: 1px solid var(--border); background: #0f1a33; color: var(--text); }
+
+/* Image previews */
+.image-previews { display: flex; flex-wrap: wrap; gap: 8px; }
+.image-preview { position: relative; width: 56px; height: 56px; border-radius: 10px; overflow: hidden; border: 1px solid var(--border); background: #111a30; }
+.image-preview img { width: 100%; height: 100%; object-fit: cover; display: block; }
+.image-preview .remove { position: absolute; top: 2px; right: 2px; background: rgba(0,0,0,0.6); color: white; border: none; border-radius: 6px; font-size: 12px; padding: 2px 4px; cursor: pointer; }
+
+/* Images inside bubbles */
+.bubble-images { display: flex; flex-wrap: wrap; gap: 8px; margin-top: 8px; }
+.bubble-images img { max-width: 180px; max-height: 180px; border-radius: 8px; border: 1px solid #23324f; }
+
+

utils.py

@@ -0,0 +1,66 @@
+import base64
+import io
+from PIL import Image
+from typing import Optional, Dict, Any
+
+def validate_base64_image(base64_string: str) -> bool:
+    """
+    Validate if a string is a valid base64 encoded image
+    """
+    try:
+        if base64_string.startswith('data:'):
+            # Extract base64 part from data URL
+            base64_part = base64_string.split(',', 1)[1] if ',' in base64_string else base64_string
+        else:
+            base64_part = base64_string
+            
+        # Try to decode and open as image
+        image_data = base64.b64decode(base64_part)
+        image = Image.open(io.BytesIO(image_data))
+        image.verify()  # Verify it's a valid image
+        return True
+        
+    except Exception:
+        return False
+
+def get_image_mime_type(base64_string: str) -> str:
+    """
+    Detect MIME type of base64 encoded image
+    """
+    try:
+        if base64_string.startswith('data:'):
+            # Extract MIME type from data URL
+            if ';' in base64_string:
+                return base64_string.split(';')[0].replace('data:', '')
+        
+        # Try to detect from image data
+        if base64_string.startswith('data:'):
+            base64_part = base64_string.split(',', 1)[1] if ',' in base64_string else base64_string
+        else:
+            base64_part = base64_string
+            
+        image_data = base64.b64decode(base64_part)
+        image = Image.open(io.BytesIO(image_data))
+        
+        format_mapping = {
+            'JPEG': 'image/jpeg',
+            'PNG': 'image/png',
+            'GIF': 'image/gif',
+            'WEBP': 'image/webp'
+        }
+        
+        return format_mapping.get(image.format, 'image/jpeg')
+        
+    except Exception:
+        return 'image/jpeg'  # Default fallback
+
+def ensure_data_url_format(image_data: str) -> str:
+    """
+    Ensure image data is in proper data URL format
+    """
+    if image_data.startswith('data:'):
+        return image_data
+    
+    # Detect MIME type and add data URL prefix
+    mime_type = get_image_mime_type(image_data)
+    return f"data:{mime_type};base64,{image_data}"