diff --git a/application/api/answer/services/compression/token_counter.py b/application/api/answer/services/compression/token_counter.py index ac676cf0..d2cbfc34 100644 --- a/application/api/answer/services/compression/token_counter.py +++ b/application/api/answer/services/compression/token_counter.py @@ -12,6 +12,12 @@ logger = logging.getLogger(__name__) class TokenCounter: """Centralized token counting for conversations and messages.""" + # Per-image token estimate. Provider tokenizers vary widely + # (Gemini ~258, GPT-4o 85-1500, Claude ~1500) and the actual cost + # depends on resolution/detail we can't see here. Errs slightly high + # so the threshold check stays conservative. + _IMAGE_PART_TOKEN_ESTIMATE = 1500 + @staticmethod def count_message_tokens(messages: List[Dict]) -> int: """ @@ -29,12 +35,36 @@ class TokenCounter: if isinstance(content, str): total_tokens += num_tokens_from_string(content) elif isinstance(content, list): - # Handle structured content (tool calls, etc.) + # Handle structured content (tool calls, image parts, etc.) for item in content: if isinstance(item, dict): - total_tokens += num_tokens_from_string(str(item)) + total_tokens += TokenCounter._count_content_part(item) return total_tokens + @staticmethod + def _count_content_part(item: Dict) -> int: + # Image/file attachments are billed by the provider per image, + # not proportional to the inline bytes/base64 string. + # ``str(item)`` on a 1MB image inflates the count by ~10000x, + # which trips spurious compression and overflows downstream + # input limits. + item_type = item.get("type") + + if "files" in item: + files = item.get("files") + count = len(files) if isinstance(files, list) and files else 1 + return TokenCounter._IMAGE_PART_TOKEN_ESTIMATE * count + + if "image_url" in item or item_type in { + "image", + "image_url", + "input_image", + "file", + }: + return TokenCounter._IMAGE_PART_TOKEN_ESTIMATE + + return num_tokens_from_string(str(item)) + @staticmethod def count_query_tokens( queries: List[Dict[str, Any]], include_tool_calls: bool = True diff --git a/application/cache.py b/application/cache.py index 3fdb6b8d..740c13c4 100644 --- a/application/cache.py +++ b/application/cache.py @@ -1,3 +1,4 @@ +import hashlib import json import logging import time @@ -10,6 +11,14 @@ from application.utils import get_hash logger = logging.getLogger(__name__) + +def _cache_default(value): + # Image attachments arrive inline as bytes (see GoogleLLM.prepare_messages_with_attachments); + # hash so the cache key stays bounded in size and stable across identical content. + if isinstance(value, (bytes, bytearray, memoryview)): + return f"" + return repr(value) + _redis_instance = None _redis_creation_failed = False _instance_lock = Lock() @@ -36,7 +45,7 @@ def get_redis_instance(): def gen_cache_key(messages, model="docgpt", tools=None): if not all(isinstance(msg, dict) for msg in messages): raise ValueError("All messages must be dictionaries.") - messages_str = json.dumps(messages) + messages_str = json.dumps(messages, default=_cache_default) tools_str = json.dumps(str(tools)) if tools else "" combined = f"{model}_{messages_str}_{tools_str}" cache_key = get_hash(combined) diff --git a/application/llm/base.py b/application/llm/base.py index 0971a886..1229a229 100644 --- a/application/llm/base.py +++ b/application/llm/base.py @@ -166,7 +166,7 @@ class BaseLLM(ABC): if is_stream: return self._stream_with_fallback( - decorated_method, method_name, *args, **kwargs + decorated_method, method_name, decorators, *args, **kwargs ) try: @@ -187,14 +187,27 @@ class BaseLLM(ABC): f"{fallback.model_id}. Error: {str(e)}" ) - fallback_method = getattr( - fallback, method_name.replace("_raw_", "") - ) + # Apply decorators to fallback's raw method directly — calling + # fallback.gen() would re-enter the orchestrator and recurse via + # fallback.fallback_llm. + fallback_method = getattr(fallback, method_name) + for decorator in decorators: + fallback_method = decorator(fallback_method) fallback_kwargs = {**kwargs, "model": fallback.model_id} - return fallback_method(*args, **fallback_kwargs) + try: + return fallback_method(fallback, *args, **fallback_kwargs) + except Exception as e2: + if self._is_non_retriable_client_error(e2): + logger.error( + f"Fallback LLM failed with non-retriable client " + f"error; giving up: {str(e2)}" + ) + else: + logger.error(f"Fallback LLM also failed; giving up: {str(e2)}") + raise def _stream_with_fallback( - self, decorated_method, method_name, *args, **kwargs + self, decorated_method, method_name, decorators, *args, **kwargs ): """ Wrapper generator that catches mid-stream errors and falls back. @@ -223,11 +236,37 @@ class BaseLLM(ABC): f"Primary LLM failed mid-stream. Falling back to " f"{fallback.model_id}. Error: {str(e)}" ) - fallback_method = getattr( - fallback, method_name.replace("_raw_", "") + # Apply decorators to fallback's raw stream method directly — + # calling fallback.gen_stream() would re-enter the orchestrator + # and recurse via fallback.fallback_llm. Emit the stream-start + # event manually so dashboards still see the fallback's + # provider/model when the response actually comes from it. + fallback._emit_stream_start_log( + fallback.model_id, + kwargs.get("messages"), + kwargs.get("tools"), + bool( + kwargs.get("_usage_attachments") + or kwargs.get("attachments") + ), ) + fallback_method = getattr(fallback, method_name) + for decorator in decorators: + fallback_method = decorator(fallback_method) fallback_kwargs = {**kwargs, "model": fallback.model_id} - yield from fallback_method(*args, **fallback_kwargs) + try: + yield from fallback_method(fallback, *args, **fallback_kwargs) + except Exception as e2: + if self._is_non_retriable_client_error(e2): + logger.error( + f"Fallback LLM failed mid-stream with non-retriable " + f"client error; giving up: {str(e2)}" + ) + else: + logger.error( + f"Fallback LLM also failed mid-stream; giving up: {str(e2)}" + ) + raise def gen(self, model, messages, stream=False, tools=None, *args, **kwargs): decorators = [gen_token_usage, gen_cache] @@ -242,22 +281,29 @@ class BaseLLM(ABC): **kwargs, ) - def gen_stream(self, model, messages, stream=True, tools=None, *args, **kwargs): - # Attachments arrive as ``_usage_attachments`` from ``Agent._llm_gen``; - # the ``stream_token_usage`` decorator pops that key, but the log - # fires before the decorator runs so it's still in ``kwargs`` here. + def _emit_stream_start_log(self, model, messages, tools, has_attachments): + # Stamped with ``self.provider_name`` so dashboards can group calls + # by vendor; the fallback path emits its own copy on the fallback + # instance so the actual responding provider is recorded. logging.info( "llm_stream_start", extra={ "model": model, "provider": self.provider_name, "message_count": len(messages) if messages is not None else 0, - "has_attachments": bool( - kwargs.get("_usage_attachments") or kwargs.get("attachments") - ), + "has_attachments": bool(has_attachments), "has_tools": bool(tools), }, ) + + def gen_stream(self, model, messages, stream=True, tools=None, *args, **kwargs): + # Attachments arrive as ``_usage_attachments`` from ``Agent._llm_gen``; + # the ``stream_token_usage`` decorator pops that key, but the log + # fires before the decorator runs so it's still in ``kwargs`` here. + has_attachments = bool( + kwargs.get("_usage_attachments") or kwargs.get("attachments") + ) + self._emit_stream_start_log(model, messages, tools, has_attachments) decorators = [stream_cache, stream_token_usage] return self._execute_with_fallback( "_raw_gen_stream", diff --git a/application/llm/handlers/base.py b/application/llm/handlers/base.py index 9a0f1f40..907b1622 100644 --- a/application/llm/handlers/base.py +++ b/application/llm/handlers/base.py @@ -280,7 +280,26 @@ class LLMHandler(ABC): # Keep serialized function calls/responses so the compressor sees actions parts_text.append(str(item)) elif "files" in item: - parts_text.append(str(item)) + # Image attachments arrive with raw bytes / base64 + # inline (see GoogleLLM.prepare_messages_with_attachments). + # ``str(item)`` would dump the whole byte/base64 + # blob into the compression prompt and bust the + # compression LLM's input limit. + files = item.get("files") or [] + descriptors = [] + if isinstance(files, list): + for f in files: + if isinstance(f, dict): + descriptors.append( + f.get("mime_type") or "file" + ) + elif isinstance(f, str): + descriptors.append(f) + if not descriptors: + descriptors = ["file"] + parts_text.append( + f"[attachment: {', '.join(descriptors)}]" + ) return "\n".join(parts_text) return "" diff --git a/tests/llm/handlers/test_llm_handlers.py b/tests/llm/handlers/test_llm_handlers.py index 3e9470a7..38f7a65c 100644 --- a/tests/llm/handlers/test_llm_handlers.py +++ b/tests/llm/handlers/test_llm_handlers.py @@ -360,7 +360,38 @@ class TestExtractTextFromContent: handler = ConcreteHandler() content = [{"files": ["/tmp/a.txt"]}] result = handler._extract_text_from_content(content) - assert "files" in result + assert result == "[attachment: /tmp/a.txt]" + + def test_list_with_inline_image_bytes(self): + # Google attaches images as inline bytes; stringifying them into + # the compression prompt would bust the compression LLM's input + # limit. The placeholder must describe the attachment without + # embedding the bytes. + handler = ConcreteHandler() + content = [ + { + "files": [ + {"file_bytes": b"\x89PNG" + b"\x00" * 1000, "mime_type": "image/png"} + ] + } + ] + result = handler._extract_text_from_content(content) + assert result == "[attachment: image/png]" + assert "PNG" not in result + assert "\\x" not in result + + def test_list_with_multiple_files(self): + handler = ConcreteHandler() + content = [ + { + "files": [ + {"file_bytes": b"a", "mime_type": "image/png"}, + {"file_uri": "https://x", "mime_type": "image/jpeg"}, + ] + } + ] + result = handler._extract_text_from_content(content) + assert result == "[attachment: image/png, image/jpeg]" def test_list_with_none_text(self): handler = ConcreteHandler() diff --git a/tests/llm/test_base.py b/tests/llm/test_base.py index 56a3caa4..aa817966 100644 --- a/tests/llm/test_base.py +++ b/tests/llm/test_base.py @@ -49,6 +49,9 @@ class FailingLLM(BaseLLM): class FallbackLLM(BaseLLM): + # _execute_with_fallback applies decorators to the fallback's raw method + # directly and never calls .gen() / .gen_stream() on it, so + # tracking lives on the raw methods. def __init__(self, **kwargs): super().__init__(**kwargs) self.gen_called = False @@ -62,14 +65,6 @@ class FallbackLLM(BaseLLM): self.gen_stream_called = True yield "fallback_chunk" - def gen(self, *args, **kwargs): - self.gen_called = True - return "fallback_gen_result" - - def gen_stream(self, *args, **kwargs): - self.gen_stream_called = True - yield "fallback_stream_chunk" - # --------------------------------------------------------------------------- # gen / gen_stream decorator application @@ -230,7 +225,7 @@ class TestExecuteWithFallbackNonStreaming: llm._fallback_llm = fallback result = llm.gen(model="m", messages=[]) - assert result == "fallback_gen_result" + assert result == "fallback_result" assert fallback.gen_called @@ -257,7 +252,7 @@ class TestStreamWithFallback: llm._fallback_llm = fallback result = list(llm.gen_stream(model="m", messages=[])) - assert "fallback_stream_chunk" in result + assert "fallback_chunk" in result assert fallback.gen_stream_called @@ -344,7 +339,7 @@ class TestNonRetriableClientError: llm._fallback_llm = fallback result = llm.gen(model="m", messages=[]) - assert result == "fallback_gen_result" + assert result == "fallback_result" assert fallback.gen_called diff --git a/tests/llm/test_fallback.py b/tests/llm/test_fallback.py index 1606ccff..1dab8af0 100644 --- a/tests/llm/test_fallback.py +++ b/tests/llm/test_fallback.py @@ -31,12 +31,20 @@ class FakeLLM(BaseLLM): self.gen_stream_called = False self.last_model_received = None # tracks the model kwarg passed to gen/gen_stream + # Track at the raw-method level. _execute_with_fallback applies + # decorators to the fallback's raw method directly and + # never calls .gen() / .gen_stream() on it, so a public-method + # override would not register fallback hops. def _raw_gen(self, baseself, model, messages, stream, tools=None, **kwargs): + self.gen_called = True + self.last_model_received = model if self.fail_at is not None: raise RuntimeError("primary model unavailable") return self.responses[0] def _raw_gen_stream(self, baseself, model, messages, stream, tools=None, **kwargs): + self.gen_stream_called = True + self.last_model_received = model yielded = 0 for chunk in self.stream_chunks: if self.fail_at is not None and yielded >= self.fail_at: @@ -44,18 +52,6 @@ class FakeLLM(BaseLLM): yield chunk yielded += 1 - # Wrap gen/gen_stream so we can track whether the fallback instance was used - # and which model kwarg it received - def gen(self, *args, **kwargs): - self.gen_called = True - self.last_model_received = kwargs.get("model") - return super().gen(*args, **kwargs) - - def gen_stream(self, *args, **kwargs): - self.gen_stream_called = True - self.last_model_received = kwargs.get("model") - return super().gen_stream(*args, **kwargs) - # Helpers @@ -294,6 +290,140 @@ class TestStreamingFallback: with pytest.raises(RuntimeError, match="mid-stream failure"): list(primary.gen_stream(**CALL_ARGS)) + def test_fallback_emits_stream_start_with_fallback_provider( + self, patch_model_utils, caplog + ): + # The fallback raw-stream path bypasses ``gen_stream``, so it must + # emit its own ``llm_stream_start`` event tagged with the fallback + # vendor — otherwise dashboards record only the failed primary + # even when the response came from the backup. + import logging as _logging + + class FallbackProvider(FakeLLM): + provider_name = "fallback-vendor" + + backup = FallbackProvider( + stream_chunks=["b1"], model_id="backup-model-id" + ) + patch_model_utils( + get_provider=lambda m, **_kwargs: "openai", + get_api_key=lambda p: "k", + create_llm=lambda type, **kw: backup, + ) + + class PrimaryProvider(FakeLLM): + provider_name = "primary-vendor" + + primary = PrimaryProvider( + stream_chunks=["x"], + fail_at=0, + backup_models=["backup-model-id"], + ) + + with caplog.at_level(_logging.INFO, logger="root"): + list( + primary.gen_stream( + model="primary-model", + messages=[{"role": "user", "content": "hi"}], + ) + ) + + starts = [r for r in caplog.records if r.message == "llm_stream_start"] + assert len(starts) == 2 + assert starts[0].provider == "primary-vendor" + assert starts[0].model == "primary-model" + assert starts[1].provider == "fallback-vendor" + assert starts[1].model == "backup-model-id" + + +# Tests — fallback never re-enters the orchestrator (Option B regression) + + +@pytest.mark.integration +class TestFallbackNoRecursion: + """When the primary fails, _execute_with_fallback applies decorators to + the fallback's raw method directly. The fallback's own ``fallback_llm`` + property must never be accessed — otherwise a fallback failure would + re-enter the orchestrator and walk the global FALLBACK_LLM_* chain + unboundedly.""" + + def test_backup_fallback_llm_property_never_accessed_on_gen_failure( + self, monkeypatch, patch_model_utils + ): + backup = FakeLLM(fail_at=0) # backup also fails + + accessed_on = [] + original_property = BaseLLM.fallback_llm + + def tracked_fallback_llm(self_llm): + accessed_on.append(self_llm) + return original_property.fget(self_llm) + + monkeypatch.setattr( + BaseLLM, "fallback_llm", property(tracked_fallback_llm) + ) + + patch_model_utils( + get_provider=lambda m, **_kwargs: "openai", + get_api_key=lambda p: "k", + create_llm=lambda type, **kw: backup, + ) + + primary = FakeLLM(fail_at=0, backup_models=["backup-model"]) + with pytest.raises(RuntimeError, match="primary model unavailable"): + primary.gen(**CALL_ARGS) + + assert primary in accessed_on # primary lazy-loaded its fallback + assert backup not in accessed_on # backup's chain was never walked + + def test_backup_fallback_llm_property_never_accessed_on_stream_failure( + self, monkeypatch, patch_model_utils + ): + backup = FakeLLM(stream_chunks=["x"], fail_at=0) + + accessed_on = [] + original_property = BaseLLM.fallback_llm + + def tracked_fallback_llm(self_llm): + accessed_on.append(self_llm) + return original_property.fget(self_llm) + + monkeypatch.setattr( + BaseLLM, "fallback_llm", property(tracked_fallback_llm) + ) + + patch_model_utils( + get_provider=lambda m, **_kwargs: "openai", + get_api_key=lambda p: "k", + create_llm=lambda type, **kw: backup, + ) + + primary = FakeLLM( + stream_chunks=["y"], fail_at=0, backup_models=["backup-model"] + ) + with pytest.raises(RuntimeError, match="mid-stream failure"): + list(primary.gen_stream(**CALL_ARGS)) + + assert primary in accessed_on + assert backup not in accessed_on + + def test_fallback_failure_propagates_without_chain(self, patch_model_utils): + """When both primary and fallback fail, the fallback's exception + propagates cleanly — no third hop, no extra retries.""" + backup = FakeLLM(fail_at=0) + + patch_model_utils( + get_provider=lambda m, **_kwargs: "openai", + get_api_key=lambda p: "k", + create_llm=lambda type, **kw: backup, + ) + + primary = FakeLLM(fail_at=0, backup_models=["backup-model"]) + with pytest.raises(RuntimeError, match="primary model unavailable"): + primary.gen(**CALL_ARGS) + + assert backup.gen_called # confirms fallback raw method WAS invoked + # Tests — backup model priority over global fallback diff --git a/tests/test_cache.py b/tests/test_cache.py index 15597133..411db4fe 100644 --- a/tests/test_cache.py +++ b/tests/test_cache.py @@ -438,3 +438,94 @@ def test_stream_cache_key_generation_failure_yields(mock_make_redis): messages = ["not_a_dict"] result = list(mock_function(None, "model", messages, stream=True, tools=None)) assert result == ["fallback_chunk"] + + +# ===================================================================== +# gen_cache_key with inline bytes (Google attachments) +# ===================================================================== + + +@pytest.mark.unit +def test_gen_cache_key_handles_inline_bytes(): + """Image attachments arrive in messages as raw bytes (see + GoogleLLM.prepare_messages_with_attachments). gen_cache_key must not + crash on json.dumps of bytes.""" + msgs = [ + { + "role": "user", + "content": [{"file_bytes": b"\x00\x01\x02", "mime_type": "image/png"}], + } + ] + key = gen_cache_key(msgs, model="x") + assert isinstance(key, str) + assert len(key) == 32 + + +@pytest.mark.unit +def test_gen_cache_key_stable_for_same_bytes(): + """Two requests with identical image bytes must produce the same key + — otherwise we'd never get cache hits on image-bearing prompts.""" + a = [ + { + "role": "user", + "content": [{"file_bytes": b"abc", "mime_type": "image/png"}], + } + ] + b = [ + { + "role": "user", + "content": [{"file_bytes": b"abc", "mime_type": "image/png"}], + } + ] + assert gen_cache_key(a, "m") == gen_cache_key(b, "m") + + +@pytest.mark.unit +def test_gen_cache_key_differs_for_different_bytes(): + """Different image bytes must produce different keys — otherwise two + different images would collide in cache.""" + a = [ + { + "role": "user", + "content": [{"file_bytes": b"abc", "mime_type": "image/png"}], + } + ] + b = [ + { + "role": "user", + "content": [{"file_bytes": b"xyz", "mime_type": "image/png"}], + } + ] + assert gen_cache_key(a, "m") != gen_cache_key(b, "m") + + +@pytest.mark.unit +def test_gen_cache_key_handles_bytearray_and_memoryview(): + """The default helper covers all bytes-like types so refactors that + swap bytes for bytearray/memoryview don't silently re-introduce the + TypeError.""" + msgs_ba = [ + { + "role": "user", + "content": [ + {"file_bytes": bytearray(b"abc"), "mime_type": "image/png"} + ], + } + ] + msgs_mv = [ + { + "role": "user", + "content": [ + {"file_bytes": memoryview(b"abc"), "mime_type": "image/png"} + ], + } + ] + msgs_b = [ + { + "role": "user", + "content": [{"file_bytes": b"abc", "mime_type": "image/png"}], + } + ] + # All three should hash the same content to the same key. + assert gen_cache_key(msgs_ba, "m") == gen_cache_key(msgs_b, "m") + assert gen_cache_key(msgs_mv, "m") == gen_cache_key(msgs_b, "m") diff --git a/tests/test_compression_service.py b/tests/test_compression_service.py index 83f27f8d..80efc7b0 100644 --- a/tests/test_compression_service.py +++ b/tests/test_compression_service.py @@ -427,6 +427,109 @@ class TestCompressionService: assert token_count_with_tools > token_count_without_tools + def test_count_message_tokens_skips_inline_image_bytes(self): + # Google attaches images as raw bytes inline. Stringifying the + # part would tokenize the byte repr (~2M tokens for a 1MB image), + # trigger spurious compression, and overflow downstream input + # limits. The estimate must stay bounded. + big_bytes = b"\x00" * 1_000_000 + messages = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "describe this"}, + { + "files": [ + {"file_bytes": big_bytes, "mime_type": "image/png"} + ] + }, + ], + } + ] + tokens = TokenCounter.count_message_tokens(messages) + assert tokens < 5000 + + def test_count_message_tokens_per_image_estimate_scales(self): + msgs_one = [ + { + "role": "user", + "content": [ + { + "files": [ + {"file_bytes": b"x", "mime_type": "image/png"} + ] + } + ], + } + ] + msgs_two = [ + { + "role": "user", + "content": [ + { + "files": [ + {"file_bytes": b"x", "mime_type": "image/png"}, + {"file_uri": "https://x", "mime_type": "image/jpeg"}, + ] + } + ], + } + ] + assert TokenCounter.count_message_tokens( + msgs_two + ) > TokenCounter.count_message_tokens(msgs_one) + + def test_count_message_tokens_skips_openai_image_url(self): + # OpenAI puts images inline as base64 data URLs. ``str(item)`` on + # the dict would tokenize the entire base64 payload. + big_b64 = "A" * 1_000_000 + messages = [ + { + "role": "user", + "content": [ + { + "type": "image_url", + "image_url": { + "url": f"data:image/png;base64,{big_b64}" + }, + }, + ], + } + ] + tokens = TokenCounter.count_message_tokens(messages) + assert tokens < 5000 + + def test_count_message_tokens_skips_anthropic_image(self): + # Anthropic puts images inline as ``source.data`` base64. + big_b64 = "A" * 1_000_000 + messages = [ + { + "role": "user", + "content": [ + { + "type": "image", + "source": { + "type": "base64", + "media_type": "image/png", + "data": big_b64, + }, + }, + ], + } + ] + tokens = TokenCounter.count_message_tokens(messages) + assert tokens < 5000 + + def test_count_message_tokens_still_counts_text_parts(self): + # The image bypass must not regress regular text-part counting. + messages = [ + { + "role": "user", + "content": [{"type": "text", "text": "hello world"}], + } + ] + assert TokenCounter.count_message_tokens(messages) > 0 + def test_format_conversation_for_compression( self, prompt_builder, sample_conversation ):