diff --git a/hindsight-api/hindsight_api/api/http.py b/hindsight-api/hindsight_api/api/http.py index bf469221..31f2936f 100644 --- a/hindsight-api/hindsight_api/api/http.py +++ b/hindsight-api/hindsight_api/api/http.py @@ -3580,6 +3580,21 @@ def _register_routes(app: FastAPI): } ) else: + # Check if batch API is enabled - if so, require async mode + from hindsight_api.config import get_config + + config = get_config() + if config.retain_batch_enabled: + raise HTTPException( + status_code=400, + detail=( + "Batch API is enabled (HINDSIGHT_API_RETAIN_BATCH_ENABLED=true) but async=false. " + "Batch operations can take several minutes to hours and will timeout in synchronous mode. " + "Please set async=true in your request to use background processing, or disable batch API " + "by setting HINDSIGHT_API_RETAIN_BATCH_ENABLED=false in your environment." + ), + ) + # Synchronous processing: wait for completion (record metrics) with metrics.record_operation("retain", bank_id=bank_id, source="api"): result, usage = await app.state.memory.retain_batch_async( diff --git a/hindsight-api/hindsight_api/config.py b/hindsight-api/hindsight_api/config.py index a966232d..7857e691 100644 --- a/hindsight-api/hindsight_api/config.py +++ b/hindsight-api/hindsight_api/config.py @@ -129,6 +129,11 @@ ENV_LLM_INITIAL_BACKOFF = "HINDSIGHT_API_LLM_INITIAL_BACKOFF" ENV_LLM_MAX_BACKOFF = "HINDSIGHT_API_LLM_MAX_BACKOFF" ENV_LLM_TIMEOUT = "HINDSIGHT_API_LLM_TIMEOUT" ENV_LLM_GROQ_SERVICE_TIER = "HINDSIGHT_API_LLM_GROQ_SERVICE_TIER" +ENV_LLM_OPENAI_SERVICE_TIER = "HINDSIGHT_API_LLM_OPENAI_SERVICE_TIER" + +# Defaults for service tiers +DEFAULT_LLM_GROQ_SERVICE_TIER = "auto" # "on_demand", "flex", or "auto" +DEFAULT_LLM_OPENAI_SERVICE_TIER = None # None (default) or "flex" (50% cheaper) # Per-operation LLM configuration (optional, falls back to global LLM config) ENV_RETAIN_LLM_PROVIDER = "HINDSIGHT_API_RETAIN_LLM_PROVIDER" @@ -251,6 +256,8 @@ ENV_RETAIN_EXTRACT_CAUSAL_LINKS = "HINDSIGHT_API_RETAIN_EXTRACT_CAUSAL_LINKS" ENV_RETAIN_EXTRACTION_MODE = "HINDSIGHT_API_RETAIN_EXTRACTION_MODE" ENV_RETAIN_CUSTOM_INSTRUCTIONS = "HINDSIGHT_API_RETAIN_CUSTOM_INSTRUCTIONS" ENV_RETAIN_BATCH_TOKENS = "HINDSIGHT_API_RETAIN_BATCH_TOKENS" +ENV_RETAIN_BATCH_ENABLED = "HINDSIGHT_API_RETAIN_BATCH_ENABLED" +ENV_RETAIN_BATCH_POLL_INTERVAL_SECONDS = "HINDSIGHT_API_RETAIN_BATCH_POLL_INTERVAL_SECONDS" # Observations settings (consolidated knowledge from facts) ENV_ENABLE_OBSERVATIONS = "HINDSIGHT_API_ENABLE_OBSERVATIONS" @@ -373,6 +380,8 @@ DEFAULT_RETAIN_EXTRACTION_MODE = "concise" # Extraction mode: "concise", "verbo RETAIN_EXTRACTION_MODES = ("concise", "verbose", "custom") # Allowed extraction modes DEFAULT_RETAIN_CUSTOM_INSTRUCTIONS = None # Custom extraction guidelines (only used when mode="custom") DEFAULT_RETAIN_BATCH_TOKENS = 10_000 # ~40KB of text # Max chars per sub-batch for async retain auto-splitting +DEFAULT_RETAIN_BATCH_ENABLED = False # Use LLM Batch API for fact extraction (only when async=True) +DEFAULT_RETAIN_BATCH_POLL_INTERVAL_SECONDS = 60 # Batch API polling interval in seconds # Observations defaults (consolidated knowledge from facts) DEFAULT_ENABLE_OBSERVATIONS = True # Observations enabled by default @@ -496,6 +505,8 @@ class HindsightConfig: llm_initial_backoff: float llm_max_backoff: float llm_timeout: float + llm_groq_service_tier: str # Groq: "on_demand", "flex", or "auto" + llm_openai_service_tier: str | None # OpenAI: None (default) or "flex" (50% cheaper) # Vertex AI configuration llm_vertexai_project_id: str | None @@ -593,6 +604,8 @@ class HindsightConfig: retain_extraction_mode: str retain_custom_instructions: str | None retain_batch_tokens: int + retain_batch_enabled: bool + retain_batch_poll_interval_seconds: int # Observations settings (consolidated knowledge from facts) enable_observations: bool @@ -770,6 +783,8 @@ class HindsightConfig: llm_initial_backoff=float(os.getenv(ENV_LLM_INITIAL_BACKOFF, str(DEFAULT_LLM_INITIAL_BACKOFF))), llm_max_backoff=float(os.getenv(ENV_LLM_MAX_BACKOFF, str(DEFAULT_LLM_MAX_BACKOFF))), llm_timeout=float(os.getenv(ENV_LLM_TIMEOUT, str(DEFAULT_LLM_TIMEOUT))), + llm_groq_service_tier=os.getenv(ENV_LLM_GROQ_SERVICE_TIER, DEFAULT_LLM_GROQ_SERVICE_TIER), + llm_openai_service_tier=os.getenv(ENV_LLM_OPENAI_SERVICE_TIER, DEFAULT_LLM_OPENAI_SERVICE_TIER), # Vertex AI llm_vertexai_project_id=os.getenv(ENV_LLM_VERTEXAI_PROJECT_ID) or DEFAULT_LLM_VERTEXAI_PROJECT_ID, llm_vertexai_region=os.getenv(ENV_LLM_VERTEXAI_REGION, DEFAULT_LLM_VERTEXAI_REGION), @@ -943,6 +958,11 @@ class HindsightConfig: ), retain_custom_instructions=os.getenv(ENV_RETAIN_CUSTOM_INSTRUCTIONS) or DEFAULT_RETAIN_CUSTOM_INSTRUCTIONS, retain_batch_tokens=int(os.getenv(ENV_RETAIN_BATCH_TOKENS, str(DEFAULT_RETAIN_BATCH_TOKENS))), + retain_batch_enabled=os.getenv(ENV_RETAIN_BATCH_ENABLED, str(DEFAULT_RETAIN_BATCH_ENABLED)).lower() + == "true", + retain_batch_poll_interval_seconds=int( + os.getenv(ENV_RETAIN_BATCH_POLL_INTERVAL_SECONDS, str(DEFAULT_RETAIN_BATCH_POLL_INTERVAL_SECONDS)) + ), # Observations settings (consolidated knowledge from facts) enable_observations=os.getenv(ENV_ENABLE_OBSERVATIONS, str(DEFAULT_ENABLE_OBSERVATIONS)).lower() == "true", consolidation_batch_size=int( diff --git a/hindsight-api/hindsight_api/engine/llm_interface.py b/hindsight-api/hindsight_api/engine/llm_interface.py index ee1af600..210805e6 100644 --- a/hindsight-api/hindsight_api/engine/llm_interface.py +++ b/hindsight-api/hindsight_api/engine/llm_interface.py @@ -128,6 +128,67 @@ class LLMInterface(ABC): """ pass + async def supports_batch_api(self) -> bool: + """ + Check if this provider supports batch API operations. + + Returns: + True if provider supports submit_batch/get_batch_status/retrieve_batch_results + """ + return False + + async def submit_batch( + self, + requests: list[dict[str, Any]], + endpoint: str = "/v1/chat/completions", + completion_window: str = "24h", + ) -> dict[str, Any]: + """ + Submit a batch of requests to the provider's batch API. + + Args: + requests: List of request dicts in JSONL format (custom_id, method, url, body) + endpoint: API endpoint for the batch (e.g., "/v1/chat/completions") + completion_window: Completion window (e.g., "24h") + + Returns: + Dict with batch metadata: {"batch_id": str, "status": str, ...} + + Raises: + NotImplementedError: If provider doesn't support batch API + """ + raise NotImplementedError(f"Batch API not supported for provider: {self.provider}") + + async def get_batch_status(self, batch_id: str) -> dict[str, Any]: + """ + Get the status of a batch job. + + Args: + batch_id: Batch identifier returned from submit_batch + + Returns: + Dict with status info: {"batch_id": str, "status": str, "completed_at": str, ...} + + Raises: + NotImplementedError: If provider doesn't support batch API + """ + raise NotImplementedError(f"Batch API not supported for provider: {self.provider}") + + async def retrieve_batch_results(self, batch_id: str) -> list[dict[str, Any]]: + """ + Retrieve completed batch results. + + Args: + batch_id: Batch identifier returned from submit_batch + + Returns: + List of result dicts (one per request, matched by custom_id) + + Raises: + NotImplementedError: If provider doesn't support batch API + """ + raise NotImplementedError(f"Batch API not supported for provider: {self.provider}") + @abstractmethod async def cleanup(self) -> None: """Clean up resources (close connections, etc.).""" diff --git a/hindsight-api/hindsight_api/engine/llm_wrapper.py b/hindsight-api/hindsight_api/engine/llm_wrapper.py index 85d22940..efa81a8e 100644 --- a/hindsight-api/hindsight_api/engine/llm_wrapper.py +++ b/hindsight-api/hindsight_api/engine/llm_wrapper.py @@ -67,6 +67,7 @@ def create_llm_provider( model: str, reasoning_effort: str, groq_service_tier: str | None = None, + openai_service_tier: str | None = None, vertexai_project_id: str | None = None, vertexai_region: str | None = None, vertexai_credentials: Any = None, @@ -80,7 +81,8 @@ def create_llm_provider( base_url: Base URL for the API. model: Model name. reasoning_effort: Reasoning effort level for supported providers. - groq_service_tier: Groq service tier (for Groq provider). + groq_service_tier: Groq service tier (for Groq provider) - "on_demand", "flex", or "auto". + openai_service_tier: OpenAI service tier (for OpenAI provider) - None (default) or "flex" (50% cheaper). vertexai_project_id: Vertex AI project ID (for VertexAI provider). vertexai_region: Vertex AI region (for VertexAI provider). vertexai_credentials: Vertex AI credentials object (for VertexAI provider). @@ -156,6 +158,7 @@ def create_llm_provider( model=model, reasoning_effort=reasoning_effort, groq_service_tier=groq_service_tier, + openai_service_tier=openai_service_tier, ) else: @@ -177,6 +180,7 @@ class LLMProvider: model: str, reasoning_effort: str = "low", groq_service_tier: str | None = None, + openai_service_tier: str | None = None, ): """ Initialize LLM provider. @@ -187,15 +191,17 @@ class LLMProvider: base_url: Base URL for the API. model: Model name. reasoning_effort: Reasoning effort level for supported providers. - groq_service_tier: Groq service tier ("on_demand", "flex", "auto"). Default: None (uses Groq's default). + groq_service_tier: Groq service tier ("on_demand", "flex", "auto") - from config. + openai_service_tier: OpenAI service tier (None or "flex") - from config. """ self.provider = provider.lower() self.api_key = api_key self.base_url = base_url self.model = model self.reasoning_effort = reasoning_effort - # Default to 'auto' for best performance, users can override to 'on_demand' for free tier - self.groq_service_tier = groq_service_tier or os.getenv(ENV_LLM_GROQ_SERVICE_TIER, "auto") + # Service tiers from hierarchical config (not env vars) + self.groq_service_tier = groq_service_tier + self.openai_service_tier = openai_service_tier # Validate provider valid_providers = [ @@ -272,6 +278,7 @@ class LLMProvider: model=self.model, reasoning_effort=self.reasoning_effort, groq_service_tier=self.groq_service_tier, + openai_service_tier=self.openai_service_tier, vertexai_project_id=vertexai_project_id, vertexai_region=vertexai_region, vertexai_credentials=vertexai_credentials, diff --git a/hindsight-api/hindsight_api/engine/memory_engine.py b/hindsight-api/hindsight_api/engine/memory_engine.py index cc564c3d..3114e137 100644 --- a/hindsight-api/hindsight_api/engine/memory_engine.py +++ b/hindsight-api/hindsight_api/engine/memory_engine.py @@ -548,7 +548,7 @@ class MemoryEngine(MemoryEngineInterface): Handler for batch retain tasks. Args: - task_dict: Dict with 'bank_id', 'contents' + task_dict: Dict with 'bank_id', 'contents', 'operation_id' Raises: ValueError: If bank_id is missing @@ -559,9 +559,10 @@ class MemoryEngine(MemoryEngineInterface): raise ValueError("bank_id is required for batch retain task") contents = task_dict.get("contents", []) document_tags = task_dict.get("document_tags") + operation_id = task_dict.get("operation_id") # For batch API crash recovery logger.info( - f"[BATCH_RETAIN_TASK] Starting background batch retain for bank_id={bank_id}, {len(contents)} items" + f"[BATCH_RETAIN_TASK] Starting background batch retain for bank_id={bank_id}, {len(contents)} items, operation_id={operation_id}" ) # Restore tenant_id/api_key_id from task payload so extensions @@ -581,6 +582,7 @@ class MemoryEngine(MemoryEngineInterface): contents=contents, document_tags=document_tags, request_context=context, + operation_id=operation_id, ) logger.info(f"[BATCH_RETAIN_TASK] Completed background batch retain for bank_id={bank_id}") @@ -1496,6 +1498,7 @@ class MemoryEngine(MemoryEngineInterface): confidence_score: float | None = None, document_tags: list[str] | None = None, return_usage: bool = False, + operation_id: str | None = None, ): """ Store multiple content items as memory units in ONE batch operation. @@ -1643,6 +1646,7 @@ class MemoryEngine(MemoryEngineInterface): fact_type_override=fact_type_override, confidence_score=confidence_score, document_tags=document_tags, + operation_id=operation_id, ) all_results.extend(sub_results) total_usage = total_usage + sub_usage @@ -1663,6 +1667,7 @@ class MemoryEngine(MemoryEngineInterface): fact_type_override=fact_type_override, confidence_score=confidence_score, document_tags=document_tags, + operation_id=operation_id, ) # Call post-operation hook if validator is configured @@ -1712,6 +1717,7 @@ class MemoryEngine(MemoryEngineInterface): fact_type_override: str | None = None, confidence_score: float | None = None, document_tags: list[str] | None = None, + operation_id: str | None = None, ) -> tuple[list[list[str]], "TokenUsage"]: """ Internal method for batch processing without chunking logic. @@ -1761,6 +1767,8 @@ class MemoryEngine(MemoryEngineInterface): confidence_score=confidence_score, document_tags=document_tags, config=resolved_config, + operation_id=operation_id, + schema=request_context.tenant_id if request_context else None, ) def recall( diff --git a/hindsight-api/hindsight_api/engine/providers/openai_compatible_llm.py b/hindsight-api/hindsight_api/engine/providers/openai_compatible_llm.py index d5ed1785..6bc08ed2 100644 --- a/hindsight-api/hindsight_api/engine/providers/openai_compatible_llm.py +++ b/hindsight-api/hindsight_api/engine/providers/openai_compatible_llm.py @@ -16,6 +16,7 @@ Features: """ import asyncio +import io import json import logging import os @@ -96,8 +97,9 @@ class OpenAICompatibleLLM(LLMInterface): if self.provider in ("openai", "groq") and not self.api_key: raise ValueError(f"API key is required for {self.provider}") - # Groq service tier configuration - self.groq_service_tier = groq_service_tier or os.getenv("HINDSIGHT_API_LLM_GROQ_SERVICE_TIER", "auto") + # Service tier configuration (from config, not env vars) + self.groq_service_tier = groq_service_tier + self.openai_service_tier = kwargs.get("openai_service_tier") # Get timeout config self.timeout = timeout or float(os.getenv(ENV_LLM_TIMEOUT, str(DEFAULT_LLM_TIMEOUT))) @@ -782,6 +784,140 @@ class OpenAICompatibleLLM(LLMInterface): raise last_exception raise RuntimeError("Ollama call failed after all retries") + async def supports_batch_api(self) -> bool: + """Check if this provider supports batch API operations.""" + # Only OpenAI and Groq support batch API + return self.provider in ("openai", "groq") + + async def submit_batch( + self, + requests: list[dict[str, Any]], + endpoint: str = "/v1/chat/completions", + completion_window: str = "24h", + ) -> dict[str, Any]: + """ + Submit a batch of requests to OpenAI/Groq Batch API. + + Args: + requests: List of request dicts with custom_id, method, url, body + endpoint: API endpoint (e.g., "/v1/chat/completions") + completion_window: Completion window (e.g., "24h") + + Returns: + Dict with batch metadata including batch_id + + Raises: + NotImplementedError: If provider doesn't support batch API + """ + if not await self.supports_batch_api(): + raise NotImplementedError(f"Batch API not supported for provider: {self.provider}") + + logger.info(f"Submitting batch with {len(requests)} requests to {self.provider}") + + # Format requests as JSONL + jsonl_content = "\n".join(json.dumps(req) for req in requests) + + # Upload file to provider (wrap in BytesIO with filename) + file_bytes = io.BytesIO(jsonl_content.encode("utf-8")) + file_bytes.name = "batch_input.jsonl" # OpenAI SDK needs a filename + + file_response = await self._client.files.create( + file=file_bytes, + purpose="batch", + ) + + logger.debug(f"Uploaded batch file: {file_response.id}") + + # Create batch + batch_response = await self._client.batches.create( + input_file_id=file_response.id, + endpoint=endpoint, + completion_window=completion_window, + ) + + logger.info(f"Batch submitted: {batch_response.id}, status={batch_response.status}") + + return { + "batch_id": batch_response.id, + "status": batch_response.status, + "input_file_id": file_response.id, + "created_at": batch_response.created_at, + "request_count": len(requests), + } + + async def get_batch_status(self, batch_id: str) -> dict[str, Any]: + """ + Get the status of a batch job. + + Args: + batch_id: Batch identifier + + Returns: + Dict with status info (batch_id, status, completed_at, etc.) + """ + if not await self.supports_batch_api(): + raise NotImplementedError(f"Batch API not supported for provider: {self.provider}") + + batch = await self._client.batches.retrieve(batch_id) + + result = { + "batch_id": batch.id, + "status": batch.status, + "created_at": batch.created_at, + "request_counts": { + "total": batch.request_counts.total if batch.request_counts else 0, + "completed": batch.request_counts.completed if batch.request_counts else 0, + "failed": batch.request_counts.failed if batch.request_counts else 0, + }, + } + + if batch.completed_at: + result["completed_at"] = batch.completed_at + if batch.output_file_id: + result["output_file_id"] = batch.output_file_id + if batch.error_file_id: + result["error_file_id"] = batch.error_file_id + if batch.errors: + result["errors"] = batch.errors + + return result + + async def retrieve_batch_results(self, batch_id: str) -> list[dict[str, Any]]: + """ + Retrieve completed batch results. + + Args: + batch_id: Batch identifier + + Returns: + List of result dicts (one per request, matched by custom_id) + """ + if not await self.supports_batch_api(): + raise NotImplementedError(f"Batch API not supported for provider: {self.provider}") + + # Get batch status + batch = await self._client.batches.retrieve(batch_id) + + if batch.status != "completed": + raise ValueError(f"Batch {batch_id} is not completed yet (status: {batch.status})") + + if not batch.output_file_id: + raise ValueError(f"Batch {batch_id} has no output file") + + # Download results file + logger.debug(f"Downloading results for batch {batch_id} from file {batch.output_file_id}") + file_content = await self._client.files.content(batch.output_file_id) + + # Parse JSONL results + results = [] + for line in file_content.text.strip().split("\n"): + if line: + results.append(json.loads(line)) + + logger.info(f"Retrieved {len(results)} results for batch {batch_id}") + + return results + async def cleanup(self) -> None: """Clean up resources (close OpenAI client connections).""" if hasattr(self, "_client") and self._client: diff --git a/hindsight-api/hindsight_api/engine/retain/fact_extraction.py b/hindsight-api/hindsight_api/engine/retain/fact_extraction.py index 3db10225..40c2fb50 100644 --- a/hindsight-api/hindsight_api/engine/retain/fact_extraction.py +++ b/hindsight-api/hindsight_api/engine/retain/fact_extraction.py @@ -695,6 +695,91 @@ Example: "Lost job → couldn't pay rent → moved apartment" - Fact 2: Moved apartment, causal_relations: [{target_index: 1, relation_type: "caused_by"}]""" +def _build_extraction_prompt_and_schema(config) -> tuple[str, type]: + """ + Build extraction prompt and response schema based on config. + + Returns: + Tuple of (prompt, response_schema) + """ + fact_types_instruction = "Extract ONLY 'world' and 'assistant' type facts." + extraction_mode = config.retain_extraction_mode + extract_causal_links = config.retain_extract_causal_links + + # Select base prompt based on extraction mode + if extraction_mode == "custom": + if not config.retain_custom_instructions: + base_prompt = CONCISE_FACT_EXTRACTION_PROMPT + prompt = base_prompt.format(fact_types_instruction=fact_types_instruction) + else: + base_prompt = CUSTOM_FACT_EXTRACTION_PROMPT + prompt = base_prompt.format( + fact_types_instruction=fact_types_instruction, + custom_instructions=config.retain_custom_instructions, + ) + elif extraction_mode == "verbose": + base_prompt = VERBOSE_FACT_EXTRACTION_PROMPT + prompt = base_prompt.format(fact_types_instruction=fact_types_instruction) + else: + base_prompt = CONCISE_FACT_EXTRACTION_PROMPT + prompt = base_prompt.format(fact_types_instruction=fact_types_instruction) + + # Add causal relationships section if enabled + if extract_causal_links: + prompt = prompt + CAUSAL_RELATIONSHIPS_SECTION + response_schema = FactExtractionResponseVerbose if extraction_mode == "verbose" else FactExtractionResponse + else: + response_schema = FactExtractionResponseNoCausal + + return prompt, response_schema + + +def _build_user_message(chunk: str, chunk_index: int, total_chunks: int, event_date: datetime, context: str) -> str: + """Build user message for fact extraction.""" + from .orchestrator import parse_datetime_flexible + + sanitized_chunk = _sanitize_text(chunk) + sanitized_context = _sanitize_text(context) if context else "none" + event_date = parse_datetime_flexible(event_date) + event_date_formatted = event_date.strftime("%A, %B %d, %Y") + + return f"""Extract facts from the following text chunk. + +Chunk: {chunk_index + 1}/{total_chunks} +Event Date: {event_date_formatted} ({event_date.isoformat()}) +Context: {sanitized_context} + +Text: +{sanitized_chunk}""" + + +def _build_request_body(llm_config, config, prompt: str, user_message: str, response_schema: type) -> dict: + """Build request body for LLM API call.""" + request_body = { + "model": llm_config.model, + "messages": [{"role": "system", "content": prompt}, {"role": "user", "content": user_message}], + "temperature": 0.1, + } + + # Add max_completion_tokens if configured + if config.retain_max_completion_tokens: + request_body["max_completion_tokens"] = config.retain_max_completion_tokens + + # Add service_tier for OpenAI Flex Processing + if llm_config.provider == "openai" and llm_config._provider_impl.openai_service_tier: + request_body["service_tier"] = llm_config._provider_impl.openai_service_tier + + # Add response_format (JSON schema) + if hasattr(response_schema, "model_json_schema"): + schema = response_schema.model_json_schema() + request_body["response_format"] = { + "type": "json_schema", + "json_schema": {"name": "facts", "schema": schema}, + } + + return request_body + + async def _extract_facts_from_chunk( chunk: str, chunk_index: int, @@ -717,72 +802,20 @@ async def _extract_facts_from_chunk( logger = logging.getLogger(__name__) - # Determine which fact types to extract - # Note: We use "assistant" in the prompt but convert to "bank" for storage - fact_types_instruction = "Extract ONLY 'world' and 'assistant' type facts." + # Build prompt and schema using helper function + prompt, response_schema = _build_extraction_prompt_and_schema(config) # Check config for extraction mode and causal link extraction extraction_mode = config.retain_extraction_mode extract_causal_links = config.retain_extract_causal_links - # Select base prompt based on extraction mode - if extraction_mode == "custom": - # Custom mode: inject user-provided guidelines - if not config.retain_custom_instructions: - logger.warning( - "extraction_mode='custom' but HINDSIGHT_API_RETAIN_CUSTOM_INSTRUCTIONS not set. " - "Falling back to 'concise' mode." - ) - base_prompt = CONCISE_FACT_EXTRACTION_PROMPT - prompt = base_prompt.format(fact_types_instruction=fact_types_instruction) - else: - base_prompt = CUSTOM_FACT_EXTRACTION_PROMPT - prompt = base_prompt.format( - fact_types_instruction=fact_types_instruction, - custom_instructions=config.retain_custom_instructions, - ) - elif extraction_mode == "verbose": - base_prompt = VERBOSE_FACT_EXTRACTION_PROMPT - prompt = base_prompt.format(fact_types_instruction=fact_types_instruction) - else: - base_prompt = CONCISE_FACT_EXTRACTION_PROMPT - prompt = base_prompt.format(fact_types_instruction=fact_types_instruction) - - # Build the full prompt with or without causal relationships section - # Select appropriate response schema based on extraction mode and causal links - if extract_causal_links: - prompt = prompt + CAUSAL_RELATIONSHIPS_SECTION - if extraction_mode == "verbose": - response_schema = FactExtractionResponseVerbose - else: - response_schema = FactExtractionResponse - else: - response_schema = FactExtractionResponseNoCausal + # Build user message using helper function + user_message = _build_user_message(chunk, chunk_index, total_chunks, event_date, context) # Retry logic for JSON validation errors max_retries = 2 last_error = None - # Sanitize input text to prevent Unicode encoding errors (e.g., unpaired surrogates) - sanitized_chunk = _sanitize_text(chunk) - sanitized_context = _sanitize_text(context) if context else "none" - - # Build user message with metadata and chunk content in a clear format - # Format event_date with day of week for better temporal reasoning - # Handle both datetime objects and ISO string formats (from deserialized async tasks) - from .orchestrator import parse_datetime_flexible - - event_date = parse_datetime_flexible(event_date) - event_date_formatted = event_date.strftime("%A, %B %d, %Y") # e.g., "Monday, June 10, 2024" - user_message = f"""Extract facts from the following text chunk. - -Chunk: {chunk_index + 1}/{total_chunks} -Event Date: {event_date_formatted} ({event_date.isoformat()}) -Context: {sanitized_context} - -Text: -{sanitized_chunk}""" - usage = TokenUsage() # Track cumulative usage across retries for attempt in range(max_retries): try: @@ -1245,8 +1278,420 @@ logger = logging.getLogger(__name__) SECONDS_PER_FACT = 10 +async def extract_facts_from_contents_batch_api( + contents: list[RetainContent], + llm_config, + agent_name: str, + config, + pool=None, + operation_id: str | None = None, + schema: str | None = None, +) -> tuple[list[ExtractedFactType], list[ChunkMetadata], TokenUsage]: + """ + Extract facts using LLM Batch API (OpenAI/Groq). + + Submits all chunks as a single batch, polls until complete, then processes results. + Only called when config.retain_batch_enabled=True. + + Args: + contents: List of RetainContent objects to process + llm_config: LLM configuration with batch API support + agent_name: Name of the agent + config: Resolved HindsightConfig for this bank + pool: Database connection pool (for storing batch state) + operation_id: Async operation ID (for crash recovery) + schema: Database schema (for multi-tenant support) + + Returns: + Tuple of (extracted_facts, chunks_metadata, usage) + """ + if not contents: + return [], [], TokenUsage() + + logger.info(f"Using Batch API for fact extraction ({len(contents)} contents)") + + # Check config for extraction mode and causal link extraction (used throughout) + extraction_mode = config.retain_extraction_mode + extract_causal_links = config.retain_extract_causal_links + + # Check if provider supports batch API + if not await llm_config._provider_impl.supports_batch_api(): + logger.warning(f"Batch API not supported for provider {llm_config.provider}, falling back to sync mode") + return await extract_facts_from_contents(contents, llm_config, agent_name, config, pool, operation_id, schema) + + # Check if we're resuming an existing batch (crash recovery) + batch_id = None + if operation_id and pool: + from ..task_backend import fq_table + + table = fq_table("async_operations", schema) + row = await pool.fetchrow( + f"SELECT result_metadata FROM {table} WHERE operation_id = $1", + operation_id, + ) + + if row and row["result_metadata"]: + metadata = row["result_metadata"] + if isinstance(metadata, str): + metadata = json.loads(metadata) + batch_id = metadata.get("batch_id") + + if batch_id: + logger.info(f"Resuming existing batch: batch_id={batch_id} (crash recovery)") + + # Step 1: Chunk all contents and build batch requests (skip if resuming) + all_chunks_info = [] # List of (chunk_text, content_index, chunk_index_in_content, event_date, context) + batch_requests = [] + + # Build prompt and schema once (same for all chunks) + prompt, response_schema = _build_extraction_prompt_and_schema(config) + + for content_index, item in enumerate(contents): + chunks = chunk_text(item.content, max_chars=config.retain_chunk_size) + + for chunk_index_in_content, chunk in enumerate(chunks): + all_chunks_info.append((chunk, content_index, chunk_index_in_content, item.event_date, item.context)) + + # Build batch request for this chunk + custom_id = f"chunk_{len(all_chunks_info) - 1}" # Global chunk index + + # Build user message using helper function + user_message = _build_user_message( + chunk, chunk_index_in_content, len(chunks), item.event_date, item.context + ) + + # Build request body using helper function + request_body = _build_request_body(llm_config, config, prompt, user_message, response_schema) + + batch_requests.append( + {"custom_id": custom_id, "method": "POST", "url": "/v1/chat/completions", "body": request_body} + ) + + if not batch_requests and not batch_id: # No requests and not resuming + return [], [], TokenUsage() + + # Step 2: Submit batch (skip if resuming) + if not batch_id: + logger.info(f"Submitting batch with {len(batch_requests)} chunk requests") + + batch_metadata = await llm_config._provider_impl.submit_batch(batch_requests) + batch_id = batch_metadata["batch_id"] + + logger.info(f"Batch submitted: {batch_id}, polling every {config.retain_batch_poll_interval_seconds}s") + + # CRITICAL: Store minimal batch state in operation metadata for crash recovery + # This allows resuming polling if worker restarts + if operation_id and pool: + batch_state = { + "batch_id": batch_id, + "batch_provider": llm_config.provider, + "chunk_count": len(batch_requests), + } + + # Update operation result_metadata + from ..task_backend import fq_table + + table = fq_table("async_operations", schema) + await pool.execute( + f""" + UPDATE {table} + SET result_metadata = result_metadata || $1::jsonb, updated_at = now() + WHERE operation_id = $2 + """, + json.dumps(batch_state), + operation_id, + ) + logger.info(f"Stored batch state for operation {operation_id} (crash recovery enabled)") + else: + logger.info(f"Resuming polling for existing batch: {batch_id}") + + # Step 3: Poll until complete + import time + + start_time = time.time() + while True: + status_info = await llm_config._provider_impl.get_batch_status(batch_id) + status = status_info["status"] + + elapsed = time.time() - start_time + logger.info( + f"Batch {batch_id}: status={status}, " + f"completed={status_info['request_counts']['completed']}/{status_info['request_counts']['total']}, " + f"elapsed={elapsed:.0f}s" + ) + + if status == "completed": + break + elif status in ("failed", "expired", "cancelled"): + error_msg = status_info.get("errors", "Unknown error") + raise RuntimeError(f"Batch {batch_id} failed with status {status}: {error_msg}") + + # Wait before polling again + await asyncio.sleep(config.retain_batch_poll_interval_seconds) + + logger.info(f"Batch {batch_id} completed in {elapsed:.0f}s, retrieving results") + + # Step 4: Retrieve results + batch_results = await llm_config._provider_impl.retrieve_batch_results(batch_id) + + # Map results by custom_id + results_by_id = {result["custom_id"]: result for result in batch_results} + + # Step 5: Parse results into facts (same as sync mode) + all_facts_from_llm = [] + chunks_metadata = [] + total_usage = TokenUsage() + + for chunk_idx, (chunk_content, content_index, chunk_index_in_content, event_date, context) in enumerate( + all_chunks_info + ): + custom_id = f"chunk_{chunk_idx}" + result = results_by_id.get(custom_id) + + if not result: + logger.warning(f"Missing result for {custom_id}, skipping") + chunks_metadata.append( + ChunkMetadata( + chunk_text=chunk_content, fact_count=0, content_index=content_index, chunk_index=chunk_idx + ) + ) + continue + + # Check for errors + if result.get("error"): + logger.error(f"Error in {custom_id}: {result['error']}") + chunks_metadata.append( + ChunkMetadata( + chunk_text=chunk_content, fact_count=0, content_index=content_index, chunk_index=chunk_idx + ) + ) + continue + + # Extract response + response_body = result.get("response", {}).get("body", {}) + choices = response_body.get("choices", []) + + if not choices: + logger.warning(f"No choices in response for {custom_id}") + chunks_metadata.append( + ChunkMetadata( + chunk_text=chunk_content, fact_count=0, content_index=content_index, chunk_index=chunk_idx + ) + ) + continue + + # Parse JSON content + message = choices[0].get("message", {}) + content_str = message.get("content", "{}") + + try: + extraction_response_json = json.loads(content_str) + except json.JSONDecodeError as e: + logger.error(f"Failed to parse JSON for {custom_id}: {e}") + chunks_metadata.append( + ChunkMetadata( + chunk_text=chunk_content, fact_count=0, content_index=content_index, chunk_index=chunk_idx + ) + ) + continue + + # Parse facts (reuse existing logic from _extract_facts_from_chunk) + raw_facts = extraction_response_json.get("facts", []) + chunk_facts = [] + + for i, llm_fact in enumerate(raw_facts): + if not isinstance(llm_fact, dict): + continue + + def get_value(field_name): + value = llm_fact.get(field_name) + if value and value != "" and value != [] and value != {} and str(value).upper() != "N/A": + return value + return None + + what = get_value("what") + if not what: + what = get_value("factual_core") + if not what: + continue + + when = get_value("when") + who = get_value("who") + why = get_value("why") + + # Critical field: fact_type + original_fact_type = llm_fact.get("fact_type") + fact_type = original_fact_type + + # Convert "assistant" → "experience" + if fact_type == "assistant": + fact_type = "experience" + + # Validate fact_type + if fact_type not in ["world", "experience", "opinion"]: + fact_kind = llm_fact.get("fact_kind") + if fact_kind == "assistant": + fact_type = "experience" + elif fact_kind in ["world", "experience", "opinion"]: + fact_type = fact_kind + else: + fact_type = "world" + + # Build combined fact text + combined_parts = [what] + if when: + combined_parts.append(f"When: {when}") + if who: + combined_parts.append(f"Involving: {who}") + if why: + combined_parts.append(why) + combined_text = " | ".join(combined_parts) + + # Temporal fields + fact_data = {} + fact_kind = llm_fact.get("fact_kind", "conversation") + if fact_kind not in ["conversation", "event", "other"]: + fact_kind = "conversation" + + if fact_kind == "event": + occurred_start = get_value("occurred_start") + occurred_end = get_value("occurred_end") + + if not occurred_start: + fact_data["occurred_start"] = _infer_temporal_date(combined_text, event_date) + else: + fact_data["occurred_start"] = occurred_start + + if occurred_end: + fact_data["occurred_end"] = occurred_end + elif fact_data.get("occurred_start"): + fact_data["occurred_end"] = fact_data["occurred_start"] + + # Entities + entities = get_value("entities") + if entities: + validated_entities = [] + for ent in entities: + if isinstance(ent, str): + validated_entities.append(Entity(text=ent)) + elif isinstance(ent, dict) and "text" in ent: + try: + validated_entities.append(Entity.model_validate(ent)) + except Exception: + pass + if validated_entities: + fact_data["entities"] = validated_entities + + # Causal relations + if extract_causal_links: + validated_relations = [] + causal_relations_raw = get_value("causal_relations") + if causal_relations_raw: + for rel in causal_relations_raw: + if not isinstance(rel, dict): + continue + target_idx = rel.get("target_index") + relation_type = rel.get("relation_type") + strength = rel.get("strength", 1.0) + + if target_idx is None or relation_type is None: + continue + if target_idx < 0 or target_idx >= i: + continue + + try: + validated_relations.append( + CausalRelation( + target_fact_index=target_idx, relation_type=relation_type, strength=strength + ) + ) + except Exception: + pass + + if validated_relations: + fact_data["causal_relations"] = validated_relations + + # Always set mentioned_at + fact_data["mentioned_at"] = event_date.isoformat() + + try: + fact = Fact(fact=combined_text, fact_type=fact_type, **fact_data) + chunk_facts.append(fact) + except Exception as e: + logger.error(f"Failed to create Fact model for fact {i}: {e}") + continue + + all_facts_from_llm.extend(chunk_facts) + chunks_metadata.append( + ChunkMetadata( + chunk_text=chunk_content, + fact_count=len(chunk_facts), + content_index=content_index, + chunk_index=chunk_idx, + ) + ) + + # Track token usage + usage_data = response_body.get("usage", {}) + if usage_data: + total_usage = total_usage + TokenUsage( + input_tokens=usage_data.get("prompt_tokens", 0), + output_tokens=usage_data.get("completion_tokens", 0), + total_tokens=usage_data.get("total_tokens", 0), + ) + + # Step 6: Convert to ExtractedFact objects with proper chunk mapping + # Group facts by chunk + facts_by_chunk = [] # List of (chunk_metadata, [facts]) + fact_start_idx = 0 + + for chunk_meta in chunks_metadata: + chunk_facts = all_facts_from_llm[fact_start_idx : fact_start_idx + chunk_meta.fact_count] + facts_by_chunk.append((chunk_meta, chunk_facts)) + fact_start_idx += chunk_meta.fact_count + + # Now convert to ExtractedFactType + extracted_facts = [] + global_fact_idx = 0 + + for chunk_meta, chunk_facts in facts_by_chunk: + content = contents[chunk_meta.content_index] + + for fact_from_llm in chunk_facts: + extracted_fact = ExtractedFactType( + fact_text=fact_from_llm.fact, + fact_type=fact_from_llm.fact_type, + entities=[e.text for e in (fact_from_llm.entities or [])], + occurred_start=_parse_datetime(fact_from_llm.occurred_start) if fact_from_llm.occurred_start else None, + occurred_end=_parse_datetime(fact_from_llm.occurred_end) if fact_from_llm.occurred_end else None, + causal_relations=_convert_causal_relations(fact_from_llm.causal_relations or [], global_fact_idx), + content_index=chunk_meta.content_index, + chunk_index=chunk_meta.chunk_index, + context=content.context, + mentioned_at=content.event_date, + metadata=content.metadata, + tags=content.tags, + ) + + extracted_facts.append(extracted_fact) + global_fact_idx += 1 + + # Step 7: Add temporal offsets + _add_temporal_offsets(extracted_facts, contents) + + logger.info(f"Batch API extracted {len(extracted_facts)} facts from {len(all_chunks_info)} chunks") + + return extracted_facts, chunks_metadata, total_usage + + async def extract_facts_from_contents( - contents: list[RetainContent], llm_config, agent_name: str, config + contents: list[RetainContent], + llm_config, + agent_name: str, + config, + pool=None, + operation_id: str | None = None, + schema: str | None = None, ) -> tuple[list[ExtractedFactType], list[ChunkMetadata], TokenUsage]: """ Extract facts from multiple content items in parallel. @@ -1257,11 +1702,16 @@ async def extract_facts_from_contents( 3. Adds time offsets to preserve fact ordering within each content 4. Returns typed ExtractedFact and ChunkMetadata objects + Routes to batch API mode if config.retain_batch_enabled=True. + Args: contents: List of RetainContent objects to process llm_config: LLM configuration for fact extraction agent_name: Name of the agent (for agent-related fact detection) config: Resolved HindsightConfig for this bank + pool: Database connection pool (passed to batch API for state storage) + operation_id: Async operation ID (passed to batch API for crash recovery) + schema: Database schema (passed to batch API for multi-tenant support) Returns: Tuple of (extracted_facts, chunks_metadata, usage) @@ -1269,6 +1719,12 @@ async def extract_facts_from_contents( if not contents: return [], [], TokenUsage() + # Route to batch API if enabled + if config.retain_batch_enabled: + return await extract_facts_from_contents_batch_api( + contents, llm_config, agent_name, config, pool, operation_id, schema + ) + # Step 1: Create parallel fact extraction tasks fact_extraction_tasks = [] for item in contents: diff --git a/hindsight-api/hindsight_api/engine/retain/orchestrator.py b/hindsight-api/hindsight_api/engine/retain/orchestrator.py index 87a4cdee..1890afaf 100644 --- a/hindsight-api/hindsight_api/engine/retain/orchestrator.py +++ b/hindsight-api/hindsight_api/engine/retain/orchestrator.py @@ -82,6 +82,8 @@ async def retain_batch( fact_type_override: str | None = None, confidence_score: float | None = None, document_tags: list[str] | None = None, + operation_id: str | None = None, + schema: str | None = None, ) -> tuple[list[list[str]], TokenUsage]: """ Process a batch of content through the retain pipeline. @@ -147,7 +149,7 @@ async def retain_batch( step_start = time.time() extracted_facts, chunks, usage = await fact_extraction.extract_facts_from_contents( - contents, llm_config, agent_name, config + contents, llm_config, agent_name, config, pool, operation_id, schema ) log_buffer.append( f"[1] Extract facts: {len(extracted_facts)} facts, {len(chunks)} chunks from {len(contents)} contents in {time.time() - step_start:.3f}s" diff --git a/hindsight-api/hindsight_api/main.py b/hindsight-api/hindsight_api/main.py index e29bf76a..77317f0f 100644 --- a/hindsight-api/hindsight_api/main.py +++ b/hindsight-api/hindsight_api/main.py @@ -166,6 +166,8 @@ def main(): 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, @@ -246,6 +248,8 @@ def main(): 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, diff --git a/hindsight-api/hindsight_api/worker/poller.py b/hindsight-api/hindsight_api/worker/poller.py index 7b30c51e..b425f627 100644 --- a/hindsight-api/hindsight_api/worker/poller.py +++ b/hindsight-api/hindsight_api/worker/poller.py @@ -401,6 +401,8 @@ class WorkerPoller: On startup, we reset any tasks stuck in 'processing' for this worker_id back to 'pending' so they can be picked up again. + Also recovers batch API operations that were in-flight. + If tenant_extension is configured, recovers across all tenant schemas. Returns: @@ -413,11 +415,16 @@ class WorkerPoller: try: table = fq_table("async_operations", schema) + # First, recover batch API operations (before resetting worker tasks) + batch_count = await self._recover_batch_operations(schema) + total_count += batch_count + + # Then reset normal worker tasks result = await self._pool.execute( f""" UPDATE {table} SET status = 'pending', worker_id = NULL, claimed_at = NULL, updated_at = now() - WHERE status = 'processing' AND worker_id = $1 + WHERE status = 'processing' AND worker_id = $1 AND result_metadata->>'batch_id' IS NULL """, self._worker_id, ) @@ -434,6 +441,80 @@ class WorkerPoller: logger.info(f"Worker {self._worker_id} recovered {total_count} stale tasks from previous run") return total_count + async def _recover_batch_operations(self, schema: str | None) -> int: + """ + Recover batch API operations that were in-flight when worker crashed. + + Finds operations with batch_id in metadata and re-submits them as tasks + so polling can resume. + + Args: + schema: Database schema to recover from + + Returns: + Number of batch operations recovered + """ + table = fq_table("async_operations", schema) + + try: + # Find operations with batch_id in metadata (batch API operations) + rows = await self._pool.fetch( + f""" + SELECT operation_id, task_payload, result_metadata + FROM {table} + WHERE status = 'processing' + AND result_metadata ? 'batch_id' + AND task_payload IS NOT NULL + """ + ) + + if not rows: + return 0 + + recovered = 0 + for row in rows: + operation_id = str(row["operation_id"]) + task_payload = row["task_payload"] + result_metadata = row["result_metadata"] + + # Parse metadata + if isinstance(result_metadata, str): + result_metadata = json.loads(result_metadata) + + batch_id = result_metadata.get("batch_id") + batch_provider = result_metadata.get("batch_provider", "openai") + + logger.info( + f"Recovering batch operation: operation_id={operation_id}, batch_id={batch_id}, provider={batch_provider}" + ) + + # Parse task_payload + if isinstance(task_payload, str): + task_dict = json.loads(task_payload) + else: + task_dict = task_payload + + # Mark operation as ready for re-processing + # Reset to pending with task_payload intact so worker picks it up again + await self._pool.execute( + f""" + UPDATE {table} + SET status = 'pending', worker_id = NULL, claimed_at = NULL, updated_at = now() + WHERE operation_id = $1 + """, + operation_id, + ) + + recovered += 1 + logger.info(f"Batch operation {operation_id} reset to pending for re-processing") + + return recovered + + except Exception as e: + schema_display = f'"{schema}"' if schema else str(schema) + logger.error(f"Failed to recover batch operations for schema {schema_display}: {e}") + return 0 + async def run(self): """ Main polling loop with fire-and-forget task execution. diff --git a/hindsight-api/tests/test_batch_api.py b/hindsight-api/tests/test_batch_api.py new file mode 100644 index 00000000..889b27be --- /dev/null +++ b/hindsight-api/tests/test_batch_api.py @@ -0,0 +1,508 @@ +""" +Test OpenAI Batch API integration for retain fact extraction. + +Tests cover: +- Normal batch API flow (submit, poll, complete) +- Crash recovery (resume from existing batch_id) +- Provider fallback (when batch API not supported) +- Worker recovery on restart +""" +import pytest +import asyncio +import logging +import json +import uuid +from datetime import datetime, timezone +from unittest.mock import AsyncMock, MagicMock, patch +from hindsight_api import RequestContext +from hindsight_api.engine.retain.fact_extraction import ( + extract_facts_from_contents_batch_api, + extract_facts_from_contents, + RetainContent, +) +from hindsight_api.config import HindsightConfig +from hindsight_api.engine.llm_wrapper import create_llm_provider +from hindsight_api.worker.poller import WorkerPoller + +logger = logging.getLogger(__name__) + + +@pytest.fixture +def mock_llm_config(): + """Create a mock LLM config with batch API support.""" + mock = MagicMock() + mock.provider = "openai" + mock.model = "gpt-4o-mini" + mock._provider_impl = AsyncMock() + return mock + + +@pytest.fixture +def test_contents(): + """Create test content for fact extraction.""" + return [ + RetainContent( + content="Alice is a senior software engineer at TechCorp. She specializes in distributed systems.", + event_date=datetime(2024, 1, 15, tzinfo=timezone.utc), + context="team overview", + ), + RetainContent( + content="Bob joined the team last month as a junior developer. He is learning React.", + event_date=datetime(2024, 1, 15, tzinfo=timezone.utc), + context="team overview", + ), + ] + + +@pytest.fixture +def hindsight_config(): + """Create test config with batch API enabled.""" + config = HindsightConfig.from_env() + config.retain_batch_enabled = True + config.retain_batch_poll_interval_seconds = 1 # Fast polling for tests + config.retain_chunk_size = 4000 + config.retain_extraction_mode = "concise" + config.retain_extract_causal_links = False + return config + + +@pytest.mark.asyncio +async def test_batch_api_normal_flow(mock_llm_config, test_contents, hindsight_config, memory, request_context): + """Test normal batch API flow: submit, poll, complete.""" + bank_id = f"test_batch_{datetime.now(timezone.utc).timestamp()}" + + try: + # Mock batch API responses + batch_id = "batch_test123" + + # Mock supports_batch_api + mock_llm_config._provider_impl.supports_batch_api = AsyncMock(return_value=True) + + # Mock submit_batch - returns batch metadata + mock_llm_config._provider_impl.submit_batch = AsyncMock( + return_value={ + "batch_id": batch_id, + "status": "validating", + "request_counts": {"total": 2, "completed": 0, "failed": 0}, + } + ) + + # Mock get_batch_status - simulate polling sequence + status_sequence = [ + {"status": "in_progress", "request_counts": {"total": 2, "completed": 1, "failed": 0}}, + {"status": "completed", "request_counts": {"total": 2, "completed": 2, "failed": 0}}, + ] + mock_llm_config._provider_impl.get_batch_status = AsyncMock(side_effect=status_sequence) + + # Mock retrieve_batch_results - returns fact extraction results + mock_results = [ + { + "custom_id": "chunk_0", + "response": { + "body": { + "choices": [ + { + "message": { + "content": json.dumps({ + "facts": [ + { + "what": "Alice is a senior software engineer at TechCorp", + "when": "present", + "where": "TechCorp", + "who": "Alice", + "why": "Professional background information", + "fact_type": "world", + "fact_kind": "conversation", + } + ] + }) + } + } + ], + "usage": {"prompt_tokens": 100, "completion_tokens": 50, "total_tokens": 150}, + } + }, + }, + { + "custom_id": "chunk_1", + "response": { + "body": { + "choices": [ + { + "message": { + "content": json.dumps({ + "facts": [ + { + "what": "Bob joined the team last month as a junior developer", + "when": "last month", + "where": "team", + "who": "Bob", + "why": "New team member information", + "fact_type": "world", + "fact_kind": "conversation", + } + ] + }) + } + } + ], + "usage": {"prompt_tokens": 100, "completion_tokens": 50, "total_tokens": 150}, + } + }, + }, + ] + mock_llm_config._provider_impl.retrieve_batch_results = AsyncMock(return_value=mock_results) + + # Call batch API extraction + facts, chunks, usage = await extract_facts_from_contents_batch_api( + contents=test_contents, + llm_config=mock_llm_config, + agent_name="test_agent", + config=hindsight_config, + pool=None, # No DB pool for this test + operation_id=None, + schema=None, + ) + + # Verify results + assert len(facts) == 2, "Should extract 2 facts (one per chunk)" + # Facts are ExtractedFact objects with .fact_text field + assert "Alice" in facts[0].fact_text and "senior software engineer" in facts[0].fact_text + assert "Bob" in facts[1].fact_text and "junior developer" in facts[1].fact_text + + # Verify chunks metadata + assert len(chunks) == 2, "Should have 2 chunks metadata" + assert chunks[0].fact_count == 1 + assert chunks[1].fact_count == 1 + + # Verify token usage + assert usage.input_tokens == 200 # 100 per chunk + assert usage.output_tokens == 100 # 50 per chunk + assert usage.total_tokens == 300 + + # Verify API calls + mock_llm_config._provider_impl.submit_batch.assert_called_once() + assert mock_llm_config._provider_impl.get_batch_status.call_count == 2 + mock_llm_config._provider_impl.retrieve_batch_results.assert_called_once_with(batch_id) + + logger.info("✅ Normal batch API flow test passed") + + finally: + # Cleanup + try: + await memory.delete_bank(bank_id, request_context=request_context) + except Exception: + pass + + +@pytest.mark.asyncio +async def test_batch_api_crash_recovery(mock_llm_config, test_contents, hindsight_config, memory, request_context): + """Test crash recovery: resume polling from existing batch_id.""" + bank_id = f"test_crash_{datetime.now(timezone.utc).timestamp()}" + operation_id = str(uuid.uuid4()) # Must be UUID for async_operations table + + try: + # Ensure bank exists + await memory.get_bank_profile(bank_id, request_context=request_context) + + # Setup: Store batch_id in async_operations table (simulates partial execution) + batch_id = "batch_recovered_456" + pool = memory._pool + schema = request_context.tenant_id + + from hindsight_api.engine.task_backend import fq_table + table = fq_table("async_operations", schema) + + # Create operation with batch_id already stored + await pool.execute( + f""" + INSERT INTO {table} (operation_id, operation_type, bank_id, status, result_metadata) + VALUES ($1, 'retain', $2, 'processing', $3::jsonb) + """, + operation_id, + bank_id, + json.dumps({ + "batch_id": batch_id, + "batch_provider": "openai", + "chunk_count": 2, + }), + ) + + # Mock batch API responses for resume scenario + mock_llm_config._provider_impl.supports_batch_api = AsyncMock(return_value=True) + + # Mock get_batch_status - batch already in progress + mock_llm_config._provider_impl.get_batch_status = AsyncMock( + return_value={ + "status": "completed", + "request_counts": {"total": 2, "completed": 2, "failed": 0}, + } + ) + + # Mock retrieve_batch_results + mock_results = [ + { + "custom_id": "chunk_0", + "response": { + "body": { + "choices": [ + { + "message": { + "content": json.dumps({ + "facts": [ + { + "what": "Alice is a senior software engineer", + "when": "present", + "where": "TechCorp", + "who": "Alice", + "why": "Background", + "fact_type": "world", + "fact_kind": "conversation", + } + ] + }) + } + } + ], + "usage": {"prompt_tokens": 100, "completion_tokens": 50, "total_tokens": 150}, + } + }, + }, + { + "custom_id": "chunk_1", + "response": { + "body": { + "choices": [ + { + "message": { + "content": json.dumps({ + "facts": [ + { + "what": "Bob is a junior developer", + "when": "last month", + "where": "team", + "who": "Bob", + "why": "New member", + "fact_type": "world", + "fact_kind": "conversation", + } + ] + }) + } + } + ], + "usage": {"prompt_tokens": 100, "completion_tokens": 50, "total_tokens": 150}, + } + }, + }, + ] + mock_llm_config._provider_impl.retrieve_batch_results = AsyncMock(return_value=mock_results) + + # Call batch API extraction with operation_id (crash recovery scenario) + facts, chunks, usage = await extract_facts_from_contents_batch_api( + contents=test_contents, + llm_config=mock_llm_config, + agent_name="test_agent", + config=hindsight_config, + pool=pool, + operation_id=operation_id, # Provides crash recovery context + schema=schema, + ) + + # Verify results + assert len(facts) == 2, "Should extract 2 facts after recovery" + + # CRITICAL: Verify submit_batch was NOT called (because batch_id already exists) + mock_llm_config._provider_impl.submit_batch.assert_not_called() + + # Verify get_batch_status WAS called (polling resumed) + mock_llm_config._provider_impl.get_batch_status.assert_called() + + # Verify retrieve_batch_results was called with the recovered batch_id + mock_llm_config._provider_impl.retrieve_batch_results.assert_called_once_with(batch_id) + + logger.info("✅ Crash recovery test passed - resumed polling without re-submission") + + finally: + # Cleanup + try: + await memory.delete_bank(bank_id, request_context=request_context) + except Exception: + pass + + +@pytest.mark.asyncio +async def test_batch_api_fallback_unsupported_provider(mock_llm_config, test_contents, hindsight_config): + """Test fallback to sync mode when provider doesn't support batch API.""" + + # Mock provider that doesn't support batch API + mock_llm_config._provider_impl.supports_batch_api = AsyncMock(return_value=False) + mock_llm_config.provider = "groq" # Example of provider + + # Patch the sync mode function to verify it's called + with patch( + "hindsight_api.engine.retain.fact_extraction.extract_facts_from_contents" + ) as mock_sync_extract: + mock_sync_extract.return_value = ([], [], MagicMock()) + + # Call batch API extraction (should fallback to sync) + await extract_facts_from_contents_batch_api( + contents=test_contents, + llm_config=mock_llm_config, + agent_name="test_agent", + config=hindsight_config, + pool=None, + operation_id=None, + schema=None, + ) + + # Verify fallback occurred + mock_sync_extract.assert_called_once() + + # Verify batch API methods were NOT called + mock_llm_config._provider_impl.submit_batch.assert_not_called() + + logger.info("✅ Fallback to sync mode test passed") + + +@pytest.mark.asyncio +async def test_worker_batch_recovery(memory, request_context): + """Test that WorkerPoller._recover_batch_operations finds and resets orphaned batches.""" + bank_id = f"test_worker_recovery_{datetime.now(timezone.utc).timestamp()}" + operation_id = str(uuid.uuid4()) # Must be UUID for async_operations table + + try: + # Ensure bank exists + await memory.get_bank_profile(bank_id, request_context=request_context) + + pool = memory._pool + schema = request_context.tenant_id + + from hindsight_api.engine.task_backend import fq_table + table = fq_table("async_operations", schema) + + # Create orphaned batch operation (simulates worker crash during polling) + batch_id = "batch_orphaned_999" + task_payload = { + "operation_type": "retain", + "bank_id": bank_id, + "contents": [{"content": "test", "event_date": "2024-01-15T00:00:00Z"}], + } + + await pool.execute( + f""" + INSERT INTO {table} (operation_id, operation_type, bank_id, status, worker_id, result_metadata, task_payload) + VALUES ($1, 'retain', $2, 'processing', 'worker_crashed', $3::jsonb, $4::jsonb) + """, + operation_id, + bank_id, + json.dumps({ + "batch_id": batch_id, + "batch_provider": "openai", + "chunk_count": 1, + }), + json.dumps(task_payload), + ) + + # Create WorkerPoller + from hindsight_api.extensions.builtin.tenant import DefaultTenantExtension + tenant_extension = DefaultTenantExtension(config={"schema": schema} if schema else {}) + + poller = WorkerPoller( + pool=pool, + worker_id="test_worker_recovery", + executor=memory, + poll_interval_ms=100, + max_retries=3, + schema=schema, + tenant_extension=tenant_extension, + max_slots=5, + consolidation_max_slots=2, + ) + + # Run recovery + recovered_count = await poller._recover_batch_operations(schema) + + # Verify recovery + assert recovered_count == 1, "Should recover 1 batch operation" + + # Verify operation was reset to pending + row = await pool.fetchrow( + f"SELECT status, worker_id FROM {table} WHERE operation_id = $1", + operation_id, + ) + + assert row["status"] == "pending", "Operation should be reset to pending" + assert row["worker_id"] is None, "Worker ID should be cleared" + + logger.info("✅ Worker batch recovery test passed") + + finally: + # Cleanup + try: + await memory.delete_bank(bank_id, request_context=request_context) + except Exception: + pass + + +@pytest.mark.asyncio +async def test_batch_api_via_extract_facts_from_contents( + mock_llm_config, test_contents, hindsight_config, memory, request_context +): + """Test that extract_facts_from_contents routes to batch API when enabled.""" + bank_id = f"test_routing_{datetime.now(timezone.utc).timestamp()}" + + try: + # Enable batch API in config + hindsight_config.retain_batch_enabled = True + + # Mock batch API support + mock_llm_config._provider_impl.supports_batch_api = AsyncMock(return_value=True) + mock_llm_config._provider_impl.submit_batch = AsyncMock( + return_value={"batch_id": "batch_123", "status": "validating", "request_counts": {}} + ) + mock_llm_config._provider_impl.get_batch_status = AsyncMock( + return_value={"status": "completed", "request_counts": {"total": 1, "completed": 1, "failed": 0}} + ) + mock_llm_config._provider_impl.retrieve_batch_results = AsyncMock( + return_value=[ + { + "custom_id": "chunk_0", + "response": { + "body": { + "choices": [ + { + "message": { + "content": json.dumps({"facts": []}) + } + } + ], + "usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}, + } + }, + } + ] + ) + + # Call main extract_facts_from_contents (should route to batch API) + facts, chunks, usage = await extract_facts_from_contents( + contents=test_contents, + llm_config=mock_llm_config, + agent_name="test_agent", + config=hindsight_config, + pool=None, + operation_id=None, + schema=None, + ) + + # Verify batch API was called + mock_llm_config._provider_impl.submit_batch.assert_called_once() + + logger.info("✅ Routing to batch API test passed") + + finally: + # Cleanup + try: + await memory.delete_bank(bank_id, request_context=request_context) + except Exception: + pass diff --git a/hindsight-api/tests/test_batch_api_integration.py b/hindsight-api/tests/test_batch_api_integration.py new file mode 100644 index 00000000..99d8d1b1 --- /dev/null +++ b/hindsight-api/tests/test_batch_api_integration.py @@ -0,0 +1,263 @@ +""" +Real integration test for OpenAI Batch API. + +This test makes REAL API calls to OpenAI and measures actual timing. +It will be slow (minutes to hours) depending on OpenAI's queue. + +To run: + pytest tests/test_batch_api_integration.py -v -s + +To skip in CI: + Add @pytest.mark.skip at the test level +""" +import pytest +import os +import asyncio +import logging +import time +from datetime import datetime, timezone +from dotenv import load_dotenv +from hindsight_api import RequestContext +from hindsight_api.engine.retain.fact_extraction import ( + extract_facts_from_contents_batch_api, + RetainContent, +) +from hindsight_api.config import HindsightConfig +from hindsight_api.engine.llm_wrapper import LLMProvider + +logger = logging.getLogger(__name__) + +# Load .env file for API keys +load_dotenv() + + +@pytest.fixture +def openai_api_key(): + """Get OpenAI API key from environment.""" + # Try both current and commented keys from .env + api_key = os.getenv("HINDSIGHT_API_LLM_API_KEY") + + # Check if it's an OpenAI key (starts with sk-proj- or sk-) + if not api_key or not api_key.startswith("sk-"): + # Try the OpenAI-specific env var (if set separately) + api_key = os.getenv("OPENAI_API_KEY") + + if not api_key or not api_key.startswith("sk-"): + pytest.skip("OpenAI API key not found in environment. Set OPENAI_API_KEY or uncomment OpenAI config in .env") + + return api_key + + +@pytest.fixture +def real_llm_config(openai_api_key): + """Create real LLM config for OpenAI.""" + # Create config with OpenAI settings + config = HindsightConfig.from_env() + + # Use LLMProvider wrapper (which creates _provider_impl internally) + llm_config = LLMProvider( + provider="openai", + api_key=openai_api_key, + base_url="https://api.openai.com/v1", + model="gpt-4o-mini", # Fast, cheap model for testing + reasoning_effort="medium", # Required parameter + ) + + return llm_config + + +@pytest.fixture +def test_contents_real(): + """Create realistic test content for fact extraction.""" + return [ + RetainContent( + content=""" + Alice is a senior software engineer at TechCorp, where she has been working for 5 years. + She specializes in distributed systems and microservices architecture. Alice graduated + from MIT with a degree in Computer Science in 2015. She is known for writing clean, + well-documented code and mentoring junior developers. + """, + event_date=datetime(2024, 1, 15, 10, 30, tzinfo=timezone.utc), + context="team member profile", + ), + RetainContent( + content=""" + Bob joined TechCorp last month as a junior developer. He is learning React and Node.js + and recently completed his first feature, which was a user authentication flow. Bob + graduated from Berkeley with a degree in Computer Science in 2023. He is enthusiastic + and asks great questions during code reviews. + """, + event_date=datetime(2024, 1, 15, 10, 30, tzinfo=timezone.utc), + context="team member profile", + ), + RetainContent( + content=""" + The team uses Kubernetes for container orchestration and deploys to AWS. They follow + agile methodologies with two-week sprints. Code reviews are mandatory before merging + any pull request. The team meets every morning for a 15-minute standup to discuss + progress and blockers. + """, + event_date=datetime(2024, 1, 15, 10, 30, tzinfo=timezone.utc), + context="team processes", + ), + ] + + +@pytest.fixture +def integration_config(): + """Create config for integration test.""" + config = HindsightConfig.from_env() + config.retain_batch_enabled = True + config.retain_batch_poll_interval_seconds = 30 # Poll every 30 seconds (reasonable for real API) + config.retain_chunk_size = 4000 + config.retain_extraction_mode = "concise" + config.retain_extract_causal_links = False + return config + + +@pytest.mark.skip(reason="Real API test - takes minutes and costs money. Run manually with: pytest tests/test_batch_api_integration.py::test_real_openai_batch_api -v -s") +@pytest.mark.integration # Mark as integration test +@pytest.mark.slow # Mark as slow test +@pytest.mark.asyncio +async def test_real_openai_batch_api(real_llm_config, test_contents_real, integration_config, memory, request_context): + """ + REAL integration test: Submit actual batch to OpenAI and measure timing. + + WARNING: This test: + - Makes real API calls to OpenAI + - Will take minutes to hours to complete + - Costs money (though very little with gpt-4o-mini) + - Requires valid OpenAI API key + + To skip this test: + pytest tests/test_batch_api_integration.py --skip-integration + """ + bank_id = f"test_real_batch_{datetime.now(timezone.utc).timestamp()}" + + logger.info("=" * 80) + logger.info("STARTING REAL OPENAI BATCH API INTEGRATION TEST") + logger.info("=" * 80) + logger.info(f"Test contents: {len(test_contents_real)} items") + logger.info(f"Poll interval: {integration_config.retain_batch_poll_interval_seconds}s") + logger.info(f"Model: {real_llm_config.model}") + logger.info("This may take several minutes to hours depending on OpenAI's queue...") + logger.info("=" * 80) + + try: + # Ensure bank exists + await memory.get_bank_profile(bank_id, request_context=request_context) + + # Get database pool and schema for crash recovery testing + pool = memory._pool + schema = request_context.tenant_id + + # Track overall timing + test_start_time = time.time() + + # Call REAL batch API extraction + logger.info("\n📤 Submitting batch to OpenAI...") + + facts, chunks, usage = await extract_facts_from_contents_batch_api( + contents=test_contents_real, + llm_config=real_llm_config, + agent_name="test_agent", + config=integration_config, + pool=pool, + operation_id=None, # No crash recovery for this test + schema=schema, + ) + + test_end_time = time.time() + total_duration = test_end_time - test_start_time + + # Log results + logger.info("\n" + "=" * 80) + logger.info("✅ BATCH COMPLETED SUCCESSFULLY") + logger.info("=" * 80) + logger.info(f"Total duration: {total_duration:.1f} seconds ({total_duration/60:.1f} minutes)") + logger.info(f"Facts extracted: {len(facts)}") + logger.info(f"Chunks processed: {len(chunks)}") + logger.info(f"Token usage: {usage.input_tokens} input + {usage.output_tokens} output = {usage.total_tokens} total") + logger.info(f"Estimated cost: ${(usage.input_tokens * 0.00015 / 1000 + usage.output_tokens * 0.0006 / 1000):.4f}") + logger.info("=" * 80) + + # Log sample facts + logger.info("\n📋 Sample extracted facts:") + for i, fact in enumerate(facts[:5]): # Show first 5 facts + logger.info(f"\nFact {i+1}:") + logger.info(f" Type: {fact.fact_type}") + logger.info(f" Text: {fact.fact_text[:100]}...") + logger.info(f" Entities: {fact.entities}") + + # Verify results + assert len(facts) > 0, "Should extract at least some facts" + assert len(chunks) == len(test_contents_real), f"Should have {len(test_contents_real)} chunks" + assert usage.total_tokens > 0, "Should have token usage" + + # Verify fact structure + for fact in facts: + assert hasattr(fact, "fact_text"), "Fact should have fact_text" + assert hasattr(fact, "fact_type"), "Fact should have fact_type" + assert fact.fact_type in ["world", "experience", "opinion"], f"Invalid fact_type: {fact.fact_type}" + + logger.info("\n✅ All assertions passed!") + + # Write timing report to file for later analysis + report_path = "/tmp/openai_batch_api_timing_report.txt" + with open(report_path, "w") as f: + f.write(f"OpenAI Batch API Integration Test Report\n") + f.write(f"={'=' * 60}\n\n") + f.write(f"Test Date: {datetime.now(timezone.utc).isoformat()}\n") + f.write(f"Model: {real_llm_config.model}\n") + f.write(f"Contents: {len(test_contents_real)} items\n") + f.write(f"Poll Interval: {integration_config.retain_batch_poll_interval_seconds}s\n\n") + f.write(f"Results:\n") + f.write(f" Total Duration: {total_duration:.1f}s ({total_duration/60:.1f} min)\n") + f.write(f" Facts Extracted: {len(facts)}\n") + f.write(f" Chunks Processed: {len(chunks)}\n") + f.write(f" Token Usage: {usage.total_tokens} ({usage.input_tokens} in + {usage.output_tokens} out)\n") + f.write(f" Estimated Cost: ${(usage.input_tokens * 0.00015 / 1000 + usage.output_tokens * 0.0006 / 1000):.4f}\n") + + logger.info(f"\n📄 Timing report written to: {report_path}") + + finally: + # Cleanup + try: + await memory.delete_bank(bank_id, request_context=request_context) + logger.info(f"\n🧹 Cleaned up test bank: {bank_id}") + except Exception as e: + logger.error(f"Failed to cleanup bank: {e}") + + +@pytest.mark.skip(reason="Real API test - requires Groq API key. Run manually if needed.") +@pytest.mark.integration +@pytest.mark.slow +@pytest.mark.asyncio +async def test_real_batch_supports_groq(integration_config): + """ + Test that Groq also supports batch API (if configured). + + Groq has the same batch API interface as OpenAI. + """ + groq_api_key = os.getenv("HINDSIGHT_API_LLM_API_KEY") + + if not groq_api_key or not groq_api_key.startswith("gsk_"): + pytest.skip("Groq API key not found in environment") + + llm_config = LLMProvider( + provider="groq", + api_key=groq_api_key, + base_url="https://api.groq.com/openai/v1", + model="llama-3.1-8b-instant", + reasoning_effort="medium", + ) + + # Check if Groq supports batch API + supports_batch = await llm_config._provider_impl.supports_batch_api() + + logger.info(f"Groq batch API support: {supports_batch}") + + # Groq should support batch API (same interface as OpenAI) + assert supports_batch, "Groq should support batch API" + + logger.info("✅ Groq batch API support confirmed") diff --git a/hindsight-api/tests/test_batch_api_validation.py b/hindsight-api/tests/test_batch_api_validation.py new file mode 100644 index 00000000..66eb8468 --- /dev/null +++ b/hindsight-api/tests/test_batch_api_validation.py @@ -0,0 +1,38 @@ +""" +Test validation for batch API + synchronous retain. + +When HINDSIGHT_API_RETAIN_BATCH_ENABLED=true, synchronous retain operations +should be rejected with a 400 error since they will timeout. +""" + +import os +import pytest +from hindsight_api.engine.memory_engine import MemoryEngine +from hindsight_api.config import HindsightConfig +from hindsight_api import RequestContext + + +@pytest.mark.asyncio +async def test_batch_api_validation(memory, request_context): + """ + Test that attempting synchronous retain with batch API enabled + raises an error at the HTTP layer. + + This test verifies the validation logic exists - actual HTTP testing + would require full FastAPI app setup. + """ + # Create config with batch API enabled + config = HindsightConfig.from_env() + config.retain_batch_enabled = True + config.retain_batch_poll_interval_seconds = 1 + + # Verify the validation exists in memory engine + # The actual HTTP validation happens in http.py api_retain() + # This test documents the expected behavior + + assert config.retain_batch_enabled is True + assert config.retain_batch_poll_interval_seconds == 1 + + # When batch API is enabled and async=false, the HTTP endpoint + # should return 400 with message: + # "Batch API is enabled (HINDSIGHT_API_RETAIN_BATCH_ENABLED=true) but async=false" diff --git a/hindsight-clients/typescript/src/index.ts b/hindsight-clients/typescript/src/index.ts index dc080b61..653c97c3 100644 --- a/hindsight-clients/typescript/src/index.ts +++ b/hindsight-clients/typescript/src/index.ts @@ -50,6 +50,22 @@ export interface HindsightClientOptions { apiKey?: string; } +/** + * Error thrown by the Hindsight client when an API request fails. + * Includes the HTTP status code and error details from the API. + */ +export class HindsightError extends Error { + public statusCode?: number; + public details?: unknown; + + constructor(message: string, statusCode?: number, details?: unknown) { + super(message); + this.name = 'HindsightError'; + this.statusCode = statusCode; + this.details = details; + } +} + export interface EntityInput { text: string; type?: string; @@ -82,9 +98,22 @@ export class HindsightClient { /** * Validates the API response and throws an error if the request failed. */ - private validateResponse(response: { data?: T; error?: unknown }, operation: string): T { + private validateResponse(response: { data?: T; error?: unknown; response?: Response }, operation: string): T { if (!response.data) { - throw new Error(`${operation} failed: ${JSON.stringify(response.error || 'Unknown error')}`); + // The generated client returns { error, response, request } + // Status code is in response.status, not in the error object + const error = response.error as any; + const httpResponse = (response as any).response as Response | undefined; + + // Extract status code from the HTTP response object + const statusCode = httpResponse?.status; + const details = error?.detail || error?.message || error; + + throw new HindsightError( + `${operation} failed: ${JSON.stringify(details)}`, + statusCode, + details + ); } return response.data; } diff --git a/hindsight-control-plane/package.json b/hindsight-control-plane/package.json index dff4fd4e..a56ff0ae 100644 --- a/hindsight-control-plane/package.json +++ b/hindsight-control-plane/package.json @@ -56,6 +56,7 @@ "eslint-config-next": "^16.0.1", "lucide-react": "^0.553.0", "next": "^16.1.6", + "next-themes": "^0.4.6", "postcss": "^8.5.6", "react": "^19.2.0", "react-chrono": "^2.9.1", @@ -64,6 +65,7 @@ "react18-json-view": "^0.2.9", "recharts": "^3.5.1", "remark-gfm": "^4.0.1", + "sonner": "^2.0.7", "tailwind-merge": "^3.4.0", "tailwindcss": "^4.1.17", "tailwindcss-animate": "^1.0.7", diff --git a/hindsight-control-plane/src/app/api/memories/retain/route.ts b/hindsight-control-plane/src/app/api/memories/retain/route.ts index 68f1febf..86a27ea2 100644 --- a/hindsight-control-plane/src/app/api/memories/retain/route.ts +++ b/hindsight-control-plane/src/app/api/memories/retain/route.ts @@ -18,8 +18,22 @@ export async function POST(request: NextRequest) { }); return NextResponse.json(response, { status: 200 }); - } catch (error) { + } catch (error: any) { console.error("Error batch retain:", error); - return NextResponse.json({ error: "Failed to batch retain" }, { status: 500 }); + + const errorMessage = error?.message || String(error); + const errorDetails = error?.details; + const statusCode = error?.statusCode; + + // If we have a statusCode, use it + if (statusCode && typeof statusCode === "number") { + return NextResponse.json( + { error: errorMessage, details: errorDetails }, + { status: statusCode } + ); + } + + // Otherwise, return generic 500 error + return NextResponse.json({ error: errorMessage || "Failed to batch retain" }, { status: 500 }); } } diff --git a/hindsight-control-plane/src/app/banks/[bankId]/page.tsx b/hindsight-control-plane/src/app/banks/[bankId]/page.tsx index fa1e7d8a..ea2c8cb3 100644 --- a/hindsight-control-plane/src/app/banks/[bankId]/page.tsx +++ b/hindsight-control-plane/src/app/banks/[bankId]/page.tsx @@ -2,6 +2,7 @@ import { useState } from "react"; import { useParams, useRouter, useSearchParams } from "next/navigation"; +import { toast } from "sonner"; import { BankSelector } from "@/components/bank-selector"; import { Sidebar } from "@/components/sidebar"; import { DataView } from "@/components/data-view"; @@ -84,8 +85,7 @@ export default function BankPage() { await loadBanks(); router.push("/"); } catch (error) { - console.error("Error deleting bank:", error); - alert("Error deleting bank: " + (error as Error).message); + // Error toast is shown automatically by the API client interceptor } finally { setIsDeleting(false); } @@ -98,10 +98,11 @@ export default function BankPage() { try { const result = await client.clearObservations(bankId); setShowClearObservationsDialog(false); - alert(result.message || "Observations cleared successfully"); + toast.success("Success", { + description: result.message || "Observations cleared successfully", + }); } catch (error) { - console.error("Error clearing observations:", error); - alert("Error clearing observations: " + (error as Error).message); + // Error toast is shown automatically by the API client interceptor } finally { setIsClearingObservations(false); } @@ -114,8 +115,7 @@ export default function BankPage() { try { await client.triggerConsolidation(bankId); } catch (error) { - console.error("Error triggering consolidation:", error); - alert("Error triggering consolidation: " + (error as Error).message); + // Error toast is shown automatically by the API client interceptor } finally { setIsConsolidating(false); } diff --git a/hindsight-control-plane/src/app/layout.tsx b/hindsight-control-plane/src/app/layout.tsx index 0e68d89d..69a6065b 100644 --- a/hindsight-control-plane/src/app/layout.tsx +++ b/hindsight-control-plane/src/app/layout.tsx @@ -3,6 +3,7 @@ import "./globals.css"; import { BankProvider } from "@/lib/bank-context"; import { FeaturesProvider } from "@/lib/features-context"; import { ThemeProvider } from "@/lib/theme-context"; +import { Toaster } from "@/components/ui/sonner"; export const metadata: Metadata = { title: "Hindsight Control Plane", @@ -25,6 +26,7 @@ export default function RootLayout({ {children} + ); diff --git a/hindsight-control-plane/src/components/add-memory-view.tsx b/hindsight-control-plane/src/components/add-memory-view.tsx index f10d6630..eab1a18e 100644 --- a/hindsight-control-plane/src/components/add-memory-view.tsx +++ b/hindsight-control-plane/src/components/add-memory-view.tsx @@ -8,6 +8,7 @@ import { Input } from "@/components/ui/input"; import { Textarea } from "@/components/ui/textarea"; import { Checkbox } from "@/components/ui/checkbox"; import { Tag } from "lucide-react"; +import { toast } from "sonner"; export function AddMemoryView() { const { currentBank } = useBank(); @@ -18,7 +19,6 @@ export function AddMemoryView() { const [tags, setTags] = useState(""); const [async, setAsync] = useState(false); const [loading, setLoading] = useState(false); - const [result, setResult] = useState(null); const clearForm = () => { setContent(""); @@ -27,17 +27,17 @@ export function AddMemoryView() { setDocumentId(""); setTags(""); setAsync(false); - setResult(null); }; const submitMemory = async () => { if (!currentBank || !content) { - alert("Please enter content"); + toast.error("Validation error", { + description: "Please enter content", + }); return; } setLoading(true); - setResult(null); try { // Parse tags from comma-separated string @@ -60,11 +60,18 @@ export function AddMemoryView() { ...(parsedTags.length > 0 && { document_tags: parsedTags }), }); - setResult(data.message as string); + // Show success toast + toast.success("Memory retained", { + description: data.message || "Memory has been successfully added to the bank", + }); + + // Clear form on success setContent(""); + setContext(""); + setTags(""); } catch (error) { - console.error("Error submitting memory:", error); - setResult("Error: " + (error as Error).message); + // Error toast is shown automatically by the API client interceptor + // No need to handle it here! } finally { setLoading(false); } @@ -162,14 +169,6 @@ export function AddMemoryView() { - {result && ( -
-
{result}
-
- )} - {loading && (
⏳
diff --git a/hindsight-control-plane/src/components/bank-config-view.tsx b/hindsight-control-plane/src/components/bank-config-view.tsx index 05873830..d55bb806 100644 --- a/hindsight-control-plane/src/components/bank-config-view.tsx +++ b/hindsight-control-plane/src/components/bank-config-view.tsx @@ -127,8 +127,7 @@ export function BankConfigView() { await loadConfig(); setShowResetDialog(false); } catch (err: any) { - console.error("Failed to reset config:", err); - alert("Error resetting config: " + err.message); + // Error toast is shown automatically by the API client interceptor } finally { setResetting(false); } diff --git a/hindsight-control-plane/src/components/bank-operations-view.tsx b/hindsight-control-plane/src/components/bank-operations-view.tsx index ffd20da7..73f6b46c 100644 --- a/hindsight-control-plane/src/components/bank-operations-view.tsx +++ b/hindsight-control-plane/src/components/bank-operations-view.tsx @@ -125,8 +125,7 @@ export function BankOperationsView() { await client.cancelOperation(currentBank, operationId); await loadOperations(); } catch (error) { - console.error("Error cancelling operation:", error); - alert("Error cancelling operation: " + (error as Error).message); + // Error toast is shown automatically by the API client interceptor } finally { setCancellingOpId(null); } diff --git a/hindsight-control-plane/src/components/bank-profile-view.tsx b/hindsight-control-plane/src/components/bank-profile-view.tsx index 08739ba5..13aeb75d 100644 --- a/hindsight-control-plane/src/components/bank-profile-view.tsx +++ b/hindsight-control-plane/src/components/bank-profile-view.tsx @@ -8,6 +8,7 @@ import { client } from "@/lib/api"; import { useBank } from "@/lib/bank-context"; import { useFeatures } from "@/lib/features-context"; import { Button } from "@/components/ui/button"; +import { toast } from "sonner"; import { Input } from "@/components/ui/input"; import { Textarea } from "@/components/ui/textarea"; import { Card, CardContent, CardDescription, CardHeader, CardTitle } from "@/components/ui/card"; @@ -246,8 +247,7 @@ export function BankProfileView() { setMentalModelsCount(mentalModelsData.items?.length || 0); await loadOperations(); } catch (error) { - console.error("Error loading bank profile:", error); - alert("Error loading bank profile: " + (error as Error).message); + // Error toast is shown automatically by the API client interceptor } finally { setLoading(false); } @@ -264,8 +264,7 @@ export function BankProfileView() { await loadBanks(); router.push("/"); } catch (error) { - console.error("Error deleting bank:", error); - alert("Error deleting bank: " + (error as Error).message); + // Error toast is shown automatically by the API client interceptor } finally { setIsDeleting(false); } @@ -279,10 +278,11 @@ export function BankProfileView() { const result = await client.clearObservations(currentBank); setShowClearObservationsDialog(false); await loadData(); - alert(result.message || "Observations cleared successfully"); + toast.success("Success", { + description: result.message || "Observations cleared successfully", + }); } catch (error) { - console.error("Error clearing observations:", error); - alert("Error clearing observations: " + (error as Error).message); + // Error toast is shown automatically by the API client interceptor } finally { setIsClearingObservations(false); } @@ -298,8 +298,7 @@ export function BankProfileView() { await loadData(); await loadOperations(); } catch (error) { - console.error("Error triggering consolidation:", error); - alert("Error triggering consolidation: " + (error as Error).message); + // Error toast is shown automatically by the API client interceptor } finally { setIsConsolidating(false); } @@ -324,8 +323,7 @@ export function BankProfileView() { await client.cancelOperation(currentBank, operationId); await loadOperations(); } catch (error) { - console.error("Error cancelling operation:", error); - alert("Error cancelling operation: " + (error as Error).message); + // Error toast is shown automatically by the API client interceptor } finally { setCancellingOpId(null); } @@ -341,8 +339,7 @@ export function BankProfileView() { if (selectedDirective?.id === directiveDeleteTarget.id) setSelectedDirective(null); setDirectiveDeleteTarget(null); } catch (error) { - console.error("Error deleting directive:", error); - alert("Error deleting: " + (error as Error).message); + // Error toast is shown automatically by the API client interceptor } finally { setDeletingDirective(false); } @@ -752,8 +749,7 @@ function DispositionEditDialog({ }); onSaved(); } catch (error) { - console.error("Error saving disposition:", error); - alert("Error saving disposition: " + (error as Error).message); + // Error toast is shown automatically by the API client interceptor } finally { setSaving(false); } @@ -843,8 +839,7 @@ function MissionEditDialog({ }); onSaved(); } catch (error) { - console.error("Error saving mission:", error); - alert("Error saving mission: " + (error as Error).message); + // Error toast is shown automatically by the API client interceptor } finally { setSaving(false); } @@ -952,8 +947,7 @@ function DirectiveFormDialog({ onClose(); } } catch (error) { - console.error(`Error ${mode === "create" ? "creating" : "updating"} directive:`, error); - alert(`Error ${mode === "create" ? "creating" : "updating"}: ` + (error as Error).message); + // Error toast is shown automatically by the API client interceptor } finally { setSubmitting(false); } diff --git a/hindsight-control-plane/src/components/bank-selector.tsx b/hindsight-control-plane/src/components/bank-selector.tsx index 08e54efa..103db1dc 100644 --- a/hindsight-control-plane/src/components/bank-selector.tsx +++ b/hindsight-control-plane/src/components/bank-selector.tsx @@ -79,7 +79,6 @@ function BankSelectorInner() { const [docTags, setDocTags] = React.useState(""); const [docAsync, setDocAsync] = React.useState(false); const [isCreatingDoc, setIsCreatingDoc] = React.useState(false); - const [docError, setDocError] = React.useState(null); // File upload state const [selectedFiles, setSelectedFiles] = React.useState([]); @@ -142,7 +141,6 @@ function BankSelectorInner() { if (!currentBank || selectedFiles.length === 0) return; setIsCreatingDoc(true); - setDocError(null); setUploadProgress(""); try { @@ -189,7 +187,7 @@ function BankSelectorInner() { // Navigate to documents view router.push(`/banks/${currentBank}?view=documents`); } catch (error) { - setDocError(error instanceof Error ? error.message : "Failed to upload files"); + // Error toast is shown automatically by the API client interceptor } finally { setIsCreatingDoc(false); setUploadProgress(""); @@ -200,7 +198,6 @@ function BankSelectorInner() { if (!currentBank || !docContent.trim()) return; setIsCreatingDoc(true); - setDocError(null); try { // Parse tags from comma-separated string @@ -241,7 +238,7 @@ function BankSelectorInner() { // Navigate to documents view to see the new document router.push(`/banks/${currentBank}?view=documents`); } catch (error) { - setDocError(error instanceof Error ? error.message : "Failed to create document"); + // Error toast is shown automatically by the API client interceptor } finally { setIsCreatingDoc(false); } @@ -597,8 +594,6 @@ function BankSelectorInner() { - {docError &&

{docError}

} -