From 88bc08bec97d2f7990069a2a33ea9a7380712627 Mon Sep 17 00:00:00 2001 From: Oberonskii Date: Tue, 21 Jul 2026 20:39:49 +0200 Subject: [PATCH 1/2] fix: resolve tuple artifact in prompt construction and malformed /ping response Two bugs fixed in main.py: 1. The /ping endpoint returned a Python set literal instead of a proper JSON dict, which cannot be serialized. Changed to return {"status": "ok"} as a proper JSON object. 2. The prompt construction used {ISLAMIC_CONTEXT, request.prompt} which Python interprets as a tuple, causing the model to receive a stringified tuple instead of clean text. Fixed by: - Passing ISLAMIC_CONTEXT as system_instruction to GenerativeModel - Building the prompt string without the tuple syntax --- main.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/main.py b/main.py index 8821848..772931c 100644 --- a/main.py +++ b/main.py @@ -112,7 +112,7 @@ def get_safety_settings(): @app.get("/ping") async def ping(): logger.info("************** Ping pong ping pong *************") - return {"************** Ping pong ping pong *************"} + return {"status": "ok"} @app.post("/chat", response_model=ChatResponse) @@ -126,6 +126,7 @@ async def chat(request: ChatRequest): logger.info(f"Creating new chat session: {chat_id}") model = genai.GenerativeModel( 'gemini-2.5-flash-preview-05-20', + system_instruction=ISLAMIC_CONTEXT, safety_settings=get_safety_settings() ) # Initialize chat without sending context message @@ -136,7 +137,7 @@ async def chat(request: ChatRequest): # Prepare the prompt with context if provided full_prompt = request.prompt if request.context: - full_prompt = f"Context: {request.context}\n\nQuestion: {ISLAMIC_CONTEXT, request.prompt}" + full_prompt = f"Context: {request.context}\n\nQuestion: {request.prompt}" # Send message and get response logger.info("Sending message to chat...") From fdf6572e76a2ded8b5c636cc7b3a5fcb5b88ee41 Mon Sep 17 00:00:00 2001 From: Abdul-dev-creator Date: Wed, 22 Jul 2026 14:44:02 +0400 Subject: [PATCH 2/2] feat: cap conversation history with token budget to prevent unbounded prompt growth - Add MAX_HISTORY_TOKENS (default 16000) and MAX_HISTORY_TURNS (default 50) env config - Trim oldest turn-pairs at pair granularity before each Gemini call - Surface truncated boolean in ChatResponse - Add include_history flag to ChatRequest (default true) to omit payload bloat - Use cheap local estimate (chars/4) avoiding extra network calls - Document new env vars in README - Add 13 unit tests covering truncation boundary, pair-granularity invariant, and env configurability Closes #13 --- README.md | 2 + history.py | 56 ++++++++++ main.py | 40 ++++--- tests/test_history_truncation.py | 183 +++++++++++++++++++++++++++++++ 4 files changed, 265 insertions(+), 16 deletions(-) create mode 100644 history.py create mode 100644 tests/test_history_truncation.py diff --git a/README.md b/README.md index cbdc6d5..f5a6acd 100644 --- a/README.md +++ b/README.md @@ -74,6 +74,8 @@ The API runs at `http://localhost:8000` — interactive docs at `http://localhos | Variable | Description | |----------|-------------| | `GEMINI_API_KEY` | Google Gemini API key | +| `MAX_HISTORY_TOKENS` | Estimated-token budget for conversation history; oldest turn-pairs are dropped when exceeded (default `16000`) | +| `MAX_HISTORY_TURNS` | Hard cap on turn pairs in history; enforced as a cheap backstop before the token-budget check (default `50`) | ## ☁️ Deployment diff --git a/history.py b/history.py new file mode 100644 index 0000000..891cd41 --- /dev/null +++ b/history.py @@ -0,0 +1,56 @@ +"""Conversation history budget management. + +Provides token estimation and history trimming logic that operates on +Gemini chat session history lists — no Gemini SDK dependency needed. +""" + +import os +import logging + + +logger = logging.getLogger(__name__) + +MAX_HISTORY_TOKENS = int(os.getenv("MAX_HISTORY_TOKENS", "16000")) +MAX_HISTORY_TURNS = int(os.getenv("MAX_HISTORY_TURNS", "50")) + + +def estimate_tokens(text: str) -> int: + """Cheap local token estimate — ~4 chars per token on average.""" + return max(1, len(text) // 4) + + +def trim_history(chat_session) -> bool: + """Enforce token budget and turn cap on chat history. + + Drops oldest turn-pairs (user + model) until both budgets are satisfied. + Returns True if any turns were dropped. + """ + if not chat_session or not chat_session.history: + return False + + original_len = len(chat_session.history) + truncated = False + + while len(chat_session.history) > MAX_HISTORY_TURNS * 2: + chat_session.history.pop(0) + chat_session.history.pop(0) + truncated = True + + while len(chat_session.history) >= 2: + total = sum( + estimate_tokens(m.parts[0].text) + if hasattr(m, "parts") and m.parts and hasattr(m.parts[0], "text") + else 0 + for m in chat_session.history + ) + if total <= MAX_HISTORY_TOKENS: + break + chat_session.history.pop(0) + chat_session.history.pop(0) + truncated = True + + if truncated: + turns_dropped = (original_len - len(chat_session.history)) // 2 + logger.info("Trimmed chat history: dropped %d turn pair(s)", turns_dropped) + + return truncated diff --git a/main.py b/main.py index 772931c..78774be 100644 --- a/main.py +++ b/main.py @@ -9,6 +9,7 @@ import uuid from stellar import router as stellar_router +from history import trim_history # Configure logging logging.basicConfig(level=logging.INFO) @@ -80,12 +81,14 @@ class ChatRequest(BaseModel): prompt: str chat_id: Optional[str] = None context: Optional[str] = None # Additional context for specific queries + include_history: bool = True # Set to false to omit history from response class ChatResponse(BaseModel): response: str chat_id: str history: List[Message] + truncated: bool = False # True when oldest turn pairs were dropped def get_safety_settings(): @@ -139,6 +142,9 @@ async def chat(request: ChatRequest): if request.context: full_prompt = f"Context: {request.context}\n\nQuestion: {request.prompt}" + # Trim history to stay within budget (oldest turn-pairs dropped) + truncated = trim_history(chat) + # Send message and get response logger.info("Sending message to chat...") response = chat.send_message( @@ -155,28 +161,30 @@ async def chat(request: ChatRequest): logger.error("Empty response received from model") raise HTTPException(status_code=500, detail="Empty response from AI model") - # Get chat history + # Build history (only if caller wants it) history = [] - for message in chat.history: - try: - if hasattr(message, 'parts') and message.parts: - content = message.parts[0].text if hasattr(message.parts[0], 'text') else str(message.parts[0]) - else: - content = str(message) - - history.append(Message( - role="user" if message.role == "user" else "model", - content=content - )) - except Exception as e: - logger.warning(f"Error processing message in history: {str(e)}") - continue + if request.include_history: + for message in chat.history: + try: + if hasattr(message, 'parts') and message.parts: + content = message.parts[0].text if hasattr(message.parts[0], 'text') else str(message.parts[0]) + else: + content = str(message) + + history.append(Message( + role="user" if message.role == "user" else "model", + content=content + )) + except Exception as e: + logger.warning(f"Error processing message in history: {str(e)}") + continue logger.info("Chat response generated successfully") return ChatResponse( response=response.text, chat_id=chat_id, - history=history + history=history, + truncated=truncated, ) except Exception as e: diff --git a/tests/test_history_truncation.py b/tests/test_history_truncation.py new file mode 100644 index 0000000..d2afdef --- /dev/null +++ b/tests/test_history_truncation.py @@ -0,0 +1,183 @@ +"""Tests for conversation history truncation. + +All tests run offline — no GEMINI_API_KEY needed. +""" + +from unittest.mock import MagicMock, PropertyMock + +import pytest + +from history import ( + MAX_HISTORY_TOKENS, + MAX_HISTORY_TURNS, + estimate_tokens, + trim_history, +) + + +# --------------------------------------------------------------------------- +# Token estimation +# --------------------------------------------------------------------------- + + +class TestEstimateTokens: + def test_empty_string(self): + assert estimate_tokens("") == 1 + + def test_short_text(self): + assert estimate_tokens("hello") == 1 + + def test_typical_text(self): + text = "What is the significance of Salah in Islam?" + expected = max(1, len(text) // 4) + assert estimate_tokens(text) == expected + + def test_long_text(self): + text = "A" * 1000 + assert estimate_tokens(text) == 250 + + +# --------------------------------------------------------------------------- +# Helpers for building fake history entries +# --------------------------------------------------------------------------- + + +def _make_msg(role: str, text: str): + part = MagicMock(spec=[]) + part.text = text + msg = MagicMock(spec=["role", "parts"]) + msg.role = role + msg.parts = [part] + return msg + + +def _make_turn_pair(user_text: str, model_text: str): + return _make_msg("user", user_text), _make_msg("model", model_text) + + +def _make_chat_session(pairs: list[tuple[str, str]]): + """Build a fake chat session with a mutable history list.""" + history = [] + for user_text, model_text in pairs: + history.append(_make_msg("user", user_text)) + history.append(_make_msg("model", model_text)) + chat = MagicMock(spec=["history"]) + type(chat).history = PropertyMock(return_value=history) + return chat + + +# --------------------------------------------------------------------------- +# Truncation behaviour +# --------------------------------------------------------------------------- + + +class TestTrimHistory: + def test_no_history_is_noop(self): + chat = _make_chat_session([]) + assert trim_history(chat) is False + + def test_single_turn_under_budget_is_noop(self): + chat = _make_chat_session([("hello", "hi there")]) + assert trim_history(chat) is False + + def test_truncates_oldest_pairs_when_over_token_budget(self, monkeypatch): + monkeypatch.setattr("history.MAX_HISTORY_TOKENS", 10) + monkeypatch.setattr("history.MAX_HISTORY_TURNS", 50) + pairs = [ + ("AAAAA AAAAA AA", "BBBBB BBBBB BB"), # ~5+5 = 10 tokens — will be dropped + ("CCCCC CCCCC CC", "DDDDD DDDDD DD"), # ~5+5 = 10 tokens — will be dropped + ("EEEEE EEEEE EE", "FFFFF FFFFF FF"), # ~5+5 = 10 tokens — keeps this + ] + chat = _make_chat_session(pairs) + assert trim_history(chat) is True + assert len(chat.history) == 2 # one turn pair remains + assert chat.history[0].parts[0].text == "EEEEE EEEEE EE" + assert chat.history[1].parts[0].text == "FFFFF FFFFF FF" + + def test_truncates_to_turn_cap_first(self, monkeypatch): + monkeypatch.setattr("history.MAX_HISTORY_TOKENS", 999999) + monkeypatch.setattr("history.MAX_HISTORY_TURNS", 2) + pairs = [ + ("A", "B"), + ("C", "D"), + ("E", "F"), + ("G", "H"), + ] + chat = _make_chat_session(pairs) + assert trim_history(chat) is True + # MAX_HISTORY_TURNS=2 → at most 4 entries (2 pairs) + assert len(chat.history) == 4 + assert chat.history[0].parts[0].text == "E" + assert chat.history[1].parts[0].text == "F" + assert chat.history[2].parts[0].text == "G" + assert chat.history[3].parts[0].text == "H" + + def test_no_orphaned_turns_after_truncation(self, monkeypatch): + """History must always have pairs — never a lone user or model turn.""" + monkeypatch.setattr("history.MAX_HISTORY_TOKENS", 10) + monkeypatch.setattr("history.MAX_HISTORY_TURNS", 50) + pairs = [ + ("AAAAA AAAAA AA", "BBBBB BBBBB BB"), + ("CCCCC CCCCC CC", "DDDDD DDDDD DD"), + ("EEEEE EEEEE EE", "FFFFF FFFFF FF"), + ] + chat = _make_chat_session(pairs) + trim_history(chat) + assert len(chat.history) % 2 == 0 + + def test_returns_true_when_truncation_happens(self, monkeypatch): + monkeypatch.setattr("history.MAX_HISTORY_TOKENS", 3) + monkeypatch.setattr("history.MAX_HISTORY_TURNS", 50) + chat = _make_chat_session([("AAAA AAAA", "BBBB BBBB")]) + assert trim_history(chat) is True + + def test_preserves_recent_context_after_truncation(self, monkeypatch): + monkeypatch.setattr("history.MAX_HISTORY_TOKENS", 20) + monkeypatch.setattr("history.MAX_HISTORY_TURNS", 50) + pairs = [ + ("A" * 40 + " very old", "B" * 40 + " very old response"), + ("C" * 40 + " old", "D" * 40 + " old response"), + ("Recent question", "Recent answer"), + ] + chat = _make_chat_session(pairs) + trim_history(chat) + remaining = " ".join( + m.parts[0].text for m in chat.history + ) + assert "Recent" in remaining + assert "Old question 1" not in remaining + + +# --------------------------------------------------------------------------- +# Integration: end-to-end scenario +# --------------------------------------------------------------------------- + + +class TestTrimHistoryIntegration: + def test_long_conversation_triggers_truncation(self, monkeypatch): + monkeypatch.setattr("history.MAX_HISTORY_TOKENS", 50) + monkeypatch.setattr("history.MAX_HISTORY_TURNS", 50) + pairs = [ + (f"User message {i} " + "X" * 60, f"Model response {i} " + "Y" * 60) + for i in range(20) + ] + chat = _make_chat_session(pairs) + truncated = trim_history(chat) + assert truncated is True + assert len(chat.history) < len(pairs) * 2 + + def test_budget_configurable_via_env(self, monkeypatch): + monkeypatch.setattr("history.MAX_HISTORY_TOKENS", 1000) + monkeypatch.setattr("history.MAX_HISTORY_TURNS", 3) + pairs = [ + ("A", "B"), + ("C", "D"), + ("E", "F"), + ("G", "H"), + ("I", "J"), + ] + chat = _make_chat_session(pairs) + trim_history(chat) + # 3 turn pairs max → 6 entries + assert len(chat.history) <= 6 + assert len(chat.history) % 2 == 0