:sparkles: Update the /v1/responses endpoint (#75)

* Remove the unused auto-refresh functionality and related imports.

They are no longer needed since the underlying library issue has been resolved.

* Enhance error handling in client initialization and message sending

* Refactor link handling to extract file paths and simplify Google search links

* Fix regex pattern for Google search link matching

* Fix regex patterns for Markdown escaping, code fence and Google search link matching

* Increase timeout value in configuration files from 60 to 120 seconds to better handle heavy tasks

* Fix Image generation

* Refactor tool handling to support standard and image generation tools separately

* Fix: use "ascii" decoding for base64-encoded image data consistency

* Fix: replace `running` with `_running` for internal client status checks

* Refactor: replace direct `_running` access with `running()` method in client status checks

* Extend models with new fields for annotations, reasoning, audio, log probabilities, and token details; adjust response handling accordingly.

* Extend models with new fields (annotations, error), add `normalize_output_text` validator, rename `created` to `created_at`, and update response handling accordingly.

* Extend response models to support tool choices, image output, and improved streaming of response items. Refactor image generation handling for consistency and add compatibility with output content.

* Set default `text` value to an empty string for `ResponseOutputContent` and ensure consistent initialization in image output handling.

* feat: Add /images endpoint with dedicated router and improved image management

Add dedicated router for /images endpoint and refactor image handling logic for better modularity. Enhance temporary image management with secure naming, token verification, and cleanup functionality.

* feat: Add token-based verification for image access

* Refactor: rename image store directory to `ai_generated_images` for clarity

* fix: Update create_response to use FastAPI Request object for base_url and refactor variable handling

* fix: Correct attribute access in request_data handling within `chat.py` for tools, tool_choice, and streaming settings

* fix: Save generated images to persistent storage

* fix: Remove unused `output_image` type from `ResponseOutputContent` and update response handling for consistency

* fix: Update image URL generation in chat response to use Markdown format for compatibility

app/main.py

@@ -6,7 +6,12 @@ from loguru import logger
 
 from .server.chat import router as chat_router
 from .server.health import router as health_router
-from .server.middleware import add_cors_middleware, add_exception_handler
+from .server.images import router as images_router
+from .server.middleware import (
+    add_cors_middleware,
+    add_exception_handler,
+    cleanup_expired_images,
+)
 from .services import GeminiClientPool, LMDBConversationStore
 
 RETENTION_CLEANUP_INTERVAL_SECONDS = 6 * 60 * 60  # 6 hours
@@ -28,6 +33,7 @@ async def _run_retention_cleanup(stop_event: asyncio.Event) -> None:
     while not stop_event.is_set():
         try:
             store.cleanup_expired()
+            cleanup_expired_images(store.retention_days)
         except Exception:
             logger.exception("LMDB retention cleanup task failed.")
 
@@ -93,5 +99,6 @@ def create_app() -> FastAPI:
 
     app.include_router(health_router, tags=["Health"])
     app.include_router(chat_router, tags=["Chat"])
+    app.include_router(images_router, tags=["Images"])
 
     return app

app/models/models.py

@@ -3,7 +3,7 @@ from __future__ import annotations
 from datetime import datetime
 from typing import Any, Dict, List, Literal, Optional, Union
 
-from pydantic import BaseModel, Field
+from pydantic import BaseModel, Field, model_validator
 
 
 class ContentItem(BaseModel):
@@ -12,7 +12,9 @@ class ContentItem(BaseModel):
     type: Literal["text", "image_url", "file", "input_audio"]
     text: Optional[str] = None
     image_url: Optional[Dict[str, str]] = None
+    input_audio: Optional[Dict[str, Any]] = None
     file: Optional[Dict[str, str]] = None
+    annotations: List[Dict[str, Any]] = Field(default_factory=list)
 
 
 class Message(BaseModel):
@@ -22,6 +24,10 @@ class Message(BaseModel):
     content: Union[str, List[ContentItem], None] = None
     name: Optional[str] = None
     tool_calls: Optional[List["ToolCall"]] = None
+    refusal: Optional[str] = None
+    reasoning_content: Optional[str] = None
+    audio: Optional[Dict[str, Any]] = None
+    annotations: List[Dict[str, Any]] = Field(default_factory=list)
 
 
 class Choice(BaseModel):
@@ -30,6 +36,7 @@ class Choice(BaseModel):
     index: int
     message: Message
     finish_reason: str
+    logprobs: Optional[Dict[str, Any]] = None
 
 
 class FunctionCall(BaseModel):
@@ -81,6 +88,8 @@ class Usage(BaseModel):
     prompt_tokens: int
     completion_tokens: int
     total_tokens: int
+    prompt_tokens_details: Optional[Dict[str, int]] = None
+    completion_tokens_details: Optional[Dict[str, int]] = None
 
 
 class ModelData(BaseModel):
@@ -161,6 +170,15 @@ class ResponseInputContent(BaseModel):
     file_url: Optional[str] = None
     file_data: Optional[str] = None
     filename: Optional[str] = None
+    annotations: List[Dict[str, Any]] = Field(default_factory=list)
+
+    @model_validator(mode="before")
+    @classmethod
+    def normalize_output_text(cls, data: Any) -> Any:
+        """Allow output_text (from previous turns) to be treated as input_text."""
+        if isinstance(data, dict) and data.get("type") == "output_text":
+            data["type"] = "input_text"
+        return data
 
 
 class ResponseInputItem(BaseModel):
@@ -216,7 +234,8 @@ class ResponseOutputContent(BaseModel):
     """Content item for Responses API output."""
 
     type: Literal["output_text"]
-    text: Optional[str] = None
+    text: Optional[str] = ""
+    annotations: List[Dict[str, Any]] = Field(default_factory=list)
 
 
 class ResponseOutputMessage(BaseModel):
@@ -254,20 +273,22 @@ class ResponseCreateResponse(BaseModel):
 
     id: str
     object: Literal["response"] = "response"
-    created: int
+    created_at: int
     model: str
     output: List[Union[ResponseOutputMessage, ResponseImageGenerationCall, ResponseToolCall]]
-    output_text: Optional[str] = None
     status: Literal[
         "in_progress",
         "completed",
         "failed",
         "incomplete",
+        "cancelled",
         "requires_action",
     ] = "completed"
+    tool_choice: Optional[Union[str, ResponseToolChoice]] = None
+    tools: Optional[List[Union[Tool, ResponseImageTool]]] = None
     usage: ResponseUsage
+    error: Optional[Dict[str, Any]] = None
     metadata: Optional[Dict[str, Any]] = None
-    system_fingerprint: Optional[str] = None
     input: Optional[Union[str, List[ResponseInputItem]]] = None
 
 

app/server/chat.py

@@ -9,7 +9,7 @@ from pathlib import Path
 from typing import Any, Iterator
 
 import orjson
-from fastapi import APIRouter, Depends, HTTPException, status
+from fastapi import APIRouter, Depends, HTTPException, Request, status
 from fastapi.responses import StreamingResponse
 from gemini_webapi.client import ChatSession
 from gemini_webapi.constants import Model
@@ -44,7 +44,7 @@ from ..services import GeminiClientPool, GeminiClientWrapper, LMDBConversationSt
 from ..services.client import CODE_BLOCK_HINT, XML_WRAP_HINT
 from ..utils import g_config
 from ..utils.helper import estimate_tokens
-from .middleware import get_temp_dir, verify_api_key
+from .middleware import get_image_store_dir, get_image_token, get_temp_dir, verify_api_key
 
 # Maximum characters Gemini Web can accept in a single request (configurable)
 MAX_CHARS_PER_REQUEST = int(g_config.gemini.max_chars_per_request * 0.9)
@@ -588,6 +588,7 @@ async def create_chat_completion(
     request: ChatCompletionRequest,
     api_key: str = Depends(verify_api_key),
     tmp_dir: Path = Depends(get_temp_dir),
+    image_store: Path = Depends(get_image_store_dir),
 ):
     pool = GeminiClientPool()
     db = LMDBConversationStore()
@@ -772,18 +773,15 @@ async def create_chat_completion(
 
 @router.post("/v1/responses")
 async def create_response(
-    request: ResponseCreateRequest,
+    request_data: ResponseCreateRequest,
+    request: Request,
     api_key: str = Depends(verify_api_key),
     tmp_dir: Path = Depends(get_temp_dir),
+    image_store: Path = Depends(get_image_store_dir),
 ):
-    base_messages, normalized_input = _response_items_to_messages(request.input)
-    if not base_messages:
-        raise HTTPException(
-            status_code=status.HTTP_400_BAD_REQUEST, detail="No message input provided."
-        )
-
-    structured_requirement = _build_structured_requirement(request.response_format)
-    if structured_requirement and request.stream:
+    base_messages, normalized_input = _response_items_to_messages(request_data.input)
+    structured_requirement = _build_structured_requirement(request_data.response_format)
+    if structured_requirement and request_data.stream:
         logger.debug(
             "Structured response requested with streaming enabled; streaming not supported for Responses."
         )
@@ -799,8 +797,8 @@ async def create_response(
     standard_tools: list[Tool] = []
     image_tools: list[ResponseImageTool] = []
 
-    if request.tools:
-        for t in request.tools:
+    if request_data.tools:
+        for t in request_data.tools:
             if isinstance(t, Tool):
                 standard_tools.append(t)
             elif isinstance(t, ResponseImageTool):
@@ -815,13 +813,15 @@ async def create_response(
 
     image_instruction = _build_image_generation_instruction(
         image_tools,
-        request.tool_choice if isinstance(request.tool_choice, ResponseToolChoice) else None,
+        request_data.tool_choice
+        if isinstance(request_data.tool_choice, ResponseToolChoice)
+        else None,
     )
     if image_instruction:
         extra_instructions.append(image_instruction)
         logger.debug("Image generation support enabled for /v1/responses request.")
 
-    preface_messages = _instructions_to_messages(request.instructions)
+    preface_messages = _instructions_to_messages(request_data.instructions)
     conversation_messages = base_messages
     if preface_messages:
         conversation_messages = [*preface_messages, *base_messages]
@@ -832,10 +832,10 @@ async def create_response(
     # Pass standard tools to the prompt builder
     # Determine tool_choice for standard tools (ignore image_generation choice here as it is handled via instruction)
     model_tool_choice = None
-    if isinstance(request.tool_choice, str):
-        model_tool_choice = request.tool_choice
-    elif isinstance(request.tool_choice, ToolChoiceFunction):
-        model_tool_choice = request.tool_choice
+    if isinstance(request_data.tool_choice, str):
+        model_tool_choice = request_data.tool_choice
+    elif isinstance(request_data.tool_choice, ToolChoiceFunction):
+        model_tool_choice = request_data.tool_choice
     # If tool_choice is ResponseToolChoice (image_generation), we don't pass it as a function tool choice.
 
     messages = _prepare_messages_for_model(
@@ -849,7 +849,7 @@ async def create_response(
     db = LMDBConversationStore()
 
     try:
-        model = Model.from_name(request.model)
+        model = Model.from_name(request_data.model)
     except ValueError as exc:
         raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail=str(exc)) from exc
 
@@ -969,7 +969,7 @@ async def create_response(
         )
 
     expects_image = (
-        request.tool_choice is not None and request.tool_choice.type == "image_generation"
+        request_data.tool_choice is not None and request_data.tool_choice.type == "image_generation"
     )
     images = model_output.images or []
     logger.debug(
@@ -984,7 +984,7 @@ async def create_response(
                 summary = f"{summary[:197]}..."
         logger.warning(
             "Image generation requested but Gemini produced no images. "
-            f"client_id={client_id}, forced_tool_choice={request.tool_choice is not None}, "
+            f"client_id={client_id}, forced_tool_choice={request_data.tool_choice is not None}, "
             f"instruction_applied={bool(image_instruction)}, assistant_preview='{summary}'"
         )
         detail = "LLM returned no images for the requested image_generation tool."
@@ -992,22 +992,35 @@ async def create_response(
             detail = f"{detail} Assistant response: {summary}"
         raise HTTPException(status_code=status.HTTP_502_BAD_GATEWAY, detail=detail)
 
+    response_contents: list[ResponseOutputContent] = []
     image_call_items: list[ResponseImageGenerationCall] = []
     for image in images:
         try:
-            image_base64, width, height = await _image_to_base64(image, tmp_dir)
+            image_base64, width, height, filename = await _image_to_base64(image, image_store)
         except Exception as exc:
             logger.warning(f"Failed to download generated image: {exc}")
             continue
+
+        img_format = "png" if isinstance(image, GeneratedImage) else "jpeg"
+
+        # Use static URL for compatibility
+        image_url = (
+            f"![{filename}]({request.base_url}images/{filename}?token={get_image_token(filename)})"
+        )
+
         image_call_items.append(
             ResponseImageGenerationCall(
                 id=f"img_{uuid.uuid4().hex}",
                 status="completed",
                 result=image_base64,
-                output_format="png" if isinstance(image, GeneratedImage) else "jpeg",
+                output_format=img_format,
                 size=f"{width}x{height}" if width and height else None,
             )
         )
+        # Add as output_text content for compatibility
+        response_contents.append(
+            ResponseOutputContent(type="output_text", text=image_url, annotations=[])
+        )
 
     tool_call_items: list[ResponseToolCall] = []
     if detected_tool_calls:
@@ -1020,12 +1033,13 @@ async def create_response(
             for call in detected_tool_calls
         ]
 
-    response_contents: list[ResponseOutputContent] = []
     if assistant_text:
-        response_contents.append(ResponseOutputContent(type="output_text", text=assistant_text))
+        response_contents.append(
+            ResponseOutputContent(type="output_text", text=assistant_text, annotations=[])
+        )
 
     if not response_contents:
-        response_contents.append(ResponseOutputContent(type="output_text", text=""))
+        response_contents.append(ResponseOutputContent(type="output_text", text="", annotations=[]))
 
     created_time = int(datetime.now(tz=timezone.utc).timestamp())
     response_id = f"resp_{uuid.uuid4().hex}"
@@ -1047,8 +1061,8 @@ async def create_response(
 
     response_payload = ResponseCreateResponse(
         id=response_id,
-        created=created_time,
-        model=request.model,
+        created_at=created_time,
+        model=request_data.model,
         output=[
             ResponseOutputMessage(
                 id=message_id,
@@ -1059,11 +1073,12 @@ async def create_response(
             *tool_call_items,
             *image_call_items,
         ],
-        output_text=assistant_text or None,
         status="completed",
         usage=usage,
         input=normalized_input or None,
-        metadata=request.metadata or None,
+        metadata=request_data.metadata or None,
+        tools=request_data.tools,
+        tool_choice=request_data.tool_choice,
     )
 
     try:
@@ -1084,7 +1099,7 @@ async def create_response(
     except Exception as exc:
         logger.warning(f"Failed to save Responses conversation to LMDB: {exc}")
 
-    if request.stream:
+    if request_data.stream:
         logger.debug(
             f"Streaming Responses API payload (response_id={response_payload.id}, text_chunks={bool(assistant_text)})."
         )
@@ -1333,7 +1348,7 @@ def _create_responses_streaming_response(
 
     response_dict = response_payload.model_dump(mode="json")
     response_id = response_payload.id
-    created_time = response_payload.created
+    created_time = response_payload.created_at
     model = response_payload.model
 
     logger.debug(
@@ -1343,14 +1358,14 @@ def _create_responses_streaming_response(
     base_event = {
         "id": response_id,
         "object": "response",
-        "created": created_time,
+        "created_at": created_time,
         "model": model,
     }
 
     created_snapshot: dict[str, Any] = {
         "id": response_id,
         "object": "response",
-        "created": created_time,
+        "created_at": created_time,
         "model": model,
         "status": "in_progress",
     }
@@ -1358,6 +1373,10 @@ def _create_responses_streaming_response(
         created_snapshot["metadata"] = response_dict["metadata"]
     if response_dict.get("input") is not None:
         created_snapshot["input"] = response_dict["input"]
+    if response_dict.get("tools") is not None:
+        created_snapshot["tools"] = response_dict["tools"]
+    if response_dict.get("tool_choice") is not None:
+        created_snapshot["tool_choice"] = response_dict["tool_choice"]
 
     async def generate_stream():
         # Emit creation event
@@ -1368,30 +1387,53 @@ def _create_responses_streaming_response(
         }
         yield f"data: {orjson.dumps(data).decode('utf-8')}\n\n"
 
-        # Stream textual content, if any
-        if assistant_text:
-            for chunk in _iter_stream_segments(assistant_text):
-                delta_event = {
-                    **base_event,
-                    "type": "response.output_text.delta",
-                    "output_index": 0,
-                    "delta": chunk,
-                }
-                yield f"data: {orjson.dumps(delta_event).decode('utf-8')}\n\n"
-
-            done_event = {
+        # Stream output items (Message/Text, Tool Calls, Images)
+        for i, item in enumerate(response_payload.output):
+            item_json = item.model_dump(mode="json", exclude_none=True)
+
+            added_event = {
                 **base_event,
-                "type": "response.output_text.done",
-                "output_index": 0,
+                "type": "response.output_item.added",
+                "output_index": i,
+                "item": item_json,
             }
-            yield f"data: {orjson.dumps(done_event).decode('utf-8')}\n\n"
-        else:
-            done_event = {
+            yield f"data: {orjson.dumps(added_event).decode('utf-8')}\n\n"
+
+            # 2. Stream content if it's a message (text)
+            if item.type == "message":
+                content_text = ""
+                # Aggregate text content to stream
+                for c in item.content:
+                    if c.type == "output_text" and c.text:
+                        content_text += c.text
+
+                if content_text:
+                    for chunk in _iter_stream_segments(content_text):
+                        delta_event = {
+                            **base_event,
+                            "type": "response.output_text.delta",
+                            "output_index": i,
+                            "delta": chunk,
+                        }
+                        yield f"data: {orjson.dumps(delta_event).decode('utf-8')}\n\n"
+
+                    # Text done
+                    done_event = {
+                        **base_event,
+                        "type": "response.output_text.done",
+                        "output_index": i,
+                    }
+                    yield f"data: {orjson.dumps(done_event).decode('utf-8')}\n\n"
+
+            # 3. Emit output_item.done for all types
+            # This confirms the item is fully transferred.
+            item_done_event = {
                 **base_event,
-                "type": "response.output_text.done",
-                "output_index": 0,
+                "type": "response.output_item.done",
+                "output_index": i,
+                "item": item_json,
             }
-            yield f"data: {orjson.dumps(done_event).decode('utf-8')}\n\n"
+            yield f"data: {orjson.dumps(item_done_event).decode('utf-8')}\n\n"
 
         # Emit completed event with full payload
         completed_event = {
@@ -1512,8 +1554,8 @@ def _extract_image_dimensions(data: bytes) -> tuple[int | None, int | None]:
     return None, None
 
 
-async def _image_to_base64(image: Image, temp_dir: Path) -> tuple[str, int | None, int | None]:
-    """Persist an image provided by gemini_webapi and return base64 plus dimensions."""
+async def _image_to_base64(image: Image, temp_dir: Path) -> tuple[str, int | None, int | None, str]:
+    """Persist an image provided by gemini_webapi and return base64 plus dimensions and filename."""
     if isinstance(image, GeneratedImage):
         saved_path = await image.save(path=str(temp_dir), full_size=True)
     else:
@@ -1522,6 +1564,13 @@ async def _image_to_base64(image: Image, temp_dir: Path) -> tuple[str, int | Non
     if not saved_path:
         raise ValueError("Failed to save generated image")
 
-    data = Path(saved_path).read_bytes()
+    # Rename file to a random UUID to ensure uniqueness and unpredictability
+    original_path = Path(saved_path)
+    random_name = f"img_{uuid.uuid4().hex}{original_path.suffix}"
+    new_path = temp_dir / random_name
+    original_path.rename(new_path)
+
+    data = new_path.read_bytes()
     width, height = _extract_image_dimensions(data)
-    return base64.b64encode(data).decode("ascii"), width, height
+    filename = random_name
+    return base64.b64encode(data).decode("ascii"), width, height, filename

app/server/images.py

@@ -0,0 +1,18 @@
+from fastapi import APIRouter, HTTPException, Query
+from fastapi.responses import FileResponse
+
+from ..server.middleware import get_image_store_dir, verify_image_token
+
+router = APIRouter()
+
+
+@router.get("/images/{filename}", tags=["Images"])
+async def get_image(filename: str, token: str | None = Query(default=None)):
+    if not verify_image_token(filename, token):
+        raise HTTPException(status_code=403, detail="Invalid token")
+
+    image_store = get_image_store_dir()
+    file_path = image_store / filename
+    if not file_path.exists():
+        raise HTTPException(status_code=404, detail="Image not found")
+    return FileResponse(file_path)

app/server/middleware.py

@@ -1,13 +1,72 @@
+import hashlib
+import hmac
 import tempfile
+import time
 from pathlib import Path
 
 from fastapi import Depends, FastAPI, HTTPException, Request, status
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.responses import ORJSONResponse
 from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer
+from loguru import logger
 
 from ..utils import g_config
 
+# Persistent directory for storing generated images
+IMAGE_STORE_DIR = Path(tempfile.gettempdir()) / "ai_generated_images"
+IMAGE_STORE_DIR.mkdir(parents=True, exist_ok=True)
+
+
+def get_image_store_dir() -> Path:
+    """Returns a persistent directory for storing images."""
+    return IMAGE_STORE_DIR
+
+
+def get_image_token(filename: str) -> str:
+    """Generate a HMAC-SHA256 token for a filename using the API key."""
+    secret = g_config.server.api_key
+    if not secret:
+        return ""
+
+    msg = filename.encode("utf-8")
+    secret_bytes = secret.encode("utf-8")
+    return hmac.new(secret_bytes, msg, hashlib.sha256).hexdigest()
+
+
+def verify_image_token(filename: str, token: str | None) -> bool:
+    """Verify the provided token against the filename."""
+    expected = get_image_token(filename)
+    if not expected:
+        return True  # No auth required
+    if not token:
+        return False
+    return hmac.compare_digest(token, expected)
+
+
+def cleanup_expired_images(retention_days: int) -> int:
+    """Delete images in IMAGE_STORE_DIR older than retention_days."""
+    if retention_days <= 0:
+        return 0
+
+    now = time.time()
+    retention_seconds = retention_days * 24 * 60 * 60
+    cutoff = now - retention_seconds
+
+    count = 0
+    for file_path in IMAGE_STORE_DIR.iterdir():
+        if not file_path.is_file():
+            continue
+        try:
+            if file_path.stat().st_mtime < cutoff:
+                file_path.unlink()
+                count += 1
+        except Exception as e:
+            logger.warning(f"Failed to delete expired image {file_path}: {e}")
+
+    if count > 0:
+        logger.info(f"Cleaned up {count} expired images.")
+    return count
+
 
 def global_exception_handler(request: Request, exc: Exception):
     if isinstance(exc, HTTPException):

app/services/client.py

@@ -76,6 +76,9 @@ class GeminiClientWrapper(GeminiClient):
             logger.exception(f"Failed to initialize GeminiClient {self.id}")
             raise
 
+    def running(self) -> bool:
+        return self._running
+
     @staticmethod
     async def process_message(
         message: Message, tempdir: Path | None = None, tagged: bool = True

app/services/pool.py

@@ -37,7 +37,7 @@ class GeminiClientPool(metaclass=Singleton):
         """Initialize all clients in the pool."""
         success_count = 0
         for client in self._clients:
-            if not client._running:
+            if not client.running():
                 try:
                     await client.init(
                         timeout=g_config.gemini.timeout,
@@ -48,7 +48,7 @@ class GeminiClientPool(metaclass=Singleton):
                 except Exception:
                     logger.exception(f"Failed to initialize client {client.id}")
 
-            if client._running:
+            if client.running():
                 success_count += 1
 
         if success_count == 0:
@@ -79,7 +79,7 @@ class GeminiClientPool(metaclass=Singleton):
 
     async def _ensure_client_ready(self, client: GeminiClientWrapper) -> bool:
         """Make sure the client is running, attempting a restart if needed."""
-        if client._running:
+        if client.running():
             return True
 
         lock = self._restart_locks.get(client.id)
@@ -87,7 +87,7 @@ class GeminiClientPool(metaclass=Singleton):
             return False  # Should not happen
 
         async with lock:
-            if client._running:
+            if client.running():
                 return True
 
             try:
@@ -110,4 +110,4 @@ class GeminiClientPool(metaclass=Singleton):
 
     def status(self) -> Dict[str, bool]:
         """Return running status for each client."""
-        return {client.id: client._running for client in self._clients}
+        return {client.id: client.running() for client in self._clients}