add simple compacting

app.py

@@ -114,6 +114,70 @@ LOG_REQUEST_BODY = os.getenv("LOG_REQUEST_BODY", "1") == "1"
 LOG_REQUEST_BODY_MAX_CHARS = int(os.getenv("LOG_REQUEST_BODY_MAX_CHARS", "2000"))
 CONTEXT_SIZE = int(os.getenv("CONTEXT_SIZE", "2048"))
 
+PROMPT_MARGIN_TOKENS = int(os.getenv("PROMPT_MARGIN_TOKENS", "256"))
+CHARS_PER_TOKEN_EST = float(os.getenv("CHARS_PER_TOKEN_EST", "4.0"))
+
+
+def _estimate_tokens(text: str) -> int:
+    """Rough token estimate based on character count."""
+    if not text:
+        return 0
+    return int(len(text) / CHARS_PER_TOKEN_EST) + 1
+
+
+def _truncate_text_to_tokens(text: str, max_tokens: int) -> str:
+    """Truncate text to an approximate token budget."""
+    if not text or max_tokens <= 0:
+        return ""
+    max_chars = int(max_tokens * CHARS_PER_TOKEN_EST)
+    if len(text) <= max_chars:
+        return text
+    return text[:max_chars] + "...[truncated]"
+
+
+def _compact_messages(messages: list[dict], max_tokens: int) -> list[dict]:
+    """
+    Compact messages to fit within the prompt budget.
+    Strategy:
+    - Cap system message content size.
+    - Drop oldest non-system messages until within budget.
+    - As a last resort, truncate the oldest remaining non-system message.
+    """
+    if not messages:
+        return messages
+
+    prompt_budget = CONTEXT_SIZE - max_tokens - PROMPT_MARGIN_TOKENS
+    if prompt_budget <= 0:
+        return messages
+
+    # Work on a copy to avoid mutating caller input
+    compacted = [dict(m) for m in messages]
+
+    # Cap system messages
+    system_cap = min(1024, max(256, prompt_budget // 3))
+    for msg in compacted:
+        if msg.get("role") == "system" and "content" in msg:
+            msg["content"] = _truncate_text_to_tokens(str(msg["content"]), system_cap)
+
+    def total_tokens(msgs: list[dict]) -> int:
+        return sum(_estimate_tokens(str(m.get("content", ""))) for m in msgs)
+
+    # Drop oldest non-system messages until under budget
+    while total_tokens(compacted) > prompt_budget:
+        idx = next((i for i, m in enumerate(compacted) if m.get("role") != "system"), None)
+        if idx is None:
+            break
+        compacted.pop(idx)
+
+    # Last resort: truncate oldest non-system content
+    if total_tokens(compacted) > prompt_budget:
+        idx = next((i for i, m in enumerate(compacted) if m.get("role") != "system"), None)
+        if idx is not None:
+            remaining_budget = max(1, prompt_budget - (total_tokens(compacted) - _estimate_tokens(str(compacted[idx].get("content", "")))))
+            compacted[idx]["content"] = _truncate_text_to_tokens(str(compacted[idx].get("content", "")), remaining_budget)
+
+    return compacted
+
 
 @dataclass
 class CachedModel:
@@ -901,8 +965,9 @@ async def chat_completions(request: ChatCompletionRequest, raw_request: Request)
             raise HTTPException(status_code=500, detail="Current model not loaded")
 
         # Forward to llama-server using aiohttp
+        compacted_messages = _compact_messages(request.messages, request.max_tokens)
         payload = {
-            "messages": request.messages,
+            "messages": compacted_messages,
             "max_tokens": request.max_tokens,
             "temperature": request.temperature,
         }
@@ -1043,6 +1108,9 @@ Always cite sources when using information from the search results."""
         if not http_session or http_session.closed:
             raise HTTPException(status_code=500, detail="HTTP session not initialized")
 
+        # Compact messages to fit within context
+        augmented_messages = _compact_messages(augmented_messages, request.max_tokens)
+
         # Get current model from cache
         cached_model = model_cache.get(current_model)
         if not cached_model: