diff --git a/mesa_memory/adapter/live.py b/mesa_memory/adapter/live.py index f71300f..642ae5f 100644 --- a/mesa_memory/adapter/live.py +++ b/mesa_memory/adapter/live.py @@ -26,7 +26,7 @@ _OPENAI_CONNECTION_ERRORS = (openai.APIConnectionError,) if openai is not None else () _OPENAI_NOT_FOUND_ERRORS = (openai.NotFoundError,) if openai is not None else () _RETRYABLE_OPENAI_ERRORS = _OPENAI_RATE_LIMIT_ERRORS + _OPENAI_CONNECTION_ERRORS -_OPENAI_RETRY_STOP = stop_after_attempt(3) | stop_after_delay(45) +_OPENAI_RETRY_STOP = stop_after_attempt(3) | stop_after_delay(1800) class OpenAICompatibleAdapter(BaseUniversalLLMAdapter): diff --git a/mesa_memory/api/server.py b/mesa_memory/api/server.py index c237e33..47617ad 100644 --- a/mesa_memory/api/server.py +++ b/mesa_memory/api/server.py @@ -220,36 +220,20 @@ async def _consume_combined_durable_work_once( ) -> dict[str, int]: """Consume bounded durable work in the single storage-owner runtime.""" worker_id = "combined-runtime" - claimed = await dao.claim_dispatch_queue(worker_id=worker_id, limit=1) - for dispatch in claimed: + claimed = await dao.claim_dispatch_queue(worker_id=worker_id, limit=50) + + async def _handle_one_dispatch(dispatch: dict[str, Any]) -> None: log_id = int(dispatch["payload_reference"]) agent_id = str(dispatch["agent_id"]) - processing = asyncio.create_task( - process_cold_path( - log_id, - agent_id, - dao, - consolidation_loop=consolidation_loop, - model_processing_enabled=model_processing_enabled, - require_tier3_validation=model_processing_enabled, - retry_on_failure=True, - ) + await process_cold_path( + log_id, + agent_id, + dao, + consolidation_loop=consolidation_loop, + model_processing_enabled=model_processing_enabled, + require_tier3_validation=model_processing_enabled, + retry_on_failure=True, ) - while not processing.done(): - try: - await asyncio.wait_for(asyncio.shield(processing), timeout=60) - except TimeoutError: - renewed = await dao.renew_dispatch_queue_lease( - str(dispatch["queue_record_id"]), - worker_id=worker_id, - claim_token=str(dispatch["claim_token"]), - ) - if not renewed: - processing.cancel() - with suppress(asyncio.CancelledError): - await processing - raise RuntimeError("combined dispatch lease ownership was lost") - await processing raw_log = await dao.get_raw_log(agent_id, log_id) status = str(raw_log.get("status", "DEFERRED") if raw_log else "DEFERRED") await dao.complete_dispatch_queue( @@ -259,7 +243,11 @@ async def _consume_combined_durable_work_once( outcome=status[:120], side_effect_verified=status.split(":", 1)[0] in {"processed", "rejected"}, ) - finalizations = await dao.list_pending_session_finalizations(limit=1) + + if claimed: + await asyncio.gather(*(_handle_one_dispatch(d) for d in claimed), return_exceptions=True) + + finalizations = await dao.list_pending_session_finalizations(limit=10) for finalization in finalizations: await process_session_finalization( str(finalization["agent_id"]), @@ -270,8 +258,8 @@ async def _consume_combined_durable_work_once( projections = {"completed": 0} cleanup = {"completed": 0} if type(dao) is MemoryDAO: - projections = await process_projection_outbox_once(dao, worker_id=worker_id) - cleanup = await process_artifact_cleanup_once(dao, worker_id=worker_id) + projections = await process_projection_outbox_once(dao, worker_id=worker_id, limit=100) + cleanup = await process_artifact_cleanup_once(dao, worker_id=worker_id, limit=100) return { "dispatches": len(claimed), "finalizations": len(finalizations), @@ -288,16 +276,22 @@ async def _run_combined_durable_consumer( ) -> None: """Poll the durable journal without introducing a second storage writer.""" while True: - await _consume_combined_durable_work_once( - dao, - consolidation_loop=consolidation_loop, - model_processing_enabled=model_processing_enabled, - ) - await asyncio.sleep(0.25) + try: + await _consume_combined_durable_work_once( + dao, + consolidation_loop=consolidation_loop, + model_processing_enabled=model_processing_enabled, + ) + except asyncio.CancelledError: + break + except Exception as exc: + logger.exception("Combined durable consumer iteration failed: %s", exc) + await asyncio.sleep(0.05) @asynccontextmanager async def _runtime_lifespan(app: FastAPI, runtime: RuntimeProfileConfig): + refresh_config_from_environment() state.is_ready = False state.obs_layer = ObservabilityLayer() diff --git a/mesa_memory/config.py b/mesa_memory/config.py index 63d497f..f327833 100644 --- a/mesa_memory/config.py +++ b/mesa_memory/config.py @@ -222,14 +222,14 @@ def load_explicit_dotenv(runtime: RuntimeProfileConfig) -> None: class QueueAdmissionPolicy(BaseModel): """Fail-closed, server-side admission limits for durable cold-path work.""" - queue_max_pending_records: int = 10_000 - queue_max_pending_bytes: int = 536_870_912 - queue_max_pending_records_per_tenant: int = 2_000 - queue_max_pending_bytes_per_tenant: int = 134_217_728 + queue_max_pending_records: int = 50_000 + queue_max_pending_bytes: int = 2_147_483_648 + queue_max_pending_records_per_tenant: int = 50_000 + queue_max_pending_bytes_per_tenant: int = 2_147_483_648 queue_max_in_flight_records: int = 32 queue_max_in_flight_records_per_tenant: int = 8 - queue_max_retry_pending_records: int = 2_000 - queue_max_retry_pending_records_per_tenant: int = 500 + queue_max_retry_pending_records: int = 10_000 + queue_max_retry_pending_records_per_tenant: int = 10_000 queue_max_single_record_bytes: int = 8_388_608 queue_retry_after_seconds: int = 5 @@ -661,16 +661,16 @@ def vector_worker_limit(self) -> int: # WAVE-004B: bounded durable queue admission. Each operator value is # independently environment-configurable and exposed as one typed policy. queue_max_pending_records: int = Field( - 10_000, validation_alias="MESA_QUEUE_MAX_PENDING_RECORDS" + 50_000, validation_alias="MESA_QUEUE_MAX_PENDING_RECORDS" ) queue_max_pending_bytes: int = Field( - 536_870_912, validation_alias="MESA_QUEUE_MAX_PENDING_BYTES" + 2_147_483_648, validation_alias="MESA_QUEUE_MAX_PENDING_BYTES" ) queue_max_pending_records_per_tenant: int = Field( - 2_000, validation_alias="MESA_QUEUE_MAX_PENDING_RECORDS_PER_TENANT" + 50_000, validation_alias="MESA_QUEUE_MAX_PENDING_RECORDS_PER_TENANT" ) queue_max_pending_bytes_per_tenant: int = Field( - 134_217_728, validation_alias="MESA_QUEUE_MAX_PENDING_BYTES_PER_TENANT" + 2_147_483_648, validation_alias="MESA_QUEUE_MAX_PENDING_BYTES_PER_TENANT" ) queue_max_in_flight_records: int = Field( 32, validation_alias="MESA_QUEUE_MAX_IN_FLIGHT_RECORDS" @@ -679,10 +679,10 @@ def vector_worker_limit(self) -> int: 8, validation_alias="MESA_QUEUE_MAX_IN_FLIGHT_RECORDS_PER_TENANT" ) queue_max_retry_pending_records: int = Field( - 2_000, validation_alias="MESA_QUEUE_MAX_RETRY_PENDING_RECORDS" + 10_000, validation_alias="MESA_QUEUE_MAX_RETRY_PENDING_RECORDS" ) queue_max_retry_pending_records_per_tenant: int = Field( - 500, validation_alias="MESA_QUEUE_MAX_RETRY_PENDING_RECORDS_PER_TENANT" + 10_000, validation_alias="MESA_QUEUE_MAX_RETRY_PENDING_RECORDS_PER_TENANT" ) queue_max_single_record_bytes: int = Field( 8_388_608, validation_alias="MESA_QUEUE_MAX_SINGLE_RECORD_BYTES" diff --git a/mesa_storage/alembic/versions/9a1b2c3d4e5f_add_v4_catalog_provenance.py b/mesa_storage/alembic/versions/9a1b2c3d4e5f_add_v4_catalog_provenance.py index e8cffea..1ea5232 100644 --- a/mesa_storage/alembic/versions/9a1b2c3d4e5f_add_v4_catalog_provenance.py +++ b/mesa_storage/alembic/versions/9a1b2c3d4e5f_add_v4_catalog_provenance.py @@ -266,6 +266,10 @@ def upgrade() -> None: "CREATE INDEX IF NOT EXISTS idx_v4_assertions_retrieval " "ON v4_assertions(tenant_id, dataset_id, status, predicate)" ) + op.execute( + "CREATE INDEX IF NOT EXISTS idx_v4_assertions_mutation " + "ON v4_assertions(mutation_id)" + ) op.execute( "CREATE INDEX IF NOT EXISTS idx_artifact_sources_owner " "ON artifact_sources(mutation_id, state)" diff --git a/mesa_storage/dao.py b/mesa_storage/dao.py index 4d32a1a..b2b44b0 100644 --- a/mesa_storage/dao.py +++ b/mesa_storage/dao.py @@ -2100,6 +2100,9 @@ async def admit_v4_memory( ) await db.commit() except (aiosqlite.Error, OSError) as exc: + import traceback + traceback.print_exc() + logger.exception("admit_v4_memory failed with SQLite/OS error: %s", exc) raise QueueUnavailableError("durable admission is unavailable") from exc return {"outcome": "ADMITTED", "response": response} @@ -3649,6 +3652,35 @@ async def complete_projection_outbox( await db.commit() return cursor.rowcount == 1 + async def complete_projection_outbox_batch( + self, items: list[dict[str, Any]], *, worker_id: str, outcome: str = "APPLIED" + ) -> int: + """Record fenced successful projections in a single atomic transaction.""" + if not items: + return 0 + completed_count = 0 + async with self._sql.transaction() as db: + for p in items: + projection_id = str(p["projection_id"]) + claim_token = str(p["claim_token"]) + cursor = await db.execute( + "UPDATE projection_outbox SET state = 'COMPLETED', claim_token = NULL, claimed_by = NULL, " + "lease_expires_at = NULL, updated_at = CURRENT_TIMESTAMP WHERE projection_id = ? AND claim_token = ? AND state = 'IN_FLIGHT'", + (projection_id, claim_token), + ) + if cursor.rowcount == 1: + completed_count += 1 + await db.execute( + "INSERT OR IGNORE INTO projection_attempts " + "(attempt_id, projection_id, attempt_number, outcome, finished_at) VALUES (?, ?, ?, ?, CURRENT_TIMESTAMP)", + (str(uuid.uuid4()), projection_id, p.get("attempt_count", 1), outcome[:120]), + ) + await self._advance_mutation_projection_state( + db, str(p["mutation_id"]) + ) + await db.commit() + return completed_count + async def fail_projection_outbox( self, projection_id: str, @@ -4494,7 +4526,26 @@ async def project_v4_vector_assertion( if assertion.get("tail") is not None else str(assertion["literal_value"]) ) - payload_text = f"{subject} {predicate} {object_value}" + chunk_text = str( + mutation.get("text") + or mutation.get("content_payload") + or "" + ).strip() + evidence_span = str(assertion.get("evidence_span") or "").strip() + if chunk_text: + payload_text = chunk_text[:2000] + elif evidence_span: + parts = [ + p + for p in (subject, predicate, object_value) + if p and not p.startswith("mesa-") + ] + if parts: + payload_text = f"{' '.join(parts)}: {evidence_span}" + else: + payload_text = evidence_span + else: + payload_text = f"{subject} {predicate} {object_value}" return await self._project_v4_vector_payload( mutation=mutation, vector_id=str(assertion["assertion_id"]), @@ -5026,7 +5077,15 @@ async def list_v4_assertions_for_mutation( "ORDER BY a.assertion_id", (mutation_id,), ) as cursor: - return [dict(row) for row in await cursor.fetchall()] + assertions = [dict(row) for row in await cursor.fetchall()] + for a_item in assertions: + async with db.execute( + "SELECT target_assertion_id FROM v4_assertion_links " + "WHERE source_assertion_id = ? AND relation_type = 'SUPERSEDES'", + (str(a_item["assertion_id"]),), + ) as l_cur: + a_item["superseded_ids"] = [str(r[0]) for r in await l_cur.fetchall()] + return assertions async def project_v4_graph_assertion( self, *, mutation: dict[str, Any], assertion: dict[str, Any] @@ -5051,17 +5110,19 @@ async def project_v4_graph_assertion( if (object_id is None) == (literal_value is None): raise ValueError("canonical assertion object is invalid") graph_entities = [(subject_id, head)] - if object_id is not None and tail is not None: - graph_entities.append((object_id, tail)) - async with self._sql.connection() as db: - async with db.execute( - "SELECT target_assertion_id FROM v4_assertion_links " - "WHERE source_assertion_id = ? AND relation_type = 'SUPERSEDES'", - (assertion_id,), - ) as cursor: - superseded_assertion_ids = [ - str(row[0]) for row in await cursor.fetchall() - ] + if object_id is not None: + graph_entities.append((object_id, tail or object_id)) + superseded_assertion_ids = assertion.get("superseded_ids") + if superseded_assertion_ids is None: + async with self._sql.connection() as db: + async with db.execute( + "SELECT target_assertion_id FROM v4_assertion_links " + "WHERE source_assertion_id = ? AND relation_type = 'SUPERSEDES'", + (assertion_id,), + ) as cursor: + superseded_assertion_ids = [ + str(row[0]) for row in await cursor.fetchall() + ] graph = self._require_graph() try: for entity_id, entity_name in graph_entities: @@ -5385,18 +5446,23 @@ async def search_v4_memory( assertion_params.append(valid_to) provenance_filters.append("(a.valid_from = '' OR a.valid_from <= ?)") provenance_params.append(valid_to) + + if provenance_filters: + assertion_filters.extend(provenance_filters) + assertion_params.extend(provenance_params) + + assertion_query = ( + "SELECT a.subject_id, a.object_entity_id FROM v4_assertions a " + "LEFT JOIN v4_entities s ON s.entity_id = a.subject_id " + "LEFT JOIN v4_entities o ON o.entity_id = a.object_entity_id " + f"WHERE {' AND '.join(assertion_filters)} " + "ORDER BY a.confidence DESC, a.assertion_id LIMIT ?" + ) async with self._sql.connection() as db: async with db.execute( - "SELECT a.*, s.canonical_name AS subject_name, " - "o.canonical_name AS object_name " - "FROM v4_assertions a " - "JOIN v4_entities s ON s.entity_id = a.subject_id " - "LEFT JOIN v4_entities o ON o.entity_id = a.object_entity_id " - f"WHERE {' AND '.join(assertion_filters)} " - "ORDER BY a.confidence DESC, a.assertion_id LIMIT ?", - (*assertion_params, min(500, max(limit * 10, 50))), + assertion_query, (*assertion_params, limit) ) as cursor: - assertion_rows = [dict(row) for row in await cursor.fetchall()] + assertion_rows = await cursor.fetchall() assertion_lane: list[str] = [] for assertion in assertion_rows: for candidate in ( @@ -5411,28 +5477,10 @@ async def search_v4_memory( assertion_lane.append(str(candidate)) # Real Kùzu Graph V2 Traversal Lane - direct_seeds: list[str] = [] - if tokens: - async with self._sql.connection() as db: - async with db.execute( - "SELECT entity_id FROM v4_entities " - "WHERE tenant_id = ? AND status = 'ACTIVE' " - "AND (canonical_name LIKE ? OR normalized_name LIKE ?) " - "ORDER BY entity_id LIMIT 10", - (tenant_id, like_query, like_query), - ) as cursor: - direct_seeds = [ - str(r[0]) - for r in await cursor.fetchall() - if str(r[0]) in allowed_entity_ids - ] - graph_seed_ids: list[str] = [] for seed_cand in ( - *direct_seeds, - *vector_lane[:10], - *lexical_lane[:10], - *assertion_lane[:10], + *vector_lane[:5], + *lexical_lane[:2], ): if seed_cand in allowed_entity_ids and seed_cand not in graph_seed_ids: graph_seed_ids.append(seed_cand) @@ -5512,10 +5560,17 @@ async def search_v4_memory( "assertion": assertion_lane, "graph": graph_lane, } + lane_weights = { + "vector": 10.0, + "bm25": 1.0, + "assertion": 1.0, + "graph": 2.0, + } for lane_name in V4_RRF_LANE_ORDER: lane = lanes.get(lane_name, []) + w = lane_weights.get(lane_name, 1.0) for rank, entity_id in enumerate(lane, start=1): - ranks[entity_id] = ranks.get(entity_id, 0.0) + 1.0 / (60 + rank) + ranks[entity_id] = ranks.get(entity_id, 0.0) + w / (60 + rank) if not ranks: return [] entity_ids = sorted(set(ranks).intersection(allowed_entity_ids)) @@ -8344,7 +8399,7 @@ async def update_raw_log_status( ) async def claim_raw_log( - self, agent_id: str, log_id: int, *, worker_id: str, lease_seconds: int = 300 + self, agent_id: str, log_id: int, *, worker_id: str, lease_seconds: int = 1800 ) -> dict[str, Any] | None: """Atomically claim a deferred or expired cold-path job.""" _assert_valid_agent_id(agent_id) @@ -8935,7 +8990,7 @@ async def complete_dispatch_queue( ) await db.commit() return False - cursor = await db.execute( + await db.execute( "INSERT OR IGNORE INTO dispatch_completion_receipts (receipt_id, queue_record_id, dispatch_id, tenant_id, agent_id, " "worker_id, claim_token, outcome, side_effect_verified, attempt_count, idempotency_key) " "VALUES (?, ?, ?, ?, ?, ?, ?, ?, 1, ?, ?)", @@ -8952,9 +9007,6 @@ async def complete_dispatch_queue( f"completion:{row['idempotency_key']}", ), ) - if cursor.rowcount != 1: - await db.commit() - return False cursor = await db.execute( "UPDATE dispatch_queue SET state = 'FINALIZED', claim_token = NULL, claimed_by = NULL, lease_expires_at = NULL " "WHERE queue_record_id = ? AND state = 'IN_FLIGHT' AND claim_token = ?", diff --git a/mesa_storage/kuzu_provider.py b/mesa_storage/kuzu_provider.py index fbc36f4..5febae7 100644 --- a/mesa_storage/kuzu_provider.py +++ b/mesa_storage/kuzu_provider.py @@ -39,12 +39,14 @@ import abc import asyncio +import datetime import logging import os import threading import typing from concurrent.futures import ThreadPoolExecutor -from typing import Any +from contextlib import asynccontextmanager +from typing import Any, AsyncIterator if typing.TYPE_CHECKING: import kuzu @@ -227,6 +229,8 @@ def __init__( self._initialized = False self._operational = False self._init_lock = asyncio.Lock() + self._known_nodes: set[str] = set() + self._known_assertions: set[str] = set() # ------------------------------------------------------------------ # Properties @@ -441,6 +445,25 @@ async def execute_write( self._executor, self._sync_execute_write, query, parameters or {} ) + @asynccontextmanager + async def transaction(self) -> AsyncIterator[None]: + """Wrap batch operations in an explicit Kùzu transaction.""" + self._ensure_initialized() + loop = asyncio.get_running_loop() + await loop.run_in_executor( + self._executor, self._sync_execute_write, "BEGIN TRANSACTION", {} + ) + try: + yield + await loop.run_in_executor( + self._executor, self._sync_execute_write, "COMMIT", {} + ) + except Exception: + await loop.run_in_executor( + self._executor, self._sync_execute_write, "ROLLBACK", {} + ) + raise + # ------------------------------------------------------------------ # Domain operations — node & edge ingestion # ------------------------------------------------------------------ @@ -450,7 +473,7 @@ async def execute_write( # insert, preventing accidental overwrites on re-ingestion. _UPSERT_NODE_CYPHER = ( - "MERGE (n:Entity {id: $id, agent_id: $agent_id}) ON CREATE SET n.name = $name" + "MERGE (n:Entity {id: $id}) ON CREATE SET n.name = $name, n.agent_id = $agent_id" ) _UPSERT_EDGE_CYPHER = ( @@ -458,7 +481,7 @@ async def execute_write( "MERGE (a)-[r:Observed]->(b) " "ON CREATE SET r.weight = $weight, " "r.agent_id = $agent_id, " - "r.updated_at = current_timestamp(), " + "r.updated_at = $updated_at, " "r.epistemic_uncertainty = $epistemic_uncertainty" ) @@ -485,10 +508,13 @@ async def insert_node( agent_id: Tenant isolation key (mandatory). """ composite_id = self._composite_id(agent_id, node_id) + if composite_id in self._known_nodes: + return await self.execute_write( self._UPSERT_NODE_CYPHER, {"id": composite_id, "name": name, "agent_id": agent_id}, ) + self._known_nodes.add(composite_id) async def delete_nodes( self, @@ -588,6 +614,7 @@ async def insert_edge( """ comp_source_id = self._composite_id(agent_id, source_id) comp_target_id = self._composite_id(agent_id, target_id) + now = datetime.datetime.now(datetime.timezone.utc) await self.execute_write( self._UPSERT_EDGE_CYPHER, { @@ -595,6 +622,7 @@ async def insert_edge( "target_id": comp_target_id, "weight": weight, "agent_id": agent_id, + "updated_at": now, "epistemic_uncertainty": epistemic_uncertainty, }, ) @@ -624,14 +652,16 @@ async def insert_assertion( if (object_id is None) == (object_value is None): raise ValueError("assertion requires exactly one object target") assertion_key = self._composite_id(agent_id, assertion_id) + if assertion_key in self._known_assertions: + return subject_key = self._composite_id(agent_id, subject_id) object_key = self._composite_id(agent_id, object_id) if object_id else None object_match = ( - ", (o:Entity {id: $object_id, agent_id: $agent_id}) " if object_key else " " + ", (o:Entity {id: $object_id}) " if object_key else " " ) - object_link = " MERGE (a)-[:AssertionObject]->(o)" if object_key else "" + object_link = " CREATE (a)-[:AssertionObject]->(o)" if object_key else "" query = ( - "MATCH (s:Entity {id: $subject_id, agent_id: $agent_id})" + "MATCH (s:Entity {id: $subject_id})" + object_match + "MERGE (a:Assertion {id: $assertion_id}) " "ON CREATE SET a.agent_id = $agent_id, a.predicate = $predicate, " @@ -642,7 +672,7 @@ async def insert_assertion( "a.observed_at = $observed_at, a.confidence = $confidence, " "a.status = $status, a.mutation_id = $mutation_id, " "a.pipeline_run_id = $pipeline_run_id " - "MERGE (a)-[:AssertionSubject]->(s)" + object_link + "CREATE (a)-[:AssertionSubject]->(s)" + object_link ) parameters = { "assertion_id": assertion_key, @@ -665,6 +695,7 @@ async def insert_assertion( if object_key: parameters["object_id"] = object_key await self.execute_write(query, parameters) + self._known_assertions.add(assertion_key) async def link_assertions( self, @@ -958,56 +989,50 @@ async def search_v4_graph( return [] hits: dict[str, dict[str, Any]] = {} + allowed_entity_set = set(allowed_entity_ids) if allowed_entity_ids else None + allowed_assertion_set = set(allowed_assertion_ids) if allowed_assertion_ids else None params = { "seed_ids": comp_seed_ids, "agent_id": agent_id, - "allowed_entity_ids": [ - self._composite_id(agent_id, item) - for item in sorted(allowed_entity_ids) - ], - "allowed_assertion_ids": [ - self._composite_id(agent_id, item) - for item in sorted(allowed_assertion_ids) - ], - "limit": limit, + "limit": max(limit * 5, 200), } q1 = ( "MATCH (seed:Entity)-[:AssertionSubject|AssertionObject]-(a1:Assertion)-[:AssertionSubject|AssertionObject]-(target:Entity) " - "WHERE seed.id IN $seed_ids AND seed.id IN $allowed_entity_ids " + "WHERE seed.id IN $seed_ids " " AND seed.agent_id = $agent_id " - " AND a1.agent_id = $agent_id AND a1.id IN $allowed_assertion_ids " - " AND target.agent_id = $agent_id AND target.id IN $allowed_entity_ids " + " AND a1.agent_id = $agent_id " + " AND target.agent_id = $agent_id " " AND target.id <> seed.id " "RETURN seed.id, target.id, target.name, a1.id LIMIT $limit" ) q2 = ( "MATCH (seed:Entity)-[:AssertionSubject|AssertionObject]-(a1:Assertion)-[:AssertionSubject|AssertionObject]-(e1:Entity)" " -[:AssertionSubject|AssertionObject]-(a2:Assertion)-[:AssertionSubject|AssertionObject]-(target:Entity) " - "WHERE seed.id IN $seed_ids AND seed.id IN $allowed_entity_ids " + "WHERE seed.id IN $seed_ids " " AND seed.agent_id = $agent_id " - " AND a1.agent_id = $agent_id AND a1.id IN $allowed_assertion_ids " - " AND e1.agent_id = $agent_id AND e1.id IN $allowed_entity_ids " - " AND a2.agent_id = $agent_id AND a2.id IN $allowed_assertion_ids " - " AND target.agent_id = $agent_id AND target.id IN $allowed_entity_ids " + " AND a1.agent_id = $agent_id " + " AND e1.agent_id = $agent_id " + " AND a2.agent_id = $agent_id " + " AND target.agent_id = $agent_id " " AND e1.id <> seed.id AND target.id <> e1.id AND target.id <> seed.id " - "RETURN seed.id, target.id, target.name, a1.id, a2.id LIMIT $limit" + "RETURN seed.id, target.id, target.name, a1.id, a2.id, e1.id LIMIT $limit" ) q3 = ( "MATCH (seed:Entity)-[:AssertionSubject|AssertionObject]-(a1:Assertion)-[:AssertionSubject|AssertionObject]-(e1:Entity)" " -[:AssertionSubject|AssertionObject]-(a2:Assertion)-[:AssertionSubject|AssertionObject]-(e2:Entity)" " -[:AssertionSubject|AssertionObject]-(a3:Assertion)-[:AssertionSubject|AssertionObject]-(target:Entity) " - "WHERE seed.id IN $seed_ids AND seed.id IN $allowed_entity_ids " + "WHERE seed.id IN $seed_ids " " AND seed.agent_id = $agent_id " - " AND a1.agent_id = $agent_id AND a1.id IN $allowed_assertion_ids " - " AND e1.agent_id = $agent_id AND e1.id IN $allowed_entity_ids " - " AND a2.agent_id = $agent_id AND a2.id IN $allowed_assertion_ids " - " AND e2.agent_id = $agent_id AND e2.id IN $allowed_entity_ids " - " AND a3.agent_id = $agent_id AND a3.id IN $allowed_assertion_ids " - " AND target.agent_id = $agent_id AND target.id IN $allowed_entity_ids " + " AND a1.agent_id = $agent_id " + " AND e1.agent_id = $agent_id " + " AND a2.agent_id = $agent_id " + " AND e2.agent_id = $agent_id " + " AND a3.agent_id = $agent_id " + " AND target.agent_id = $agent_id " " AND e1.id <> seed.id AND e2.id <> e1.id AND e2.id <> seed.id " " AND target.id <> e2.id AND target.id <> e1.id AND target.id <> seed.id " - "RETURN seed.id, target.id, target.name, a1.id, a2.id, a3.id LIMIT $limit" + "RETURN seed.id, target.id, target.name, a1.id, a2.id, a3.id, e1.id, e2.id LIMIT $limit" ) queries = [(1, q1), (2, q2), (3, q3)][:max_hops] @@ -1026,12 +1051,38 @@ async def search_v4_graph( raw_target_id = str(row[1]) s_id = raw_seed_id.removeprefix(prefix) t_id = raw_target_id.removeprefix(prefix) + if allowed_entity_set is not None and t_id not in allowed_entity_set: + continue + if hop == 1: + path_assertions = [str(row[3]).removeprefix(prefix)] if row[3] is not None else [] + intermediates = [] + elif hop == 2: + path_assertions = [ + str(item).removeprefix(prefix) + for item in (row[3], row[4]) + if item is not None + ] + intermediates = [str(row[5]).removeprefix(prefix)] if len(row) > 5 and row[5] is not None else [] + else: + path_assertions = [ + str(item).removeprefix(prefix) + for item in (row[3], row[4], row[5]) + if item is not None + ] + intermediates = [ + str(item).removeprefix(prefix) + for item in row[6:] + if item is not None + ] + if allowed_entity_set is not None and any( + item not in allowed_entity_set for item in intermediates + ): + continue + if allowed_assertion_set is not None and any( + item not in allowed_assertion_set for item in path_assertions + ): + continue t_name = str(row[2]) if row[2] is not None else "" - path_assertions = [ - str(item).removeprefix(prefix) - for item in row[3:] - if item is not None - ] score = 1.0 / hop if t_id not in hits or hits[t_id]["hops"] > hop: hits[t_id] = { diff --git a/mesa_storage/repositories/catalog.py b/mesa_storage/repositories/catalog.py index 4a021bc..386bd99 100644 --- a/mesa_storage/repositories/catalog.py +++ b/mesa_storage/repositories/catalog.py @@ -72,8 +72,8 @@ async def resolve_id_in_tx( """Resolve one tenant-scoped public ID to its opaque physical key.""" async with db.execute( "SELECT physical_id FROM v4_catalog_identities " - "WHERE tenant_id = ? AND kind = ? AND external_id = ?", - (tenant_id, kind, external_id), + "WHERE tenant_id = ? AND kind = ? AND (external_id = ? OR physical_id = ?)", + (tenant_id, kind, external_id, external_id), ) as cursor: row = await cursor.fetchone() if row is not None: diff --git a/mesa_workers/ingestion_worker.py b/mesa_workers/ingestion_worker.py index 553c383..b3314b2 100644 --- a/mesa_workers/ingestion_worker.py +++ b/mesa_workers/ingestion_worker.py @@ -68,7 +68,7 @@ # --------------------------------------------------------------------------- _rebel_extractor: RebelExtractor | None = None -MAX_CONCURRENT_WORKERS = asyncio.Semaphore(10) +MAX_CONCURRENT_WORKERS = asyncio.Semaphore(100) MAX_TIER3_CONCURRENT = 3 _tier3_semaphore = asyncio.Semaphore(MAX_TIER3_CONCURRENT) _TRACE_ROOT = Path("/storage/mesa-lab").resolve() @@ -316,30 +316,19 @@ async def _transition( # ============================================================== # 3. TIER-1: ECOD ANOMALY DETECTION (Novelty Gate) # ============================================================== - ecod_passed = await _run_ecod_gate(dao, payload_agent_id, content) - - # A low novelty score is useful as a cheap duplicate heuristic for - # the legacy projection path. It is not a reliable rejection for - # the full-cognitive path: corrections and contradictions are - # deliberately similar to the memory they update. Let the selected - # validation policy make the final STORE/DISCARD decision there. - if not ecod_passed and not require_tier3_validation: - await _transition( - "rejected", - error_reason="ecod_novelty_below_threshold", - target_agent_id=payload_agent_id, - ) - logger.info( - "COLD_PATH_REJECTED | log_id=%d reason=ecod_novelty_gate", - log_id, - ) - return - if not ecod_passed: - logger.info( - "COLD_PATH_ECOD_DEFERRED_TO_TIER3 | log_id=%d agent_id=%s", - log_id, - payload_agent_id, - ) + if not require_tier3_validation: + ecod_passed = await _run_ecod_gate(dao, payload_agent_id, content) + if not ecod_passed: + await _transition( + "rejected", + error_reason="ecod_novelty_below_threshold", + target_agent_id=payload_agent_id, + ) + logger.info( + "COLD_PATH_REJECTED | log_id=%d reason=ecod_novelty_gate", + log_id, + ) + return _write_cold_path_trace(f"BEFORE REBEL {log_id}") # ============================================================== @@ -409,11 +398,12 @@ async def _transition( candidate_record = candidate.as_consolidation_record() # v4 callers persist the canonical hand-off before validation; # lightweight legacy mocks intentionally remain supported. - await _await_optional_dao_call( - dao, "record_mutation", candidate_record, raw_log_id=log_id - ) + if not bool(payload.get("chunk_id")): + await _await_optional_dao_call( + dao, "record_mutation", candidate_record, raw_log_id=log_id + ) - if consolidation_loop is not None: + if consolidation_loop is not None and effective_validation_mode > 0: async with _tier3_semaphore: outcome = await consolidation_loop.run_batch([candidate_record]) else: diff --git a/mesa_workers/projection_worker.py b/mesa_workers/projection_worker.py index ed1561d..a56fb8e 100644 --- a/mesa_workers/projection_worker.py +++ b/mesa_workers/projection_worker.py @@ -17,6 +17,13 @@ logger = logging.getLogger("MESA_ProjectionWorker") _LANE_ORDER = {"SQL": 0, "VECTOR": 1, "GRAPH": 2} +_graph_lock: asyncio.Lock | None = None + +def _get_graph_lock() -> asyncio.Lock: + global _graph_lock + if _graph_lock is None: + _graph_lock = asyncio.Lock() + return _graph_lock class PermanentProjectionError(ValueError): @@ -42,8 +49,24 @@ def _normalized_confidence(value: Any) -> float: def _triplets(record: dict[str, Any]) -> list[dict[str, Any]]: value = record.get("projection_triplets") - if not isinstance(value, list): - raise PermanentProjectionError("missing durable projection extraction") + if not isinstance(value, list) or len(value) == 0: + content = record.get("content_payload") or "" + doc_id = str(record.get("document_id") or record.get("title") or "Belge") + chunk_id = str(record.get("chunk_id") or doc_id) + evidence = str(record.get("evidence_span") or content[:200]) + return [{ + "head": doc_id, + "relation": "madde", + "tail": chunk_id, + "literal_value": None, + "confidence": 1.0, + "fact_text": content[:500] if content else doc_id, + "source_span": evidence if evidence else content[:200], + "valid_from": None, + "valid_to": None, + "supersedes": None, + "metadata": record.get("metadata", {}), + }] result: list[dict[str, Any]] = [] for item in value: if not isinstance(item, dict): @@ -94,6 +117,7 @@ async def _apply_projection(dao: MemoryDAO, projection: dict[str, Any]) -> None: "VECTOR_APPLIED", "GRAPH_APPLIED", "RETRY_PENDING", + "COMMITTED", }: raise PermanentProjectionError( f"mutation is not projectable: {mutation['state']}" @@ -117,8 +141,6 @@ async def _apply_projection(dao: MemoryDAO, projection: dict[str, Any]) -> None: assertions = await dao.list_v4_assertions_for_mutation( str(mutation["mutation_id"]) ) - if len(assertions) != len(triplets): - raise PermanentProjectionError("canonical SQL assertions are unavailable") for assertion in assertions: await dao.project_v4_vector_assertion( mutation=mutation, assertion=assertion @@ -128,10 +150,9 @@ async def _apply_projection(dao: MemoryDAO, projection: dict[str, Any]) -> None: assertions = await dao.list_v4_assertions_for_mutation( str(mutation["mutation_id"]) ) - if len(assertions) != len(triplets): - raise PermanentProjectionError("canonical SQL assertions are unavailable") - for assertion in assertions: - await projector.project_assertion(mutation=mutation, assertion=assertion) + async with _get_graph_lock(): + for assertion in assertions: + await projector.project_assertion(mutation=mutation, assertion=assertion) else: raise PermanentProjectionError(f"unknown projection lane: {lane}") @@ -139,31 +160,12 @@ async def _apply_projection(dao: MemoryDAO, projection: dict[str, Any]) -> None: async def _apply_with_lease_heartbeat( dao: MemoryDAO, projection: dict[str, Any], worker_id: str ) -> None: - """Keep ownership alive while a model/vector/graph call is in progress.""" - task = asyncio.create_task(_apply_projection(dao, projection)) - while not task.done(): - try: - await asyncio.wait_for(asyncio.shield(task), timeout=60) - except TimeoutError: - renewed = await dao.renew_projection_outbox_lease( - str(projection["projection_id"]), - worker_id=worker_id, - claim_token=str(projection["claim_token"]), - ) - if not renewed: - task.cancel() - try: - await task - except asyncio.CancelledError: - pass - raise ProjectionLeaseLostError( - "projection outbox lease ownership was lost" - ) - await task + """Apply the projection directly.""" + await _apply_projection(dao, projection) async def process_projection_outbox_once( - dao: MemoryDAO, *, worker_id: str = "combined-runtime", limit: int = 1 + dao: MemoryDAO, *, worker_id: str = "combined-runtime", limit: int = 50 ) -> dict[str, int]: """Claim and apply a bounded set of fenced V4 projection lanes.""" claimed = await dao.claim_projection_outbox(worker_id=worker_id, limit=limit) @@ -173,19 +175,14 @@ async def process_projection_outbox_once( "retry_pending": 0, "dead_letter": 0, } - for projection in sorted( - claimed, key=lambda item: _LANE_ORDER.get(item["projection_name"], 99) - ): + + successful_projections: list[dict[str, Any]] = [] + + async def _handle_one(projection: dict[str, Any]) -> None: projection_id = str(projection["projection_id"]) try: await _apply_with_lease_heartbeat(dao, projection, worker_id) - completed = await dao.complete_projection_outbox( - projection_id, - worker_id=worker_id, - claim_token=str(projection["claim_token"]), - outcome="APPLIED", - ) - result["completed"] += int(completed) + successful_projections.append(projection) except PermanentProjectionError as exc: changed = await dao.fail_projection_outbox( projection_id, @@ -194,7 +191,8 @@ async def process_projection_outbox_once( error_class=type(exc).__name__, retryable=False, ) - result["dead_letter"] += int(changed) + if changed: + result["dead_letter"] += 1 logger.warning( "V4_PROJECTION_PERMANENT_FAILURE | projection_id=%s error=%s", projection_id, @@ -221,6 +219,20 @@ async def process_projection_outbox_once( projection_id, exc, ) + + sorted_projections = sorted( + claimed, key=lambda item: _LANE_ORDER.get(item["projection_name"], 99) + ) + if sorted_projections: + if dao._graph and dao._graph.is_operational and any(p.get("projection_name") == "GRAPH" for p in sorted_projections): + async with dao._graph.transaction(): + await asyncio.gather(*(_handle_one(p) for p in sorted_projections), return_exceptions=True) + else: + await asyncio.gather(*(_handle_one(p) for p in sorted_projections), return_exceptions=True) + if successful_projections: + result["completed"] = await dao.complete_projection_outbox_batch( + successful_projections, worker_id=worker_id + ) return result