8114ef440e
* Fix async method parity and server keepalive timeout
The Python client's async methods were missing parameters available in
their sync counterparts, and the server's default keepalive timeout was
shorter than the client's, causing ServerDisconnectedError on reused
connections.
Server:
- Set uvicorn timeout_keep_alive to 30s (default was 5s). The Python
client (aiohttp) has a 15s client-side keepalive, so the server must
hold connections longer to prevent the client from writing to a
closed socket.
Python client - async method parity:
- arecall(): add trace, query_timestamp, include_entities,
include_chunks, max_entity_tokens, max_chunk_tokens. Return
RecallResponse instead of list[RecallResult].
- areflect(): add max_tokens and response_schema.
- acreate_bank(): new async method.
- aset_mission(): new async method.
- adelete_bank(): new async method.
Tests:
- Add test verifying uvicorn keepalive timeout exceeds client default.
- Add async tests for arecall (include_chunks, include_entities, trace,
full params), areflect (max_tokens, structured output), and
adelete_bank.
* Fix flaky tag tests by using entity-rich content and asserting on tags
The tag tests were unreliable because:
- Generic content ("Project X meeting notes") was frequently collapsed
during fact extraction, leaving no memories to recall
- Assertions checked LLM-rewritten text for literal substrings instead
of checking tags, which is what the tests are actually verifying
Fix: use distinctive, entity-rich content (named people with specific
actions) that reliably survives fact extraction, and assert on tag
membership rather than text content.
410 lines
17 KiB
Python
410 lines
17 KiB
Python
"""
|
|
Command-line interface for Hindsight API.
|
|
|
|
Run the server with:
|
|
hindsight-api
|
|
|
|
Run as background daemon:
|
|
hindsight-api --daemon
|
|
|
|
Stop with Ctrl+C.
|
|
"""
|
|
|
|
import argparse
|
|
import asyncio
|
|
import atexit
|
|
import os
|
|
import signal
|
|
import sys
|
|
import warnings
|
|
|
|
import uvicorn
|
|
|
|
from . import MemoryEngine, __version__
|
|
from .api import create_app
|
|
from .banner import print_banner
|
|
from .config import DEFAULT_WORKERS, ENV_WORKERS, HindsightConfig, _get_raw_config
|
|
from .daemon import (
|
|
DEFAULT_DAEMON_PORT,
|
|
DEFAULT_IDLE_TIMEOUT,
|
|
IdleTimeoutMiddleware,
|
|
daemonize,
|
|
)
|
|
from .extensions import DefaultExtensionContext, OperationValidatorExtension, TenantExtension, load_extension
|
|
|
|
# Filter deprecation warnings from third-party libraries
|
|
warnings.filterwarnings("ignore", message="websockets.legacy is deprecated")
|
|
warnings.filterwarnings("ignore", message="websockets.server.WebSocketServerProtocol is deprecated")
|
|
|
|
# Disable tokenizers parallelism to avoid warnings
|
|
os.environ["TOKENIZERS_PARALLELISM"] = "false"
|
|
|
|
# Global reference for cleanup
|
|
_memory: MemoryEngine | None = None
|
|
|
|
|
|
def _cleanup():
|
|
"""Synchronous cleanup function to stop resources on exit."""
|
|
global _memory
|
|
if _memory is not None and _memory._pg0 is not None:
|
|
try:
|
|
loop = asyncio.new_event_loop()
|
|
loop.run_until_complete(_memory._pg0.stop())
|
|
loop.close()
|
|
print("\npg0 stopped.")
|
|
except Exception as e:
|
|
print(f"\nError stopping pg0: {e}")
|
|
|
|
|
|
def _signal_handler(signum, frame):
|
|
"""Handle SIGINT/SIGTERM to ensure cleanup."""
|
|
print(f"\nReceived signal {signum}, shutting down...")
|
|
_cleanup()
|
|
sys.exit(0)
|
|
|
|
|
|
def main():
|
|
"""Main entry point for the CLI."""
|
|
global _memory
|
|
|
|
# Load configuration from environment (for CLI args defaults)
|
|
config = _get_raw_config()
|
|
|
|
parser = argparse.ArgumentParser(
|
|
prog="hindsight-api",
|
|
description="Hindsight API Server",
|
|
)
|
|
|
|
# Server options
|
|
parser.add_argument(
|
|
"--host", default=config.host, help=f"Host to bind to (default: {config.host}, env: HINDSIGHT_API_HOST)"
|
|
)
|
|
parser.add_argument(
|
|
"--port",
|
|
type=int,
|
|
default=config.port,
|
|
help=f"Port to bind to (default: {config.port}, env: HINDSIGHT_API_PORT)",
|
|
)
|
|
parser.add_argument(
|
|
"--log-level",
|
|
default=config.log_level,
|
|
choices=["critical", "error", "warning", "info", "debug", "trace"],
|
|
help=f"Log level (default: {config.log_level}, env: HINDSIGHT_API_LOG_LEVEL)",
|
|
)
|
|
|
|
# Development options
|
|
parser.add_argument("--reload", action="store_true", help="Enable auto-reload on code changes (development only)")
|
|
parser.add_argument(
|
|
"--workers",
|
|
type=int,
|
|
default=int(os.getenv(ENV_WORKERS, str(DEFAULT_WORKERS))),
|
|
help=f"Number of worker processes (env: {ENV_WORKERS}, default: {DEFAULT_WORKERS})",
|
|
)
|
|
|
|
# Access log options
|
|
parser.add_argument("--access-log", action="store_true", help="Enable access log")
|
|
parser.add_argument("--no-access-log", dest="access_log", action="store_false", help="Disable access log (default)")
|
|
parser.set_defaults(access_log=False)
|
|
|
|
# Proxy options
|
|
parser.add_argument(
|
|
"--proxy-headers", action="store_true", help="Enable X-Forwarded-Proto, X-Forwarded-For headers"
|
|
)
|
|
parser.add_argument(
|
|
"--forwarded-allow-ips", default=None, help="Comma separated list of IPs to trust with proxy headers"
|
|
)
|
|
|
|
# SSL options
|
|
parser.add_argument("--ssl-keyfile", default=None, help="SSL key file")
|
|
parser.add_argument("--ssl-certfile", default=None, help="SSL certificate file")
|
|
|
|
# Daemon mode options
|
|
parser.add_argument(
|
|
"--daemon",
|
|
action="store_true",
|
|
help=f"Run as background daemon (uses port {DEFAULT_DAEMON_PORT}, auto-exits after idle)",
|
|
)
|
|
parser.add_argument(
|
|
"--idle-timeout",
|
|
type=int,
|
|
default=DEFAULT_IDLE_TIMEOUT,
|
|
help=f"Idle timeout in seconds before auto-exit in daemon mode (default: {DEFAULT_IDLE_TIMEOUT})",
|
|
)
|
|
|
|
args = parser.parse_args()
|
|
|
|
# Daemon mode handling
|
|
if args.daemon:
|
|
# Use port from args (may be custom for profiles)
|
|
if args.port == config.port: # No custom port specified
|
|
args.port = DEFAULT_DAEMON_PORT
|
|
args.host = "127.0.0.1" # Only bind to localhost for security
|
|
|
|
# Fork into background
|
|
# No lockfile needed - port binding prevents duplicate daemons
|
|
daemonize()
|
|
|
|
# Print banner (not in daemon mode)
|
|
if not args.daemon:
|
|
print()
|
|
print_banner()
|
|
|
|
# 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,
|
|
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_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_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,
|
|
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,
|
|
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,
|
|
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_custom_instructions=config.retain_custom_instructions,
|
|
retain_batch_tokens=config.retain_batch_tokens,
|
|
retain_batch_enabled=config.retain_batch_enabled,
|
|
retain_batch_poll_interval_seconds=config.retain_batch_poll_interval_seconds,
|
|
enable_observations=config.enable_observations,
|
|
consolidation_batch_size=config.consolidation_batch_size,
|
|
consolidation_max_tokens=config.consolidation_max_tokens,
|
|
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,
|
|
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,
|
|
)
|
|
config.configure_logging()
|
|
if not args.daemon:
|
|
config.log_config()
|
|
|
|
# Register cleanup handlers
|
|
atexit.register(_cleanup)
|
|
signal.signal(signal.SIGINT, _signal_handler)
|
|
signal.signal(signal.SIGTERM, _signal_handler)
|
|
|
|
# Load operation validator extension if configured
|
|
operation_validator = load_extension("OPERATION_VALIDATOR", OperationValidatorExtension)
|
|
if operation_validator:
|
|
import logging
|
|
|
|
logging.info(f"Loaded operation validator: {operation_validator.__class__.__name__}")
|
|
|
|
# Load tenant extension if configured
|
|
tenant_extension = load_extension("TENANT", TenantExtension)
|
|
if tenant_extension:
|
|
import logging
|
|
|
|
logging.info(f"Loaded tenant extension: {tenant_extension.__class__.__name__}")
|
|
|
|
# Create MemoryEngine (reads configuration from environment)
|
|
_memory = MemoryEngine(
|
|
operation_validator=operation_validator,
|
|
tenant_extension=tenant_extension,
|
|
run_migrations=config.run_migrations_on_startup,
|
|
)
|
|
|
|
# Set extension context on tenant extension (needed for schema provisioning)
|
|
if tenant_extension:
|
|
extension_context = DefaultExtensionContext(
|
|
database_url=config.database_url,
|
|
memory_engine=_memory,
|
|
)
|
|
tenant_extension.set_context(extension_context)
|
|
logging.info("Extension context set on tenant extension")
|
|
|
|
# Create FastAPI app
|
|
app = create_app(
|
|
memory=_memory,
|
|
http_api_enabled=True,
|
|
mcp_api_enabled=config.mcp_enabled,
|
|
mcp_mount_path="/mcp",
|
|
initialize_memory=True,
|
|
)
|
|
|
|
# Wrap with idle timeout middleware in daemon mode
|
|
idle_middleware = None
|
|
if args.daemon:
|
|
idle_middleware = IdleTimeoutMiddleware(app, idle_timeout=args.idle_timeout)
|
|
app = idle_middleware
|
|
|
|
# Prepare uvicorn config
|
|
# When using workers or reload, we must use import string so each worker can import the app
|
|
use_import_string = args.workers > 1 or args.reload
|
|
# Check for uvloop availability
|
|
try:
|
|
import uvloop # noqa: F401
|
|
|
|
loop_impl = "uvloop"
|
|
print("uvloop available, will use for event loop")
|
|
except ImportError:
|
|
loop_impl = "asyncio"
|
|
print("uvloop not installed, using default asyncio event loop")
|
|
|
|
uvicorn_config = {
|
|
"app": "hindsight_api.server:app" if use_import_string else app,
|
|
"host": args.host,
|
|
"port": args.port,
|
|
"log_level": args.log_level,
|
|
"access_log": args.access_log,
|
|
"proxy_headers": args.proxy_headers,
|
|
"ws": "wsproto", # Use wsproto instead of websockets to avoid deprecation warnings
|
|
"loop": loop_impl, # Explicitly set event loop implementation
|
|
"timeout_keep_alive": 30, # Exceed aiohttp's 15s client timeout so the client always closes first
|
|
}
|
|
|
|
# Add optional parameters if provided
|
|
if args.reload:
|
|
uvicorn_config["reload"] = True
|
|
if args.workers > 1:
|
|
uvicorn_config["workers"] = args.workers
|
|
if args.forwarded_allow_ips:
|
|
uvicorn_config["forwarded_allow_ips"] = args.forwarded_allow_ips
|
|
if args.ssl_keyfile:
|
|
uvicorn_config["ssl_keyfile"] = args.ssl_keyfile
|
|
if args.ssl_certfile:
|
|
uvicorn_config["ssl_certfile"] = args.ssl_certfile
|
|
|
|
# Print startup info (not in daemon mode)
|
|
if not args.daemon:
|
|
from .banner import print_startup_info
|
|
|
|
print_startup_info(
|
|
host=args.host,
|
|
port=args.port,
|
|
database_url=config.database_url,
|
|
llm_provider=config.llm_provider,
|
|
llm_model=config.llm_model,
|
|
embeddings_provider=config.embeddings_provider,
|
|
reranker_provider=config.reranker_provider,
|
|
mcp_enabled=config.mcp_enabled,
|
|
version=__version__,
|
|
vector_extension=config.vector_extension,
|
|
text_search_extension=config.text_search_extension,
|
|
)
|
|
|
|
# Start idle checker in daemon mode
|
|
if idle_middleware is not None:
|
|
# Start the idle checker in a background thread with its own event loop
|
|
import logging
|
|
import threading
|
|
|
|
def run_idle_checker():
|
|
import time
|
|
|
|
time.sleep(2) # Wait for uvicorn to start
|
|
try:
|
|
loop = asyncio.new_event_loop()
|
|
asyncio.set_event_loop(loop)
|
|
loop.run_until_complete(idle_middleware._check_idle())
|
|
except Exception as e:
|
|
logging.error(f"Idle checker error: {e}", exc_info=True)
|
|
|
|
threading.Thread(target=run_idle_checker, daemon=True).start()
|
|
|
|
uvicorn.run(**uvicorn_config)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|