hindsight-api 0.4.2__py3-none-any.whl → 0.4.3__py3-none-any.whl

hindsight_api/mcp_tools.py

@@ -32,6 +32,9 @@ class MCPToolsConfig:
     # How to resolve bank_id for operations
     bank_id_resolver: Callable[[], str | None]
 
+    # How to resolve API key for tenant auth (optional)
+    api_key_resolver: Callable[[], str | None] | None = None
+
     # Whether to include bank_id as a parameter on tools (for multi-bank support)
     include_bank_id_param: bool = False
 
@@ -46,6 +49,16 @@ class MCPToolsConfig:
     retain_fire_and_forget: bool = False  # If True, use asyncio.create_task pattern
 
 
+def _get_request_context(config: MCPToolsConfig) -> RequestContext:
+    """Create RequestContext with API key from resolver if available.
+
+    This enables tenant auth to work with MCP tools by propagating
+    the Bearer token from the MCP middleware to the memory engine.
+    """
+    api_key = config.api_key_resolver() if config.api_key_resolver else None
+    return RequestContext(api_key=api_key)
+
+
 def parse_timestamp(timestamp: str) -> datetime | None:
     """Parse an ISO format timestamp string.
 
@@ -155,12 +168,14 @@ def _register_retain(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfig)
                 if error:
                     return {"status": "error", "message": error}
 
+                request_context = _get_request_context(config)
+
                 async def _retain():
                     try:
                         await memory.retain_batch_async(
                             bank_id=target_bank,
                             contents=[content_dict],
-                            request_context=RequestContext(),
+                            request_context=request_context,
                         )
                     except Exception as e:
                         logger.error(f"Error storing memory: {e}", exc_info=True)
@@ -196,16 +211,17 @@ def _register_retain(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfig)
                         return f"Error: {error}"
 
                     contents = [content_dict]
+                    request_context = _get_request_context(config)
                     if async_processing:
                         result = await memory.submit_async_retain(
-                            bank_id=target_bank, contents=contents, request_context=RequestContext()
+                            bank_id=target_bank, contents=contents, request_context=request_context
                         )
                         return f"Memory queued for background processing (operation_id: {result.get('operation_id', 'N/A')})"
                     else:
                         await memory.retain_batch_async(
                             bank_id=target_bank,
                             contents=contents,
-                            request_context=RequestContext(),
+                            request_context=request_context,
                         )
                         return f"Memory stored successfully in bank '{target_bank}'"
                 except Exception as e:
@@ -237,12 +253,14 @@ def _register_retain(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfig)
             if error:
                 return {"status": "error", "message": error}
 
+            request_context = _get_request_context(config)
+
             async def _retain():
                 try:
                     await memory.retain_batch_async(
                         bank_id=target_bank,
                         contents=[content_dict],
-                        request_context=RequestContext(),
+                        request_context=request_context,
                     )
                 except Exception as e:
                     logger.error(f"Error storing memory: {e}", exc_info=True)
@@ -280,7 +298,7 @@ def _register_recall(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfig)
                     fact_type=list(VALID_RECALL_FACT_TYPES),
                     budget=Budget.HIGH,
                     max_tokens=max_tokens,
-                    request_context=RequestContext(),
+                    request_context=_get_request_context(config),
                 )
 
                 return recall_result.model_dump_json(indent=2)
@@ -311,7 +329,7 @@ def _register_recall(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfig)
                     fact_type=list(VALID_RECALL_FACT_TYPES),
                     budget=Budget.HIGH,
                     max_tokens=max_tokens,
-                    request_context=RequestContext(),
+                    request_context=_get_request_context(config),
                 )
 
                 return recall_result.model_dump()
@@ -370,7 +388,7 @@ def _register_reflect(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfig
                     query=query,
                     budget=budget_enum,
                     context=context,
-                    request_context=RequestContext(),
+                    request_context=_get_request_context(config),
                 )
 
                 return reflect_result.model_dump_json(indent=2)
@@ -423,7 +441,7 @@ def _register_reflect(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfig
                     query=query,
                     budget=budget_enum,
                     context=context,
-                    request_context=RequestContext(),
+                    request_context=_get_request_context(config),
                 )
 
                 return reflect_result.model_dump()
@@ -447,7 +465,7 @@ def _register_list_banks(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsCon
             JSON list of banks with their IDs, names, dispositions, and missions.
         """
         try:
-            banks = await memory.list_banks(request_context=RequestContext())
+            banks = await memory.list_banks(request_context=_get_request_context(config))
             return json.dumps({"banks": banks}, indent=2)
         except Exception as e:
             logger.error(f"Error listing banks: {e}", exc_info=True)
@@ -471,8 +489,9 @@ def _register_create_bank(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsCo
             mission: Optional mission describing who the agent is and what they're trying to accomplish
         """
         try:
+            request_context = _get_request_context(config)
             # get_bank_profile auto-creates bank if it doesn't exist
-            profile = await memory.get_bank_profile(bank_id, request_context=RequestContext())
+            profile = await memory.get_bank_profile(bank_id, request_context=request_context)
 
             # Update name/mission if provided
             if name is not None or mission is not None:
@@ -480,10 +499,10 @@ def _register_create_bank(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsCo
                     bank_id,
                     name=name,
                     mission=mission,
-                    request_context=RequestContext(),
+                    request_context=request_context,
                 )
                 # Fetch updated profile
-                profile = await memory.get_bank_profile(bank_id, request_context=RequestContext())
+                profile = await memory.get_bank_profile(bank_id, request_context=request_context)
 
             # Serialize disposition if it's a Pydantic model
             if "disposition" in profile and hasattr(profile["disposition"], "model_dump"):

hindsight_api/metrics.py

@@ -189,7 +189,7 @@ class MetricsCollectorBase:
         Args:
             provider: LLM provider name (openai, anthropic, gemini, groq, ollama, lmstudio)
             model: Model name
-            scope: Scope identifier (e.g., "memory", "reflect", "entity_observation")
+            scope: Scope identifier (e.g., "memory", "reflect", "consolidation")
             duration: Call duration in seconds
             input_tokens: Number of input/prompt tokens
             output_tokens: Number of output/completion tokens
@@ -321,7 +321,7 @@ class MetricsCollector(MetricsCollectorBase):
                 pass
 
         Args:
-            operation: Operation name (retain, recall, reflect, entity_observation)
+            operation: Operation name (retain, recall, reflect, consolidation)
             bank_id: Memory bank ID
             source: Source of the operation (api, reflect, internal)
             budget: Optional budget level (low, mid, high)
@@ -371,7 +371,7 @@ class MetricsCollector(MetricsCollectorBase):
         Args:
             provider: LLM provider name (openai, anthropic, gemini, groq, ollama, lmstudio)
             model: Model name
-            scope: Scope identifier (e.g., "memory", "reflect", "entity_observation")
+            scope: Scope identifier (e.g., "memory", "reflect", "consolidation")
             duration: Call duration in seconds
             input_tokens: Number of input/prompt tokens
             output_tokens: Number of output/completion tokens

hindsight_api/pg0.py

@@ -40,7 +40,7 @@ class EmbeddedPostgres:
             # Only set port if explicitly specified
             if self.port is not None:
                 kwargs["port"] = self.port
-            self._pg0 = Pg0(**kwargs)  # type: ignore[invalid-argument-type] - dict kwargs
+            self._pg0 = Pg0(**kwargs)
         return self._pg0
 
     async def start(self, max_retries: int = 5, retry_delay: float = 4.0) -> str:

hindsight_api/worker/main.py

@@ -124,12 +124,6 @@ def main():
         default=config.worker_poll_interval_ms,
         help=f"Poll interval in milliseconds (default: {config.worker_poll_interval_ms}, env: HINDSIGHT_API_WORKER_POLL_INTERVAL_MS)",
     )
-    parser.add_argument(
-        "--batch-size",
-        type=int,
-        default=config.worker_batch_size,
-        help=f"Tasks to claim per poll (default: {config.worker_batch_size}, env: HINDSIGHT_API_WORKER_BATCH_SIZE)",
-    )
     parser.add_argument(
         "--max-retries",
         type=int,
@@ -168,8 +162,9 @@ def main():
 
     print(f"Starting Hindsight Worker: {args.worker_id}")
     print(f"  Poll interval: {args.poll_interval}ms")
-    print(f"  Batch size: {args.batch_size}")
     print(f"  Max retries: {args.max_retries}")
+    print(f"  Max slots: {config.worker_max_slots}")
+    print(f"  Consolidation max slots: {config.worker_consolidation_max_slots}")
     print(f"  HTTP server: {args.http_host}:{args.http_port}")
     print()
 
@@ -183,21 +178,25 @@ def main():
 
         from ..extensions import TenantExtension, load_extension
 
+        # Load tenant extension BEFORE creating MemoryEngine so it can
+        # set correct schema context during task execution. Without this,
+        # _authenticate_tenant sees no extension and resets schema to "public",
+        # causing worker writes to land in the wrong schema.
+        tenant_extension = load_extension("TENANT", TenantExtension)
+
         # Initialize MemoryEngine
         # Workers use SyncTaskBackend because they execute tasks directly,
         # they don't need to store tasks (they poll from DB)
         memory = MemoryEngine(
             run_migrations=False,  # Workers don't run migrations
             task_backend=SyncTaskBackend(),
+            tenant_extension=tenant_extension,
         )
 
         await memory.initialize()
 
         print(f"Database connected: {config.database_url}")
 
-        # 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:
@@ -209,9 +208,10 @@ def main():
             worker_id=args.worker_id,
             executor=memory.execute_task,
             poll_interval_ms=args.poll_interval,
-            batch_size=args.batch_size,
             max_retries=args.max_retries,
             tenant_extension=tenant_extension,
+            max_slots=config.worker_max_slots,
+            consolidation_max_slots=config.worker_consolidation_max_slots,
         )
 
         # Create the HTTP app for metrics/health