Add wall-clock timeout to reflect operations (#643)
* Initial plan * feat: add wall-clock timeout to reflect operations (fixes vectorize-io/hindsight#642) Add a configurable wall-clock timeout (default: 300s / 5 minutes) for the entire reflect operation. This prevents reflect calls from hanging for up to 40 minutes when LLM calls are slow or iteration counts are high. Changes: - Add DEFAULT_REFLECT_WALL_TIMEOUT (300s) config constant - Add HINDSIGHT_API_REFLECT_WALL_TIMEOUT env variable support - Wrap run_reflect_agent() with asyncio.wait_for() in reflect_async() - Return HTTP 504 on timeout in the reflect HTTP endpoint - Add unit test for wall-clock timeout enforcement Co-authored-by: ThePlenkov <6381507+ThePlenkov@users.noreply.github.com> Agent-Logs-Url: https://github.com/ThePlenkov/hindsight/sessions/a123d68b-aca1-4040-8bba-8c4f0fab2e2c * fix: address PR review findings (OpenAPI 504, docs, type hints, main.py TypeError, overlapping exceptions, lazy logging) Co-authored-by: ThePlenkov <6381507+ThePlenkov@users.noreply.github.com> Agent-Logs-Url: https://github.com/ThePlenkov/hindsight/sessions/dd574a88-53a3-4f9e-bba7-5a40b0eddb99 --------- Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com> Co-authored-by: ThePlenkov <6381507+ThePlenkov@users.noreply.github.com>
This commit is contained in:
Petr Plenkov
GitHub
parent
365fa3ce50
commit
8ce06e3e7c
6 changed files with 80 additions and 194 deletions
|
|
@ -2527,6 +2527,7 @@ def _register_routes(app: FastAPI):
|
|||
"5. Returns plain text answer and the facts used",
|
||||
operation_id="reflect",
|
||||
tags=["Memory"],
|
||||
responses={504: {"description": "Reflect operation timed out"}},
|
||||
)
|
||||
async def api_reflect(
|
||||
bank_id: str, request: ReflectRequest, request_context: RequestContext = Depends(get_request_context)
|
||||
|
|
@ -2630,6 +2631,12 @@ def _register_routes(app: FastAPI):
|
|||
raise HTTPException(status_code=e.status_code, detail=e.reason)
|
||||
except (AuthenticationError, HTTPException):
|
||||
raise
|
||||
except TimeoutError as e:
|
||||
logger.error("Timeout in /v1/default/banks/%s/reflect: %s", bank_id, e)
|
||||
raise HTTPException(
|
||||
status_code=504,
|
||||
detail=str(e) or "Reflect operation timed out. Consider reducing the budget or simplifying the query.",
|
||||
)
|
||||
except Exception as e:
|
||||
import traceback
|
||||
|
||||
|
|
|
|||
|
|
@ -338,6 +338,7 @@ ENV_WORKER_CONSOLIDATION_MAX_SLOTS = "HINDSIGHT_API_WORKER_CONSOLIDATION_MAX_SLO
|
|||
# Reflect agent settings
|
||||
ENV_REFLECT_MAX_ITERATIONS = "HINDSIGHT_API_REFLECT_MAX_ITERATIONS"
|
||||
ENV_REFLECT_MAX_CONTEXT_TOKENS = "HINDSIGHT_API_REFLECT_MAX_CONTEXT_TOKENS"
|
||||
ENV_REFLECT_WALL_TIMEOUT = "HINDSIGHT_API_REFLECT_WALL_TIMEOUT"
|
||||
ENV_REFLECT_MISSION = "HINDSIGHT_API_REFLECT_MISSION"
|
||||
|
||||
# Disposition settings
|
||||
|
|
@ -499,6 +500,7 @@ DEFAULT_WORKER_CONSOLIDATION_MAX_SLOTS = 2 # Max concurrent consolidation tasks
|
|||
# Reflect agent settings
|
||||
DEFAULT_REFLECT_MAX_ITERATIONS = 10 # Max tool call iterations before forcing response
|
||||
DEFAULT_REFLECT_MAX_CONTEXT_TOKENS = 100_000 # Max accumulated context tokens before forcing final prompt
|
||||
DEFAULT_REFLECT_WALL_TIMEOUT = 300 # Wall-clock timeout in seconds for the entire reflect operation (5 minutes)
|
||||
|
||||
# Disposition defaults (None = not set, fall back to bank DB value or 3)
|
||||
DEFAULT_DISPOSITION_SKEPTICISM = None
|
||||
|
|
@ -801,6 +803,7 @@ class HindsightConfig:
|
|||
# Reflect agent settings
|
||||
reflect_max_iterations: int
|
||||
reflect_max_context_tokens: int
|
||||
reflect_wall_timeout: int
|
||||
|
||||
# OpenTelemetry tracing configuration
|
||||
otel_traces_enabled: bool
|
||||
|
|
@ -1274,6 +1277,7 @@ class HindsightConfig:
|
|||
reflect_max_context_tokens=int(
|
||||
os.getenv(ENV_REFLECT_MAX_CONTEXT_TOKENS, str(DEFAULT_REFLECT_MAX_CONTEXT_TOKENS))
|
||||
),
|
||||
reflect_wall_timeout=int(os.getenv(ENV_REFLECT_WALL_TIMEOUT, str(DEFAULT_REFLECT_WALL_TIMEOUT))),
|
||||
reflect_mission=os.getenv(ENV_REFLECT_MISSION) or None,
|
||||
# Disposition settings (None = fall back to DB value)
|
||||
disposition_skepticism=int(os.getenv(ENV_DISPOSITION_SKEPTICISM))
|
||||
|
|
|
|||
|
|
@ -5228,6 +5228,7 @@ class MemoryEngine(MemoryEngineInterface):
|
|||
effective_budget = budget or Budget.LOW
|
||||
max_iterations = max(1, int(base_max_iterations * budget_multipliers.get(effective_budget, 1.0)))
|
||||
max_context_tokens = config.reflect_max_context_tokens
|
||||
wall_timeout = config.reflect_wall_timeout
|
||||
|
||||
# Run agentic loop - acquire connections only when needed for DB operations
|
||||
# (not held during LLM calls which can be slow)
|
||||
|
|
@ -5332,7 +5333,9 @@ class MemoryEngine(MemoryEngineInterface):
|
|||
span_context = None
|
||||
|
||||
try:
|
||||
agent_result = await run_reflect_agent(
|
||||
try:
|
||||
agent_result = await asyncio.wait_for(
|
||||
run_reflect_agent(
|
||||
llm_config=self._reflect_llm_config.with_config(resolved_reflect_config),
|
||||
bank_id=bank_id,
|
||||
query=query,
|
||||
|
|
@ -5351,12 +5354,25 @@ class MemoryEngine(MemoryEngineInterface):
|
|||
include_recall=include_recall,
|
||||
budget=effective_budget,
|
||||
max_context_tokens=max_context_tokens,
|
||||
),
|
||||
timeout=wall_timeout,
|
||||
)
|
||||
except asyncio.TimeoutError:
|
||||
total_time = time.time() - reflect_start
|
||||
logger.error(
|
||||
"[REFLECT %s] Wall-clock timeout after %.1fs (limit: %ss) for query: %.50s...",
|
||||
reflect_id, total_time, wall_timeout, query,
|
||||
)
|
||||
raise TimeoutError(
|
||||
f"Reflect operation timed out after {wall_timeout} seconds. "
|
||||
f"Consider reducing the budget or simplifying the query."
|
||||
)
|
||||
|
||||
total_time = time.time() - reflect_start
|
||||
logger.info(
|
||||
f"[REFLECT {reflect_id}] Complete: {len(agent_result.text)} chars, "
|
||||
f"{agent_result.iterations} iterations, {agent_result.tools_called} tool calls | {total_time:.3f}s"
|
||||
"[REFLECT %s] Complete: %d chars, %d iterations, %d tool calls | %.3fs",
|
||||
reflect_id, len(agent_result.text), agent_result.iterations,
|
||||
agent_result.tools_called, total_time,
|
||||
)
|
||||
|
||||
# Convert agent tool trace to ToolCallTrace objects
|
||||
|
|
|
|||
|
|
@ -13,6 +13,7 @@ Stop with Ctrl+C.
|
|||
import argparse
|
||||
import asyncio
|
||||
import atexit
|
||||
import dataclasses
|
||||
import os
|
||||
import signal
|
||||
import sys
|
||||
|
|
@ -152,178 +153,7 @@ def main():
|
|||
# Configure Python logging based on log level
|
||||
# Update config with CLI override if provided
|
||||
if args.log_level != config.log_level:
|
||||
config = HindsightConfig(
|
||||
database_url=config.database_url,
|
||||
database_schema=config.database_schema,
|
||||
vector_extension=config.vector_extension,
|
||||
text_search_extension=config.text_search_extension,
|
||||
llm_provider=config.llm_provider,
|
||||
llm_api_key=config.llm_api_key,
|
||||
llm_model=config.llm_model,
|
||||
llm_base_url=config.llm_base_url,
|
||||
llm_max_concurrent=config.llm_max_concurrent,
|
||||
llm_max_retries=config.llm_max_retries,
|
||||
llm_initial_backoff=config.llm_initial_backoff,
|
||||
llm_max_backoff=config.llm_max_backoff,
|
||||
llm_timeout=config.llm_timeout,
|
||||
llm_groq_service_tier=config.llm_groq_service_tier,
|
||||
llm_openai_service_tier=config.llm_openai_service_tier,
|
||||
llm_vertexai_project_id=config.llm_vertexai_project_id,
|
||||
llm_vertexai_region=config.llm_vertexai_region,
|
||||
llm_vertexai_service_account_key=config.llm_vertexai_service_account_key,
|
||||
llm_gemini_safety_settings=config.llm_gemini_safety_settings,
|
||||
retain_llm_provider=config.retain_llm_provider,
|
||||
retain_llm_api_key=config.retain_llm_api_key,
|
||||
retain_llm_model=config.retain_llm_model,
|
||||
retain_llm_base_url=config.retain_llm_base_url,
|
||||
retain_llm_max_concurrent=config.retain_llm_max_concurrent,
|
||||
retain_llm_max_retries=config.retain_llm_max_retries,
|
||||
retain_llm_initial_backoff=config.retain_llm_initial_backoff,
|
||||
retain_llm_max_backoff=config.retain_llm_max_backoff,
|
||||
retain_llm_timeout=config.retain_llm_timeout,
|
||||
reflect_llm_provider=config.reflect_llm_provider,
|
||||
reflect_llm_api_key=config.reflect_llm_api_key,
|
||||
reflect_llm_model=config.reflect_llm_model,
|
||||
reflect_llm_base_url=config.reflect_llm_base_url,
|
||||
reflect_llm_max_concurrent=config.reflect_llm_max_concurrent,
|
||||
reflect_llm_max_retries=config.reflect_llm_max_retries,
|
||||
reflect_llm_initial_backoff=config.reflect_llm_initial_backoff,
|
||||
reflect_llm_max_backoff=config.reflect_llm_max_backoff,
|
||||
reflect_llm_timeout=config.reflect_llm_timeout,
|
||||
consolidation_llm_provider=config.consolidation_llm_provider,
|
||||
consolidation_llm_api_key=config.consolidation_llm_api_key,
|
||||
consolidation_llm_model=config.consolidation_llm_model,
|
||||
consolidation_llm_base_url=config.consolidation_llm_base_url,
|
||||
consolidation_llm_max_concurrent=config.consolidation_llm_max_concurrent,
|
||||
consolidation_llm_max_retries=config.consolidation_llm_max_retries,
|
||||
consolidation_llm_initial_backoff=config.consolidation_llm_initial_backoff,
|
||||
consolidation_llm_max_backoff=config.consolidation_llm_max_backoff,
|
||||
consolidation_llm_timeout=config.consolidation_llm_timeout,
|
||||
embeddings_provider=config.embeddings_provider,
|
||||
embeddings_local_model=config.embeddings_local_model,
|
||||
embeddings_local_force_cpu=config.embeddings_local_force_cpu,
|
||||
embeddings_local_trust_remote_code=config.embeddings_local_trust_remote_code,
|
||||
embeddings_tei_url=config.embeddings_tei_url,
|
||||
embeddings_openai_base_url=config.embeddings_openai_base_url,
|
||||
embeddings_cohere_api_key=config.embeddings_cohere_api_key,
|
||||
embeddings_cohere_model=config.embeddings_cohere_model,
|
||||
embeddings_cohere_base_url=config.embeddings_cohere_base_url,
|
||||
embeddings_litellm_api_base=config.embeddings_litellm_api_base,
|
||||
embeddings_litellm_api_key=config.embeddings_litellm_api_key,
|
||||
embeddings_litellm_model=config.embeddings_litellm_model,
|
||||
embeddings_litellm_sdk_api_key=config.embeddings_litellm_sdk_api_key,
|
||||
embeddings_litellm_sdk_model=config.embeddings_litellm_sdk_model,
|
||||
embeddings_litellm_sdk_api_base=config.embeddings_litellm_sdk_api_base,
|
||||
reranker_provider=config.reranker_provider,
|
||||
reranker_local_model=config.reranker_local_model,
|
||||
reranker_local_force_cpu=config.reranker_local_force_cpu,
|
||||
reranker_local_max_concurrent=config.reranker_local_max_concurrent,
|
||||
reranker_local_trust_remote_code=config.reranker_local_trust_remote_code,
|
||||
reranker_local_fp16=config.reranker_local_fp16,
|
||||
reranker_local_bucket_batching=config.reranker_local_bucket_batching,
|
||||
reranker_local_batch_size=config.reranker_local_batch_size,
|
||||
reranker_tei_url=config.reranker_tei_url,
|
||||
reranker_tei_batch_size=config.reranker_tei_batch_size,
|
||||
reranker_tei_max_concurrent=config.reranker_tei_max_concurrent,
|
||||
reranker_max_candidates=config.reranker_max_candidates,
|
||||
reranker_cohere_api_key=config.reranker_cohere_api_key,
|
||||
reranker_cohere_model=config.reranker_cohere_model,
|
||||
reranker_cohere_base_url=config.reranker_cohere_base_url,
|
||||
reranker_litellm_api_base=config.reranker_litellm_api_base,
|
||||
reranker_litellm_api_key=config.reranker_litellm_api_key,
|
||||
reranker_litellm_model=config.reranker_litellm_model,
|
||||
reranker_litellm_max_tokens_per_doc=config.reranker_litellm_max_tokens_per_doc,
|
||||
reranker_litellm_sdk_api_key=config.reranker_litellm_sdk_api_key,
|
||||
reranker_litellm_sdk_model=config.reranker_litellm_sdk_model,
|
||||
reranker_litellm_sdk_api_base=config.reranker_litellm_sdk_api_base,
|
||||
reranker_zeroentropy_api_key=config.reranker_zeroentropy_api_key,
|
||||
reranker_zeroentropy_model=config.reranker_zeroentropy_model,
|
||||
host=args.host,
|
||||
port=args.port,
|
||||
base_path=config.base_path,
|
||||
log_level=args.log_level,
|
||||
log_format=config.log_format,
|
||||
mcp_enabled=config.mcp_enabled,
|
||||
mcp_enabled_tools=config.mcp_enabled_tools,
|
||||
enable_bank_config_api=config.enable_bank_config_api,
|
||||
graph_retriever=config.graph_retriever,
|
||||
mpfp_top_k_neighbors=config.mpfp_top_k_neighbors,
|
||||
recall_max_concurrent=config.recall_max_concurrent,
|
||||
recall_connection_budget=config.recall_connection_budget,
|
||||
recall_max_query_tokens=config.recall_max_query_tokens,
|
||||
retain_max_completion_tokens=config.retain_max_completion_tokens,
|
||||
retain_chunk_size=config.retain_chunk_size,
|
||||
retain_extract_causal_links=config.retain_extract_causal_links,
|
||||
retain_extraction_mode=config.retain_extraction_mode,
|
||||
retain_mission=config.retain_mission,
|
||||
retain_custom_instructions=config.retain_custom_instructions,
|
||||
retain_default_strategy=config.retain_default_strategy,
|
||||
retain_strategies=config.retain_strategies,
|
||||
retain_batch_tokens=config.retain_batch_tokens,
|
||||
retain_entity_lookup=config.retain_entity_lookup,
|
||||
retain_batch_enabled=config.retain_batch_enabled,
|
||||
retain_batch_poll_interval_seconds=config.retain_batch_poll_interval_seconds,
|
||||
file_storage_type=config.file_storage_type,
|
||||
file_storage_s3_bucket=config.file_storage_s3_bucket,
|
||||
file_storage_s3_region=config.file_storage_s3_region,
|
||||
file_storage_s3_endpoint=config.file_storage_s3_endpoint,
|
||||
file_storage_s3_access_key_id=config.file_storage_s3_access_key_id,
|
||||
file_storage_s3_secret_access_key=config.file_storage_s3_secret_access_key,
|
||||
file_storage_gcs_bucket=config.file_storage_gcs_bucket,
|
||||
file_storage_gcs_service_account_key=config.file_storage_gcs_service_account_key,
|
||||
file_storage_azure_container=config.file_storage_azure_container,
|
||||
file_storage_azure_account_name=config.file_storage_azure_account_name,
|
||||
file_storage_azure_account_key=config.file_storage_azure_account_key,
|
||||
file_parser=config.file_parser,
|
||||
file_parser_allowlist=config.file_parser_allowlist,
|
||||
file_parser_iris_token=config.file_parser_iris_token,
|
||||
file_parser_iris_org_id=config.file_parser_iris_org_id,
|
||||
file_conversion_max_batch_size_mb=config.file_conversion_max_batch_size_mb,
|
||||
file_conversion_max_batch_size=config.file_conversion_max_batch_size,
|
||||
enable_file_upload_api=config.enable_file_upload_api,
|
||||
file_delete_after_retain=config.file_delete_after_retain,
|
||||
enable_observations=config.enable_observations,
|
||||
enable_observation_history=config.enable_observation_history,
|
||||
enable_mental_model_history=config.enable_mental_model_history,
|
||||
consolidation_batch_size=config.consolidation_batch_size,
|
||||
consolidation_llm_batch_size=config.consolidation_llm_batch_size,
|
||||
consolidation_max_tokens=config.consolidation_max_tokens,
|
||||
consolidation_source_facts_max_tokens=config.consolidation_source_facts_max_tokens,
|
||||
consolidation_source_facts_max_tokens_per_observation=config.consolidation_source_facts_max_tokens_per_observation,
|
||||
observations_mission=config.observations_mission,
|
||||
entity_labels=config.entity_labels,
|
||||
entities_allow_free_form=config.entities_allow_free_form,
|
||||
skip_llm_verification=config.skip_llm_verification,
|
||||
lazy_reranker=config.lazy_reranker,
|
||||
run_migrations_on_startup=config.run_migrations_on_startup,
|
||||
db_pool_min_size=config.db_pool_min_size,
|
||||
db_pool_max_size=config.db_pool_max_size,
|
||||
db_command_timeout=config.db_command_timeout,
|
||||
db_acquire_timeout=config.db_acquire_timeout,
|
||||
worker_enabled=config.worker_enabled,
|
||||
worker_id=config.worker_id,
|
||||
worker_poll_interval_ms=config.worker_poll_interval_ms,
|
||||
worker_max_retries=config.worker_max_retries,
|
||||
worker_http_port=config.worker_http_port,
|
||||
worker_max_slots=config.worker_max_slots,
|
||||
worker_consolidation_max_slots=config.worker_consolidation_max_slots,
|
||||
reflect_max_iterations=config.reflect_max_iterations,
|
||||
reflect_max_context_tokens=config.reflect_max_context_tokens,
|
||||
reflect_mission=config.reflect_mission,
|
||||
disposition_skepticism=config.disposition_skepticism,
|
||||
disposition_literalism=config.disposition_literalism,
|
||||
disposition_empathy=config.disposition_empathy,
|
||||
mental_model_refresh_concurrency=config.mental_model_refresh_concurrency,
|
||||
otel_traces_enabled=config.otel_traces_enabled,
|
||||
otel_exporter_otlp_endpoint=config.otel_exporter_otlp_endpoint,
|
||||
otel_exporter_otlp_headers=config.otel_exporter_otlp_headers,
|
||||
otel_service_name=config.otel_service_name,
|
||||
otel_deployment_environment=config.otel_deployment_environment,
|
||||
webhook_url=config.webhook_url,
|
||||
webhook_secret=config.webhook_secret,
|
||||
webhook_event_types=config.webhook_event_types,
|
||||
webhook_delivery_poll_interval_seconds=config.webhook_delivery_poll_interval_seconds,
|
||||
)
|
||||
config = dataclasses.replace(config, host=args.host, port=args.port, log_level=args.log_level)
|
||||
config.configure_logging()
|
||||
if not args.daemon:
|
||||
config.log_config()
|
||||
|
|
|
|||
|
|
@ -5,8 +5,10 @@ These tests verify:
|
|||
1. Tool name normalization for various LLM output formats
|
||||
2. Recovery from unknown tool calls
|
||||
3. Recovery from tool execution errors
|
||||
4. Wall-clock timeout enforcement
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
import pytest
|
||||
|
|
@ -416,6 +418,32 @@ class TestReflectAgentMocked:
|
|||
assert result is not None
|
||||
assert result.iterations == 3
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_wall_clock_timeout(self, mock_llm: MagicMock, mock_functions: dict[str, AsyncMock]) -> None:
|
||||
"""Test that asyncio.wait_for can enforce a wall-clock timeout on run_reflect_agent."""
|
||||
|
||||
async def slow_llm_call(*args: object, **kwargs: object) -> LLMToolCallResult:
|
||||
await asyncio.sleep(10) # Simulate a slow LLM call
|
||||
return LLMToolCallResult(
|
||||
tool_calls=[LLMToolCall(id="1", name="recall", arguments={"query": "test"})],
|
||||
finish_reason="tool_calls",
|
||||
)
|
||||
|
||||
mock_llm.call_with_tools.side_effect = slow_llm_call
|
||||
|
||||
with pytest.raises(asyncio.TimeoutError):
|
||||
await asyncio.wait_for(
|
||||
run_reflect_agent(
|
||||
llm_config=mock_llm,
|
||||
bank_id="test-bank",
|
||||
query="test query",
|
||||
bank_profile={"name": "Test", "mission": "Testing"},
|
||||
max_iterations=5,
|
||||
**mock_functions,
|
||||
),
|
||||
timeout=0.1, # Very short timeout to trigger quickly
|
||||
)
|
||||
|
||||
|
||||
class TestContextOverflowHelpers:
|
||||
"""Unit tests for context-overflow detection helpers."""
|
||||
|
|
|
|||
|
|
@ -892,6 +892,7 @@ export HINDSIGHT_API_OBSERVATIONS_MISSION="Observations are recurring patterns i
|
|||
|----------|-------------|---------|
|
||||
| `HINDSIGHT_API_REFLECT_MAX_ITERATIONS` | Max tool call iterations before forcing a response | `10` |
|
||||
| `HINDSIGHT_API_REFLECT_MAX_CONTEXT_TOKENS` | Max accumulated context tokens in the reflect loop before forcing final synthesis. Prevents `context_length_exceeded` errors on large banks. Lower this if your LLM has a context window smaller than 128K. | `100000` |
|
||||
| `HINDSIGHT_API_REFLECT_WALL_TIMEOUT` | Wall-clock timeout in seconds for the entire reflect operation. If exceeded, the request returns HTTP 504. | `300` |
|
||||
| `HINDSIGHT_API_REFLECT_MISSION` | Global reflect mission (identity and reasoning framing). Overridden per bank via config API. | - |
|
||||
|
||||
#### Disposition
|
||||
|
|
|
|||
Loading…
Reference in a new issue