✨ Multi-segment text merging, ultra-long message splitting, xml prompt, tag, session, and token calculation optimization. (#16)

* ✨ Enhance text input handling by appending multiple fragments

* 💡 Add guidance for wrapping XML blocks in responses

* 💡 Skip tagging for the first message in a new session.

* ✨ Implement message splitting for Gemini requests exceeding the character limit

* ✨ try to reuse session as much as possible

* 💡 fix response token calculation to include the whole history

* 💡 Enhance streaming response to include token usage calculation

Author: WAcry

app/server/chat.py

@@ -20,6 +20,7 @@
     GeminiClientWrapper,
     LMDBConversationStore,
 )
+from ..utils import g_config
 from ..utils.helper import estimate_tokens
 from .middleware import get_temp_dir, verify_api_key
 
@@ -64,24 +65,23 @@ async def create_chat_completion(
         )
 
     # Check if conversation is reusable
-    session = None
-    client = None
-    if _check_reusable(request.messages):
-        try:
-            # Exclude the last message from user
-            if old_conv := db.find(model.model_name, request.messages[:-1]):
-                client = pool.acquire(old_conv.client_id)
-                session = client.start_chat(metadata=old_conv.metadata, model=model)
-        except Exception as e:
-            session = None
-            logger.warning(f"Error checking LMDB for reusable session: {e}")
+    session, client, remaining_messages = _find_reusable_session(
+        db, pool, model, request.messages
+    )
 
     if session:
-        # Just send the last message to the existing session
-        model_input, files = await GeminiClientWrapper.process_message(
-            request.messages[-1], tmp_dir, tagged=False
+        # Prepare the model input depending on how many turns are missing.
+        if len(remaining_messages) == 1:
+            model_input, files = await GeminiClientWrapper.process_message(
+                remaining_messages[0], tmp_dir, tagged=False
+            )
+        else:
+            model_input, files = await GeminiClientWrapper.process_conversation(
+                remaining_messages, tmp_dir
+            )
+        logger.debug(
+            f"Reused session {session.metadata} – sending {len(remaining_messages)} new messages."
         )
-        logger.debug(f"Found reusable session: {session.metadata}")
     else:
         # Start a new session and concat messages into a single string
         try:
@@ -97,13 +97,53 @@ async def create_chat_completion(
             raise
         logger.debug("New session started.")
 
+    # 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)
+
+    async def _send_with_split(session, text: str, files: list[Path | str] | None = None):
+        """Send text to Gemini, automatically splitting into multiple batches if it is
+        longer than ``MAX_CHARS_PER_REQUEST``.
+
+        Every intermediate batch (that is **not** the last one) is suffixed with a hint
+        telling Gemini that more content will come, and it should simply reply with
+        "ok". The final batch carries any file uploads and the real user prompt so
+        that Gemini can produce the actual answer.
+        """
+        if len(text) <= MAX_CHARS_PER_REQUEST:
+            # No need to split – a single request is fine.
+            return await session.send_message(text, files=files)
+
+        chunks: list[str] = []
+        pos = 0
+        total = len(text)
+        while pos < total:
+            end = min(pos + MAX_CHARS_PER_REQUEST, total)
+            chunk = text[pos:end]
+            pos = end
+
+            # If this is NOT the last chunk, add the continuation hint.
+            if end < total:
+                chunk += "\n(More messages to come, please reply with just 'ok'.)"
+            chunks.append(chunk)
+
+        # Fire off all but the last chunk, discarding the interim "ok" replies.
+        for chk in chunks[:-1]:
+            try:
+                await session.send_message(chk)
+            except Exception as e:
+                logger.exception(f"Error sending chunk to Gemini: {e}")
+                raise
+
+        # The last chunk carries the files (if any) and we return its response.
+        return await session.send_message(chunks[-1], files=files)
+
     # Generate response
     try:
         assert session and client, "Session and client not available"
         logger.debug(
             f"Client ID: {client.id}, Input length: {len(model_input)}, files count: {len(files)}"
         )
-        response = await session.send_message(model_input, files=files)
+        response = await _send_with_split(session, model_input, files=files)
     except Exception as e:
         logger.exception(f"Error generating content from Gemini API: {e}")
         raise
@@ -132,35 +172,81 @@ async def create_chat_completion(
     completion_id = f"chatcmpl-{uuid.uuid4()}"
     timestamp = int(datetime.now(tz=timezone.utc).timestamp())
     if request.stream:
-        return _create_streaming_response(model_output, completion_id, timestamp, request.model)
+        return _create_streaming_response(
+            model_output,
+            completion_id,
+            timestamp,
+            request.model,
+            request.messages,
+        )
     else:
         return _create_standard_response(
-            model_output, completion_id, timestamp, request.model, model_input
+            model_output, completion_id, timestamp, request.model, request.messages
         )
 
-
-def _check_reusable(messages: list[Message]) -> bool:
-    """
-    Check if the conversation is reusable based on the message history.
+# --- Helper to find reusable session with partial history ---
+def _find_reusable_session(
+    db: LMDBConversationStore,
+    pool: GeminiClientPool,
+    model: Model,
+    messages: list[Message],
+):
+    """Find an existing chat session that matches the *longest* prefix of
+    ``messages`` **whose last element is an assistant/system reply**.
+
+    Rationale
+    ---------
+    When a reply was generated by *another* server instance, the local LMDB may
+    only contain an older part of the conversation.  However, as long as we can
+    line-up **any** earlier assistant/system response, we can restore the
+    corresponding Gemini session and replay the *remaining* turns locally
+    (including that missing assistant reply and the subsequent user prompts).
+
+    The algorithm therefore walks backwards through the history **one message at
+    a time**, each time requiring the current tail to be assistant/system before
+    querying LMDB.  As soon as a match is found we recreate the session and
+    return the untouched suffix as ``remaining_messages``.
     """
-    if not messages or len(messages) < 2:
-        return False
 
-    # Last message must from the user
-    if messages[-1].role != "user" or not messages[-1].content:
-        return False
+    if len(messages) < 2:
+        return None, None, messages
+
+    # Start with the full history and iteratively trim from the end.
+    search_end = len(messages)
+    while search_end >= 2:
+        search_history = messages[:search_end]
+
+        # Only try to match if the last stored message would be assistant/system.
+        if search_history[-1].role in {"assistant", "system"}:
+            try:
+                if conv := db.find(model.model_name, search_history):
+                    client = pool.acquire(conv.client_id)
+                    session = client.start_chat(metadata=conv.metadata, model=model)
+                    remain = messages[search_end:]
+                    return session, client, remain
+            except Exception as e:
+                logger.warning(f"Error checking LMDB for reusable session: {e}")
+                break
 
-    # The second last message must be from the assistant or system
-    if messages[-2].role not in ["assistant", "system"]:
-        return False
+        # Trim one message and try again.
+        search_end -= 1
 
-    return True
+    return None, None, messages
 
 
 def _create_streaming_response(
-    model_output: str, completion_id: str, created_time: int, model: str
+    model_output: str,
+    completion_id: str,
+    created_time: int,
+    model: str,
+    messages: list[Message],
 ) -> StreamingResponse:
-    """Create streaming response"""
+    """Create streaming response with `usage` calculation included in the final chunk."""
+
+    # Calculate token usage
+    prompt_tokens = sum(estimate_tokens(msg.content) for msg in messages)
+    completion_tokens = estimate_tokens(model_output)
+    total_tokens = prompt_tokens + completion_tokens
 
     async def generate_stream():
         # Send start event
@@ -193,6 +279,11 @@ async def generate_stream():
             "created": created_time,
             "model": model,
             "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}],
+            "usage": {
+                "prompt_tokens": prompt_tokens,
+                "completion_tokens": completion_tokens,
+                "total_tokens": total_tokens,
+            },
         }
         yield f"data: {orjson.dumps(data).decode('utf-8')}\n\n"
         yield "data: [DONE]\n\n"
@@ -201,11 +292,15 @@ async def generate_stream():
 
 
 def _create_standard_response(
-    model_output: str, completion_id: str, created_time: int, model: str, model_input: str
+    model_output: str,
+    completion_id: str,
+    created_time: int,
+    model: str,
+    messages: list[Message],
 ) -> dict:
     """Create standard response"""
     # Calculate token usage
-    prompt_tokens = estimate_tokens(model_input)
+    prompt_tokens = sum(estimate_tokens(msg.content) for msg in messages)
     completion_tokens = estimate_tokens(model_output)
     total_tokens = prompt_tokens + completion_tokens
 

app/services/client.py

@@ -43,7 +43,12 @@ async def process_message(
             # TODO: Use Pydantic to enforce the value checking
             for item in message.content:
                 if item.type == "text":
-                    model_input = item.text or ""
+                    # Append multiple text fragments
+                    if item.text:
+                        if model_input:
+                            model_input += "\n" + item.text
+                        else:
+                            model_input = item.text
 
                 elif item.type == "image_url":
                     if not item.image_url:
@@ -66,6 +71,8 @@ async def process_message(
         if model_input and tagged:
             model_input = add_tag(message.role, model_input)
 
+        if "<" in model_input and ">" in model_input:
+            model_input += "\nFor any xml block, e.g. tool call, always wrap it by: \n`````xml\n...\n`````\n"
         return model_input, files
 
     @staticmethod
@@ -76,16 +83,25 @@ async def process_conversation(
         Process the entire conversation and return a formatted string and list of
         files. The last message is assumed to be the assistant's response.
         """
+        # Determine once whether we need to wrap messages with role tags: only required
+        # if the history already contains assistant/system messages. When every message
+        # so far is from the user, we can skip tagging entirely.
+        need_tag = any(m.role not in ("user", "system") for m in messages)
+
         conversation: list[str] = []
         files: list[Path | str] = []
 
         for msg in messages:
-            input_part, files_part = await GeminiClientWrapper.process_message(msg, tempdir)
+            input_part, files_part = await GeminiClientWrapper.process_message(
+                msg, tempdir, tagged=need_tag
+            )
             conversation.append(input_part)
             files.extend(files_part)
 
-        # Left with the last message as the assistant's response
-        conversation.append(add_tag("assistant", "", unclose=True))
+        # Append an opening assistant tag only when we used tags above so that Gemini
+        # knows where to start its reply.
+        if need_tag:
+            conversation.append(add_tag("assistant", "", unclose=True))
 
         return "\n".join(conversation), files
 

app/utils/config.py

@@ -44,6 +44,11 @@ class GeminiConfig(BaseModel):
         default=540, ge=1, description="Interval in seconds to refresh Gemini cookies"
     )
     verbose: bool = Field(False, description="Enable verbose logging for Gemini API requests")
+    max_chars_per_request: int = Field(
+        default=1_000_000,
+        ge=1,
+        description="Maximum characters Gemini Web can accept per request",
+    )
 
 
 class CORSConfig(BaseModel):

config/config.yaml

@@ -21,6 +21,7 @@ gemini:
   auto_refresh: true       # Auto-refresh session cookies
   refresh_interval: 540    # Refresh interval in seconds
   verbose: false           # Enable verbose logging for Gemini requests
+  max_chars_per_request: 1000000     # Maximum characters Gemini Web accepts per request. Non-pro users might have a lower limit
 
 storage:
   path: "data/lmdb"        # Database storage path