From 8e39cb7bc827a4548e119c4a1f01fdef5e224914 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Nicol=C3=B2=20Boschi?= Date: Mon, 26 Jan 2026 09:54:25 +0100 Subject: [PATCH] fix: improve mental model consolidation (#197) * fix: improve mental model consolidation * fix skill names * fixes * fix: add missing list_tenants to test mocks and update CLI for async refresh --- ...8r9_add_consolidated_at_to_memory_units.py | 53 ++++ hindsight-api/hindsight_api/api/http.py | 74 ++--- hindsight-api/hindsight_api/config.py | 18 ++ .../engine/consolidation/consolidator.py | 294 ++++++++---------- .../hindsight_api/engine/llm_wrapper.py | 32 +- .../hindsight_api/engine/memory_engine.py | 209 ++++++++++++- .../hindsight_api/engine/reflect/agent.py | 156 ++++++++-- .../hindsight_api/engine/reflect/models.py | 13 + .../hindsight_api/engine/response_models.py | 2 + .../hindsight_api/extensions/__init__.py | 11 +- .../extensions/builtin/tenant.py | 6 +- .../extensions/operation_validator.py | 85 ++++- .../hindsight_api/extensions/tenant.py | 26 ++ hindsight-api/hindsight_api/main.py | 4 + hindsight-api/hindsight_api/worker/main.py | 13 +- hindsight-api/hindsight_api/worker/poller.py | 255 +++++++++------ .../tests/test_custom_embedding_dimension.py | 17 +- hindsight-api/tests/test_extensions.py | 19 +- .../tests/test_http_api_integration.py | 24 +- hindsight-api/tests/test_main_module.py | 15 +- hindsight-api/tests/test_reflect_agent.py | 51 ++- hindsight-api/tests/test_schema_isolation.py | 7 +- hindsight-api/tests/test_server_module.py | 15 +- hindsight-api/tests/test_worker.py | 210 ++++++++++++- hindsight-cli/src/api.rs | 2 +- hindsight-cli/src/commands/reflection.rs | 14 +- .../python/.openapi-generator/FILES | 1 + .../python/hindsight_client_api/__init__.py | 1 + .../api/reflections_api.py | 17 +- .../hindsight_client_api/models/__init__.py | 1 + .../models/async_operation_submit_response.py | 89 ++++++ .../models/consolidation_response.py | 20 +- .../typescript/generated/sdk.gen.ts | 2 +- .../typescript/generated/types.gen.ts | 48 ++- .../src/app/banks/[bankId]/page.tsx | 43 ++- hindsight-control-plane/src/app/layout.tsx | 5 +- .../src/components/bank-profile-view.tsx | 49 ++- .../src/components/data-view.tsx | 114 ++++--- .../src/components/memory-detail-panel.tsx | 69 +++- .../src/components/reflections-view.tsx | 39 ++- hindsight-control-plane/src/lib/api.ts | 23 +- .../src/lib/features-context.tsx | 63 ++++ hindsight-docs/static/openapi.json | 63 ++-- 43 files changed, 1719 insertions(+), 553 deletions(-) create mode 100644 hindsight-api/hindsight_api/alembic/versions/s4n5o6p7q8r9_add_consolidated_at_to_memory_units.py create mode 100644 hindsight-clients/python/hindsight_client_api/models/async_operation_submit_response.py create mode 100644 hindsight-control-plane/src/lib/features-context.tsx diff --git a/hindsight-api/hindsight_api/alembic/versions/s4n5o6p7q8r9_add_consolidated_at_to_memory_units.py b/hindsight-api/hindsight_api/alembic/versions/s4n5o6p7q8r9_add_consolidated_at_to_memory_units.py new file mode 100644 index 00000000..98b78239 --- /dev/null +++ b/hindsight-api/hindsight_api/alembic/versions/s4n5o6p7q8r9_add_consolidated_at_to_memory_units.py @@ -0,0 +1,53 @@ +"""Add consolidated_at column to memory_units for incremental consolidation tracking. + +This allows consolidation to track progress at the memory level rather than +using a bank-level watermark. If consolidation crashes, already-processed +memories won't be reprocessed. + +Revision ID: s4n5o6p7q8r9 +Revises: r3m4n5o6p7q8 +Create Date: 2025-01-22 +""" + +from collections.abc import Sequence + +from alembic import context, op + +revision: str = "s4n5o6p7q8r9" +down_revision: str | Sequence[str] | None = "r3m4n5o6p7q8" +branch_labels: str | Sequence[str] | None = None +depends_on: str | Sequence[str] | None = None + + +def _get_schema_prefix() -> str: + """Get schema prefix for table names (required for multi-tenant support).""" + schema = context.config.get_main_option("target_schema") + return f'"{schema}".' if schema else "" + + +def upgrade() -> None: + schema = _get_schema_prefix() + + # Add consolidated_at column to memory_units + op.execute( + f""" + ALTER TABLE {schema}memory_units + ADD COLUMN IF NOT EXISTS consolidated_at TIMESTAMPTZ DEFAULT NULL + """ + ) + + # Create index for efficient querying of unconsolidated memories + op.execute( + f""" + CREATE INDEX IF NOT EXISTS idx_memory_units_unconsolidated + ON {schema}memory_units (bank_id, created_at) + WHERE consolidated_at IS NULL AND fact_type IN ('experience', 'world') + """ + ) + + +def downgrade() -> None: + schema = _get_schema_prefix() + + op.execute(f"DROP INDEX IF EXISTS {schema}idx_memory_units_unconsolidated") + op.execute(f"ALTER TABLE {schema}memory_units DROP COLUMN IF EXISTS consolidated_at") diff --git a/hindsight-api/hindsight_api/api/http.py b/hindsight-api/hindsight_api/api/http.py index 42dde613..02b14765 100644 --- a/hindsight-api/hindsight_api/api/http.py +++ b/hindsight-api/hindsight_api/api/http.py @@ -1191,11 +1191,8 @@ class OperationResponse(BaseModel): class ConsolidationResponse(BaseModel): """Response model for consolidation trigger endpoint.""" - status: str = Field(description="Status of the consolidation (completed or queued)") - processed: int = Field(description="Number of memories processed") - created: int = Field(description="Number of mental models created") - updated: int = Field(description="Number of mental models updated") - message: str = Field(description="Human-readable summary") + operation_id: str = Field(description="ID of the async consolidation operation") + deduplicated: bool = Field(default=False, description="True if an existing pending task was reused") class OperationsListResponse(BaseModel): @@ -2058,42 +2055,19 @@ def _register_routes(app: FastAPI): ) total_documents = doc_count_result["count"] if doc_count_result else 0 - # Get consolidation stats - bank_row = await conn.fetchrow( + # Get consolidation stats from memory-level tracking + consolidation_stats = await conn.fetchrow( f""" - SELECT last_consolidated_at - FROM {fq_table("banks")} + SELECT + MAX(consolidated_at) as last_consolidated_at, + COUNT(*) FILTER (WHERE consolidated_at IS NULL AND fact_type IN ('experience', 'world')) as pending + FROM {fq_table("memory_units")} WHERE bank_id = $1 """, bank_id, ) - last_consolidated_at = bank_row["last_consolidated_at"] if bank_row else None - - # Count memories pending consolidation (created after last_consolidated_at) - if last_consolidated_at: - pending_consolidation_result = await conn.fetchrow( - f""" - SELECT COUNT(*) as count - FROM {fq_table("memory_units")} - WHERE bank_id = $1 - AND created_at > $2 - AND fact_type IN ('experience', 'world') - """, - bank_id, - last_consolidated_at, - ) - else: - # If never consolidated, count all experience/world memories - pending_consolidation_result = await conn.fetchrow( - f""" - SELECT COUNT(*) as count - FROM {fq_table("memory_units")} - WHERE bank_id = $1 - AND fact_type IN ('experience', 'world') - """, - bank_id, - ) - pending_consolidation = pending_consolidation_result["count"] if pending_consolidation_result else 0 + last_consolidated_at = consolidation_stats["last_consolidated_at"] if consolidation_stats else None + pending_consolidation = consolidation_stats["pending"] if consolidation_stats else 0 # Count total mental models mental_model_count_result = await conn.fetchrow( @@ -2354,9 +2328,9 @@ def _register_routes(app: FastAPI): @app.post( "/v1/default/banks/{bank_id}/reflections/{reflection_id}/refresh", - response_model=ReflectionResponse, + response_model=AsyncOperationSubmitResponse, summary="Refresh reflection", - description="Re-run the source query through reflect and update the content.", + description="Submit an async task to re-run the source query through reflect and update the content.", operation_id="refresh_reflection", tags=["Reflections"], ) @@ -2365,16 +2339,16 @@ def _register_routes(app: FastAPI): reflection_id: str, request_context: RequestContext = Depends(get_request_context), ): - """Refresh a reflection by re-running its source query.""" + """Refresh a reflection by re-running its source query (async).""" try: - reflection = await app.state.memory.refresh_reflection( + result = await app.state.memory.submit_async_refresh_reflection( bank_id=bank_id, reflection_id=reflection_id, request_context=request_context, ) - if reflection is None: - raise HTTPException(status_code=404, detail=f"Reflection '{reflection_id}' not found") - return ReflectionResponse(**reflection) + return AsyncOperationSubmitResponse(operation_id=result["operation_id"], status="queued") + except ValueError as e: + raise HTTPException(status_code=404, detail=str(e)) except (AuthenticationError, HTTPException): raise except Exception as e: @@ -3206,18 +3180,12 @@ def _register_routes(app: FastAPI): tags=["Banks"], ) async def api_trigger_consolidation(bank_id: str, request_context: RequestContext = Depends(get_request_context)): - """Trigger consolidation for a bank.""" + """Trigger consolidation for a bank (async).""" try: - result = await app.state.memory.run_consolidation(bank_id, request_context=request_context) - processed = result.get("processed", 0) - created = result.get("created", 0) - updated = result.get("updated", 0) + result = await app.state.memory.submit_async_consolidation(bank_id=bank_id, request_context=request_context) return ConsolidationResponse( - status="completed", - processed=processed, - created=created, - updated=updated, - message=f"Consolidation completed: {processed} memories processed, {created} mental models created, {updated} updated", + operation_id=result["operation_id"], + deduplicated=result.get("deduplicated", False), ) except (AuthenticationError, HTTPException): raise diff --git a/hindsight-api/hindsight_api/config.py b/hindsight-api/hindsight_api/config.py index 44df44f2..c7bd9f37 100644 --- a/hindsight-api/hindsight_api/config.py +++ b/hindsight-api/hindsight_api/config.py @@ -39,6 +39,11 @@ ENV_REFLECT_LLM_API_KEY = "HINDSIGHT_API_REFLECT_LLM_API_KEY" ENV_REFLECT_LLM_MODEL = "HINDSIGHT_API_REFLECT_LLM_MODEL" ENV_REFLECT_LLM_BASE_URL = "HINDSIGHT_API_REFLECT_LLM_BASE_URL" +ENV_CONSOLIDATION_LLM_PROVIDER = "HINDSIGHT_API_CONSOLIDATION_LLM_PROVIDER" +ENV_CONSOLIDATION_LLM_API_KEY = "HINDSIGHT_API_CONSOLIDATION_LLM_API_KEY" +ENV_CONSOLIDATION_LLM_MODEL = "HINDSIGHT_API_CONSOLIDATION_LLM_MODEL" +ENV_CONSOLIDATION_LLM_BASE_URL = "HINDSIGHT_API_CONSOLIDATION_LLM_BASE_URL" + ENV_EMBEDDINGS_PROVIDER = "HINDSIGHT_API_EMBEDDINGS_PROVIDER" ENV_EMBEDDINGS_LOCAL_MODEL = "HINDSIGHT_API_EMBEDDINGS_LOCAL_MODEL" ENV_EMBEDDINGS_TEI_URL = "HINDSIGHT_API_EMBEDDINGS_TEI_URL" @@ -293,6 +298,11 @@ class HindsightConfig: reflect_llm_model: str | None reflect_llm_base_url: str | None + consolidation_llm_provider: str | None + consolidation_llm_api_key: str | None + consolidation_llm_model: str | None + consolidation_llm_base_url: str | None + # Embeddings embeddings_provider: str embeddings_local_model: str @@ -385,6 +395,10 @@ class HindsightConfig: reflect_llm_api_key=os.getenv(ENV_REFLECT_LLM_API_KEY) or None, reflect_llm_model=os.getenv(ENV_REFLECT_LLM_MODEL) or None, reflect_llm_base_url=os.getenv(ENV_REFLECT_LLM_BASE_URL) or None, + consolidation_llm_provider=os.getenv(ENV_CONSOLIDATION_LLM_PROVIDER) or None, + consolidation_llm_api_key=os.getenv(ENV_CONSOLIDATION_LLM_API_KEY) or None, + consolidation_llm_model=os.getenv(ENV_CONSOLIDATION_LLM_MODEL) or None, + consolidation_llm_base_url=os.getenv(ENV_CONSOLIDATION_LLM_BASE_URL) or None, # Embeddings embeddings_provider=os.getenv(ENV_EMBEDDINGS_PROVIDER, DEFAULT_EMBEDDINGS_PROVIDER), embeddings_local_model=os.getenv(ENV_EMBEDDINGS_LOCAL_MODEL, DEFAULT_EMBEDDINGS_LOCAL_MODEL), @@ -531,6 +545,10 @@ class HindsightConfig: reflect_provider = self.reflect_llm_provider or self.llm_provider reflect_model = self.reflect_llm_model or self.llm_model logger.info(f"LLM (reflect): provider={reflect_provider}, model={reflect_model}") + if self.consolidation_llm_provider or self.consolidation_llm_model: + consolidation_provider = self.consolidation_llm_provider or self.llm_provider + consolidation_model = self.consolidation_llm_model or self.llm_model + logger.info(f"LLM (consolidation): provider={consolidation_provider}, model={consolidation_model}") logger.info(f"Embeddings: provider={self.embeddings_provider}") logger.info(f"Reranker: provider={self.reranker_provider}") logger.info(f"Graph retriever: {self.graph_retriever}") diff --git a/hindsight-api/hindsight_api/engine/consolidation/consolidator.py b/hindsight-api/hindsight_api/engine/consolidation/consolidator.py index 9b7ec6e3..d80b23f3 100644 --- a/hindsight-api/hindsight_api/engine/consolidation/consolidator.py +++ b/hindsight-api/hindsight_api/engine/consolidation/consolidator.py @@ -93,12 +93,14 @@ async def run_consolidation_job( logger.debug(f"Consolidation disabled for bank {bank_id}") return {"status": "disabled", "bank_id": bank_id} - async with memory_engine._pool.acquire() as conn: - # Get bank profile and last_consolidated_at + pool = memory_engine._pool + + # Get bank profile + async with pool.acquire() as conn: t0 = time.time() bank_row = await conn.fetchrow( f""" - SELECT bank_id, name, mission, last_consolidated_at + SELECT bank_id, name, mission FROM {fq_table("banks")} WHERE bank_id = $1 """, @@ -110,31 +112,51 @@ async def run_consolidation_job( return {"status": "bank_not_found", "bank_id": bank_id} mission = bank_row["mission"] or "General memory consolidation" - last_consolidated_at = bank_row["last_consolidated_at"] perf.record_timing("fetch_bank", time.time() - t0) - # Fetch memories created after last_consolidated_at (exclude mental_model type) - t0 = time.time() - if last_consolidated_at: + # Count total unconsolidated memories for progress logging + total_count = await conn.fetchval( + f""" + SELECT COUNT(*) + FROM {fq_table("memory_units")} + WHERE bank_id = $1 + AND consolidated_at IS NULL + AND fact_type IN ('experience', 'world') + """, + bank_id, + ) + + if total_count == 0: + logger.debug(f"No new memories to consolidate for bank {bank_id}") + return {"status": "no_new_memories", "bank_id": bank_id, "memories_processed": 0} + + logger.info(f"[CONSOLIDATION] bank={bank_id} total_unconsolidated={total_count}") + perf.log(f"[1] Found {total_count} pending memories to consolidate") + + # Process each memory with individual commits for crash recovery + stats = { + "memories_processed": 0, + "mental_models_created": 0, + "mental_models_updated": 0, + "mental_models_merged": 0, + "actions_executed": 0, + "skipped": 0, + } + + batch_num = 0 + while True: + batch_num += 1 + batch_start = time.time() + + # Fetch next batch of unconsolidated memories + async with pool.acquire() as conn: + t0 = time.time() memories = await conn.fetch( f""" - SELECT id, text, fact_type, occurred_start, event_date, tags - FROM {fq_table("memory_units")} - WHERE bank_id = $1 AND created_at > $2 - AND fact_type IN ('experience', 'world') - ORDER BY created_at ASC - LIMIT $3 - """, - bank_id, - last_consolidated_at, - max_memories_per_batch, - ) - else: - memories = await conn.fetch( - f""" - SELECT id, text, fact_type, occurred_start, event_date, tags + SELECT id, text, fact_type, occurred_start, event_date, tags, mentioned_at FROM {fq_table("memory_units")} WHERE bank_id = $1 + AND consolidated_at IS NULL AND fact_type IN ('experience', 'world') ORDER BY created_at ASC LIMIT $2 @@ -142,45 +164,16 @@ async def run_consolidation_job( bank_id, max_memories_per_batch, ) - perf.record_timing("fetch_memories", time.time() - t0) + perf.record_timing("fetch_memories", time.time() - t0) if not memories: - logger.debug(f"No new memories to consolidate for bank {bank_id}") - # Update timestamp anyway to prevent reprocessing - await _update_last_consolidated_at(conn, bank_id) - return {"status": "no_new_memories", "bank_id": bank_id, "memories_processed": 0} + break # No more unconsolidated memories - logger.info( - f"[CONSOLIDATION] bank={bank_id} memories={len(memories)} " - f"batch_size={max_memories_per_batch} since={last_consolidated_at or 'beginning'}" - ) - perf.log(f"[1] Found {len(memories)} pending memories to consolidate") + for memory in memories: + mem_start = time.time() - # Process each memory sequentially - # Important: We process ALL pending memories before updating the watermark - # to avoid losing memories when many have the same timestamp - stats = { - "memories_processed": 0, - "mental_models_created": 0, - "mental_models_updated": 0, - "mental_models_merged": 0, - "actions_executed": 0, # Total actions (can be > memories_processed due to multiple actions per fact) - "skipped": 0, - } - - # Track processed memory IDs to avoid reprocessing - processed_ids: set[uuid.UUID] = set() - batch_num = 0 - - while memories: - batch_num += 1 - batch_start = time.time() - - for memory in memories: - if memory["id"] in processed_ids: - continue - - mem_start = time.time() + # Process the memory (uses its own connection internally) + async with pool.acquire() as conn: result = await _process_memory( conn=conn, memory_engine=memory_engine, @@ -190,118 +183,81 @@ async def run_consolidation_job( request_context=request_context, perf=perf, ) - mem_time = time.time() - mem_start - perf.record_timing("process_memory_total", mem_time) - processed_ids.add(memory["id"]) - stats["memories_processed"] += 1 - - action = result.get("action") - if action == "created": - stats["mental_models_created"] += 1 - stats["actions_executed"] += 1 - elif action == "updated": - stats["mental_models_updated"] += 1 - stats["actions_executed"] += 1 - elif action == "merged": - stats["mental_models_merged"] += 1 - stats["actions_executed"] += 1 - elif action == "multiple": - # Multiple actions from one fact (tag routing) - stats["mental_models_created"] += result.get("created", 0) - stats["mental_models_updated"] += result.get("updated", 0) - stats["mental_models_merged"] += result.get("merged", 0) - stats["actions_executed"] += result.get("total_actions", 0) - elif action == "skipped": - stats["skipped"] += 1 - - batch_time = time.time() - batch_start - perf.log( - f"[2] Batch {batch_num}: {len(memories)} memories in {batch_time:.3f}s " - f"(avg {batch_time / len(memories):.3f}s/memory)" - ) - - # Fetch next batch of memories (excluding already processed) - t0 = time.time() - if last_consolidated_at: - memories = await conn.fetch( + # Mark memory as consolidated (committed immediately) + await conn.execute( f""" - SELECT id, text, fact_type, occurred_start, event_date, tags - FROM {fq_table("memory_units")} - WHERE bank_id = $1 AND created_at > $2 - AND fact_type IN ('experience', 'world') - AND id != ALL($4) - ORDER BY created_at ASC - LIMIT $3 + UPDATE {fq_table("memory_units")} + SET consolidated_at = NOW() + WHERE id = $1 """, - bank_id, - last_consolidated_at, - max_memories_per_batch, - list(processed_ids), + memory["id"], ) - else: - memories = await conn.fetch( - f""" - SELECT id, text, fact_type, occurred_start, event_date, tags - FROM {fq_table("memory_units")} - WHERE bank_id = $1 - AND fact_type IN ('experience', 'world') - AND id != ALL($3) - ORDER BY created_at ASC - LIMIT $2 - """, - bank_id, - max_memories_per_batch, - list(processed_ids), + + mem_time = time.time() - mem_start + perf.record_timing("process_memory_total", mem_time) + + stats["memories_processed"] += 1 + + action = result.get("action") + if action == "created": + stats["mental_models_created"] += 1 + stats["actions_executed"] += 1 + elif action == "updated": + stats["mental_models_updated"] += 1 + stats["actions_executed"] += 1 + elif action == "merged": + stats["mental_models_merged"] += 1 + stats["actions_executed"] += 1 + elif action == "multiple": + stats["mental_models_created"] += result.get("created", 0) + stats["mental_models_updated"] += result.get("updated", 0) + stats["mental_models_merged"] += result.get("merged", 0) + stats["actions_executed"] += result.get("total_actions", 0) + elif action == "skipped": + stats["skipped"] += 1 + + # Log progress periodically + if stats["memories_processed"] % 10 == 0: + logger.info( + f"[CONSOLIDATION] bank={bank_id} progress: " + f"{stats['memories_processed']}/{total_count} memories processed" ) - perf.record_timing("fetch_memories", time.time() - t0) - # Update last_consolidated_at only after ALL memories are processed - t0 = time.time() - await _update_last_consolidated_at(conn, bank_id) - perf.record_timing("update_watermark", time.time() - t0) - - # Build summary + batch_time = time.time() - batch_start perf.log( - f"[3] Results: {stats['memories_processed']} memories → " - f"{stats['actions_executed']} actions " - f"({stats['mental_models_created']} created, " - f"{stats['mental_models_updated']} updated, " - f"{stats['mental_models_merged']} merged, " - f"{stats['skipped']} skipped)" + f"[2] Batch {batch_num}: {len(memories)} memories in {batch_time:.3f}s " + f"(avg {batch_time / len(memories):.3f}s/memory)" ) - # Add timing breakdown - timing_parts = [] - if "recall" in perf.timings: - timing_parts.append(f"recall={perf.timings['recall']:.3f}s") - if "llm" in perf.timings: - timing_parts.append(f"llm={perf.timings['llm']:.3f}s") - if "embedding" in perf.timings: - timing_parts.append(f"embedding={perf.timings['embedding']:.3f}s") - if "db_write" in perf.timings: - timing_parts.append(f"db_write={perf.timings['db_write']:.3f}s") - - if timing_parts: - perf.log(f"[4] Timing breakdown: {', '.join(timing_parts)}") - - perf.flush() - - return {"status": "completed", "bank_id": bank_id, **stats} - - -async def _update_last_consolidated_at(conn: "Connection", bank_id: str) -> None: - """Update the bank's last_consolidated_at timestamp.""" - await conn.execute( - f""" - UPDATE {fq_table("banks")} - SET last_consolidated_at = $1 - WHERE bank_id = $2 - """, - datetime.now(timezone.utc), - bank_id, + # Build summary + perf.log( + f"[3] Results: {stats['memories_processed']} memories -> " + f"{stats['actions_executed']} actions " + f"({stats['mental_models_created']} created, " + f"{stats['mental_models_updated']} updated, " + f"{stats['mental_models_merged']} merged, " + f"{stats['skipped']} skipped)" ) + # Add timing breakdown + timing_parts = [] + if "recall" in perf.timings: + timing_parts.append(f"recall={perf.timings['recall']:.3f}s") + if "llm" in perf.timings: + timing_parts.append(f"llm={perf.timings['llm']:.3f}s") + if "embedding" in perf.timings: + timing_parts.append(f"embedding={perf.timings['embedding']:.3f}s") + if "db_write" in perf.timings: + timing_parts.append(f"db_write={perf.timings['db_write']:.3f}s") + + if timing_parts: + perf.log(f"[4] Timing breakdown: {', '.join(timing_parts)}") + + perf.flush() + + return {"status": "completed", "bank_id": bank_id, **stats} + async def _process_memory( conn: "Connection", @@ -372,6 +328,7 @@ async def _process_memory( memory_id=memory_id, action=action, mental_models=related_mental_models, + source_mentioned_at=memory.get("mentioned_at"), perf=perf, ) results.append(result) @@ -384,6 +341,7 @@ async def _process_memory( action=action, event_date=memory.get("event_date"), occurred_start=memory.get("occurred_start"), + mentioned_at=memory.get("mentioned_at"), perf=perf, ) results.append(result) @@ -416,12 +374,14 @@ async def _execute_update_action( memory_id: uuid.UUID, action: dict[str, Any], mental_models: list[dict[str, Any]], + source_mentioned_at: datetime | None = None, perf: ConsolidationPerfLog | None = None, ) -> dict[str, Any]: """ Execute an update action on an existing mental model. - Updates the mental model text, adds to history, and increments proof_count. + Updates the mental model text, adds to history, increments proof_count, + and updates mentioned_at if the new source memory has a more recent date. """ learning_id = action.get("learning_id") new_text = action.get("text") @@ -458,6 +418,7 @@ async def _execute_update_action( perf.record_timing("embedding", time.time() - t0) # Update the mental model + # Update mentioned_at if source memory has a more recent date t0 = time.time() await conn.execute( f""" @@ -467,7 +428,8 @@ async def _execute_update_action( history = $3, source_memory_ids = $4, proof_count = $5, - updated_at = now() + updated_at = now(), + mentioned_at = GREATEST(mentioned_at, COALESCE($7, mentioned_at)) WHERE id = $6 """, new_text, @@ -476,6 +438,7 @@ async def _execute_update_action( source_ids, len(source_ids), uuid.UUID(learning_id), + source_mentioned_at, ) # Create links from memory to mental model @@ -496,6 +459,7 @@ async def _execute_create_action( action: dict[str, Any], event_date: datetime | None = None, occurred_start: datetime | None = None, + mentioned_at: datetime | None = None, perf: ConsolidationPerfLog | None = None, ) -> dict[str, Any]: """ @@ -520,6 +484,7 @@ async def _execute_create_action( tags=tags, event_date=event_date, occurred_start=occurred_start, + mentioned_at=mentioned_at, perf=perf, ) @@ -751,7 +716,7 @@ Focus on DURABLE knowledge that serves this mission, not ephemeral state. ] try: - result = await memory_engine._llm_config.call( + result = await memory_engine._consolidation_llm_config.call( messages=messages, skip_validation=True, # Raw JSON response scope="consolidation", @@ -790,6 +755,7 @@ async def _create_mental_model_directly( tags: list[str] | None = None, event_date: datetime | None = None, occurred_start: datetime | None = None, + mentioned_at: datetime | None = None, perf: ConsolidationPerfLog | None = None, ) -> dict[str, Any]: """ @@ -809,6 +775,7 @@ async def _create_mental_model_directly( now = datetime.now(timezone.utc) mm_event_date = event_date or now mm_occurred_start = occurred_start or now + mm_mentioned_at = mentioned_at or now mm_tags = tags or [] t0 = time.time() @@ -817,9 +784,9 @@ async def _create_mental_model_directly( f""" INSERT INTO {fq_table("memory_units")} ( id, bank_id, text, fact_type, embedding, proof_count, source_memory_ids, history, - tags, event_date, occurred_start + tags, event_date, occurred_start, mentioned_at ) - VALUES ($1, $2, $3, 'mental_model', $4::vector, 1, $5, '[]'::jsonb, $6, $7, $8) + VALUES ($1, $2, $3, 'mental_model', $4::vector, 1, $5, '[]'::jsonb, $6, $7, $8, $9) RETURNING id """, mental_model_id, @@ -830,6 +797,7 @@ async def _create_mental_model_directly( mm_tags, mm_event_date, mm_occurred_start, + mm_mentioned_at, ) # Create links between memory and mental model (includes entity links, memory_links) diff --git a/hindsight-api/hindsight_api/engine/llm_wrapper.py b/hindsight-api/hindsight_api/engine/llm_wrapper.py index 0eb89add..d89015ba 100644 --- a/hindsight-api/hindsight_api/engine/llm_wrapper.py +++ b/hindsight-api/hindsight_api/engine/llm_wrapper.py @@ -647,7 +647,13 @@ class LLMProvider: success=True, ) - return LLMToolCallResult(content=content, tool_calls=tool_calls, finish_reason=finish_reason) + return LLMToolCallResult( + content=content, + tool_calls=tool_calls, + finish_reason=finish_reason, + input_tokens=input_tokens, + output_tokens=output_tokens, + ) except APIConnectionError as e: last_exception = e @@ -797,6 +803,10 @@ class LLMProvider: content = "".join(content_parts) if content_parts else None finish_reason = "tool_calls" if tool_calls else "stop" + # Extract token usage + input_tokens = response.usage.input_tokens or 0 + output_tokens = response.usage.output_tokens or 0 + # Record metrics metrics = get_metrics_collector() metrics.record_llm_call( @@ -804,12 +814,18 @@ class LLMProvider: model=self.model, scope=scope, duration=time.time() - start_time, - input_tokens=response.usage.input_tokens or 0, - output_tokens=response.usage.output_tokens or 0, + input_tokens=input_tokens, + output_tokens=output_tokens, success=True, ) - return LLMToolCallResult(content=content, tool_calls=tool_calls, finish_reason=finish_reason) + return LLMToolCallResult( + content=content, + tool_calls=tool_calls, + finish_reason=finish_reason, + input_tokens=input_tokens, + output_tokens=output_tokens, + ) except (APIConnectionError, APIStatusError) as e: if isinstance(e, APIStatusError) and e.status_code in (401, 403): @@ -930,7 +946,13 @@ class LLMProvider: success=True, ) - return LLMToolCallResult(content=content, tool_calls=tool_calls, finish_reason=finish_reason) + return LLMToolCallResult( + content=content, + tool_calls=tool_calls, + finish_reason=finish_reason, + input_tokens=input_tokens, + output_tokens=output_tokens, + ) except genai_errors.APIError as e: if e.code in (401, 403): diff --git a/hindsight-api/hindsight_api/engine/memory_engine.py b/hindsight-api/hindsight_api/engine/memory_engine.py index ecf0f6a4..1be8cd6e 100644 --- a/hindsight-api/hindsight_api/engine/memory_engine.py +++ b/hindsight-api/hindsight_api/engine/memory_engine.py @@ -222,6 +222,10 @@ class MemoryEngine(MemoryEngineInterface): reflect_llm_api_key: str | None = None, reflect_llm_model: str | None = None, reflect_llm_base_url: str | None = None, + consolidation_llm_provider: str | None = None, + consolidation_llm_api_key: str | None = None, + consolidation_llm_model: str | None = None, + consolidation_llm_base_url: str | None = None, embeddings: Embeddings | None = None, cross_encoder: CrossEncoderModel | None = None, query_analyzer: QueryAnalyzer | None = None, @@ -257,6 +261,10 @@ class MemoryEngine(MemoryEngineInterface): reflect_llm_api_key: API key for reflect LLM. Falls back to memory_llm_api_key. reflect_llm_model: Model for reflect operations. Falls back to memory_llm_model. reflect_llm_base_url: Base URL for reflect LLM. Falls back to memory_llm_base_url. + consolidation_llm_provider: LLM provider for consolidation operations. Falls back to memory_llm_provider. + consolidation_llm_api_key: API key for consolidation LLM. Falls back to memory_llm_api_key. + consolidation_llm_model: Model for consolidation operations. Falls back to memory_llm_model. + consolidation_llm_base_url: Base URL for consolidation LLM. Falls back to memory_llm_base_url. embeddings: Embeddings implementation. If not provided, created from env vars. cross_encoder: Cross-encoder model. If not provided, created from env vars. query_analyzer: Query analyzer implementation. If not provided, uses DateparserQueryAnalyzer. @@ -398,6 +406,27 @@ class MemoryEngine(MemoryEngineInterface): model=reflect_model, ) + # Consolidation LLM config - for mental model consolidation (can use efficient models) + consolidation_provider = consolidation_llm_provider or config.consolidation_llm_provider or memory_llm_provider + consolidation_api_key = consolidation_llm_api_key or config.consolidation_llm_api_key or memory_llm_api_key + consolidation_model = consolidation_llm_model or config.consolidation_llm_model or memory_llm_model + consolidation_base_url = consolidation_llm_base_url or config.consolidation_llm_base_url or memory_llm_base_url + # Apply provider-specific base URL defaults for consolidation + if consolidation_base_url is None: + if consolidation_provider.lower() == "groq": + consolidation_base_url = "https://api.groq.com/openai/v1" + elif consolidation_provider.lower() == "ollama": + consolidation_base_url = "http://localhost:11434/v1" + else: + consolidation_base_url = "" + + self._consolidation_llm_config = LLMConfig( + provider=consolidation_provider, + api_key=consolidation_api_key, + base_url=consolidation_base_url, + model=consolidation_model, + ) + # Initialize cross-encoder reranker (cached for performance) self._cross_encoder_reranker = CrossEncoderReranker(cross_encoder=cross_encoder) @@ -584,7 +613,10 @@ class MemoryEngine(MemoryEngineInterface): ] for fact_type, facts in reflect_result.based_on.items() }, - "mental_models": [], # Mental models are included in based_on["mental-models"] + # Extract mental models from based_on["mental-models"] for easy UI access + "mental_models": [ + {"id": str(fact.id), "text": fact.text} for fact in reflect_result.based_on.get("mental-models", []) + ], } # Update the reflection with the generated content and reflect_response @@ -598,6 +630,79 @@ class MemoryEngine(MemoryEngineInterface): logger.info(f"[CREATE_REFLECTION_TASK] Completed for bank_id={bank_id}, reflection_id={reflection_id}") + async def _handle_refresh_reflection(self, task_dict: dict[str, Any]): + """ + Handler for refresh_reflection tasks. + + Re-runs the source query through reflect and updates the reflection content. + + Args: + task_dict: Dict with 'bank_id', 'reflection_id', 'operation_id' + + Raises: + ValueError: If required fields are missing + Exception: Any exception from reflect/update (propagates to execute_task for retry) + """ + bank_id = task_dict.get("bank_id") + reflection_id = task_dict.get("reflection_id") + + if not bank_id or not reflection_id: + raise ValueError("bank_id and reflection_id are required for refresh_reflection task") + + logger.info(f"[REFRESH_REFLECTION_TASK] Starting for bank_id={bank_id}, reflection_id={reflection_id}") + + from hindsight_api.models import RequestContext + + internal_context = RequestContext() + + # Get the current reflection to get source_query + reflection = await self.get_reflection(bank_id, reflection_id, request_context=internal_context) + if not reflection: + raise ValueError(f"Reflection {reflection_id} not found in bank {bank_id}") + + source_query = reflection["source_query"] + + # Run reflect to generate new content, excluding the reflection being refreshed + reflect_result = await self.reflect_async( + bank_id=bank_id, + query=source_query, + request_context=internal_context, + exclude_reflection_ids=[reflection_id], + ) + + generated_content = reflect_result.text or "No content generated" + + # Build reflect_response payload to store + reflect_response = { + "text": reflect_result.text, + "based_on": { + fact_type: [ + { + "id": str(fact.id), + "text": fact.text, + "type": fact_type, + } + for fact in facts + ] + for fact_type, facts in reflect_result.based_on.items() + }, + # Extract mental models from based_on["mental-models"] for easy UI access + "mental_models": [ + {"id": str(fact.id), "text": fact.text} for fact in reflect_result.based_on.get("mental-models", []) + ], + } + + # Update the reflection with the generated content and reflect_response + await self.update_reflection( + bank_id=bank_id, + reflection_id=reflection_id, + content=generated_content, + reflect_response=reflect_response, + request_context=internal_context, + ) + + logger.info(f"[REFRESH_REFLECTION_TASK] Completed for bank_id={bank_id}, reflection_id={reflection_id}") + async def execute_task(self, task_dict: dict[str, Any]): """ Execute a task by routing it to the appropriate handler. @@ -638,6 +743,8 @@ class MemoryEngine(MemoryEngineInterface): await self._handle_consolidation(task_dict) elif task_type == "create_reflection": await self._handle_create_reflection(task_dict) + elif task_type == "refresh_reflection": + await self._handle_refresh_reflection(task_dict) else: logger.error(f"Unknown task type: {task_type}") # Don't retry unknown task types @@ -792,6 +899,23 @@ class MemoryEngine(MemoryEngineInterface): ) if reflect_is_different: await self._reflect_llm_config.verify_connection() + # Verify consolidation config if different from all others + consolidation_is_different = ( + ( + self._consolidation_llm_config.provider != self._llm_config.provider + or self._consolidation_llm_config.model != self._llm_config.model + ) + and ( + self._consolidation_llm_config.provider != self._retain_llm_config.provider + or self._consolidation_llm_config.model != self._retain_llm_config.model + ) + and ( + self._consolidation_llm_config.provider != self._reflect_llm_config.provider + or self._consolidation_llm_config.model != self._reflect_llm_config.model + ) + ) + if consolidation_is_different: + await self._consolidation_llm_config.verify_connection() # Build list of initialization tasks init_tasks = [ @@ -2995,11 +3119,11 @@ class MemoryEngine(MemoryEngineInterface): await self._authenticate_tenant(request_context) pool = await self._get_pool() async with acquire_with_retry(pool) as conn: - # Get the memory unit + # Get the memory unit (include source_memory_ids for mental models) row = await conn.fetchrow( f""" SELECT id, text, context, event_date, occurred_start, occurred_end, - mentioned_at, fact_type, document_id, chunk_id, tags + mentioned_at, fact_type, document_id, chunk_id, tags, source_memory_ids FROM {fq_table("memory_units")} WHERE id = $1 AND bank_id = $2 """, @@ -3022,7 +3146,7 @@ class MemoryEngine(MemoryEngineInterface): ) entities = [r["canonical_name"] for r in entities_rows] - return { + result = { "id": str(row["id"]), "text": row["text"], "context": row["context"] if row["context"] else "", @@ -3037,6 +3161,35 @@ class MemoryEngine(MemoryEngineInterface): "tags": row["tags"] if row["tags"] else [], } + # For mental models, include source_memory_ids and fetch source_memories + if row["fact_type"] == "mental_model" and row["source_memory_ids"]: + source_ids = row["source_memory_ids"] + result["source_memory_ids"] = [str(sid) for sid in source_ids] + + # Fetch source memories + source_rows = await conn.fetch( + f""" + SELECT id, text, fact_type, context, occurred_start, mentioned_at + FROM {fq_table("memory_units")} + WHERE id = ANY($1::uuid[]) + ORDER BY mentioned_at DESC NULLS LAST + """, + source_ids, + ) + result["source_memories"] = [ + { + "id": str(r["id"]), + "text": r["text"], + "type": r["fact_type"], + "context": r["context"], + "occurred_start": r["occurred_start"].isoformat() if r["occurred_start"] else None, + "mentioned_at": r["mentioned_at"].isoformat() if r["mentioned_at"] else None, + } + for r in source_rows + ] + + return result + async def list_documents( self, bank_id: str, @@ -3671,13 +3824,22 @@ class MemoryEngine(MemoryEngineInterface): DirectiveRef(id=d.id, name=d.name, rules=d.rules) for d in agent_result.directives_applied ] + # Convert agent usage to TokenUsage format + from hindsight_api.engine.response_models import TokenUsage + + usage = TokenUsage( + input_tokens=agent_result.usage.input_tokens, + output_tokens=agent_result.usage.output_tokens, + total_tokens=agent_result.usage.total_tokens, + ) + # Return response (compatible with existing API) result = ReflectResult( text=agent_result.text, based_on=based_on, new_opinions=[], # Learnings stored as mental models structured_output=agent_result.structured_output, - usage=None, # Token tracking not yet implemented for agentic loop + usage=usage, tool_trace=tool_trace_result, llm_trace=llm_trace_result, directives_applied=directives_applied_result, @@ -5650,3 +5812,40 @@ class MemoryEngine(MemoryEngineInterface): result_metadata={"reflection_id": reflection_id, "name": name, "source_query": source_query}, dedupe_by_bank=False, ) + + async def submit_async_refresh_reflection( + self, + bank_id: str, + reflection_id: str, + *, + request_context: "RequestContext", + ) -> dict[str, Any]: + """Submit an async reflection refresh operation. + + This schedules a background task to re-run the source query and update the content. + + Args: + bank_id: Bank identifier + reflection_id: Reflection UUID to refresh + request_context: Request context for authentication + + Returns: + Dict with operation_id + """ + await self._authenticate_tenant(request_context) + + # Verify reflection exists + reflection = await self.get_reflection(bank_id, reflection_id, request_context=request_context) + if not reflection: + raise ValueError(f"Reflection {reflection_id} not found in bank {bank_id}") + + return await self._submit_async_operation( + bank_id=bank_id, + operation_type="refresh_reflection", + task_type="refresh_reflection", + task_payload={ + "reflection_id": reflection_id, + }, + result_metadata={"reflection_id": reflection_id, "name": reflection["name"]}, + dedupe_by_bank=False, + ) diff --git a/hindsight-api/hindsight_api/engine/reflect/agent.py b/hindsight-api/hindsight_api/engine/reflect/agent.py index e4925eeb..603f5fdb 100644 --- a/hindsight-api/hindsight_api/engine/reflect/agent.py +++ b/hindsight-api/hindsight_api/engine/reflect/agent.py @@ -10,10 +10,11 @@ Uses hierarchical retrieval: import asyncio import json import logging +import re import time from typing import TYPE_CHECKING, Any, Awaitable, Callable -from .models import DirectiveInfo, LLMCall, ReflectAgentResult, ToolCall +from .models import DirectiveInfo, LLMCall, ReflectAgentResult, TokenUsageSummary, ToolCall from .prompts import FINAL_SYSTEM_PROMPT, _extract_directive_rules, build_final_prompt, build_system_prompt_for_tools from .tools_schema import get_reflect_tools @@ -77,12 +78,27 @@ def _is_done_tool(name: str) -> bool: return _normalize_tool_name(name) == "done" +# Pattern to match done() call as text - handles done({...}) with nested JSON +_DONE_CALL_PATTERN = re.compile(r"done\s*\(\s*\{.*$", re.DOTALL) + + +def _clean_answer_text(text: str) -> str: + """Clean up answer text by removing any done() tool call syntax. + + Some LLMs output the done() call as text instead of a proper tool call. + This strips out patterns like: done({"answer": "...", ...}) + """ + # Remove done() call pattern from the end of the text + cleaned = _DONE_CALL_PATTERN.sub("", text).strip() + return cleaned if cleaned else text + + async def _generate_structured_output( answer: str, response_schema: dict, llm_config: "LLMProvider", reflect_id: str, -) -> dict[str, Any] | None: +) -> tuple[dict[str, Any] | None, int, int]: """Generate structured output from an answer using the provided JSON schema. Args: @@ -92,7 +108,8 @@ async def _generate_structured_output( reflect_id: Reflect ID for logging Returns: - Structured output dict if successful, None otherwise + Tuple of (structured_output, input_tokens, output_tokens). + structured_output is None if generation fails. """ try: from typing import Any as TypingAny @@ -149,7 +166,7 @@ Return ONLY a valid JSON object that matches this exact schema. Pay special atte Do not include any explanation, only the JSON object.""" - structured_result = await llm_config.call( + structured_result, usage = await llm_config.call( messages=[ { "role": "system", @@ -160,6 +177,7 @@ Do not include any explanation, only the JSON object.""" response_format=DynamicModel, scope="reflect_structured", skip_validation=True, # We'll handle the dict ourselves + return_usage=True, ) # Convert to dict @@ -172,11 +190,11 @@ Do not include any explanation, only the JSON object.""" structured_output = json.loads(str(structured_result)) logger.info(f"[REFLECT {reflect_id}] Generated structured output with {len(structured_output)} fields") - return structured_output + return structured_output, usage.input_tokens, usage.output_tokens except Exception as e: logger.warning(f"[REFLECT {reflect_id}] Failed to generate structured output: {e}") - return None + return None, 0, 0 async def run_reflect_agent( @@ -246,13 +264,32 @@ async def run_reflect_agent( llm_trace: list[dict[str, Any]] = [] context_history: list[dict[str, Any]] = [] # For final prompt fallback + # Token usage tracking - accumulate across all LLM calls + total_input_tokens = 0 + total_output_tokens = 0 + # Track available IDs for validation (prevents hallucinated citations) available_memory_ids: set[str] = set() available_reflection_ids: set[str] = set() available_mental_model_ids: set[str] = set() def _get_llm_trace() -> list[LLMCall]: - return [LLMCall(scope=c["scope"], duration_ms=c["duration_ms"]) for c in llm_trace] + return [ + LLMCall( + scope=c["scope"], + duration_ms=c["duration_ms"], + input_tokens=c.get("input_tokens", 0), + output_tokens=c.get("output_tokens", 0), + ) + for c in llm_trace + ] + + def _get_usage() -> TokenUsageSummary: + return TokenUsageSummary( + input_tokens=total_input_tokens, + output_tokens=total_output_tokens, + total_tokens=total_input_tokens + total_output_tokens, + ) def _log_completion(answer: str, iterations: int, forced: bool = False): elapsed_ms = int((time.time() - start_time) * 1000) @@ -286,21 +323,36 @@ async def run_reflect_agent( # Force text response on last iteration - no tools prompt = build_final_prompt(query, context_history, bank_profile, context) llm_start = time.time() - response = await llm_config.call( + response, usage = await llm_config.call( messages=[ {"role": "system", "content": FINAL_SYSTEM_PROMPT}, {"role": "user", "content": prompt}, ], scope="reflect_agent_final", max_completion_tokens=max_tokens, + return_usage=True, ) - llm_trace.append({"scope": "final", "duration_ms": int((time.time() - llm_start) * 1000)}) - answer = response.strip() + llm_duration = int((time.time() - llm_start) * 1000) + total_input_tokens += usage.input_tokens + total_output_tokens += usage.output_tokens + llm_trace.append( + { + "scope": "final", + "duration_ms": llm_duration, + "input_tokens": usage.input_tokens, + "output_tokens": usage.output_tokens, + } + ) + answer = _clean_answer_text(response.strip()) # Generate structured output if schema provided structured_output = None if response_schema and answer: - structured_output = await _generate_structured_output(answer, response_schema, llm_config, reflect_id) + structured_output, struct_in, struct_out = await _generate_structured_output( + answer, response_schema, llm_config, reflect_id + ) + total_input_tokens += struct_in + total_output_tokens += struct_out _log_completion(answer, iteration + 1, forced=True) return ReflectAgentResult( @@ -310,6 +362,7 @@ async def run_reflect_agent( tools_called=total_tools_called, tool_trace=tool_trace, llm_trace=_get_llm_trace(), + usage=_get_usage(), directives_applied=directives_applied, ) @@ -324,7 +377,16 @@ async def run_reflect_agent( tool_choice="required" if iteration == 0 else "auto", # Force tool use on first iteration ) llm_duration = int((time.time() - llm_start) * 1000) - llm_trace.append({"scope": f"agent_{iteration + 1}", "duration_ms": llm_duration}) + total_input_tokens += result.input_tokens + total_output_tokens += result.output_tokens + llm_trace.append( + { + "scope": f"agent_{iteration + 1}", + "duration_ms": llm_duration, + "input_tokens": result.input_tokens, + "output_tokens": result.output_tokens, + } + ) except Exception as e: err_duration = int((time.time() - llm_start) * 1000) @@ -338,21 +400,36 @@ async def run_reflect_agent( continue prompt = build_final_prompt(query, context_history, bank_profile, context) llm_start = time.time() - response = await llm_config.call( + response, usage = await llm_config.call( messages=[ {"role": "system", "content": FINAL_SYSTEM_PROMPT}, {"role": "user", "content": prompt}, ], scope="reflect_agent_final", max_completion_tokens=max_tokens, + return_usage=True, ) - llm_trace.append({"scope": "final", "duration_ms": int((time.time() - llm_start) * 1000)}) - answer = response.strip() + llm_duration = int((time.time() - llm_start) * 1000) + total_input_tokens += usage.input_tokens + total_output_tokens += usage.output_tokens + llm_trace.append( + { + "scope": "final", + "duration_ms": llm_duration, + "input_tokens": usage.input_tokens, + "output_tokens": usage.output_tokens, + } + ) + answer = _clean_answer_text(response.strip()) # Generate structured output if schema provided structured_output = None if response_schema and answer: - structured_output = await _generate_structured_output(answer, response_schema, llm_config, reflect_id) + structured_output, struct_in, struct_out = await _generate_structured_output( + answer, response_schema, llm_config, reflect_id + ) + total_input_tokens += struct_in + total_output_tokens += struct_out _log_completion(answer, iteration + 1, forced=True) return ReflectAgentResult( @@ -362,20 +439,23 @@ async def run_reflect_agent( tools_called=total_tools_called, tool_trace=tool_trace, llm_trace=_get_llm_trace(), + usage=_get_usage(), directives_applied=directives_applied, ) # No tool calls - LLM wants to respond with text if not result.tool_calls: if result.content: - answer = result.content.strip() + answer = _clean_answer_text(result.content.strip()) # Generate structured output if schema provided structured_output = None if response_schema and answer: - structured_output = await _generate_structured_output( + structured_output, struct_in, struct_out = await _generate_structured_output( answer, response_schema, llm_config, reflect_id ) + total_input_tokens += struct_in + total_output_tokens += struct_out _log_completion(answer, iteration + 1) return ReflectAgentResult( @@ -385,26 +465,42 @@ async def run_reflect_agent( tools_called=total_tools_called, tool_trace=tool_trace, llm_trace=_get_llm_trace(), + usage=_get_usage(), directives_applied=directives_applied, ) # Empty response, force final prompt = build_final_prompt(query, context_history, bank_profile, context) llm_start = time.time() - response = await llm_config.call( + response, usage = await llm_config.call( messages=[ {"role": "system", "content": FINAL_SYSTEM_PROMPT}, {"role": "user", "content": prompt}, ], scope="reflect_agent_final", max_completion_tokens=max_tokens, + return_usage=True, ) - llm_trace.append({"scope": "final", "duration_ms": int((time.time() - llm_start) * 1000)}) - answer = response.strip() + llm_duration = int((time.time() - llm_start) * 1000) + total_input_tokens += usage.input_tokens + total_output_tokens += usage.output_tokens + llm_trace.append( + { + "scope": "final", + "duration_ms": llm_duration, + "input_tokens": usage.input_tokens, + "output_tokens": usage.output_tokens, + } + ) + answer = _clean_answer_text(response.strip()) # Generate structured output if schema provided structured_output = None if response_schema and answer: - structured_output = await _generate_structured_output(answer, response_schema, llm_config, reflect_id) + structured_output, struct_in, struct_out = await _generate_structured_output( + answer, response_schema, llm_config, reflect_id + ) + total_input_tokens += struct_in + total_output_tokens += struct_out _log_completion(answer, iteration + 1, forced=True) return ReflectAgentResult( @@ -414,6 +510,7 @@ async def run_reflect_agent( tools_called=total_tools_called, tool_trace=tool_trace, llm_trace=_get_llm_trace(), + usage=_get_usage(), directives_applied=directives_applied, ) @@ -456,6 +553,7 @@ async def run_reflect_agent( total_tools_called, tool_trace, _get_llm_trace(), + _get_usage(), _log_completion, reflect_id, directives_applied=directives_applied, @@ -577,6 +675,7 @@ async def run_reflect_agent( tools_called=total_tools_called, tool_trace=tool_trace, llm_trace=_get_llm_trace(), + usage=_get_usage(), directives_applied=directives_applied, ) @@ -602,6 +701,7 @@ async def _process_done_tool( total_tools_called: int, tool_trace: list[ToolCall], llm_trace: list[LLMCall], + usage: TokenUsageSummary, log_completion: Callable, reflect_id: str, directives_applied: list[DirectiveInfo], @@ -622,8 +722,17 @@ async def _process_done_tool( # Generate structured output if schema provided structured_output = None + final_usage = usage if response_schema and llm_config and answer: - structured_output = await _generate_structured_output(answer, response_schema, llm_config, reflect_id) + structured_output, struct_in, struct_out = await _generate_structured_output( + answer, response_schema, llm_config, reflect_id + ) + # Add structured output tokens to usage + final_usage = TokenUsageSummary( + input_tokens=usage.input_tokens + struct_in, + output_tokens=usage.output_tokens + struct_out, + total_tokens=usage.total_tokens + struct_in + struct_out, + ) log_completion(answer, iterations) return ReflectAgentResult( @@ -633,6 +742,7 @@ async def _process_done_tool( tools_called=total_tools_called, tool_trace=tool_trace, llm_trace=llm_trace, + usage=final_usage, used_memory_ids=used_memory_ids, used_reflection_ids=used_reflection_ids, used_mental_model_ids=used_mental_model_ids, diff --git a/hindsight-api/hindsight_api/engine/reflect/models.py b/hindsight-api/hindsight_api/engine/reflect/models.py index 71e6d044..6c1ad452 100644 --- a/hindsight-api/hindsight_api/engine/reflect/models.py +++ b/hindsight-api/hindsight_api/engine/reflect/models.py @@ -85,6 +85,8 @@ class LLMCall(BaseModel): scope: str = Field(description="Call scope: agent_1, agent_2, final, etc.") duration_ms: int = Field(description="Execution time in milliseconds") + input_tokens: int = Field(default=0, description="Input tokens used") + output_tokens: int = Field(default=0, description="Output tokens used") class DirectiveInfo(BaseModel): @@ -95,6 +97,14 @@ class DirectiveInfo(BaseModel): rules: list[str] = Field(default_factory=list, description="Directive rules/observations that were applied") +class TokenUsageSummary(BaseModel): + """Total token usage across all LLM calls.""" + + input_tokens: int = Field(default=0, description="Total input tokens used") + output_tokens: int = Field(default=0, description="Total output tokens used") + total_tokens: int = Field(default=0, description="Total tokens (input + output)") + + class ReflectAgentResult(BaseModel): """Result from the reflect agent.""" @@ -106,6 +116,9 @@ class ReflectAgentResult(BaseModel): tools_called: int = Field(default=0, description="Total number of tool calls made") tool_trace: list[ToolCall] = Field(default_factory=list, description="Trace of all tool calls made") llm_trace: list[LLMCall] = Field(default_factory=list, description="Trace of all LLM calls made") + usage: TokenUsageSummary = Field( + default_factory=TokenUsageSummary, description="Total token usage across all LLM calls" + ) used_memory_ids: list[str] = Field(default_factory=list, description="Validated memory IDs actually used in answer") used_reflection_ids: list[str] = Field( default_factory=list, description="Validated reflection IDs actually used in answer" diff --git a/hindsight-api/hindsight_api/engine/response_models.py b/hindsight-api/hindsight_api/engine/response_models.py index 12c827ee..4f24e50d 100644 --- a/hindsight-api/hindsight_api/engine/response_models.py +++ b/hindsight-api/hindsight_api/engine/response_models.py @@ -28,6 +28,8 @@ class LLMToolCallResult(BaseModel): content: str | None = Field(default=None, description="Text content if any") tool_calls: list[LLMToolCall] = Field(default_factory=list, description="Tool calls requested by the LLM") finish_reason: str | None = Field(default=None, description="Reason the LLM stopped: 'stop', 'tool_calls', etc.") + input_tokens: int = Field(default=0, description="Input tokens used in this call") + output_tokens: int = Field(default=0, description="Output tokens used in this call") class ToolCallTrace(BaseModel): diff --git a/hindsight-api/hindsight_api/extensions/__init__.py b/hindsight-api/hindsight_api/extensions/__init__.py index fd84a272..592d0f8d 100644 --- a/hindsight-api/hindsight_api/extensions/__init__.py +++ b/hindsight-api/hindsight_api/extensions/__init__.py @@ -21,6 +21,10 @@ from hindsight_api.extensions.context import DefaultExtensionContext, ExtensionC from hindsight_api.extensions.http import HttpExtension from hindsight_api.extensions.loader import load_extension from hindsight_api.extensions.operation_validator import ( + # Consolidation operation + ConsolidateContext, + ConsolidateResult, + # Core operations OperationValidationError, OperationValidatorExtension, RecallContext, @@ -33,6 +37,7 @@ from hindsight_api.extensions.operation_validator import ( ) from hindsight_api.extensions.tenant import ( AuthenticationError, + Tenant, TenantContext, TenantExtension, ) @@ -47,7 +52,7 @@ __all__ = [ "DefaultExtensionContext", # HTTP Extension "HttpExtension", - # Operation Validator + # Operation Validator - Core "OperationValidationError", "OperationValidatorExtension", "RecallContext", @@ -57,10 +62,14 @@ __all__ = [ "RetainContext", "RetainResult", "ValidationResult", + # Operation Validator - Consolidation + "ConsolidateContext", + "ConsolidateResult", # Tenant/Auth "ApiKeyTenantExtension", "AuthenticationError", "RequestContext", + "Tenant", "TenantContext", "TenantExtension", ] diff --git a/hindsight-api/hindsight_api/extensions/builtin/tenant.py b/hindsight-api/hindsight_api/extensions/builtin/tenant.py index a5a1b80a..d56266c5 100644 --- a/hindsight-api/hindsight_api/extensions/builtin/tenant.py +++ b/hindsight-api/hindsight_api/extensions/builtin/tenant.py @@ -1,6 +1,6 @@ """Built-in tenant extension implementations.""" -from hindsight_api.extensions.tenant import AuthenticationError, TenantContext, TenantExtension +from hindsight_api.extensions.tenant import AuthenticationError, Tenant, TenantContext, TenantExtension from hindsight_api.models import RequestContext @@ -31,3 +31,7 @@ class ApiKeyTenantExtension(TenantExtension): if context.api_key != self.expected_api_key: raise AuthenticationError("Invalid API key") return TenantContext(schema_name="public") + + async def list_tenants(self) -> list[Tenant]: + """Return public schema for single-tenant setup.""" + return [Tenant(schema="public")] diff --git a/hindsight-api/hindsight_api/extensions/operation_validator.py b/hindsight-api/hindsight_api/extensions/operation_validator.py index a1dec0eb..de4fb814 100644 --- a/hindsight-api/hindsight_api/extensions/operation_validator.py +++ b/hindsight-api/hindsight_api/extensions/operation_validator.py @@ -1,4 +1,4 @@ -"""Operation Validator Extension for validating retain/recall/reflect operations.""" +"""Operation Validator Extension for validating retain/recall/reflect/consolidate operations.""" from abc import ABC, abstractmethod from dataclasses import dataclass, field @@ -97,6 +97,19 @@ class ReflectContext: context: str | None = None +# ============================================================================= +# Consolidation Pre-operation Context +# ============================================================================= + + +@dataclass +class ConsolidateContext: + """Context for a consolidation operation validation (pre-operation).""" + + bank_id: str + request_context: "RequestContext" + + # ============================================================================= # Post-operation Contexts (includes results) # ============================================================================= @@ -164,9 +177,28 @@ class ReflectResultContext: error: str | None = None +# ============================================================================= +# Consolidation Post-operation Context +# ============================================================================= + + +@dataclass +class ConsolidateResult: + """Result context for post-consolidation hook.""" + + bank_id: str + request_context: "RequestContext" + # Result + processed: int = 0 + created: int = 0 + updated: int = 0 + success: bool = True + error: str | None = None + + class OperationValidatorExtension(Extension, ABC): """ - Validates and hooks into retain/recall/reflect operations. + Validates and hooks into retain/recall/reflect/consolidate operations. This extension allows implementing custom logic such as: - Rate limiting (pre-operation) @@ -185,9 +217,13 @@ class OperationValidatorExtension(Extension, ABC): -> config = {"max_requests": "100"} Hook execution order: - 1. validate_retain/validate_recall/validate_reflect (pre-operation) + 1. validate_* (pre-operation) 2. [operation executes] - 3. on_retain_complete/on_recall_complete/on_reflect_complete (post-operation) + 3. on_*_complete (post-operation) + + Supported operations: + - retain, recall, reflect (core memory operations) + - consolidate (mental models consolidation) """ # ========================================================================= @@ -325,3 +361,44 @@ class OperationValidatorExtension(Extension, ABC): - error: Error message (if failed) """ pass + + # ========================================================================= + # Consolidation - Pre-operation validation hook (optional - override to implement) + # ========================================================================= + + async def validate_consolidate(self, ctx: ConsolidateContext) -> ValidationResult: + """ + Validate a consolidation operation before execution. + + Override to implement custom validation logic for consolidation. + + Args: + ctx: Context containing: + - bank_id: Bank identifier + - request_context: Request context with auth info + + Returns: + ValidationResult indicating whether the operation is allowed. + """ + return ValidationResult.accept() + + # ========================================================================= + # Consolidation - Post-operation hook (optional - override to implement) + # ========================================================================= + + async def on_consolidate_complete(self, result: ConsolidateResult) -> None: + """ + Called after a consolidation operation completes (success or failure). + + Override to implement post-operation logic such as usage tracking or audit logging. + + Args: + result: Result context containing: + - bank_id: Bank identifier + - processed: Number of memories processed + - created: Number of mental models created + - updated: Number of mental models updated + - success: Whether the operation succeeded + - error: Error message (if failed) + """ + pass diff --git a/hindsight-api/hindsight_api/extensions/tenant.py b/hindsight-api/hindsight_api/extensions/tenant.py index 6b178a32..75dc0dfa 100644 --- a/hindsight-api/hindsight_api/extensions/tenant.py +++ b/hindsight-api/hindsight_api/extensions/tenant.py @@ -28,6 +28,18 @@ class TenantContext: schema_name: str +@dataclass +class Tenant: + """ + Represents a tenant for worker discovery. + + Used by list_tenants() to return tenant information including + the PostgreSQL schema name for database operations. + """ + + schema: str + + class TenantExtension(Extension, ABC): """ Extension for multi-tenancy and API key authentication. @@ -61,3 +73,17 @@ class TenantExtension(Extension, ABC): AuthenticationError: If authentication fails. """ ... + + @abstractmethod + async def list_tenants(self) -> list[Tenant]: + """ + List all tenants that should be processed by workers. + + This method is used by the worker to discover all tenants that need + task polling. Workers will poll for pending tasks in each tenant's schema. + + Returns: + List of Tenant objects containing schema information. + For single-tenant setups, return [Tenant(schema="public")]. + """ + ... diff --git a/hindsight-api/hindsight_api/main.py b/hindsight-api/hindsight_api/main.py index 728dba9e..f9a07ae9 100644 --- a/hindsight-api/hindsight_api/main.py +++ b/hindsight-api/hindsight_api/main.py @@ -184,6 +184,10 @@ def main(): 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, + 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, embeddings_provider=config.embeddings_provider, embeddings_local_model=config.embeddings_local_model, embeddings_tei_url=config.embeddings_tei_url, diff --git a/hindsight-api/hindsight_api/worker/main.py b/hindsight-api/hindsight_api/worker/main.py index a230b73e..d5f6fb47 100644 --- a/hindsight-api/hindsight_api/worker/main.py +++ b/hindsight-api/hindsight_api/worker/main.py @@ -181,6 +181,8 @@ def main(): nonlocal memory, poller import uvicorn + from ..extensions import TenantExtension, load_extension + # Initialize MemoryEngine # Workers use SyncTaskBackend because they execute tasks directly, # they don't need to store tasks (they poll from DB) @@ -193,7 +195,15 @@ def main(): print(f"Database connected: {config.database_url}") - # Create and start the poller + # Load tenant extension for dynamic schema discovery + tenant_extension = load_extension("TENANT", TenantExtension) + + if tenant_extension: + print("Tenant extension loaded - schemas will be discovered dynamically on each poll") + else: + print("No tenant extension configured, using public schema only") + + # Create a single poller that handles all schemas dynamically poller = WorkerPoller( pool=memory._pool, worker_id=args.worker_id, @@ -201,6 +211,7 @@ def main(): poll_interval_ms=args.poll_interval, batch_size=args.batch_size, max_retries=args.max_retries, + tenant_extension=tenant_extension, ) # Create the HTTP app for metrics/health diff --git a/hindsight-api/hindsight_api/worker/poller.py b/hindsight-api/hindsight_api/worker/poller.py index 52cbecbc..3cd307d5 100644 --- a/hindsight-api/hindsight_api/worker/poller.py +++ b/hindsight-api/hindsight_api/worker/poller.py @@ -11,11 +11,14 @@ import logging import time import traceback from collections.abc import Awaitable, Callable +from dataclasses import dataclass from typing import TYPE_CHECKING, Any if TYPE_CHECKING: import asyncpg + from hindsight_api.extensions.tenant import TenantExtension + logger = logging.getLogger(__name__) # Progress logging interval in seconds @@ -29,12 +32,23 @@ def fq_table(table: str, schema: str | None = None) -> str: return table +@dataclass +class ClaimedTask: + """A task claimed from the database with its schema context.""" + + operation_id: str + task_dict: dict[str, Any] + schema: str | None + + class WorkerPoller: """ Polls PostgreSQL for pending tasks and executes them. Uses FOR UPDATE SKIP LOCKED for safe distributed claiming, allowing multiple workers to process tasks without conflicts. + + Supports dynamic multi-tenant discovery via tenant_extension. """ def __init__( @@ -46,6 +60,7 @@ class WorkerPoller: batch_size: int = 10, max_retries: int = 3, schema: str | None = None, + tenant_extension: "TenantExtension | None" = None, ): """ Initialize the worker poller. @@ -57,7 +72,9 @@ class WorkerPoller: poll_interval_ms: Interval between polls when no tasks found (milliseconds) batch_size: Maximum number of tasks to claim per poll cycle max_retries: Maximum retry attempts before marking task as failed - schema: Database schema for multi-tenant support (optional) + schema: Database schema for single-tenant support (ignored if tenant_extension is set) + tenant_extension: Extension for dynamic multi-tenant discovery. If set, list_tenants() + is called on each poll cycle to discover schemas dynamically. """ self._pool = pool self._worker_id = worker_id @@ -66,27 +83,56 @@ class WorkerPoller: self._batch_size = batch_size self._max_retries = max_retries self._schema = schema + self._tenant_extension = tenant_extension self._shutdown = asyncio.Event() self._current_tasks: set[asyncio.Task] = set() self._in_flight_count = 0 self._in_flight_lock = asyncio.Lock() self._last_progress_log = 0.0 self._tasks_completed_since_log = 0 - self._active_banks: set[str] = set() + # Track active tasks locally: operation_id -> (op_type, bank_id, schema) + self._active_tasks: dict[str, tuple[str, str, str | None]] = {} - async def claim_batch(self) -> list[tuple[str, dict[str, Any]]]: + async def _get_schemas(self) -> list[str | None]: + """Get list of schemas to poll. Returns [None] for public schema.""" + if self._tenant_extension is not None: + tenants = await self._tenant_extension.list_tenants() + # Convert "public" to None for SQL compatibility, keep others as-is + return [t.schema if t.schema != "public" else None for t in tenants] + # Single schema mode + return [self._schema] + + async def claim_batch(self) -> list[ClaimedTask]: """ - Claim up to batch_size pending tasks atomically. + Claim up to batch_size pending tasks atomically across all tenant schemas. Uses FOR UPDATE SKIP LOCKED to ensure no conflicts with other workers. For consolidation tasks specifically, skips pending tasks if there's already a processing consolidation for the same bank (to avoid duplicate work). + If tenant_extension is configured, dynamically discovers schemas on each call. + Returns: - List of tuples (operation_id, task_dict) + List of ClaimedTask objects containing operation_id, task_dict, and schema """ - table = fq_table("async_operations", self._schema) + schemas = await self._get_schemas() + all_tasks: list[ClaimedTask] = [] + remaining_batch = self._batch_size + + for schema in schemas: + if remaining_batch <= 0: + break + + tasks = await self._claim_batch_for_schema(schema, remaining_batch) + all_tasks.extend(tasks) + remaining_batch -= len(tasks) + + return all_tasks + + async def _claim_batch_for_schema(self, schema: str | None, limit: int) -> list[ClaimedTask]: + """Claim tasks from a specific schema.""" + table = fq_table("async_operations", schema) async with self._pool.acquire() as conn: async with conn.transaction(): @@ -113,7 +159,7 @@ class WorkerPoller: LIMIT $1 FOR UPDATE SKIP LOCKED """, - self._batch_size, + limit, ) if not rows: @@ -131,12 +177,19 @@ class WorkerPoller: operation_ids, ) - # Parse and return task payloads - return [(str(row["operation_id"]), json.loads(row["task_payload"])) for row in rows] + # Parse and return task payloads with schema context + return [ + ClaimedTask( + operation_id=str(row["operation_id"]), + task_dict=json.loads(row["task_payload"]), + schema=schema, + ) + for row in rows + ] - async def _mark_completed(self, operation_id: str): + async def _mark_completed(self, operation_id: str, schema: str | None): """Mark a task as completed.""" - table = fq_table("async_operations", self._schema) + table = fq_table("async_operations", schema) await self._pool.execute( f""" UPDATE {table} @@ -146,9 +199,9 @@ class WorkerPoller: operation_id, ) - async def _mark_failed(self, operation_id: str, error_message: str): + async def _mark_failed(self, operation_id: str, error_message: str, schema: str | None): """Mark a task as failed with error message.""" - table = fq_table("async_operations", self._schema) + table = fq_table("async_operations", schema) # Truncate error message if too long (max 5000 chars in schema) error_message = error_message[:5000] if len(error_message) > 5000 else error_message await self._pool.execute( @@ -161,9 +214,9 @@ class WorkerPoller: error_message, ) - async def _retry_or_fail(self, operation_id: str, error_message: str): + async def _retry_or_fail(self, operation_id: str, error_message: str, schema: str | None): """Increment retry count or mark as failed if max retries exceeded.""" - table = fq_table("async_operations", self._schema) + table = fq_table("async_operations", schema) # Get current retry count row = await self._pool.fetchrow( @@ -180,7 +233,7 @@ class WorkerPoller: if retry_count >= self._max_retries: # Max retries exceeded, mark as failed await self._mark_failed( - operation_id, f"Max retries ({self._max_retries}) exceeded. Last error: {error_message}" + operation_id, f"Max retries ({self._max_retries}) exceeded. Last error: {error_message}", schema ) logger.error(f"Task {operation_id} failed after {retry_count} retries") else: @@ -196,20 +249,29 @@ class WorkerPoller: ) logger.warning(f"Task {operation_id} failed, will retry (attempt {retry_count + 1}/{self._max_retries})") - async def execute_task(self, operation_id: str, task_dict: dict[str, Any]): + async def execute_task(self, task: ClaimedTask): """Execute a single task and update its status.""" - task_type = task_dict.get("type", "unknown") - bank_id = task_dict.get("bank_id", "unknown") + task_type = task.task_dict.get("type", "unknown") + bank_id = task.task_dict.get("bank_id", "unknown") + + # Track this task as active + async with self._in_flight_lock: + self._active_tasks[task.operation_id] = (task_type, bank_id, task.schema) try: - logger.debug(f"Executing task {operation_id} (type={task_type}, bank={bank_id})") - await self._executor(task_dict) - await self._mark_completed(operation_id) - logger.debug(f"Task {operation_id} completed successfully") + schema_info = f", schema={task.schema}" if task.schema else "" + logger.debug(f"Executing task {task.operation_id} (type={task_type}, bank={bank_id}{schema_info})") + await self._executor(task.task_dict) + await self._mark_completed(task.operation_id, task.schema) + logger.debug(f"Task {task.operation_id} completed successfully") except Exception as e: error_msg = f"{type(e).__name__}: {e}\n{traceback.format_exc()}" - logger.error(f"Task {operation_id} failed: {e}") - await self._retry_or_fail(operation_id, error_msg) + logger.error(f"Task {task.operation_id} failed: {e}") + await self._retry_or_fail(task.operation_id, error_msg, task.schema) + finally: + # Remove from active tasks + async with self._in_flight_lock: + self._active_tasks.pop(task.operation_id, None) async def recover_own_tasks(self) -> int: """ @@ -219,25 +281,33 @@ 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. + If tenant_extension is configured, recovers across all tenant schemas. + Returns: Number of tasks recovered """ - table = fq_table("async_operations", self._schema) + schemas = await self._get_schemas() + total_count = 0 - 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 - """, - self._worker_id, - ) + for schema in schemas: + table = fq_table("async_operations", schema) - # Parse "UPDATE N" to get count - count = int(result.split()[-1]) if result else 0 - if count > 0: - logger.info(f"Worker {self._worker_id} recovered {count} stale tasks from previous run") - return count + 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 + """, + self._worker_id, + ) + + # Parse "UPDATE N" to get count + count = int(result.split()[-1]) if result else 0 + total_count += count + + if total_count > 0: + logger.info(f"Worker {self._worker_id} recovered {total_count} stale tasks from previous run") + return total_count async def run(self): """ @@ -245,6 +315,8 @@ class WorkerPoller: Continuously polls for pending tasks, claims them, and executes them until shutdown is signaled. + + If tenant_extension is configured, dynamically discovers schemas on each poll. """ # Recover any tasks from a previous crash before starting await self.recover_own_tasks() @@ -253,17 +325,22 @@ class WorkerPoller: while not self._shutdown.is_set(): try: - # Claim a batch of tasks + # Claim a batch of tasks (across all tenant schemas if configured) tasks = await self.claim_batch() if tasks: # Log batch info - task_types = {} - for _, task_dict in tasks: - t = task_dict.get("type", "unknown") + task_types: dict[str, int] = {} + schemas_seen: set[str | None] = set() + for task in tasks: + t = task.task_dict.get("type", "unknown") task_types[t] = task_types.get(t, 0) + 1 + schemas_seen.add(task.schema) types_str = ", ".join(f"{k}:{v}" for k, v in task_types.items()) - logger.info(f"Worker {self._worker_id} claimed {len(tasks)} tasks: {types_str}") + schemas_str = ", ".join(s or "public" for s in schemas_seen) + logger.info( + f"Worker {self._worker_id} claimed {len(tasks)} tasks: {types_str} (schemas: {schemas_str})" + ) # Track in-flight tasks async with self._in_flight_lock: @@ -272,7 +349,7 @@ class WorkerPoller: # Execute tasks concurrently try: await asyncio.gather( - *[self.execute_task(op_id, task_dict) for op_id, task_dict in tasks], + *[self.execute_task(task) for task in tasks], return_exceptions=True, ) finally: @@ -336,58 +413,60 @@ class WorkerPoller: self._last_progress_log = now try: - table = fq_table("async_operations", self._schema) - async with self._pool.acquire() as conn: - # Get global stats by status - stats = await conn.fetch( - f""" - SELECT status, COUNT(*) as count - FROM {table} - WHERE created_at > now() - interval '24 hours' - GROUP BY status - """ - ) - - # Get currently processing tasks grouped by type and bank - processing = await conn.fetch( - f""" - SELECT operation_type, bank_id, COUNT(*) as count - FROM {table} - WHERE status = 'processing' - GROUP BY operation_type, bank_id - """ - ) - - # Build stats dict - status_counts = {row["status"]: row["count"] for row in stats} - pending = status_counts.get("pending", 0) - processing_count = status_counts.get("processing", 0) - completed = status_counts.get("completed", 0) - failed = status_counts.get("failed", 0) - - # Build processing breakdown - processing_info = [] - banks_working = set() - for row in processing: - op_type = row["operation_type"] - bank_id = row["bank_id"] - count = row["count"] - banks_working.add(bank_id) - processing_info.append(f"{op_type}:{bank_id}({count})") - - # Format log + # Get local active tasks (this worker only) async with self._in_flight_lock: in_flight = self._in_flight_count + active_tasks = dict(self._active_tasks) # Copy to avoid holding lock + # Build local processing breakdown grouped by (op_type, bank_id) + task_groups: dict[tuple[str, str], int] = {} + for op_type, bank_id, _ in active_tasks.values(): + key = (op_type, bank_id) + task_groups[key] = task_groups.get(key, 0) + 1 + + processing_info = [f"{op}:{bank}({cnt})" for (op, bank), cnt in task_groups.items()] processing_str = ", ".join(processing_info[:10]) if processing_info else "none" if len(processing_info) > 10: processing_str += f" +{len(processing_info) - 10} more" + # Get global stats from DB across all schemas + schemas = await self._get_schemas() + global_pending = 0 + all_worker_counts: dict[str, int] = {} + + async with self._pool.acquire() as conn: + for schema in schemas: + table = fq_table("async_operations", schema) + + row = await conn.fetchrow(f"SELECT COUNT(*) as count FROM {table} WHERE status = 'pending'") + global_pending += row["count"] if row else 0 + + # Get processing breakdown by worker + worker_rows = await conn.fetch( + f""" + SELECT worker_id, COUNT(*) as count + FROM {table} + WHERE status = 'processing' + GROUP BY worker_id + """ + ) + for wr in worker_rows: + wid = wr["worker_id"] or "unknown" + all_worker_counts[wid] = all_worker_counts.get(wid, 0) + wr["count"] + + # Format other workers' processing counts + other_workers = [] + for wid, cnt in all_worker_counts.items(): + if wid != self._worker_id: + other_workers.append(f"{wid}:{cnt}") + others_str = ", ".join(other_workers) if other_workers else "none" + + schemas_str = ", ".join(s or "public" for s in schemas) logger.info( f"[WORKER_STATS] worker={self._worker_id} in_flight={in_flight} | " - f"global: pending={pending} processing={processing_count} " - f"completed_24h={completed} failed_24h={failed} | " - f"active: {processing_str}" + f"global: pending={global_pending} (schemas: {schemas_str}) | " + f"others: {others_str} | " + f"my_active: {processing_str}" ) except Exception as e: diff --git a/hindsight-api/tests/test_custom_embedding_dimension.py b/hindsight-api/tests/test_custom_embedding_dimension.py index 57b62c5e..8d382481 100644 --- a/hindsight-api/tests/test_custom_embedding_dimension.py +++ b/hindsight-api/tests/test_custom_embedding_dimension.py @@ -9,18 +9,18 @@ Includes tests for: import asyncio import os -import pytest from datetime import datetime + +import pytest from sqlalchemy import create_engine, text from hindsight_api import MemoryEngine, RequestContext -from hindsight_api.engine.embeddings import LocalSTEmbeddings, OpenAIEmbeddings, CohereEmbeddings -from hindsight_api.engine.cross_encoder import LocalSTCrossEncoder, CohereCrossEncoder +from hindsight_api.engine.cross_encoder import CohereCrossEncoder, LocalSTCrossEncoder +from hindsight_api.engine.embeddings import CohereEmbeddings, LocalSTEmbeddings, OpenAIEmbeddings from hindsight_api.engine.query_analyzer import DateparserQueryAnalyzer from hindsight_api.engine.task_backend import SyncTaskBackend -from hindsight_api.extensions import TenantExtension, TenantContext -from hindsight_api.migrations import run_migrations, ensure_embedding_dimension - +from hindsight_api.extensions import TenantContext, TenantExtension +from hindsight_api.migrations import ensure_embedding_dimension, run_migrations # ============================================================================= # Shared Utilities @@ -36,6 +36,11 @@ class SchemaTenantExtension(TenantExtension): async def authenticate(self, request_context: RequestContext) -> TenantContext: return TenantContext(schema_name=self.schema_name) + async def list_tenants(self) -> list: + from hindsight_api.extensions.tenant import Tenant + + return [Tenant(schema=self.schema_name)] + def get_test_schema(prefix: str, worker_id: str) -> str: """Get unique schema name per xdist worker.""" diff --git a/hindsight-api/tests/test_extensions.py b/hindsight-api/tests/test_extensions.py index 88a2f5d3..9956a509 100644 --- a/hindsight-api/tests/test_extensions.py +++ b/hindsight-api/tests/test_extensions.py @@ -24,6 +24,9 @@ from hindsight_api.extensions import ( TenantExtension, ValidationResult, load_extension, + # Consolidation operation + ConsolidateContext, + ConsolidateResult, ) @@ -128,14 +131,18 @@ class TrackingValidator(OperationValidatorExtension): def __init__(self, config: dict): super().__init__(config) - # Pre-hook tracking + # Pre-hook tracking - Core operations self.pre_retain_calls: list[RetainContext] = [] self.pre_recall_calls: list[RecallContext] = [] self.pre_reflect_calls: list[ReflectContext] = [] - # Post-hook tracking + # Post-hook tracking - Core operations self.post_retain_calls: list[RetainResult] = [] self.post_recall_calls: list[RecallResult] = [] self.post_reflect_calls: list[ReflectResultContext] = [] + # Pre-hook tracking - Consolidation + self.pre_consolidate_calls: list[ConsolidateContext] = [] + # Post-hook tracking - Consolidation + self.post_consolidate_calls: list[ConsolidateResult] = [] async def validate_retain(self, ctx: RetainContext) -> ValidationResult: self.pre_retain_calls.append(ctx) @@ -158,6 +165,14 @@ class TrackingValidator(OperationValidatorExtension): async def on_reflect_complete(self, result: ReflectResultContext) -> None: self.post_reflect_calls.append(result) + # Consolidation hooks + async def validate_consolidate(self, ctx: ConsolidateContext) -> ValidationResult: + self.pre_consolidate_calls.append(ctx) + return ValidationResult.accept() + + async def on_consolidate_complete(self, result: ConsolidateResult) -> None: + self.post_consolidate_calls.append(result) + class TestMemoryEngineValidation: """Tests for validation integration with MemoryEngine. diff --git a/hindsight-api/tests/test_http_api_integration.py b/hindsight-api/tests/test_http_api_integration.py index c892cc10..bd51217f 100644 --- a/hindsight-api/tests/test_http_api_integration.py +++ b/hindsight-api/tests/test_http_api_integration.py @@ -969,24 +969,22 @@ async def test_reflect_returns_token_usage(api_client): assert "text" in result assert len(result["text"]) > 0 - # Verify usage field exists (may be None for agentic reflect which makes multiple LLM calls) + # Verify usage field exists and is populated (agentic reflect aggregates all LLM calls) assert "usage" in result, "Response should include 'usage' field" usage = result["usage"] - # Usage is optional - agentic reflect doesn't aggregate multiple LLM call usages - if usage is not None: - assert "input_tokens" in usage, "Usage should have 'input_tokens'" - assert "output_tokens" in usage, "Usage should have 'output_tokens'" - assert "total_tokens" in usage, "Usage should have 'total_tokens'" + # Usage must be present - agentic reflect now aggregates token usage from all LLM calls + assert usage is not None, "Usage should not be None - reflect aggregates all LLM call usages" + assert "input_tokens" in usage, "Usage should have 'input_tokens'" + assert "output_tokens" in usage, "Usage should have 'output_tokens'" + assert "total_tokens" in usage, "Usage should have 'total_tokens'" - # Verify token counts are valid - assert usage["input_tokens"] > 0, f"Expected input_tokens > 0, got {usage['input_tokens']}" - assert usage["output_tokens"] >= 0, f"Expected output_tokens >= 0, got {usage['output_tokens']}" - assert usage["total_tokens"] == usage["input_tokens"] + usage["output_tokens"] + # Verify token counts are valid + assert usage["input_tokens"] > 0, f"Expected input_tokens > 0, got {usage['input_tokens']}" + assert usage["output_tokens"] >= 0, f"Expected output_tokens >= 0, got {usage['output_tokens']}" + assert usage["total_tokens"] == usage["input_tokens"] + usage["output_tokens"] - print(f"Reflect token usage: input={usage['input_tokens']}, output={usage['output_tokens']}, total={usage['total_tokens']}") - else: - print("Reflect usage is None (expected for agentic reflect)") + print(f"Reflect token usage: input={usage['input_tokens']}, output={usage['output_tokens']}, total={usage['total_tokens']}") @pytest.mark.asyncio diff --git a/hindsight-api/tests/test_main_module.py b/hindsight-api/tests/test_main_module.py index 0923fad0..453ee563 100644 --- a/hindsight-api/tests/test_main_module.py +++ b/hindsight-api/tests/test_main_module.py @@ -355,14 +355,14 @@ class TestMainModuleExtensionLoading: # Mock extensions for testing from hindsight_api.extensions import ( - TenantExtension, - TenantContext, - RequestContext, OperationValidatorExtension, - ValidationResult, - RetainContext, RecallContext, ReflectContext, + RequestContext, + RetainContext, + TenantContext, + TenantExtension, + ValidationResult, ) @@ -376,6 +376,11 @@ class MockTenantExtension(TenantExtension): async def authenticate(self, request_context: RequestContext) -> TenantContext: return TenantContext(schema_name="public") + async def list_tenants(self) -> list: + from hindsight_api.extensions.tenant import Tenant + + return [Tenant(schema="public")] + def set_context(self, context) -> None: self._context_set = True diff --git a/hindsight-api/tests/test_reflect_agent.py b/hindsight-api/tests/test_reflect_agent.py index 5ebd79ef..76749da4 100644 --- a/hindsight-api/tests/test_reflect_agent.py +++ b/hindsight-api/tests/test_reflect_agent.py @@ -13,9 +13,52 @@ from unittest.mock import AsyncMock, MagicMock, patch from hindsight_api.engine.reflect.agent import ( _normalize_tool_name, _is_done_tool, + _clean_answer_text, run_reflect_agent, ) -from hindsight_api.engine.response_models import LLMToolCall, LLMToolCallResult +from hindsight_api.engine.response_models import LLMToolCall, LLMToolCallResult, TokenUsage + + +class TestCleanAnswerText: + """Test cleanup of answer text that includes done() tool call syntax.""" + + def test_clean_text_with_done_call(self): + """Text ending with done() call should have it stripped.""" + text = '''The team's OKRs focus on performance.done({"answer":"The team's OKRs","memory_ids":[]})''' + cleaned = _clean_answer_text(text) + assert cleaned == "The team's OKRs focus on performance." + assert "done(" not in cleaned + + def test_clean_text_with_done_call_and_whitespace(self): + """done() call with whitespace should be stripped.""" + text = '''Answer text here. done( {"answer": "short", "memory_ids": []} )''' + cleaned = _clean_answer_text(text) + assert cleaned == "Answer text here." + + def test_clean_text_without_done_call(self): + """Text without done() call should be unchanged.""" + text = "This is a normal answer without any tool calls." + cleaned = _clean_answer_text(text) + assert cleaned == text + + def test_clean_text_with_done_word_in_content(self): + """The word 'done' in regular text should not be stripped.""" + text = "The task is done and completed successfully." + cleaned = _clean_answer_text(text) + assert cleaned == text + + def test_clean_empty_text(self): + """Empty text should return empty.""" + assert _clean_answer_text("") == "" + + def test_clean_text_multiline_done(self): + """done() call spanning multiple lines should be stripped.""" + text = '''Summary of findings.done({ + "answer": "Summary", + "memory_ids": ["id1", "id2"] + })''' + cleaned = _clean_answer_text(text) + assert cleaned == "Summary of findings." class TestToolNameNormalization: @@ -70,8 +113,10 @@ class TestReflectAgentMocked: """Create a mock LLM provider.""" llm = MagicMock() llm.call_with_tools = AsyncMock() - # Also mock call() for final iteration fallback - llm.call = AsyncMock(return_value="Fallback answer from final iteration") + # Also mock call() for final iteration fallback - returns (response, usage) tuple + llm.call = AsyncMock( + return_value=("Fallback answer from final iteration", TokenUsage(input_tokens=100, output_tokens=50, total_tokens=150)) + ) return llm @pytest.fixture diff --git a/hindsight-api/tests/test_schema_isolation.py b/hindsight-api/tests/test_schema_isolation.py index cf7e767c..9eded90b 100644 --- a/hindsight-api/tests/test_schema_isolation.py +++ b/hindsight-api/tests/test_schema_isolation.py @@ -11,8 +11,8 @@ import uuid import pytest import pytest_asyncio -from hindsight_api.extensions import RequestContext, TenantContext, TenantExtension from hindsight_api.engine.memory_engine import _current_schema, fq_table +from hindsight_api.extensions import RequestContext, TenantContext, TenantExtension from hindsight_api.migrations import run_migrations @@ -52,6 +52,11 @@ class MultiSchemaTestTenantExtension(TenantExtension): raise AuthenticationError(f"Unknown API key: {context.api_key}") + async def list_tenants(self) -> list: + from hindsight_api.extensions.tenant import Tenant + + return [Tenant(schema=schema) for schema in self.valid_schemas] + async def drop_schema(conn, schema_name: str) -> None: """Drop a schema and all its contents.""" diff --git a/hindsight-api/tests/test_server_module.py b/hindsight-api/tests/test_server_module.py index bc1d29bf..5559f519 100644 --- a/hindsight-api/tests/test_server_module.py +++ b/hindsight-api/tests/test_server_module.py @@ -249,14 +249,14 @@ class TestServerModuleExtensionLoading: # Mock extensions for testing from hindsight_api.extensions import ( - TenantExtension, - TenantContext, - RequestContext, OperationValidatorExtension, - ValidationResult, - RetainContext, RecallContext, ReflectContext, + RequestContext, + RetainContext, + TenantContext, + TenantExtension, + ValidationResult, ) @@ -270,6 +270,11 @@ class MockTenantExtension(TenantExtension): async def authenticate(self, request_context: RequestContext) -> TenantContext: return TenantContext(schema_name="public") + async def list_tenants(self) -> list: + from hindsight_api.extensions.tenant import Tenant + + return [Tenant(schema="public")] + def set_context(self, context) -> None: self._context_set = True diff --git a/hindsight-api/tests/test_worker.py b/hindsight-api/tests/test_worker.py index 7302c358..659554d6 100644 --- a/hindsight-api/tests/test_worker.py +++ b/hindsight-api/tests/test_worker.py @@ -162,6 +162,11 @@ class TestWorkerPoller: claimed = await poller.claim_batch() assert len(claimed) == 3 + # ClaimedTask objects have operation_id, task_dict, schema attributes + for task in claimed: + assert task.operation_id is not None + assert task.task_dict is not None + # Verify tasks are marked as processing with worker_id rows = await pool.fetch( "SELECT status, worker_id FROM async_operations WHERE bank_id = $1", @@ -206,6 +211,7 @@ class TestWorkerPoller: async def test_execute_task_marks_completed(self, pool, clean_operations): """Test that successful task execution marks task as completed.""" from hindsight_api.worker import WorkerPoller + from hindsight_api.worker.poller import ClaimedTask # Create a pending task bank_id = f"test-worker-{uuid.uuid4().hex[:8]}" @@ -234,7 +240,8 @@ class TestWorkerPoller: # Execute the task task_dict = json.loads(payload) - await poller.execute_task(str(op_id), task_dict) + claimed_task = ClaimedTask(operation_id=str(op_id), task_dict=task_dict, schema=None) + await poller.execute_task(claimed_task) assert len(executed) == 1 @@ -250,6 +257,7 @@ class TestWorkerPoller: async def test_execute_task_retries_on_failure(self, pool, clean_operations): """Test that failed task execution triggers retry mechanism.""" from hindsight_api.worker import WorkerPoller + from hindsight_api.worker.poller import ClaimedTask # Create a pending task with retry_count=0 bank_id = f"test-worker-{uuid.uuid4().hex[:8]}" @@ -277,7 +285,8 @@ class TestWorkerPoller: # Execute (should fail and retry) task_dict = json.loads(payload) - await poller.execute_task(str(op_id), task_dict) + claimed_task = ClaimedTask(operation_id=str(op_id), task_dict=task_dict, schema=None) + await poller.execute_task(claimed_task) # Verify task is back to pending with incremented retry_count row = await pool.fetchrow( @@ -292,6 +301,7 @@ class TestWorkerPoller: async def test_execute_task_fails_after_max_retries(self, pool, clean_operations): """Test that task is marked failed after exceeding max retries.""" from hindsight_api.worker import WorkerPoller + from hindsight_api.worker.poller import ClaimedTask # Create a task that has already used all retries bank_id = f"test-worker-{uuid.uuid4().hex[:8]}" @@ -319,7 +329,8 @@ class TestWorkerPoller: # Execute (should fail permanently) task_dict = json.loads(payload) - await poller.execute_task(str(op_id), task_dict) + claimed_task = ClaimedTask(operation_id=str(op_id), task_dict=task_dict, schema=None) + await poller.execute_task(claimed_task) # Verify task is marked as failed row = await pool.fetchrow( @@ -384,9 +395,8 @@ class TestWorkerPoller: # Should only claim the consolidation for the other bank assert len(claimed) == 1 - claimed_op_id, claimed_payload = claimed[0] - assert claimed_op_id == str(other_op_id) - assert claimed_payload["bank_id"] == other_bank_id + assert claimed[0].operation_id == str(other_op_id) + assert claimed[0].task_dict["bank_id"] == other_bank_id # Verify the pending consolidation for first bank is still pending row = await pool.fetchrow( @@ -437,8 +447,7 @@ class TestWorkerPoller: # Should claim the retain task (non-consolidation tasks are unaffected) assert len(claimed) == 1 - claimed_op_id, _ = claimed[0] - assert claimed_op_id == str(retain_op_id) + assert claimed[0].operation_id == str(retain_op_id) class TestWorkerRecovery: @@ -601,7 +610,7 @@ class TestConcurrentWorkers: batch_size=5, # Each worker tries to claim 5 ) claimed = await poller.claim_batch() - workers_claimed[worker_id] = [op_id for op_id, _ in claimed] + workers_claimed[worker_id] = [task.operation_id for task in claimed] # Run all workers concurrently await asyncio.gather( @@ -825,3 +834,186 @@ class TestSyncTaskBackend: # Should not raise, error is logged await backend.submit_task({"type": "test"}) + + +class TestDynamicTenantDiscovery: + """Tests for dynamic tenant discovery via TenantExtension.""" + + @pytest.mark.asyncio + async def test_poller_discovers_tenants_dynamically(self, pool, clean_operations): + """Test that poller calls list_tenants() on each poll cycle.""" + from hindsight_api.extensions.tenant import Tenant, TenantExtension + from hindsight_api.worker import WorkerPoller + + # Create a mock tenant extension that tracks calls + class MockTenantExtension(TenantExtension): + def __init__(self): + self.list_tenants_calls = 0 + self.tenants_to_return: list[Tenant] = [Tenant(schema="public")] + + async def authenticate(self, context): + raise NotImplementedError("Not used in this test") + + async def list_tenants(self) -> list[Tenant]: + self.list_tenants_calls += 1 + return self.tenants_to_return + + mock_extension = MockTenantExtension() + + # Create pending tasks in public schema + bank_id = f"test-worker-{uuid.uuid4().hex[:8]}" + for i in range(2): + op_id = uuid.uuid4() + payload = json.dumps({"type": "test_task", "index": i, "bank_id": bank_id}) + await pool.execute( + """ + INSERT INTO async_operations (operation_id, bank_id, operation_type, status, task_payload) + VALUES ($1, $2, 'test', 'pending', $3::jsonb) + """, + op_id, + bank_id, + payload, + ) + + poller = WorkerPoller( + pool=pool, + worker_id="test-worker-1", + executor=lambda x: None, + batch_size=10, + tenant_extension=mock_extension, + ) + + # First claim_batch should call list_tenants + claimed1 = await poller.claim_batch() + assert mock_extension.list_tenants_calls == 1 + assert len(claimed1) == 2 + + # Add more tasks + for i in range(2): + op_id = uuid.uuid4() + payload = json.dumps({"type": "test_task", "index": i + 10, "bank_id": bank_id}) + await pool.execute( + """ + INSERT INTO async_operations (operation_id, bank_id, operation_type, status, task_payload) + VALUES ($1, $2, 'test', 'pending', $3::jsonb) + """, + op_id, + bank_id, + payload, + ) + + # Second claim_batch should call list_tenants again + claimed2 = await poller.claim_batch() + assert mock_extension.list_tenants_calls == 2 + assert len(claimed2) == 2 + + @pytest.mark.asyncio + async def test_poller_picks_up_new_tenants_without_restart(self, pool, clean_operations): + """Test that new tenants are discovered on subsequent poll cycles.""" + from hindsight_api.extensions.tenant import Tenant, TenantExtension + from hindsight_api.worker import WorkerPoller + + class DynamicTenantExtension(TenantExtension): + def __init__(self): + # Start with just public + self.tenants: list[Tenant] = [Tenant(schema="public")] + self.list_tenants_calls = 0 + + async def authenticate(self, context): + raise NotImplementedError("Not used in this test") + + async def list_tenants(self) -> list[Tenant]: + self.list_tenants_calls += 1 + return self.tenants + + dynamic_extension = DynamicTenantExtension() + + # Create a task in public schema + bank_id = f"test-worker-{uuid.uuid4().hex[:8]}" + op_id = uuid.uuid4() + payload = json.dumps({"type": "test_task", "bank_id": bank_id}) + await pool.execute( + """ + INSERT INTO async_operations (operation_id, bank_id, operation_type, status, task_payload) + VALUES ($1, $2, 'test', 'pending', $3::jsonb) + """, + op_id, + bank_id, + payload, + ) + + poller = WorkerPoller( + pool=pool, + worker_id="test-worker-1", + executor=lambda x: None, + batch_size=10, + tenant_extension=dynamic_extension, + ) + + # First poll - only public schema + claimed1 = await poller.claim_batch() + assert len(claimed1) == 1 + assert claimed1[0].schema is None # public is represented as None + assert dynamic_extension.list_tenants_calls == 1 + + # Simulate tenant list changing (but we won't add a non-existent schema) + # In real world, the schema would be created before list_tenants returns it + # Here we just verify that list_tenants is called again + + # Add another task to public + op_id2 = uuid.uuid4() + payload2 = json.dumps({"type": "test_task", "bank_id": bank_id}) + await pool.execute( + """ + INSERT INTO async_operations (operation_id, bank_id, operation_type, status, task_payload) + VALUES ($1, $2, 'test', 'pending', $3::jsonb) + """, + op_id2, + bank_id, + payload2, + ) + + # Second poll - list_tenants should be called again + claimed2 = await poller.claim_batch() + assert len(claimed2) == 1 + assert dynamic_extension.list_tenants_calls == 2 # Called again on second poll + + # Third poll with no tasks - still calls list_tenants + claimed3 = await poller.claim_batch() + assert len(claimed3) == 0 + assert dynamic_extension.list_tenants_calls == 3 # Called again even with no tasks + + @pytest.mark.asyncio + async def test_poller_without_tenant_extension_uses_public(self, pool, clean_operations): + """Test that poller uses public schema when no tenant extension is configured.""" + from hindsight_api.worker import WorkerPoller + + # Create pending tasks + bank_id = f"test-worker-{uuid.uuid4().hex[:8]}" + for i in range(3): + op_id = uuid.uuid4() + payload = json.dumps({"type": "test_task", "index": i, "bank_id": bank_id}) + await pool.execute( + """ + INSERT INTO async_operations (operation_id, bank_id, operation_type, status, task_payload) + VALUES ($1, $2, 'test', 'pending', $3::jsonb) + """, + op_id, + bank_id, + payload, + ) + + # No tenant_extension provided + poller = WorkerPoller( + pool=pool, + worker_id="test-worker-1", + executor=lambda x: None, + batch_size=10, + ) + + claimed = await poller.claim_batch() + assert len(claimed) == 3 + + # All tasks should have schema=None (public) + for task in claimed: + assert task.schema is None diff --git a/hindsight-cli/src/api.rs b/hindsight-cli/src/api.rs index 0acbf080..2b496ecf 100644 --- a/hindsight-cli/src/api.rs +++ b/hindsight-cli/src/api.rs @@ -485,7 +485,7 @@ impl ApiClient { }) } - pub fn refresh_reflection(&self, bank_id: &str, reflection_id: &str, _verbose: bool) -> Result { + pub fn refresh_reflection(&self, bank_id: &str, reflection_id: &str, _verbose: bool) -> Result { self.runtime.block_on(async { let response = self.client.refresh_reflection(bank_id, reflection_id, None).await?; Ok(response.into_inner()) diff --git a/hindsight-cli/src/commands/reflection.rs b/hindsight-cli/src/commands/reflection.rs index 96e50554..199c5bea 100644 --- a/hindsight-cli/src/commands/reflection.rs +++ b/hindsight-cli/src/commands/reflection.rs @@ -233,7 +233,7 @@ pub fn refresh( output_format: OutputFormat, ) -> Result<()> { let spinner = if output_format == OutputFormat::Pretty { - Some(ui::create_spinner("Refreshing reflection...")) + Some(ui::create_spinner("Submitting reflection refresh...")) } else { None }; @@ -245,13 +245,17 @@ pub fn refresh( } match response { - Ok(reflection) => { + Ok(operation) => { if output_format == OutputFormat::Pretty { - ui::print_success(&format!("Reflection '{}' refreshed successfully", reflection_id)); + ui::print_success(&format!( + "Reflection refresh submitted. Operation ID: {}", + operation.operation_id + )); + println!(" {} {}", ui::dim("Status:"), operation.status); println!(); - print_reflection_detail(&reflection); + println!("{}", ui::dim("Use 'hindsight operations get' to check the operation status.")); } else { - output::print_output(&reflection, output_format)?; + output::print_output(&operation, output_format)?; } Ok(()) } diff --git a/hindsight-clients/python/.openapi-generator/FILES b/hindsight-clients/python/.openapi-generator/FILES index b8bbaefd..afc687ce 100644 --- a/hindsight-clients/python/.openapi-generator/FILES +++ b/hindsight-clients/python/.openapi-generator/FILES @@ -14,6 +14,7 @@ hindsight_client_api/configuration.py hindsight_client_api/exceptions.py hindsight_client_api/models/__init__.py hindsight_client_api/models/add_background_request.py +hindsight_client_api/models/async_operation_submit_response.py hindsight_client_api/models/background_response.py hindsight_client_api/models/bank_list_item.py hindsight_client_api/models/bank_list_response.py diff --git a/hindsight-clients/python/hindsight_client_api/__init__.py b/hindsight-clients/python/hindsight_client_api/__init__.py index d6d3bcb9..61ed2500 100644 --- a/hindsight-clients/python/hindsight_client_api/__init__.py +++ b/hindsight-clients/python/hindsight_client_api/__init__.py @@ -39,6 +39,7 @@ from hindsight_client_api.exceptions import ApiException # import models into sdk package from hindsight_client_api.models.add_background_request import AddBackgroundRequest +from hindsight_client_api.models.async_operation_submit_response import AsyncOperationSubmitResponse from hindsight_client_api.models.background_response import BackgroundResponse from hindsight_client_api.models.bank_list_item import BankListItem from hindsight_client_api.models.bank_list_response import BankListResponse diff --git a/hindsight-clients/python/hindsight_client_api/api/reflections_api.py b/hindsight-clients/python/hindsight_client_api/api/reflections_api.py index e617b1b0..f9d47c8c 100644 --- a/hindsight-clients/python/hindsight_client_api/api/reflections_api.py +++ b/hindsight-clients/python/hindsight_client_api/api/reflections_api.py @@ -19,6 +19,7 @@ from typing_extensions import Annotated from pydantic import Field, StrictStr, field_validator from typing import Any, List, Optional from typing_extensions import Annotated +from hindsight_client_api.models.async_operation_submit_response import AsyncOperationSubmitResponse from hindsight_client_api.models.create_reflection_request import CreateReflectionRequest from hindsight_client_api.models.create_reflection_response import CreateReflectionResponse from hindsight_client_api.models.reflection_list_response import ReflectionListResponse @@ -1300,10 +1301,10 @@ class ReflectionsApi: _content_type: Optional[StrictStr] = None, _headers: Optional[Dict[StrictStr, Any]] = None, _host_index: Annotated[StrictInt, Field(ge=0, le=0)] = 0, - ) -> ReflectionResponse: + ) -> AsyncOperationSubmitResponse: """Refresh reflection - Re-run the source query through reflect and update the content. + Submit an async task to re-run the source query through reflect and update the content. :param bank_id: (required) :type bank_id: str @@ -1344,7 +1345,7 @@ class ReflectionsApi: ) _response_types_map: Dict[str, Optional[str]] = { - '200': "ReflectionResponse", + '200': "AsyncOperationSubmitResponse", '422': "HTTPValidationError", } response_data = await self.api_client.call_api( @@ -1376,10 +1377,10 @@ class ReflectionsApi: _content_type: Optional[StrictStr] = None, _headers: Optional[Dict[StrictStr, Any]] = None, _host_index: Annotated[StrictInt, Field(ge=0, le=0)] = 0, - ) -> ApiResponse[ReflectionResponse]: + ) -> ApiResponse[AsyncOperationSubmitResponse]: """Refresh reflection - Re-run the source query through reflect and update the content. + Submit an async task to re-run the source query through reflect and update the content. :param bank_id: (required) :type bank_id: str @@ -1420,7 +1421,7 @@ class ReflectionsApi: ) _response_types_map: Dict[str, Optional[str]] = { - '200': "ReflectionResponse", + '200': "AsyncOperationSubmitResponse", '422': "HTTPValidationError", } response_data = await self.api_client.call_api( @@ -1455,7 +1456,7 @@ class ReflectionsApi: ) -> RESTResponseType: """Refresh reflection - Re-run the source query through reflect and update the content. + Submit an async task to re-run the source query through reflect and update the content. :param bank_id: (required) :type bank_id: str @@ -1496,7 +1497,7 @@ class ReflectionsApi: ) _response_types_map: Dict[str, Optional[str]] = { - '200': "ReflectionResponse", + '200': "AsyncOperationSubmitResponse", '422': "HTTPValidationError", } response_data = await self.api_client.call_api( diff --git a/hindsight-clients/python/hindsight_client_api/models/__init__.py b/hindsight-clients/python/hindsight_client_api/models/__init__.py index a6eebd19..2a897d1a 100644 --- a/hindsight-clients/python/hindsight_client_api/models/__init__.py +++ b/hindsight-clients/python/hindsight_client_api/models/__init__.py @@ -15,6 +15,7 @@ # import models into model package from hindsight_client_api.models.add_background_request import AddBackgroundRequest +from hindsight_client_api.models.async_operation_submit_response import AsyncOperationSubmitResponse from hindsight_client_api.models.background_response import BackgroundResponse from hindsight_client_api.models.bank_list_item import BankListItem from hindsight_client_api.models.bank_list_response import BankListResponse diff --git a/hindsight-clients/python/hindsight_client_api/models/async_operation_submit_response.py b/hindsight-clients/python/hindsight_client_api/models/async_operation_submit_response.py new file mode 100644 index 00000000..420d1e9e --- /dev/null +++ b/hindsight-clients/python/hindsight_client_api/models/async_operation_submit_response.py @@ -0,0 +1,89 @@ +# coding: utf-8 + +""" + Hindsight HTTP API + + HTTP API for Hindsight + + The version of the OpenAPI document: 0.1.0 + Generated by OpenAPI Generator (https://openapi-generator.tech) + + Do not edit the class manually. +""" # noqa: E501 + + +from __future__ import annotations +import pprint +import re # noqa: F401 +import json + +from pydantic import BaseModel, ConfigDict, StrictStr +from typing import Any, ClassVar, Dict, List +from typing import Optional, Set +from typing_extensions import Self + +class AsyncOperationSubmitResponse(BaseModel): + """ + Response model for submitting an async operation. + """ # noqa: E501 + operation_id: StrictStr + status: StrictStr + __properties: ClassVar[List[str]] = ["operation_id", "status"] + + model_config = ConfigDict( + populate_by_name=True, + validate_assignment=True, + protected_namespaces=(), + ) + + + def to_str(self) -> str: + """Returns the string representation of the model using alias""" + return pprint.pformat(self.model_dump(by_alias=True)) + + def to_json(self) -> str: + """Returns the JSON representation of the model using alias""" + # TODO: pydantic v2: use .model_dump_json(by_alias=True, exclude_unset=True) instead + return json.dumps(self.to_dict()) + + @classmethod + def from_json(cls, json_str: str) -> Optional[Self]: + """Create an instance of AsyncOperationSubmitResponse from a JSON string""" + return cls.from_dict(json.loads(json_str)) + + def to_dict(self) -> Dict[str, Any]: + """Return the dictionary representation of the model using alias. + + This has the following differences from calling pydantic's + `self.model_dump(by_alias=True)`: + + * `None` is only added to the output dict for nullable fields that + were set at model initialization. Other fields with value `None` + are ignored. + """ + excluded_fields: Set[str] = set([ + ]) + + _dict = self.model_dump( + by_alias=True, + exclude=excluded_fields, + exclude_none=True, + ) + return _dict + + @classmethod + def from_dict(cls, obj: Optional[Dict[str, Any]]) -> Optional[Self]: + """Create an instance of AsyncOperationSubmitResponse from a dict""" + if obj is None: + return None + + if not isinstance(obj, dict): + return cls.model_validate(obj) + + _obj = cls.model_validate({ + "operation_id": obj.get("operation_id"), + "status": obj.get("status") + }) + return _obj + + diff --git a/hindsight-clients/python/hindsight_client_api/models/consolidation_response.py b/hindsight-clients/python/hindsight_client_api/models/consolidation_response.py index e1f7c511..502ceae0 100644 --- a/hindsight-clients/python/hindsight_client_api/models/consolidation_response.py +++ b/hindsight-clients/python/hindsight_client_api/models/consolidation_response.py @@ -17,8 +17,8 @@ import pprint import re # noqa: F401 import json -from pydantic import BaseModel, ConfigDict, Field, StrictInt, StrictStr -from typing import Any, ClassVar, Dict, List +from pydantic import BaseModel, ConfigDict, Field, StrictBool, StrictStr +from typing import Any, ClassVar, Dict, List, Optional from typing import Optional, Set from typing_extensions import Self @@ -26,12 +26,9 @@ class ConsolidationResponse(BaseModel): """ Response model for consolidation trigger endpoint. """ # noqa: E501 - status: StrictStr = Field(description="Status of the consolidation (completed or queued)") - processed: StrictInt = Field(description="Number of memories processed") - created: StrictInt = Field(description="Number of mental models created") - updated: StrictInt = Field(description="Number of mental models updated") - message: StrictStr = Field(description="Human-readable summary") - __properties: ClassVar[List[str]] = ["status", "processed", "created", "updated", "message"] + operation_id: StrictStr = Field(description="ID of the async consolidation operation") + deduplicated: Optional[StrictBool] = Field(default=False, description="True if an existing pending task was reused") + __properties: ClassVar[List[str]] = ["operation_id", "deduplicated"] model_config = ConfigDict( populate_by_name=True, @@ -84,11 +81,8 @@ class ConsolidationResponse(BaseModel): return cls.model_validate(obj) _obj = cls.model_validate({ - "status": obj.get("status"), - "processed": obj.get("processed"), - "created": obj.get("created"), - "updated": obj.get("updated"), - "message": obj.get("message") + "operation_id": obj.get("operation_id"), + "deduplicated": obj.get("deduplicated") if obj.get("deduplicated") is not None else False }) return _obj diff --git a/hindsight-clients/typescript/generated/sdk.gen.ts b/hindsight-clients/typescript/generated/sdk.gen.ts index d3b83066..3236dac6 100644 --- a/hindsight-clients/typescript/generated/sdk.gen.ts +++ b/hindsight-clients/typescript/generated/sdk.gen.ts @@ -455,7 +455,7 @@ export const updateReflection = ( /** * Refresh reflection * - * Re-run the source query through reflect and update the content. + * Submit an async task to re-run the source query through reflect and update the content. */ export const refreshReflection = ( options: Options, diff --git a/hindsight-clients/typescript/generated/types.gen.ts b/hindsight-clients/typescript/generated/types.gen.ts index 5d42fe45..a730f1d1 100644 --- a/hindsight-clients/typescript/generated/types.gen.ts +++ b/hindsight-clients/typescript/generated/types.gen.ts @@ -24,6 +24,22 @@ export type AddBackgroundRequest = { update_disposition?: boolean; }; +/** + * AsyncOperationSubmitResponse + * + * Response model for submitting an async operation. + */ +export type AsyncOperationSubmitResponse = { + /** + * Operation Id + */ + operation_id: string; + /** + * Status + */ + status: string; +}; + /** * BackgroundResponse * @@ -295,35 +311,17 @@ export type ChunkResponse = { */ export type ConsolidationResponse = { /** - * Status + * Operation Id * - * Status of the consolidation (completed or queued) + * ID of the async consolidation operation */ - status: string; + operation_id: string; /** - * Processed + * Deduplicated * - * Number of memories processed + * True if an existing pending task was reused */ - processed: number; - /** - * Created - * - * Number of mental models created - */ - created: number; - /** - * Updated - * - * Number of mental models updated - */ - updated: number; - /** - * Message - * - * Human-readable summary - */ - message: string; + deduplicated?: boolean; }; /** @@ -2486,7 +2484,7 @@ export type RefreshReflectionResponses = { /** * Successful Response */ - 200: ReflectionResponse; + 200: AsyncOperationSubmitResponse; }; export type RefreshReflectionResponse = diff --git a/hindsight-control-plane/src/app/banks/[bankId]/page.tsx b/hindsight-control-plane/src/app/banks/[bankId]/page.tsx index 879e6f0d..e936e761 100644 --- a/hindsight-control-plane/src/app/banks/[bankId]/page.tsx +++ b/hindsight-control-plane/src/app/banks/[bankId]/page.tsx @@ -10,6 +10,7 @@ import { ThinkView } from "@/components/think-view"; import { SearchDebugView } from "@/components/search-debug-view"; import { BankProfileView } from "@/components/bank-profile-view"; import { ReflectionsView } from "@/components/reflections-view"; +import { useFeatures } from "@/lib/features-context"; type NavItem = "recall" | "reflect" | "data" | "documents" | "entities" | "profile"; type DataSubTab = "world" | "experience" | "models" | "reflections"; @@ -18,10 +19,12 @@ export default function BankPage() { const params = useParams(); const router = useRouter(); const searchParams = useSearchParams(); + const { features } = useFeatures(); const bankId = params.bankId as string; const view = (searchParams.get("view") || "profile") as NavItem; const subTab = (searchParams.get("subTab") || "world") as DataSubTab; + const mentalModelsEnabled = features?.mental_models ?? false; const handleTabChange = (tab: NavItem) => { router.push(`/banks/${bankId}?view=${tab}`); @@ -120,6 +123,11 @@ export default function BankPage() { }`} > Mental Models + {!mentalModelsEnabled && ( + + Off + + )} {subTab === "models" && (
)} @@ -143,7 +151,40 @@ export default function BankPage() {
{subTab === "world" && } {subTab === "experience" && } - {subTab === "models" && } + {subTab === "models" && + (mentalModelsEnabled ? ( + + ) : ( +
+
+ + + + + +
+

+ Mental Models Not Enabled +

+

+ Mental models consolidation is disabled on this server. Set{" "} + + HINDSIGHT_API_ENABLE_MENTAL_MODELS=true + {" "} + to enable. +

+
+ ))} {subTab === "reflections" && }
diff --git a/hindsight-control-plane/src/app/layout.tsx b/hindsight-control-plane/src/app/layout.tsx index 23c78e9a..0e68d89d 100644 --- a/hindsight-control-plane/src/app/layout.tsx +++ b/hindsight-control-plane/src/app/layout.tsx @@ -1,6 +1,7 @@ import type { Metadata } from "next"; import "./globals.css"; import { BankProvider } from "@/lib/bank-context"; +import { FeaturesProvider } from "@/lib/features-context"; import { ThemeProvider } from "@/lib/theme-context"; export const metadata: Metadata = { @@ -20,7 +21,9 @@ export default function RootLayout({ - {children} + + {children} + diff --git a/hindsight-control-plane/src/components/bank-profile-view.tsx b/hindsight-control-plane/src/components/bank-profile-view.tsx index 9700190f..e21529e5 100644 --- a/hindsight-control-plane/src/components/bank-profile-view.tsx +++ b/hindsight-control-plane/src/components/bank-profile-view.tsx @@ -5,6 +5,7 @@ import ReactMarkdown from "react-markdown"; import { useRouter } from "next/navigation"; import { client } from "@/lib/api"; import { useBank } from "@/lib/bank-context"; +import { useFeatures } from "@/lib/features-context"; import { Button } from "@/components/ui/button"; import { Input } from "@/components/ui/input"; import { Textarea } from "@/components/ui/textarea"; @@ -212,6 +213,8 @@ function DispositionEditor({ export function BankProfileView() { const router = useRouter(); const { currentBank, setCurrentBank, loadBanks } = useBank(); + const { features } = useFeatures(); + const mentalModelsEnabled = features?.mental_models ?? false; const [profile, setProfile] = useState(null); const [stats, setStats] = useState(null); const [operations, setOperations] = useState([]); @@ -391,12 +394,10 @@ export function BankProfileView() { setIsConsolidating(true); try { - const result = await client.triggerConsolidation(currentBank); + await client.triggerConsolidation(currentBank); + // Reload to show the new operation in the list await loadData(); - alert( - result.message || - `Consolidation completed: ${result.created} created, ${result.updated} updated` - ); + await loadOperations(); } catch (error) { console.error("Error triggering consolidation:", error); alert("Error triggering consolidation: " + (error as Error).message); @@ -534,20 +535,32 @@ export function BankProfileView() { Edit Profile - + {isConsolidating ? ( ) : ( )} {isConsolidating ? "Consolidating..." : "Run Consolidation"} + {!mentalModelsEnabled && ( + Off + )} setShowClearMentalModelsDialog(true)} + disabled={!mentalModelsEnabled} className="text-amber-600 dark:text-amber-400 focus:text-amber-700 dark:focus:text-amber-300" + title={!mentalModelsEnabled ? "Mental models feature is not enabled" : undefined} > Clear Mental Models + {!mentalModelsEnabled && ( + Off + )} -
-

+

+

Mental Models + {!mentalModelsEnabled && (Off)}

-

- {stats.total_mental_models || 0} +

+ {mentalModelsEnabled ? stats.total_mental_models || 0 : "—"}

diff --git a/hindsight-control-plane/src/components/data-view.tsx b/hindsight-control-plane/src/components/data-view.tsx index f96e8ee7..0ae94d29 100644 --- a/hindsight-control-plane/src/components/data-view.tsx +++ b/hindsight-control-plane/src/components/data-view.tsx @@ -616,10 +616,24 @@ export function DataView({ factType }: DataViewProps) { - Memory - Entities - Occurred - Mentioned + + {factType === "mental_model" ? "Mental Model" : "Memory"} + + {factType === "mental_model" ? ( + <> + Sources + Created + Mentioned + + ) : ( + <> + Entities + Occurred + Mentioned + + )} @@ -637,6 +651,12 @@ export function DataView({ factType }: DataViewProps) { day: "numeric", }) : null; + const createdDisplay = row.created_at + ? new Date(row.created_at).toLocaleDateString("en-US", { + month: "short", + day: "numeric", + }) + : null; return ( )} - - {row.entities ? ( -
- {row.entities - .split(", ") - .slice(0, 2) - .map((entity: string, i: number) => ( - - {entity} - - ))} - {row.entities.split(", ").length > 2 && ( - - +{row.entities.split(", ").length - 2} + {factType === "mental_model" ? ( + <> + + {row.proof_count || 1} + + + {createdDisplay || ( + - + )} + + + {mentionedDisplay || ( + - + )} + + + ) : ( + <> + + {row.entities ? ( +
+ {row.entities + .split(", ") + .slice(0, 2) + .map((entity: string, i: number) => ( + + {entity} + + ))} + {row.entities.split(", ").length > 2 && ( + + +{row.entities.split(", ").length - 2} + + )} +
+ ) : ( + + - )} -
- ) : ( - - - )} - - - {occurredDisplay || ( - - - )} - - - {mentionedDisplay || ( - - - )} - + + + {occurredDisplay || ( + - + )} + + + {mentionedDisplay || ( + - + )} + + + )} + +

{source.text}

+ {source.context && ( +

+ Context: {source.context} +

+ )} + + ))} + + + )} )} @@ -528,6 +565,12 @@ export function MemoryDetailPanel({ {modalType && modalId && ( )} + + {/* Source Memory Modal */} + setSourceMemoryModalId(null)} + /> ); } diff --git a/hindsight-control-plane/src/components/reflections-view.tsx b/hindsight-control-plane/src/components/reflections-view.tsx index 979b91c2..bf62a1c3 100644 --- a/hindsight-control-plane/src/components/reflections-view.tsx +++ b/hindsight-control-plane/src/components/reflections-view.tsx @@ -541,13 +541,46 @@ function ReflectionDetailPanel({ if (!currentBank) return; setRefreshing(true); + const originalRefreshedAt = reflection.last_refreshed_at; + try { - const updated = await client.refreshReflection(currentBank, reflection.id); - onRefreshed(updated); + // Submit the refresh task + await client.refreshReflection(currentBank, reflection.id); + + // Poll until last_refreshed_at changes + const pollInterval = 1000; // 1 second + const maxAttempts = 120; // 2 minutes max + let attempts = 0; + + const poll = async (): Promise => { + attempts++; + try { + const updated = await client.getReflection(currentBank, reflection.id); + if (updated.last_refreshed_at !== originalRefreshedAt) { + // Refresh complete + onRefreshed(updated); + setRefreshing(false); + return; + } + if (attempts >= maxAttempts) { + // Timeout + setRefreshing(false); + alert("Refresh is taking longer than expected. Check the operations list for status."); + return; + } + // Continue polling + setTimeout(poll, pollInterval); + } catch (error) { + console.error("Error polling reflection:", error); + setRefreshing(false); + } + }; + + // Start polling after a short delay + setTimeout(poll, pollInterval); } catch (error) { console.error("Error refreshing reflection:", error); alert("Error refreshing: " + (error as Error).message); - } finally { setRefreshing(false); } }; diff --git a/hindsight-control-plane/src/lib/api.ts b/hindsight-control-plane/src/lib/api.ts index d718a675..7230fcbe 100644 --- a/hindsight-control-plane/src/lib/api.ts +++ b/hindsight-control-plane/src/lib/api.ts @@ -261,11 +261,8 @@ export class ControlPlaneClient { */ async triggerConsolidation(bankId: string) { return this.fetchApi<{ - status: string; - processed: number; - created: number; - updated: number; - message: string; + operation_id: string; + deduplicated: boolean; }>(`/api/banks/${bankId}/consolidate`, { method: "POST", }); @@ -658,23 +655,11 @@ export class ControlPlaneClient { } /** - * Refresh a reflection (re-run source query) + * Refresh a reflection (re-run source query) - async operation */ async refreshReflection(bankId: string, reflectionId: string) { return this.fetchApi<{ - id: string; - bank_id: string; - name: string; - source_query: string; - content: string; - tags: string[]; - last_refreshed_at: string; - created_at: string; - reflect_response?: { - text: string; - based_on: Record>; - mental_models?: Array<{ id: string; text: string }>; - }; + operation_id: string; }>(`/api/banks/${bankId}/reflections/${reflectionId}/refresh`, { method: "POST", }); diff --git a/hindsight-control-plane/src/lib/features-context.tsx b/hindsight-control-plane/src/lib/features-context.tsx new file mode 100644 index 00000000..9fb6b5d4 --- /dev/null +++ b/hindsight-control-plane/src/lib/features-context.tsx @@ -0,0 +1,63 @@ +"use client"; + +import React, { createContext, useContext, useState, useEffect } from "react"; +import { client } from "./api"; + +interface Features { + mental_models: boolean; + mcp: boolean; + worker: boolean; +} + +interface FeaturesContextType { + features: Features | null; + loading: boolean; + error: string | null; +} + +const defaultFeatures: Features = { + mental_models: false, + mcp: false, + worker: false, +}; + +const FeaturesContext = createContext(undefined); + +export function FeaturesProvider({ children }: { children: React.ReactNode }) { + const [features, setFeatures] = useState(null); + const [loading, setLoading] = useState(true); + const [error, setError] = useState(null); + + useEffect(() => { + const loadFeatures = async () => { + try { + const response = await client.getVersion(); + setFeatures(response.features); + setError(null); + } catch (err) { + console.error("Error loading features:", err); + setError("Failed to load feature flags"); + // Use defaults on error + setFeatures(defaultFeatures); + } finally { + setLoading(false); + } + }; + + loadFeatures(); + }, []); + + return ( + + {children} + + ); +} + +export function useFeatures() { + const context = useContext(FeaturesContext); + if (context === undefined) { + throw new Error("useFeatures must be used within a FeaturesProvider"); + } + return context; +} diff --git a/hindsight-docs/static/openapi.json b/hindsight-docs/static/openapi.json index c4cc88d6..bfb4a6b3 100644 --- a/hindsight-docs/static/openapi.json +++ b/hindsight-docs/static/openapi.json @@ -1205,7 +1205,7 @@ "Reflections" ], "summary": "Refresh reflection", - "description": "Re-run the source query through reflect and update the content.", + "description": "Submit an async task to re-run the source query through reflect and update the content.", "operationId": "refresh_reflection", "parameters": [ { @@ -1249,7 +1249,7 @@ "content": { "application/json": { "schema": { - "$ref": "#/components/schemas/ReflectionResponse" + "$ref": "#/components/schemas/AsyncOperationSubmitResponse" } } } @@ -2980,6 +2980,29 @@ "update_disposition": true } }, + "AsyncOperationSubmitResponse": { + "properties": { + "operation_id": { + "type": "string", + "title": "Operation Id" + }, + "status": { + "type": "string", + "title": "Status" + } + }, + "type": "object", + "required": [ + "operation_id", + "status" + ], + "title": "AsyncOperationSubmitResponse", + "description": "Response model for submitting an async operation.", + "example": { + "operation_id": "550e8400-e29b-41d4-a716-446655440000", + "status": "queued" + } + }, "BackgroundResponse": { "properties": { "mission": { @@ -3427,39 +3450,21 @@ }, "ConsolidationResponse": { "properties": { - "status": { + "operation_id": { "type": "string", - "title": "Status", - "description": "Status of the consolidation (completed or queued)" + "title": "Operation Id", + "description": "ID of the async consolidation operation" }, - "processed": { - "type": "integer", - "title": "Processed", - "description": "Number of memories processed" - }, - "created": { - "type": "integer", - "title": "Created", - "description": "Number of mental models created" - }, - "updated": { - "type": "integer", - "title": "Updated", - "description": "Number of mental models updated" - }, - "message": { - "type": "string", - "title": "Message", - "description": "Human-readable summary" + "deduplicated": { + "type": "boolean", + "title": "Deduplicated", + "description": "True if an existing pending task was reused", + "default": false } }, "type": "object", "required": [ - "status", - "processed", - "created", - "updated", - "message" + "operation_id" ], "title": "ConsolidationResponse", "description": "Response model for consolidation trigger endpoint."