From 98a2a3cbaf3e6ba5444a170e5993f0d18738edbb Mon Sep 17 00:00:00 2001 From: phernandez Date: Sat, 28 Mar 2026 14:42:18 -0500 Subject: [PATCH] Unify MCP telemetry spans across routers and services Signed-off-by: phernandez --- .../api/v2/routers/knowledge_router.py | 534 ++++++++++++------ .../api/v2/routers/memory_router.py | 126 +++-- .../api/v2/routers/resource_router.py | 390 +++++++------ .../api/v2/routers/search_router.py | 61 +- src/basic_memory/api/v2/utils.py | 360 ++++++------ src/basic_memory/mcp/clients/knowledge.py | 174 ++++-- src/basic_memory/mcp/clients/memory.py | 41 +- src/basic_memory/mcp/clients/resource.py | 21 +- src/basic_memory/mcp/clients/search.py | 23 +- src/basic_memory/mcp/tools/build_context.py | 2 +- src/basic_memory/mcp/tools/edit_note.py | 2 +- src/basic_memory/mcp/tools/read_note.py | 64 ++- src/basic_memory/mcp/tools/search.py | 4 +- src/basic_memory/mcp/tools/utils.py | 190 +++++-- src/basic_memory/mcp/tools/write_note.py | 2 +- src/basic_memory/services/context_service.py | 275 ++++----- src/basic_memory/services/entity_service.py | 366 +++++++++--- src/basic_memory/services/file_service.py | 145 +++-- src/basic_memory/services/search_service.py | 387 +++++++------ src/basic_memory/telemetry.py | 36 +- .../api/v2/test_knowledge_router_telemetry.py | 225 ++++++++ tests/api/v2/test_search_router_telemetry.py | 7 +- tests/api/v2/test_utils_telemetry.py | 63 +++ tests/mcp/test_client_telemetry.py | 91 +++ tests/mcp/test_tool_telemetry.py | 131 ++--- .../services/test_entity_service_telemetry.py | 131 +++++ .../services/test_search_service_telemetry.py | 24 +- tests/test_telemetry.py | 18 - 28 files changed, 2607 insertions(+), 1286 deletions(-) create mode 100644 tests/api/v2/test_knowledge_router_telemetry.py create mode 100644 tests/api/v2/test_utils_telemetry.py create mode 100644 tests/mcp/test_client_telemetry.py create mode 100644 tests/services/test_entity_service_telemetry.py diff --git a/src/basic_memory/api/v2/routers/knowledge_router.py b/src/basic_memory/api/v2/routers/knowledge_router.py index 88a53f5e..c642cb52 100644 --- a/src/basic_memory/api/v2/routers/knowledge_router.py +++ b/src/basic_memory/api/v2/routers/knowledge_router.py @@ -13,6 +13,7 @@ Key improvements: from fastapi import APIRouter, HTTPException, BackgroundTasks, Depends, Response, Path, Query from loguru import logger +from basic_memory import telemetry from basic_memory.deps import ( EntityServiceV2ExternalDep, SearchServiceV2ExternalDep, @@ -142,47 +143,66 @@ async def resolve_identifier( "resolution_method": "permalink" } """ - logger.info(f"API v2 request: resolve_identifier for '{data.identifier}'") + with telemetry.operation( + "api.request.knowledge.resolve_entity", + entrypoint="api", + domain="knowledge", + action="resolve_entity", + ): + logger.info(f"API v2 request: resolve_identifier for '{data.identifier}'") - # Try to resolve by external_id first - entity = await entity_repository.get_by_external_id(data.identifier) - resolution_method = "external_id" if entity else "search" + with telemetry.scope( + "api.knowledge.resolve_entity.lookup_entity", + domain="knowledge", + action="resolve_entity", + phase="lookup_entity", + ): + entity = await entity_repository.get_by_external_id(data.identifier) + resolution_method = "external_id" if entity else "search" - # If not found by external_id, try other resolution methods - # Pass source_path for context-aware resolution (prefers notes closer to source) - # Pass strict to control fuzzy search fallback (default False allows fuzzy matching) - if not entity: - entity = await link_resolver.resolve_link( - data.identifier, source_path=data.source_path, strict=data.strict + if not entity: + with telemetry.scope( + "api.knowledge.resolve_entity.resolve_link", + domain="knowledge", + action="resolve_entity", + phase="resolve_link", + ): + entity = await link_resolver.resolve_link( + data.identifier, source_path=data.source_path, strict=data.strict + ) + if entity: + if entity.permalink == data.identifier: + resolution_method = "permalink" + elif entity.title == data.identifier: + resolution_method = "title" + elif entity.file_path == data.identifier: + resolution_method = "path" + else: + resolution_method = "search" + + if not entity: + raise HTTPException(status_code=404, detail=f"Entity not found: '{data.identifier}'") + + with telemetry.scope( + "api.knowledge.resolve_entity.shape_response", + domain="knowledge", + action="resolve_entity", + phase="shape_response", + ): + result = EntityResolveResponse( + external_id=entity.external_id, + entity_id=entity.id, + permalink=entity.permalink, + file_path=entity.file_path, + title=entity.title, + resolution_method=resolution_method, + ) + + logger.debug( + f"API v2 response: resolved '{data.identifier}' to external_id={result.external_id} via {resolution_method}" ) - if entity: - # Determine resolution method - if entity.permalink == data.identifier: - resolution_method = "permalink" - elif entity.title == data.identifier: - resolution_method = "title" - elif entity.file_path == data.identifier: - resolution_method = "path" - else: - resolution_method = "search" - if not entity: - raise HTTPException(status_code=404, detail=f"Entity not found: '{data.identifier}'") - - result = EntityResolveResponse( - external_id=entity.external_id, - entity_id=entity.id, - permalink=entity.permalink, - file_path=entity.file_path, - title=entity.title, - resolution_method=resolution_method, - ) - - logger.debug( - f"API v2 response: resolved '{data.identifier}' to external_id={result.external_id} via {resolution_method}" - ) - - return result + return result ## Read endpoints @@ -208,18 +228,36 @@ async def get_entity_by_id( Raises: HTTPException: 404 if entity not found """ - logger.info(f"API v2 request: get_entity_by_id entity_id={entity_id}") + with telemetry.operation( + "api.request.knowledge.get_entity", + entrypoint="api", + domain="knowledge", + action="get_entity", + ): + logger.info(f"API v2 request: get_entity_by_id entity_id={entity_id}") - entity = await entity_repository.get_by_external_id(entity_id) - if not entity: - raise HTTPException( - status_code=404, detail=f"Entity with external_id '{entity_id}' not found" - ) + with telemetry.scope( + "api.knowledge.get_entity.load_entity", + domain="knowledge", + action="get_entity", + phase="load_entity", + ): + entity = await entity_repository.get_by_external_id(entity_id) + if not entity: + raise HTTPException( + status_code=404, detail=f"Entity with external_id '{entity_id}' not found" + ) - result = EntityResponseV2.model_validate(entity) - logger.info(f"API v2 response: external_id={entity_id}, title='{result.title}'") + with telemetry.scope( + "api.knowledge.get_entity.shape_response", + domain="knowledge", + action="get_entity", + phase="shape_response", + ): + result = EntityResponseV2.model_validate(entity) + logger.info(f"API v2 response: external_id={entity_id}, title='{result.title}'") - return result + return result ## Create endpoints @@ -248,39 +286,80 @@ async def create_entity( Returns: Created entity with generated external_id (UUID) and file content """ - logger.info( - "API v2 request", endpoint="create_entity", note_type=data.note_type, title=data.title - ) - - if fast: - entity = await entity_service.fast_write_entity(data) - task_scheduler.schedule( - "reindex_entity", - entity_id=entity.id, - project_id=project_id, - ) - else: - entity = await entity_service.create_entity(data) - await search_service.index_entity(entity) - _schedule_vector_sync_if_enabled( - task_scheduler=task_scheduler, - app_config=app_config, - entity_id=entity.id, - project_id=project_id, + with telemetry.operation( + "api.request.knowledge.create_entity", + entrypoint="api", + domain="knowledge", + action="create_entity", + fast=fast, + ): + logger.info( + "API v2 request", endpoint="create_entity", note_type=data.note_type, title=data.title ) - result = EntityResponseV2.model_validate(entity) - if fast: - result = result.model_copy(update={"observations": [], "relations": []}) + with telemetry.scope( + "api.knowledge.create_entity.write_entity", + domain="knowledge", + action="create_entity", + phase="write_entity", + fast=fast, + ): + if fast: + entity = await entity_service.fast_write_entity(data) + else: + entity = await entity_service.create_entity(data) - # Always read and return file content - content = await file_service.read_file_content(entity.file_path) - result = result.model_copy(update={"content": content}) + if fast: + with telemetry.scope( + "api.knowledge.create_entity.enqueue_reindex", + domain="knowledge", + action="create_entity", + phase="enqueue_reindex", + fast=fast, + ): + task_scheduler.schedule( + "reindex_entity", + entity_id=entity.id, + project_id=project_id, + ) + else: + with telemetry.scope( + "api.knowledge.create_entity.search_index", + domain="knowledge", + action="create_entity", + phase="search_index", + ): + await search_service.index_entity(entity) + with telemetry.scope( + "api.knowledge.create_entity.vector_sync", + domain="knowledge", + action="create_entity", + phase="vector_sync", + ): + _schedule_vector_sync_if_enabled( + task_scheduler=task_scheduler, + app_config=app_config, + entity_id=entity.id, + project_id=project_id, + ) - logger.info( - f"API v2 response: endpoint='create_entity' external_id={entity.external_id}, title={result.title}, permalink={result.permalink}, status_code=201" - ) - return result + result = EntityResponseV2.model_validate(entity) + if fast: + result = result.model_copy(update={"observations": [], "relations": []}) + + with telemetry.scope( + "api.knowledge.create_entity.read_content", + domain="knowledge", + action="create_entity", + phase="read_content", + ): + content = await file_service.read_file_content(entity.file_path) + result = result.model_copy(update={"content": content}) + + logger.info( + f"API v2 response: endpoint='create_entity' external_id={entity.external_id}, title={result.title}, permalink={result.permalink}, status_code=201" + ) + return result ## Update endpoints @@ -315,61 +394,104 @@ async def update_entity_by_id( Returns: Updated entity with file content """ - logger.info(f"API v2 request: update_entity_by_id entity_id={entity_id}") + with telemetry.operation( + "api.request.knowledge.update_entity", + entrypoint="api", + domain="knowledge", + action="update_entity", + fast=fast, + ): + logger.info(f"API v2 request: update_entity_by_id entity_id={entity_id}") - # Check if entity exists (external_id is the source of truth for v2) - existing = await entity_repository.get_by_external_id(entity_id) - created = existing is None + with telemetry.scope( + "api.knowledge.update_entity.load_entity", + domain="knowledge", + action="update_entity", + phase="load_entity", + ): + existing = await entity_repository.get_by_external_id(entity_id) + created = existing is None - if fast: - entity = await entity_service.fast_write_entity(data, external_id=entity_id) - response.status_code = 200 if existing else 201 - task_scheduler.schedule( - "reindex_entity", - entity_id=entity.id, - project_id=project_id, - resolve_relations=created, - ) - else: - if existing: - # Update the existing entity in-place to avoid path-based duplication - entity = await entity_service.update_entity(existing, data) - response.status_code = 200 - else: - # Create new entity, then bind external_id to the requested UUID - entity = await entity_service.create_entity(data) - if entity.external_id != entity_id: - entity = await entity_repository.update( - entity.id, - {"external_id": entity_id}, + with telemetry.scope( + "api.knowledge.update_entity.write_entity", + domain="knowledge", + action="update_entity", + phase="write_entity", + fast=fast, + ): + if fast: + entity = await entity_service.fast_write_entity(data, external_id=entity_id) + response.status_code = 200 if existing else 201 + else: + if existing: + entity = await entity_service.update_entity(existing, data) + response.status_code = 200 + else: + entity = await entity_service.create_entity(data) + if entity.external_id != entity_id: + entity = await entity_repository.update( + entity.id, + {"external_id": entity_id}, + ) + if not entity: + raise HTTPException( + status_code=404, + detail=f"Entity with external_id '{entity_id}' not found", + ) + response.status_code = 201 + + if fast: + with telemetry.scope( + "api.knowledge.update_entity.enqueue_reindex", + domain="knowledge", + action="update_entity", + phase="enqueue_reindex", + fast=fast, + ): + task_scheduler.schedule( + "reindex_entity", + entity_id=entity.id, + project_id=project_id, + resolve_relations=created, + ) + else: + with telemetry.scope( + "api.knowledge.update_entity.search_index", + domain="knowledge", + action="update_entity", + phase="search_index", + ): + await search_service.index_entity(entity) + with telemetry.scope( + "api.knowledge.update_entity.vector_sync", + domain="knowledge", + action="update_entity", + phase="vector_sync", + ): + _schedule_vector_sync_if_enabled( + task_scheduler=task_scheduler, + app_config=app_config, + entity_id=entity.id, + project_id=project_id, ) - if not entity: - raise HTTPException( - status_code=404, - detail=f"Entity with external_id '{entity_id}' not found", - ) - response.status_code = 201 - await search_service.index_entity(entity) - _schedule_vector_sync_if_enabled( - task_scheduler=task_scheduler, - app_config=app_config, - entity_id=entity.id, - project_id=project_id, + result = EntityResponseV2.model_validate(entity) + if fast: + result = result.model_copy(update={"observations": [], "relations": []}) + + with telemetry.scope( + "api.knowledge.update_entity.read_content", + domain="knowledge", + action="update_entity", + phase="read_content", + ): + content = await file_service.read_file_content(entity.file_path) + result = result.model_copy(update={"content": content}) + + logger.info( + f"API v2 response: external_id={entity_id}, created={created}, status_code={response.status_code}" ) - - result = EntityResponseV2.model_validate(entity) - if fast: - result = result.model_copy(update={"observations": [], "relations": []}) - - # Always read and return file content - content = await file_service.read_file_content(entity.file_path) - result = result.model_copy(update={"content": content}) - - logger.info( - f"API v2 response: external_id={entity_id}, created={created}, status_code={response.status_code}" - ) - return result + return result @router.patch("/entities/{entity_id}", response_model=EntityResponseV2) @@ -401,69 +523,113 @@ async def edit_entity_by_id( Raises: HTTPException: 404 if entity not found, 400 if edit fails """ - logger.info( - f"API v2 request: edit_entity_by_id entity_id={entity_id}, operation='{data.operation}'" - ) - - # Verify entity exists - entity = await entity_repository.get_by_external_id(entity_id) - if not entity: # pragma: no cover - raise HTTPException( - status_code=404, detail=f"Entity with external_id '{entity_id}' not found" - ) - - try: - if fast: - updated_entity = await entity_service.fast_edit_entity( - entity=entity, - operation=data.operation, - content=data.content, - section=data.section, - find_text=data.find_text, - expected_replacements=data.expected_replacements, - ) - task_scheduler.schedule( - "reindex_entity", - entity_id=updated_entity.id, - project_id=project_id, - ) - else: - # Edit using the entity's permalink or path - identifier = entity.permalink or entity.file_path - updated_entity = await entity_service.edit_entity( - identifier=identifier, - operation=data.operation, - content=data.content, - section=data.section, - find_text=data.find_text, - expected_replacements=data.expected_replacements, - ) - - await search_service.index_entity(updated_entity) - _schedule_vector_sync_if_enabled( - task_scheduler=task_scheduler, - app_config=app_config, - entity_id=updated_entity.id, - project_id=project_id, - ) - - result = EntityResponseV2.model_validate(updated_entity) - if fast: - result = result.model_copy(update={"observations": [], "relations": []}) - - # Always read and return file content - content = await file_service.read_file_content(updated_entity.file_path) - result = result.model_copy(update={"content": content}) - + with telemetry.operation( + "api.request.knowledge.edit_entity", + entrypoint="api", + domain="knowledge", + action="edit_entity", + fast=fast, + ): logger.info( - f"API v2 response: external_id={entity_id}, operation='{data.operation}', status_code=200" + f"API v2 request: edit_entity_by_id entity_id={entity_id}, operation='{data.operation}'" ) - return result + with telemetry.scope( + "api.knowledge.edit_entity.load_entity", + domain="knowledge", + action="edit_entity", + phase="load_entity", + ): + entity = await entity_repository.get_by_external_id(entity_id) + if not entity: # pragma: no cover + raise HTTPException( + status_code=404, detail=f"Entity with external_id '{entity_id}' not found" + ) - except Exception as e: - logger.error(f"Error editing entity {entity_id}: {e}") - raise HTTPException(status_code=400, detail=str(e)) + try: + with telemetry.scope( + "api.knowledge.edit_entity.write_entity", + domain="knowledge", + action="edit_entity", + phase="write_entity", + fast=fast, + ): + if fast: + updated_entity = await entity_service.fast_edit_entity( + entity=entity, + operation=data.operation, + content=data.content, + section=data.section, + find_text=data.find_text, + expected_replacements=data.expected_replacements, + ) + else: + identifier = entity.permalink or entity.file_path + updated_entity = await entity_service.edit_entity( + identifier=identifier, + operation=data.operation, + content=data.content, + section=data.section, + find_text=data.find_text, + expected_replacements=data.expected_replacements, + ) + + if fast: + with telemetry.scope( + "api.knowledge.edit_entity.enqueue_reindex", + domain="knowledge", + action="edit_entity", + phase="enqueue_reindex", + fast=fast, + ): + task_scheduler.schedule( + "reindex_entity", + entity_id=updated_entity.id, + project_id=project_id, + ) + else: + with telemetry.scope( + "api.knowledge.edit_entity.search_index", + domain="knowledge", + action="edit_entity", + phase="search_index", + ): + await search_service.index_entity(updated_entity) + with telemetry.scope( + "api.knowledge.edit_entity.vector_sync", + domain="knowledge", + action="edit_entity", + phase="vector_sync", + ): + _schedule_vector_sync_if_enabled( + task_scheduler=task_scheduler, + app_config=app_config, + entity_id=updated_entity.id, + project_id=project_id, + ) + + result = EntityResponseV2.model_validate(updated_entity) + if fast: + result = result.model_copy(update={"observations": [], "relations": []}) + + with telemetry.scope( + "api.knowledge.edit_entity.read_content", + domain="knowledge", + action="edit_entity", + phase="read_content", + ): + content = await file_service.read_file_content(updated_entity.file_path) + result = result.model_copy(update={"content": content}) + + logger.info( + f"API v2 response: external_id={entity_id}, operation='{data.operation}', status_code=200" + ) + + return result + + except Exception as e: + logger.error(f"Error editing entity {entity_id}: {e}") + raise HTTPException(status_code=400, detail=str(e)) ## Delete endpoints diff --git a/src/basic_memory/api/v2/routers/memory_router.py b/src/basic_memory/api/v2/routers/memory_router.py index 91de4a08..17ef1eed 100644 --- a/src/basic_memory/api/v2/routers/memory_router.py +++ b/src/basic_memory/api/v2/routers/memory_router.py @@ -9,6 +9,7 @@ from typing import Annotated, Optional from fastapi import APIRouter, Query, Path from loguru import logger +from basic_memory import telemetry from basic_memory.deps import ContextServiceV2ExternalDep, EntityRepositoryV2ExternalDep from basic_memory.schemas.base import TimeFrame, parse_timeframe from basic_memory.schemas.memory import ( @@ -50,30 +51,55 @@ async def recent( Returns: GraphContext with recent activity and related entities """ - # return all types by default - types = ( - [SearchItemType.ENTITY, SearchItemType.RELATION, SearchItemType.OBSERVATION] - if not type - else type - ) + with telemetry.operation( + "api.request.memory.recent_activity", + entrypoint="api", + domain="memory", + action="recent_activity", + page=page, + page_size=page_size, + ): + types = ( + [SearchItemType.ENTITY, SearchItemType.RELATION, SearchItemType.OBSERVATION] + if not type + else type + ) - logger.debug( - f"V2 Getting recent context for project {project_id}: `{types}` depth: `{depth}` timeframe: `{timeframe}` page: `{page}` page_size: `{page_size}` max_related: `{max_related}`" - ) - # Parse timeframe - since = parse_timeframe(timeframe) - limit = page_size - offset = (page - 1) * page_size + logger.debug( + f"V2 Getting recent context for project {project_id}: `{types}` depth: `{depth}` timeframe: `{timeframe}` page: `{page}` page_size: `{page_size}` max_related: `{max_related}`" + ) + since = parse_timeframe(timeframe) + limit = page_size + offset = (page - 1) * page_size - # Build context - context = await context_service.build_context( - types=types, depth=depth, since=since, limit=limit, offset=offset, max_related=max_related - ) - recent_context = await to_graph_context( - context, entity_repository=entity_repository, page=page, page_size=page_size - ) - logger.debug(f"V2 Recent context: {recent_context.model_dump_json()}") - return recent_context + with telemetry.scope( + "api.memory.recent_activity.build_context", + domain="memory", + action="recent_activity", + phase="build_context", + page=page, + page_size=page_size, + ): + context = await context_service.build_context( + types=types, + depth=depth, + since=since, + limit=limit, + offset=offset, + max_related=max_related, + ) + with telemetry.scope( + "api.memory.recent_activity.shape_response", + domain="memory", + action="recent_activity", + phase="shape_response", + result_count=len(context.results), + ): + recent_context = await to_graph_context( + context, entity_repository=entity_repository, page=page, page_size=page_size + ) + logger.debug(f"V2 Recent context: {recent_context.model_dump_json()}") + return recent_context # get_memory_context needs to be declared last so other paths can match @@ -111,20 +137,46 @@ async def get_memory_context( Returns: GraphContext with the entity and its related context """ - logger.debug( - f"V2 Getting context for project {project_id}, URI: `{uri}` depth: `{depth}` timeframe: `{timeframe}` page: `{page}` page_size: `{page_size}` max_related: `{max_related}`" - ) - memory_url = normalize_memory_url(uri) + with telemetry.operation( + "api.request.memory.build_context", + entrypoint="api", + domain="memory", + action="build_context", + page=page, + page_size=page_size, + ): + logger.debug( + f"V2 Getting context for project {project_id}, URI: `{uri}` depth: `{depth}` timeframe: `{timeframe}` page: `{page}` page_size: `{page_size}` max_related: `{max_related}`" + ) + memory_url = normalize_memory_url(uri) - # Parse timeframe - since = parse_timeframe(timeframe) if timeframe else None - limit = page_size - offset = (page - 1) * page_size + since = parse_timeframe(timeframe) if timeframe else None + limit = page_size + offset = (page - 1) * page_size - # Build context - context = await context_service.build_context( - memory_url, depth=depth, since=since, limit=limit, offset=offset, max_related=max_related - ) - return await to_graph_context( - context, entity_repository=entity_repository, page=page, page_size=page_size - ) + with telemetry.scope( + "api.memory.build_context.build_context", + domain="memory", + action="build_context", + phase="build_context", + page=page, + page_size=page_size, + ): + context = await context_service.build_context( + memory_url, + depth=depth, + since=since, + limit=limit, + offset=offset, + max_related=max_related, + ) + with telemetry.scope( + "api.memory.build_context.shape_response", + domain="memory", + action="build_context", + phase="shape_response", + result_count=len(context.results), + ): + return await to_graph_context( + context, entity_repository=entity_repository, page=page, page_size=page_size + ) diff --git a/src/basic_memory/api/v2/routers/resource_router.py b/src/basic_memory/api/v2/routers/resource_router.py index 0f27b412..d459bb9d 100644 --- a/src/basic_memory/api/v2/routers/resource_router.py +++ b/src/basic_memory/api/v2/routers/resource_router.py @@ -15,6 +15,7 @@ from pathlib import Path as PathLib from fastapi import APIRouter, HTTPException, Response, Path from loguru import logger +from basic_memory import telemetry from basic_memory.deps import ( ProjectConfigV2ExternalDep, FileServiceV2ExternalDep, @@ -55,36 +56,62 @@ async def get_resource_content( Raises: HTTPException: 404 if entity or file not found """ - logger.debug(f"V2 Getting content for project {project_id}, entity_id: {entity_id}") + with telemetry.operation( + "api.request.resource.get_content", + entrypoint="api", + domain="resource", + action="get_content", + ): + logger.debug(f"V2 Getting content for project {project_id}, entity_id: {entity_id}") - # Get entity by external_id - entity = await entity_repository.get_by_external_id(entity_id) - if not entity: - raise HTTPException(status_code=404, detail=f"Entity {entity_id} not found") + with telemetry.scope( + "api.resource.get_content.load_entity", + domain="resource", + action="get_content", + phase="load_entity", + ): + entity = await entity_repository.get_by_external_id(entity_id) + if not entity: + raise HTTPException(status_code=404, detail=f"Entity {entity_id} not found") - # Validate entity file path to prevent path traversal - project_path = PathLib(config.home) - if not validate_project_path(entity.file_path, project_path): - logger.error( # pragma: no cover - f"Invalid file path in entity {entity.id}: {entity.file_path}" - ) - raise HTTPException( # pragma: no cover - status_code=500, - detail="Entity contains invalid file path", - ) + with telemetry.scope( + "api.resource.get_content.validate_path", + domain="resource", + action="get_content", + phase="validate_path", + ): + project_path = PathLib(config.home) + if not validate_project_path(entity.file_path, project_path): + logger.error( # pragma: no cover + f"Invalid file path in entity {entity.id}: {entity.file_path}" + ) + raise HTTPException( # pragma: no cover + status_code=500, + detail="Entity contains invalid file path", + ) - # Check file exists via file_service (for cloud compatibility) - if not await file_service.exists(entity.file_path): - raise HTTPException( # pragma: no cover - status_code=404, - detail=f"File not found: {entity.file_path}", - ) + with telemetry.scope( + "api.resource.get_content.ensure_exists", + domain="resource", + action="get_content", + phase="ensure_exists", + ): + if not await file_service.exists(entity.file_path): + raise HTTPException( # pragma: no cover + status_code=404, + detail=f"File not found: {entity.file_path}", + ) - # Read content via file_service as bytes (works with both local and S3) - content = await file_service.read_file_bytes(entity.file_path) - content_type = file_service.content_type(entity.file_path) + with telemetry.scope( + "api.resource.get_content.read_content", + domain="resource", + action="get_content", + phase="read_content", + ): + content = await file_service.read_file_bytes(entity.file_path) + content_type = file_service.content_type(entity.file_path) - return Response(content=content, media_type=content_type) + return Response(content=content, media_type=content_type) @router.post("", response_model=ResourceResponse) @@ -112,74 +139,94 @@ async def create_resource( Raises: HTTPException: 400 for invalid file paths, 409 if file already exists """ - try: - # Validate path to prevent path traversal attacks - project_path = PathLib(config.home) - if not validate_project_path(data.file_path, project_path): - logger.warning( - f"Invalid file path attempted: {data.file_path} in project {config.name}" + with telemetry.operation( + "api.request.resource.create", + entrypoint="api", + domain="resource", + action="create", + ): + try: + # Validate path to prevent path traversal attacks + project_path = PathLib(config.home) + if not validate_project_path(data.file_path, project_path): + logger.warning( + f"Invalid file path attempted: {data.file_path} in project {config.name}" + ) + raise HTTPException( + status_code=400, + detail=f"Invalid file path: {data.file_path}. " + "Path must be relative and stay within project boundaries.", + ) + + existing_entity = await entity_repository.get_by_file_path(data.file_path) + if existing_entity: + raise HTTPException( + status_code=409, + detail=f"Resource already exists at {data.file_path} with entity_id {existing_entity.external_id}. " + f"Use PUT /resource/{existing_entity.external_id} to update it.", + ) + + with telemetry.scope( + "api.resource.create.write_file", + domain="resource", + action="create", + phase="write_file", + ): + await file_service.ensure_directory(PathLib(data.file_path).parent) + checksum = await file_service.write_file(data.file_path, data.content) + + with telemetry.scope( + "api.resource.create.read_metadata", + domain="resource", + action="create", + phase="read_metadata", + ): + file_metadata = await file_service.get_file_metadata(data.file_path) + + file_name = PathLib(data.file_path).name + content_type = file_service.content_type(data.file_path) + note_type = "canvas" if data.file_path.endswith(".canvas") else "file" + + entity = EntityModel( + external_id=str(uuid.uuid4()), + title=file_name, + note_type=note_type, + content_type=content_type, + file_path=data.file_path, + checksum=checksum, + created_at=file_metadata.created_at, + updated_at=file_metadata.modified_at, ) - raise HTTPException( - status_code=400, - detail=f"Invalid file path: {data.file_path}. " - "Path must be relative and stay within project boundaries.", + with telemetry.scope( + "api.resource.create.upsert_entity", + domain="resource", + action="create", + phase="upsert_entity", + ): + entity = await entity_repository.add(entity) + + with telemetry.scope( + "api.resource.create.search_index", + domain="resource", + action="create", + phase="search_index", + ): + await search_service.index_entity(entity) # pyright: ignore + + return ResourceResponse( + entity_id=entity.id, + external_id=entity.external_id, + file_path=data.file_path, + checksum=checksum, + size=file_metadata.size, + created_at=file_metadata.created_at.timestamp(), + modified_at=file_metadata.modified_at.timestamp(), ) - - # Check if entity already exists - existing_entity = await entity_repository.get_by_file_path(data.file_path) - if existing_entity: - raise HTTPException( - status_code=409, - detail=f"Resource already exists at {data.file_path} with entity_id {existing_entity.external_id}. " - f"Use PUT /resource/{existing_entity.external_id} to update it.", - ) - - # Cloud compatibility: avoid assuming a local filesystem path. - # Delegate directory creation + writes to FileService (local or S3). - await file_service.ensure_directory(PathLib(data.file_path).parent) - checksum = await file_service.write_file(data.file_path, data.content) - - # Get file info - file_metadata = await file_service.get_file_metadata(data.file_path) - - # Determine file details - file_name = PathLib(data.file_path).name - content_type = file_service.content_type(data.file_path) - note_type = "canvas" if data.file_path.endswith(".canvas") else "file" - - # Create a new entity model - # Explicitly set external_id to ensure NOT NULL constraint is satisfied (fixes #512) - entity = EntityModel( - external_id=str(uuid.uuid4()), - title=file_name, - note_type=note_type, - content_type=content_type, - file_path=data.file_path, - checksum=checksum, - created_at=file_metadata.created_at, - updated_at=file_metadata.modified_at, - ) - entity = await entity_repository.add(entity) - - # Index the file for search - await search_service.index_entity(entity) # pyright: ignore - - # Return success response - return ResourceResponse( - entity_id=entity.id, - external_id=entity.external_id, - file_path=data.file_path, - checksum=checksum, - size=file_metadata.size, - created_at=file_metadata.created_at.timestamp(), - modified_at=file_metadata.modified_at.timestamp(), - ) - except HTTPException: - # Re-raise HTTP exceptions without wrapping - raise - except Exception as e: # pragma: no cover - logger.error(f"Error creating resource {data.file_path}: {e}") - raise HTTPException(status_code=500, detail=f"Failed to create resource: {str(e)}") + except HTTPException: + raise + except Exception as e: # pragma: no cover + logger.error(f"Error creating resource {data.file_path}: {e}") + raise HTTPException(status_code=500, detail=f"Failed to create resource: {str(e)}") @router.put("/{entity_id}", response_model=ResourceResponse) @@ -211,79 +258,94 @@ async def update_resource( Raises: HTTPException: 404 if entity not found, 400 for invalid paths """ - try: - # Get existing entity by external_id - entity = await entity_repository.get_by_external_id(entity_id) - if not entity: - raise HTTPException(status_code=404, detail=f"Entity {entity_id} not found") + with telemetry.operation( + "api.request.resource.update", + entrypoint="api", + domain="resource", + action="update", + ): + try: + entity = await entity_repository.get_by_external_id(entity_id) + if not entity: + raise HTTPException(status_code=404, detail=f"Entity {entity_id} not found") - # Determine target file path - target_file_path = data.file_path if data.file_path else entity.file_path + target_file_path = data.file_path if data.file_path else entity.file_path - # Validate path to prevent path traversal attacks - project_path = PathLib(config.home) - if not validate_project_path(target_file_path, project_path): - logger.warning( - f"Invalid file path attempted: {target_file_path} in project {config.name}" + project_path = PathLib(config.home) + if not validate_project_path(target_file_path, project_path): + logger.warning( + f"Invalid file path attempted: {target_file_path} in project {config.name}" + ) + raise HTTPException( + status_code=400, + detail=f"Invalid file path: {target_file_path}. " + "Path must be relative and stay within project boundaries.", + ) + + with telemetry.scope( + "api.resource.update.write_file", + domain="resource", + action="update", + phase="write_file", + ): + if data.file_path and data.file_path != entity.file_path: + await file_service.ensure_directory(PathLib(target_file_path).parent) + if await file_service.exists(entity.file_path): + await file_service.delete_file(entity.file_path) + else: + await file_service.ensure_directory(PathLib(target_file_path).parent) + + checksum = await file_service.write_file(target_file_path, data.content) + + with telemetry.scope( + "api.resource.update.read_metadata", + domain="resource", + action="update", + phase="read_metadata", + ): + file_metadata = await file_service.get_file_metadata(target_file_path) + + file_name = PathLib(target_file_path).name + content_type = file_service.content_type(target_file_path) + note_type = "canvas" if target_file_path.endswith(".canvas") else "file" + + with telemetry.scope( + "api.resource.update.update_entity", + domain="resource", + action="update", + phase="update_entity", + ): + updated_entity = await entity_repository.update( + entity.id, + { + "title": file_name, + "note_type": note_type, + "content_type": content_type, + "file_path": target_file_path, + "checksum": checksum, + "updated_at": file_metadata.modified_at, + }, + ) + + with telemetry.scope( + "api.resource.update.search_index", + domain="resource", + action="update", + phase="search_index", + ): + await search_service.index_entity(updated_entity) # pyright: ignore + + return ResourceResponse( + entity_id=entity.id, + external_id=entity.external_id, + file_path=target_file_path, + checksum=checksum, + size=file_metadata.size, + created_at=file_metadata.created_at.timestamp(), + modified_at=file_metadata.modified_at.timestamp(), ) - raise HTTPException( - status_code=400, - detail=f"Invalid file path: {target_file_path}. " - "Path must be relative and stay within project boundaries.", - ) - - # If moving file, handle the move - if data.file_path and data.file_path != entity.file_path: - # Ensure new parent directory exists (no-op for S3) - await file_service.ensure_directory(PathLib(target_file_path).parent) - - # If old file exists, remove it via file_service (for cloud compatibility) - if await file_service.exists(entity.file_path): - await file_service.delete_file(entity.file_path) - else: - # Ensure directory exists for in-place update - await file_service.ensure_directory(PathLib(target_file_path).parent) - - # Write content to target file - checksum = await file_service.write_file(target_file_path, data.content) - - # Get file info - file_metadata = await file_service.get_file_metadata(target_file_path) - - # Determine file details - file_name = PathLib(target_file_path).name - content_type = file_service.content_type(target_file_path) - note_type = "canvas" if target_file_path.endswith(".canvas") else "file" - - # Update entity using internal ID - updated_entity = await entity_repository.update( - entity.id, - { - "title": file_name, - "note_type": note_type, - "content_type": content_type, - "file_path": target_file_path, - "checksum": checksum, - "updated_at": file_metadata.modified_at, - }, - ) - - # Index the updated file for search - await search_service.index_entity(updated_entity) # pyright: ignore - - # Return success response - return ResourceResponse( - entity_id=entity.id, - external_id=entity.external_id, - file_path=target_file_path, - checksum=checksum, - size=file_metadata.size, - created_at=file_metadata.created_at.timestamp(), - modified_at=file_metadata.modified_at.timestamp(), - ) - except HTTPException: - # Re-raise HTTP exceptions without wrapping - raise - except Exception as e: # pragma: no cover - logger.error(f"Error updating resource {entity_id}: {e}") - raise HTTPException(status_code=500, detail=f"Failed to update resource: {str(e)}") + except HTTPException: + raise + except Exception as e: # pragma: no cover + logger.error(f"Error updating resource {entity_id}: {e}") + raise HTTPException(status_code=500, detail=f"Failed to update resource: {str(e)}") diff --git a/src/basic_memory/api/v2/routers/search_router.py b/src/basic_memory/api/v2/routers/search_router.py index 18b99374..eb098729 100644 --- a/src/basic_memory/api/v2/routers/search_router.py +++ b/src/basic_memory/api/v2/routers/search_router.py @@ -51,18 +51,28 @@ async def search( with telemetry.operation( "api.request.search", entrypoint="api", + domain="search", + action="search", page=page, page_size=page_size, retrieval_mode=query.retrieval_mode.value, - has_text_query=bool(query.text and query.text.strip()), - has_title_query=bool(query.title), - has_permalink_query=bool(query.permalink or query.permalink_match), + has_query=bool( + (query.text and query.text.strip()) or query.title or query.permalink or query.permalink_match + ), + has_filters=bool(query.note_types or query.entity_types or query.metadata_filters), ): offset = (page - 1) * page_size - # Fetch one extra item to detect whether more pages exist (N+1 trick) fetch_limit = page_size + 1 try: - results = await search_service.search(query, limit=fetch_limit, offset=offset) + with telemetry.scope( + "api.search.search.execute_query", + domain="search", + action="search", + phase="execute_query", + page=page, + page_size=page_size, + ): + results = await search_service.search(query, limit=fetch_limit, offset=offset) except SemanticSearchDisabledError as exc: raise HTTPException(status_code=400, detail=str(exc)) from exc except SemanticDependenciesMissingError as exc: @@ -70,17 +80,38 @@ async def search( except ValueError as exc: raise HTTPException(status_code=400, detail=str(exc)) from exc - has_more = len(results) > page_size - if has_more: - results = results[:page_size] + with telemetry.scope( + "api.search.search.paginate_results", + domain="search", + action="search", + phase="paginate_results", + result_count=len(results), + ): + has_more = len(results) > page_size + if has_more: + results = results[:page_size] - search_results = await to_search_results(entity_service, results) - return SearchResponse( - results=search_results, - current_page=page, - page_size=page_size, - has_more=has_more, - ) + with telemetry.scope( + "api.search.search.hydrate_results", + domain="search", + action="search", + phase="hydrate_results", + result_count=len(results), + ): + search_results = await to_search_results(entity_service, results) + with telemetry.scope( + "api.search.search.build_response", + domain="search", + action="search", + phase="build_response", + result_count=len(search_results), + ): + return SearchResponse( + results=search_results, + current_page=page, + page_size=page_size, + has_more=has_more, + ) @router.post("/search/reindex") diff --git a/src/basic_memory/api/v2/utils.py b/src/basic_memory/api/v2/utils.py index c92248f0..c72a0c1b 100644 --- a/src/basic_memory/api/v2/utils.py +++ b/src/basic_memory/api/v2/utils.py @@ -1,5 +1,6 @@ from typing import Optional, List +from basic_memory import telemetry from basic_memory.repository import EntityRepository from basic_memory.repository.search_repository import SearchIndexRow from basic_memory.schemas.memory import ( @@ -24,169 +25,212 @@ async def to_graph_context( page: Optional[int] = None, page_size: Optional[int] = None, ): - # First pass: collect all entity IDs needed for external_id lookup - # This includes: entity primary results, observation parent entities, relation from/to entities - entity_ids_needed: set[int] = set() - for context_item in context_result.results: - for item in ( - [context_item.primary_result] + context_item.observations + context_item.related_results - ): - if item.type == SearchItemType.ENTITY: - # Entity's own ID for its external_id - entity_ids_needed.add(item.id) - elif item.type == SearchItemType.OBSERVATION: - # Parent entity ID for entity_external_id - if item.entity_id: # pyright: ignore - entity_ids_needed.add(item.entity_id) # pyright: ignore - elif item.type == SearchItemType.RELATION: - # Source and target entity IDs for external_ids - if item.from_id: # pyright: ignore - entity_ids_needed.add(item.from_id) # pyright: ignore - if item.to_id: - entity_ids_needed.add(item.to_id) - - # Batch fetch all entities at once - get both title and external_id - entity_title_lookup: dict[int, str] = {} - entity_external_id_lookup: dict[int, str] = {} - if entity_ids_needed: - entities = await entity_repository.find_by_ids(list(entity_ids_needed)) - for e in entities: - entity_title_lookup[e.id] = e.title - entity_external_id_lookup[e.id] = e.external_id - - # Helper function to convert items to summaries - def to_summary(item: SearchIndexRow | ContextResultRow): - match item.type: - case SearchItemType.ENTITY: - return EntitySummary( - external_id=entity_external_id_lookup.get(item.id, ""), - entity_id=item.id, - title=item.title, # pyright: ignore - permalink=item.permalink, - content=item.content, - file_path=item.file_path, - created_at=item.created_at, - ) - case SearchItemType.OBSERVATION: - entity_ext_id = None - if item.entity_id: # pyright: ignore - entity_ext_id = entity_external_id_lookup.get(item.entity_id) # pyright: ignore - return ObservationSummary( - observation_id=item.id, - entity_id=item.entity_id, # pyright: ignore - entity_external_id=entity_ext_id, - title=entity_title_lookup.get(item.entity_id), # pyright: ignore - file_path=item.file_path, - category=item.category, # pyright: ignore - content=item.content, # pyright: ignore - permalink=item.permalink, # pyright: ignore - created_at=item.created_at, - ) - case SearchItemType.RELATION: - from_title = entity_title_lookup.get(item.from_id) if item.from_id else None # pyright: ignore - to_title = entity_title_lookup.get(item.to_id) if item.to_id else None - from_ext_id = entity_external_id_lookup.get(item.from_id) if item.from_id else None # pyright: ignore - to_ext_id = entity_external_id_lookup.get(item.to_id) if item.to_id else None - return RelationSummary( - relation_id=item.id, - entity_id=item.entity_id, # pyright: ignore - title=item.title, # pyright: ignore - file_path=item.file_path, - permalink=item.permalink, # pyright: ignore - relation_type=item.relation_type, # pyright: ignore - from_entity=from_title, - from_entity_id=item.from_id, # pyright: ignore - from_entity_external_id=from_ext_id, - to_entity=to_title, - to_entity_id=item.to_id, - to_entity_external_id=to_ext_id, - created_at=item.created_at, - ) - case _: # pragma: no cover - raise ValueError(f"Unexpected type: {item.type}") - - # Process the hierarchical results - hierarchical_results = [] - for context_item in context_result.results: - # Process primary result - primary_result = to_summary(context_item.primary_result) - - # Process observations (always ObservationSummary, validated by context_service) - observations = [to_summary(obs) for obs in context_item.observations] - - # Process related results - related = [to_summary(rel) for rel in context_item.related_results] - - # Add to hierarchical results - hierarchical_results.append( - ContextResult( - primary_result=primary_result, - observations=observations, # pyright: ignore[reportArgumentType] - related_results=related, - ) - ) - - # Create schema metadata from service metadata - metadata = MemoryMetadata( - uri=context_result.metadata.uri, - types=context_result.metadata.types, - depth=context_result.metadata.depth, - timeframe=context_result.metadata.timeframe, - generated_at=context_result.metadata.generated_at, - primary_count=context_result.metadata.primary_count, - related_count=context_result.metadata.related_count, - total_results=context_result.metadata.primary_count + context_result.metadata.related_count, - total_relations=context_result.metadata.total_relations, - total_observations=context_result.metadata.total_observations, - ) - - # Return new GraphContext with just hierarchical results - return GraphContext( - results=hierarchical_results, - metadata=metadata, + with telemetry.scope( + "memory.hydrate_context", + domain="memory", + action="build_context", + phase="hydrate_context", page=page, page_size=page_size, - has_more=context_result.metadata.has_more, - ) + result_count=len(context_result.results), + ): + # First pass: collect all entity IDs needed for external_id lookup + # This includes: entity primary results, observation parent entities, relation from/to entities + entity_ids_needed: set[int] = set() + for context_item in context_result.results: + for item in ( + [context_item.primary_result] + context_item.observations + context_item.related_results + ): + if item.type == SearchItemType.ENTITY: + # Entity's own ID for its external_id + entity_ids_needed.add(item.id) + elif item.type == SearchItemType.OBSERVATION: + # Parent entity ID for entity_external_id + if item.entity_id: # pyright: ignore + entity_ids_needed.add(item.entity_id) # pyright: ignore + elif item.type == SearchItemType.RELATION: + # Source and target entity IDs for external_ids + if item.from_id: # pyright: ignore + entity_ids_needed.add(item.from_id) # pyright: ignore + if item.to_id: + entity_ids_needed.add(item.to_id) + + # Batch fetch all entities at once - get both title and external_id + entity_title_lookup: dict[int, str] = {} + entity_external_id_lookup: dict[int, str] = {} + if entity_ids_needed: + with telemetry.scope( + "memory.hydrate_context.lookup_entities", + domain="memory", + action="build_context", + phase="lookup_entities", + result_count=len(entity_ids_needed), + ): + entities = await entity_repository.find_by_ids(list(entity_ids_needed)) + for e in entities: + entity_title_lookup[e.id] = e.title + entity_external_id_lookup[e.id] = e.external_id + + # Helper function to convert items to summaries + def to_summary(item: SearchIndexRow | ContextResultRow): + match item.type: + case SearchItemType.ENTITY: + return EntitySummary( + external_id=entity_external_id_lookup.get(item.id, ""), + entity_id=item.id, + title=item.title, # pyright: ignore + permalink=item.permalink, + content=item.content, + file_path=item.file_path, + created_at=item.created_at, + ) + case SearchItemType.OBSERVATION: + entity_ext_id = None + if item.entity_id: # pyright: ignore + entity_ext_id = entity_external_id_lookup.get(item.entity_id) # pyright: ignore + return ObservationSummary( + observation_id=item.id, + entity_id=item.entity_id, # pyright: ignore + entity_external_id=entity_ext_id, + title=entity_title_lookup.get(item.entity_id), # pyright: ignore + file_path=item.file_path, + category=item.category, # pyright: ignore + content=item.content, # pyright: ignore + permalink=item.permalink, # pyright: ignore + created_at=item.created_at, + ) + case SearchItemType.RELATION: + from_title = ( + entity_title_lookup.get(item.from_id) if item.from_id else None + ) # pyright: ignore + to_title = entity_title_lookup.get(item.to_id) if item.to_id else None + from_ext_id = ( + entity_external_id_lookup.get(item.from_id) if item.from_id else None + ) # pyright: ignore + to_ext_id = entity_external_id_lookup.get(item.to_id) if item.to_id else None + return RelationSummary( + relation_id=item.id, + entity_id=item.entity_id, # pyright: ignore + title=item.title, # pyright: ignore + file_path=item.file_path, + permalink=item.permalink, # pyright: ignore + relation_type=item.relation_type, # pyright: ignore + from_entity=from_title, + from_entity_id=item.from_id, # pyright: ignore + from_entity_external_id=from_ext_id, + to_entity=to_title, + to_entity_id=item.to_id, + to_entity_external_id=to_ext_id, + created_at=item.created_at, + ) + case _: # pragma: no cover + raise ValueError(f"Unexpected type: {item.type}") + + with telemetry.scope( + "memory.hydrate_context.shape_results", + domain="memory", + action="build_context", + phase="shape_results", + result_count=len(context_result.results), + ): + hierarchical_results = [] + for context_item in context_result.results: + primary_result = to_summary(context_item.primary_result) + observations = [to_summary(obs) for obs in context_item.observations] + related = [to_summary(rel) for rel in context_item.related_results] + hierarchical_results.append( + ContextResult( + primary_result=primary_result, + observations=observations, # pyright: ignore[reportArgumentType] + related_results=related, + ) + ) + + metadata = MemoryMetadata( + uri=context_result.metadata.uri, + types=context_result.metadata.types, + depth=context_result.metadata.depth, + timeframe=context_result.metadata.timeframe, + generated_at=context_result.metadata.generated_at, + primary_count=context_result.metadata.primary_count, + related_count=context_result.metadata.related_count, + total_results=context_result.metadata.primary_count + context_result.metadata.related_count, + total_relations=context_result.metadata.total_relations, + total_observations=context_result.metadata.total_observations, + ) + + return GraphContext( + results=hierarchical_results, + metadata=metadata, + page=page, + page_size=page_size, + has_more=context_result.metadata.has_more, + ) async def to_search_results(entity_service: EntityService, results: List[SearchIndexRow]): - search_results = [] - for r in results: - entities = await entity_service.get_entities_by_id([r.entity_id, r.from_id, r.to_id]) # pyright: ignore + with telemetry.scope( + "search.hydrate_results", + domain="search", + action="search", + phase="hydrate_results", + result_count=len(results), + ): + entity_batches = [] + with telemetry.scope( + "search.hydrate_results.fetch_entities", + domain="search", + action="search", + phase="fetch_entities", + result_count=len(results), + ): + for result in results: + entity_batches.append( + await entity_service.get_entities_by_id( + [result.entity_id, result.from_id, result.to_id] # pyright: ignore + ) + ) - # Determine which IDs to set based on type - entity_id = None - observation_id = None - relation_id = None + search_results = [] + with telemetry.scope( + "search.hydrate_results.shape_results", + domain="search", + action="search", + phase="shape_results", + result_count=len(results), + ): + for result, entities in zip(results, entity_batches): + entity_id = None + observation_id = None + relation_id = None - if r.type == SearchItemType.ENTITY: - entity_id = r.id - elif r.type == SearchItemType.OBSERVATION: - observation_id = r.id - entity_id = r.entity_id # Parent entity - elif r.type == SearchItemType.RELATION: - relation_id = r.id - entity_id = r.entity_id # Parent entity + if result.type == SearchItemType.ENTITY: + entity_id = result.id + elif result.type == SearchItemType.OBSERVATION: + observation_id = result.id + entity_id = result.entity_id + elif result.type == SearchItemType.RELATION: + relation_id = result.id + entity_id = result.entity_id - search_results.append( - SearchResult( - title=r.title, # pyright: ignore - type=r.type, # pyright: ignore - permalink=r.permalink, - score=r.score, # pyright: ignore - entity=entities[0].permalink if entities else None, - content=r.content, - matched_chunk=r.matched_chunk_text, - file_path=r.file_path, - metadata=r.metadata, - entity_id=entity_id, - observation_id=observation_id, - relation_id=relation_id, - category=r.category, - from_entity=entities[0].permalink if entities else None, - to_entity=entities[1].permalink if len(entities) > 1 else None, - relation_type=r.relation_type, - ) - ) - return search_results + search_results.append( + SearchResult( + title=result.title, # pyright: ignore + type=result.type, # pyright: ignore + permalink=result.permalink, + score=result.score, # pyright: ignore + entity=entities[0].permalink if entities else None, + content=result.content, + matched_chunk=result.matched_chunk_text, + file_path=result.file_path, + metadata=result.metadata, + entity_id=entity_id, + observation_id=observation_id, + relation_id=relation_id, + category=result.category, + from_entity=entities[0].permalink if entities else None, + to_entity=entities[1].permalink if len(entities) > 1 else None, + relation_type=result.relation_type, + ) + ) + return search_results diff --git a/src/basic_memory/mcp/clients/knowledge.py b/src/basic_memory/mcp/clients/knowledge.py index e6800128..cb705fe0 100644 --- a/src/basic_memory/mcp/clients/knowledge.py +++ b/src/basic_memory/mcp/clients/knowledge.py @@ -7,6 +7,7 @@ from typing import Any from httpx import AsyncClient +from basic_memory import telemetry from basic_memory.mcp.tools.utils import call_get, call_post, call_put, call_patch, call_delete from basic_memory.schemas.response import ( EntityResponse, @@ -58,12 +59,21 @@ class KnowledgeClient: ToolError: If the request fails """ params = {"fast": fast} if fast is not None else None - response = await call_post( - self.http_client, - f"{self._base_path}/entities", - json=entity_data, - params=params, - ) + with telemetry.scope( + "mcp.client.knowledge.create_entity", + client_name="knowledge", + operation="create_entity", + fast=fast, + ): + response = await call_post( + self.http_client, + f"{self._base_path}/entities", + json=entity_data, + params=params, + client_name="knowledge", + operation="create_entity", + path_template="/v2/projects/{project_id}/knowledge/entities", + ) return EntityResponse.model_validate(response.json()) async def update_entity( @@ -86,12 +96,21 @@ class KnowledgeClient: ToolError: If the request fails """ params = {"fast": fast} if fast is not None else None - response = await call_put( - self.http_client, - f"{self._base_path}/entities/{entity_id}", - json=entity_data, - params=params, - ) + with telemetry.scope( + "mcp.client.knowledge.update_entity", + client_name="knowledge", + operation="update_entity", + fast=fast, + ): + response = await call_put( + self.http_client, + f"{self._base_path}/entities/{entity_id}", + json=entity_data, + params=params, + client_name="knowledge", + operation="update_entity", + path_template="/v2/projects/{project_id}/knowledge/entities/{entity_id}", + ) return EntityResponse.model_validate(response.json()) async def get_entity(self, entity_id: str) -> EntityResponse: @@ -106,10 +125,18 @@ class KnowledgeClient: Raises: ToolError: If the entity is not found or request fails """ - response = await call_get( - self.http_client, - f"{self._base_path}/entities/{entity_id}", - ) + with telemetry.scope( + "mcp.client.knowledge.get_entity", + client_name="knowledge", + operation="get_entity", + ): + response = await call_get( + self.http_client, + f"{self._base_path}/entities/{entity_id}", + client_name="knowledge", + operation="get_entity", + path_template="/v2/projects/{project_id}/knowledge/entities/{entity_id}", + ) return EntityResponse.model_validate(response.json()) async def patch_entity( @@ -132,12 +159,21 @@ class KnowledgeClient: ToolError: If the request fails """ params = {"fast": fast} if fast is not None else None - response = await call_patch( - self.http_client, - f"{self._base_path}/entities/{entity_id}", - json=patch_data, - params=params, - ) + with telemetry.scope( + "mcp.client.knowledge.patch_entity", + client_name="knowledge", + operation="patch_entity", + fast=fast, + ): + response = await call_patch( + self.http_client, + f"{self._base_path}/entities/{entity_id}", + json=patch_data, + params=params, + client_name="knowledge", + operation="patch_entity", + path_template="/v2/projects/{project_id}/knowledge/entities/{entity_id}", + ) return EntityResponse.model_validate(response.json()) async def delete_entity(self, entity_id: str) -> DeleteEntitiesResponse: @@ -152,10 +188,18 @@ class KnowledgeClient: Raises: ToolError: If the entity is not found or request fails """ - response = await call_delete( - self.http_client, - f"{self._base_path}/entities/{entity_id}", - ) + with telemetry.scope( + "mcp.client.knowledge.delete_entity", + client_name="knowledge", + operation="delete_entity", + ): + response = await call_delete( + self.http_client, + f"{self._base_path}/entities/{entity_id}", + client_name="knowledge", + operation="delete_entity", + path_template="/v2/projects/{project_id}/knowledge/entities/{entity_id}", + ) return DeleteEntitiesResponse.model_validate(response.json()) async def move_entity(self, entity_id: str, destination_path: str) -> EntityResponse: @@ -171,11 +215,19 @@ class KnowledgeClient: Raises: ToolError: If the request fails """ - response = await call_put( - self.http_client, - f"{self._base_path}/entities/{entity_id}/move", - json={"destination_path": destination_path}, - ) + with telemetry.scope( + "mcp.client.knowledge.move_entity", + client_name="knowledge", + operation="move_entity", + ): + response = await call_put( + self.http_client, + f"{self._base_path}/entities/{entity_id}/move", + json={"destination_path": destination_path}, + client_name="knowledge", + operation="move_entity", + path_template="/v2/projects/{project_id}/knowledge/entities/{entity_id}/move", + ) return EntityResponse.model_validate(response.json()) async def move_directory( @@ -193,14 +245,22 @@ class KnowledgeClient: Raises: ToolError: If the request fails """ - response = await call_post( - self.http_client, - f"{self._base_path}/move-directory", - json={ - "source_directory": source_directory, - "destination_directory": destination_directory, - }, - ) + with telemetry.scope( + "mcp.client.knowledge.move_directory", + client_name="knowledge", + operation="move_directory", + ): + response = await call_post( + self.http_client, + f"{self._base_path}/move-directory", + json={ + "source_directory": source_directory, + "destination_directory": destination_directory, + }, + client_name="knowledge", + operation="move_directory", + path_template="/v2/projects/{project_id}/knowledge/move-directory", + ) return DirectoryMoveResult.model_validate(response.json()) async def delete_directory(self, directory: str) -> DirectoryDeleteResult: @@ -215,11 +275,19 @@ class KnowledgeClient: Raises: ToolError: If the request fails """ - response = await call_post( - self.http_client, - f"{self._base_path}/delete-directory", - json={"directory": directory}, - ) + with telemetry.scope( + "mcp.client.knowledge.delete_directory", + client_name="knowledge", + operation="delete_directory", + ): + response = await call_post( + self.http_client, + f"{self._base_path}/delete-directory", + json={"directory": directory}, + client_name="knowledge", + operation="delete_directory", + path_template="/v2/projects/{project_id}/knowledge/delete-directory", + ) return DirectoryDeleteResult.model_validate(response.json()) # --- Resolution --- @@ -237,10 +305,18 @@ class KnowledgeClient: Raises: ToolError: If the identifier cannot be resolved """ - response = await call_post( - self.http_client, - f"{self._base_path}/resolve", - json={"identifier": identifier, "strict": strict}, - ) + with telemetry.scope( + "mcp.client.knowledge.resolve_entity", + client_name="knowledge", + operation="resolve_entity", + ): + response = await call_post( + self.http_client, + f"{self._base_path}/resolve", + json={"identifier": identifier, "strict": strict}, + client_name="knowledge", + operation="resolve_entity", + path_template="/v2/projects/{project_id}/knowledge/resolve", + ) data = response.json() return data["external_id"] diff --git a/src/basic_memory/mcp/clients/memory.py b/src/basic_memory/mcp/clients/memory.py index de42c63e..bc1e8b56 100644 --- a/src/basic_memory/mcp/clients/memory.py +++ b/src/basic_memory/mcp/clients/memory.py @@ -7,6 +7,7 @@ from typing import Optional from httpx import AsyncClient +from basic_memory import telemetry from basic_memory.mcp.tools.utils import call_get from basic_memory.schemas.memory import GraphContext @@ -71,11 +72,21 @@ class MemoryClient: if timeframe: params["timeframe"] = timeframe - response = await call_get( - self.http_client, - f"{self._base_path}/{path}", - params=params, - ) + with telemetry.scope( + "mcp.client.memory.build_context", + client_name="memory", + operation="build_context", + page=page, + page_size=page_size, + ): + response = await call_get( + self.http_client, + f"{self._base_path}/{path}", + params=params, + client_name="memory", + operation="build_context", + path_template="/v2/projects/{project_id}/memory/{path}", + ) return GraphContext.model_validate(response.json()) async def recent( @@ -112,9 +123,19 @@ class MemoryClient: # Join types as comma-separated string if provided params["type"] = ",".join(types) if isinstance(types, list) else types - response = await call_get( - self.http_client, - f"{self._base_path}/recent", - params=params, - ) + with telemetry.scope( + "mcp.client.memory.recent_activity", + client_name="memory", + operation="recent_activity", + page=page, + page_size=page_size, + ): + response = await call_get( + self.http_client, + f"{self._base_path}/recent", + params=params, + client_name="memory", + operation="recent_activity", + path_template="/v2/projects/{project_id}/memory/recent", + ) return GraphContext.model_validate(response.json()) diff --git a/src/basic_memory/mcp/clients/resource.py b/src/basic_memory/mcp/clients/resource.py index cc44807b..1b3a50cd 100644 --- a/src/basic_memory/mcp/clients/resource.py +++ b/src/basic_memory/mcp/clients/resource.py @@ -7,6 +7,7 @@ from typing import Optional from httpx import AsyncClient, Response +from basic_memory import telemetry from basic_memory.mcp.tools.utils import call_get @@ -64,8 +65,18 @@ class ResourceClient: if page_size is not None: params["page_size"] = page_size - return await call_get( - self.http_client, - f"{self._base_path}/{entity_id}", - params=params if params else None, - ) + with telemetry.scope( + "mcp.client.resource.read", + client_name="resource", + operation="read", + page=page, + page_size=page_size, + ): + return await call_get( + self.http_client, + f"{self._base_path}/{entity_id}", + params=params if params else None, + client_name="resource", + operation="read", + path_template="/v2/projects/{project_id}/resource/{entity_id}", + ) diff --git a/src/basic_memory/mcp/clients/search.py b/src/basic_memory/mcp/clients/search.py index a7c3ccb9..0727851b 100644 --- a/src/basic_memory/mcp/clients/search.py +++ b/src/basic_memory/mcp/clients/search.py @@ -7,6 +7,7 @@ from typing import Any from httpx import AsyncClient +from basic_memory import telemetry from basic_memory.mcp.tools.utils import call_post from basic_memory.schemas.search import SearchResponse @@ -56,10 +57,20 @@ class SearchClient: Raises: ToolError: If the request fails """ - response = await call_post( - self.http_client, - f"{self._base_path}/", - json=query, - params={"page": page, "page_size": page_size}, - ) + with telemetry.scope( + "mcp.client.search.search", + client_name="search", + operation="search", + page=page, + page_size=page_size, + ): + response = await call_post( + self.http_client, + f"{self._base_path}/", + json=query, + params={"page": page, "page_size": page_size}, + client_name="search", + operation="search", + path_template="/v2/projects/{project_id}/search/", + ) return SearchResponse.model_validate(response.json()) diff --git a/src/basic_memory/mcp/tools/build_context.py b/src/basic_memory/mcp/tools/build_context.py index 4131b365..28a7265a 100644 --- a/src/basic_memory/mcp/tools/build_context.py +++ b/src/basic_memory/mcp/tools/build_context.py @@ -206,7 +206,7 @@ async def build_context( "mcp.tool.build_context", entrypoint="mcp", tool_name="build_context", - requested_project=project, + project_name=project, workspace_id=workspace, depth=depth or 1, timeframe=timeframe, diff --git a/src/basic_memory/mcp/tools/edit_note.py b/src/basic_memory/mcp/tools/edit_note.py index c0bb3c46..49323a2d 100644 --- a/src/basic_memory/mcp/tools/edit_note.py +++ b/src/basic_memory/mcp/tools/edit_note.py @@ -275,7 +275,7 @@ async def edit_note( "mcp.tool.edit_note", entrypoint="mcp", tool_name="edit_note", - requested_project=project, + project_name=project, workspace_id=workspace, edit_operation=operation, output_format=output_format, diff --git a/src/basic_memory/mcp/tools/read_note.py b/src/basic_memory/mcp/tools/read_note.py index 69ac9f7b..3d336ba3 100644 --- a/src/basic_memory/mcp/tools/read_note.py +++ b/src/basic_memory/mcp/tools/read_note.py @@ -16,8 +16,8 @@ from basic_memory.mcp.project_context import ( resolve_project_and_path, ) from basic_memory.mcp.server import mcp -from basic_memory.mcp.tools.search import search_notes from basic_memory.schemas.memory import memory_url_path +from basic_memory.schemas.search import SearchQuery from basic_memory.utils import validate_project_path @@ -144,7 +144,7 @@ async def read_note( "mcp.tool.read_note", entrypoint="mcp", tool_name="read_note", - requested_project=project, + project_name=project, workspace_id=workspace, output_format=output_format, page=page, @@ -199,24 +199,31 @@ async def read_note( ) # Import here to avoid circular import - from basic_memory.mcp.clients import KnowledgeClient, ResourceClient + from basic_memory.mcp.clients import KnowledgeClient, ResourceClient, SearchClient # Use typed clients for API calls knowledge_client = KnowledgeClient(client, active_project.external_id) resource_client = ResourceClient(client, active_project.external_id) + search_client = SearchClient(client, active_project.external_id) async def _read_json_payload(entity_id: str) -> dict: - entity = await knowledge_client.get_entity(entity_id) - response = await resource_client.read(entity_id, page=page, page_size=page_size) - content_text = response.text - body_content, parsed_frontmatter = _parse_opening_frontmatter(content_text) - return { - "title": entity.title, - "permalink": entity.permalink, - "file_path": entity.file_path, - "content": content_text if include_frontmatter else body_content, - "frontmatter": parsed_frontmatter, - } + with telemetry.scope( + "mcp.read_note.shape_response", + domain="mcp", + action="read_note", + phase="shape_response", + ): + entity = await knowledge_client.get_entity(entity_id) + response = await resource_client.read(entity_id, page=page, page_size=page_size) + content_text = response.text + body_content, parsed_frontmatter = _parse_opening_frontmatter(content_text) + return { + "title": entity.title, + "permalink": entity.permalink, + "file_path": entity.file_path, + "content": content_text if include_frontmatter else body_content, + "frontmatter": parsed_frontmatter, + } def _empty_json_payload() -> dict: return { @@ -233,6 +240,17 @@ async def read_note( results = payload.get("results") return results if isinstance(results, list) else [] + async def _search_candidates(identifier_text: str, *, title_only: bool) -> dict: + query = SearchQuery(title=identifier_text) if title_only else SearchQuery( + text=identifier_text + ) + response = await search_client.search( + query.model_dump(mode="json", exclude_none=True), + page=page, + page_size=page_size, + ) + return response.model_dump(mode="json") + def _result_title(item: dict) -> str: return str(item.get("title") or "") @@ -265,14 +283,7 @@ async def read_note( # Fallback 1: Try title search via API logger.info(f"Search title for: {identifier}") - title_results = await search_notes( - query=identifier, - search_type="title", - project=active_project.name, - workspace=workspace, - output_format="json", - context=context, - ) + title_results = await _search_candidates(identifier, title_only=True) title_candidates = _search_results(title_results) if title_candidates: @@ -319,14 +330,7 @@ async def read_note( # Fallback 2: Text search as a last resort logger.info(f"Title search failed, trying text search for: {identifier}") - text_results = await search_notes( - query=identifier, - search_type="text", - project=active_project.name, - workspace=workspace, - output_format="json", - context=context, - ) + text_results = await _search_candidates(identifier, title_only=False) # We didn't find a direct match, construct a helpful error message text_candidates = _search_results(text_results) diff --git a/src/basic_memory/mcp/tools/search.py b/src/basic_memory/mcp/tools/search.py index c35d3676..9f33f9a9 100644 --- a/src/basic_memory/mcp/tools/search.py +++ b/src/basic_memory/mcp/tools/search.py @@ -528,7 +528,7 @@ async def search_notes( "mcp.tool.search_notes", entrypoint="mcp", tool_name="search_notes", - requested_project=project, + project_name=project, workspace_id=workspace, search_type=search_type or "default", output_format=output_format, @@ -537,7 +537,7 @@ async def search_notes( has_query=bool(query and query.strip()), note_type_filter_count=len(note_types), entity_type_filter_count=len(entity_types), - has_metadata_filters=bool(metadata_filters), + has_filters=bool(metadata_filters or tags or status or note_types or entity_types or after_date), has_tags_filter=bool(tags), has_status_filter=bool(status), ): diff --git a/src/basic_memory/mcp/tools/utils.py b/src/basic_memory/mcp/tools/utils.py index 53a76bca..6f847c10 100644 --- a/src/basic_memory/mcp/tools/utils.py +++ b/src/basic_memory/mcp/tools/utils.py @@ -23,6 +23,7 @@ from httpx._types import ( from loguru import logger from mcp.server.fastmcp.exceptions import ToolError +from basic_memory import telemetry from basic_memory.config import ConfigManager @@ -135,10 +136,35 @@ def _resolve_error_message( return get_error_message(status_code, url, method) +def _request_scope( + method: str, + *, + client_name: str | None, + operation: str | None, + path_template: str | None, + params: QueryParamTypes | None = None, + has_body: bool = False, +): + """Create the shared MCP transport span used by all HTTP helpers.""" + return telemetry.scope( + "mcp.http.request", + method=method, + client_name=client_name, + operation=operation, + path_template=path_template, + phase="request", + has_query=bool(params), + has_body=has_body, + ) + + async def call_get( client: AsyncClient, url: URL | str, *, + client_name: str | None = None, + operation: str | None = None, + path_template: str | None = None, params: QueryParamTypes | None = None, headers: HeaderTypes | None = None, cookies: CookieTypes | None = None, @@ -170,16 +196,23 @@ async def call_get( error_message = None try: - response = await client.get( - url, + with _request_scope( + "GET", + client_name=client_name, + operation=operation, + path_template=path_template, params=params, - headers=headers, - cookies=cookies, - auth=auth, - follow_redirects=follow_redirects, - timeout=timeout, - extensions=extensions, - ) + ): + response = await client.get( + url, + params=params, + headers=headers, + cookies=cookies, + auth=auth, + follow_redirects=follow_redirects, + timeout=timeout, + extensions=extensions, + ) if response.is_success: return response @@ -212,6 +245,9 @@ async def call_put( client: AsyncClient, url: URL | str, *, + client_name: str | None = None, + operation: str | None = None, + path_template: str | None = None, content: RequestContent | None = None, data: RequestData | None = None, files: RequestFiles | None = None, @@ -251,20 +287,28 @@ async def call_put( error_message = None try: - response = await client.put( - url, - content=content, - data=data, - files=files, - json=json, + with _request_scope( + "PUT", + client_name=client_name, + operation=operation, + path_template=path_template, params=params, - headers=headers, - cookies=cookies, - auth=auth, - follow_redirects=follow_redirects, - timeout=timeout, - extensions=extensions, - ) + has_body=any(value is not None for value in (content, data, files, json)), + ): + response = await client.put( + url, + content=content, + data=data, + files=files, + json=json, + params=params, + headers=headers, + cookies=cookies, + auth=auth, + follow_redirects=follow_redirects, + timeout=timeout, + extensions=extensions, + ) if response.is_success: return response @@ -298,6 +342,9 @@ async def call_patch( client: AsyncClient, url: URL | str, *, + client_name: str | None = None, + operation: str | None = None, + path_template: str | None = None, content: RequestContent | None = None, data: RequestData | None = None, files: RequestFiles | None = None, @@ -336,20 +383,28 @@ async def call_patch( logger.debug(f"Calling PATCH '{url}'") try: - response = await client.patch( - url, - content=content, - data=data, - files=files, - json=json, + with _request_scope( + "PATCH", + client_name=client_name, + operation=operation, + path_template=path_template, params=params, - headers=headers, - cookies=cookies, - auth=auth, - follow_redirects=follow_redirects, - timeout=timeout, - extensions=extensions, - ) + has_body=any(value is not None for value in (content, data, files, json)), + ): + response = await client.patch( + url, + content=content, + data=data, + files=files, + json=json, + params=params, + headers=headers, + cookies=cookies, + auth=auth, + follow_redirects=follow_redirects, + timeout=timeout, + extensions=extensions, + ) if response.is_success: return response @@ -388,6 +443,9 @@ async def call_post( client: AsyncClient, url: URL | str, *, + client_name: str | None = None, + operation: str | None = None, + path_template: str | None = None, content: RequestContent | None = None, data: RequestData | None = None, files: RequestFiles | None = None, @@ -427,20 +485,28 @@ async def call_post( error_message = None try: - response = await client.post( - url=url, - content=content, - data=data, - files=files, - json=json, + with _request_scope( + "POST", + client_name=client_name, + operation=operation, + path_template=path_template, params=params, - headers=headers, - cookies=cookies, - auth=auth, - follow_redirects=follow_redirects, - timeout=timeout, - extensions=extensions, - ) + has_body=any(value is not None for value in (content, data, files, json)), + ): + response = await client.post( + url=url, + content=content, + data=data, + files=files, + json=json, + params=params, + headers=headers, + cookies=cookies, + auth=auth, + follow_redirects=follow_redirects, + timeout=timeout, + extensions=extensions, + ) logger.debug(f"response: {response.json()}") if response.is_success: @@ -506,6 +572,9 @@ async def call_delete( client: AsyncClient, url: URL | str, *, + client_name: str | None = None, + operation: str | None = None, + path_template: str | None = None, params: QueryParamTypes | None = None, headers: HeaderTypes | None = None, cookies: CookieTypes | None = None, @@ -537,16 +606,23 @@ async def call_delete( error_message = None try: - response = await client.delete( - url=url, + with _request_scope( + "DELETE", + client_name=client_name, + operation=operation, + path_template=path_template, params=params, - headers=headers, - cookies=cookies, - auth=auth, - follow_redirects=follow_redirects, - timeout=timeout, - extensions=extensions, - ) + ): + response = await client.delete( + url=url, + params=params, + headers=headers, + cookies=cookies, + auth=auth, + follow_redirects=follow_redirects, + timeout=timeout, + extensions=extensions, + ) if response.is_success: return response diff --git a/src/basic_memory/mcp/tools/write_note.py b/src/basic_memory/mcp/tools/write_note.py index 2ec997e9..f832565d 100644 --- a/src/basic_memory/mcp/tools/write_note.py +++ b/src/basic_memory/mcp/tools/write_note.py @@ -153,7 +153,7 @@ async def write_note( "mcp.tool.write_note", entrypoint="mcp", tool_name="write_note", - requested_project=project, + project_name=project, workspace_id=workspace, note_type=note_type, overwrite=effective_overwrite, diff --git a/src/basic_memory/services/context_service.py b/src/basic_memory/services/context_service.py index ccede46b..b9fe88d5 100644 --- a/src/basic_memory/services/context_service.py +++ b/src/basic_memory/services/context_service.py @@ -10,6 +10,7 @@ from typing import List, Optional, Tuple, TYPE_CHECKING from loguru import logger from sqlalchemy import text +from basic_memory import telemetry from basic_memory.repository.entity_repository import EntityRepository from basic_memory.repository.observation_repository import ObservationRepository from basic_memory.repository.postgres_search_repository import PostgresSearchRepository @@ -110,146 +111,162 @@ class ContextService: f"Building context for URI: '{memory_url}' depth: '{depth}' since: '{since}' limit: '{limit}' offset: '{offset}' max_related: '{max_related}'" ) - # Fetch one extra item to detect whether more pages exist (N+1 trick) - fetch_limit = limit + 1 + with telemetry.scope( + "memory.build_context", + domain="memory", + action="build_context", + phase="build_context", + limit=limit, + offset=offset, + ): + fetch_limit = limit + 1 - normalized_path: Optional[str] = None - if memory_url: - path = memory_url_path(memory_url) - # Check for wildcards before normalization - has_wildcard = "*" in path + normalized_path: Optional[str] = None + with telemetry.scope( + "memory.build_context.resolve_primary", + domain="memory", + action="build_context", + phase="resolve_primary", + ): + if memory_url: + path = memory_url_path(memory_url) + has_wildcard = "*" in path - if has_wildcard: - # For wildcard patterns, normalize each segment separately to preserve the * - parts = path.split("*") - normalized_parts = [ - generate_permalink(part, split_extension=False) if part else "" - for part in parts - ] - normalized_path = "*".join(normalized_parts) - logger.debug(f"Pattern search for '{normalized_path}'") - primary = await self.search_repository.search( - permalink_match=normalized_path, limit=fetch_limit, offset=offset - ) - else: - # For exact paths, normalize the whole thing - normalized_path = generate_permalink(path, split_extension=False) - logger.debug(f"Direct lookup for '{normalized_path}'") - primary = await self.search_repository.search( - permalink=normalized_path, limit=fetch_limit, offset=offset - ) - - # Trigger: exact permalink lookup returned no results - # Why: the identifier may be valid but not an exact permalink match - # (e.g., missing project prefix, title instead of permalink) - # Outcome: use LinkResolver's multi-strategy resolution to find the entity, - # then retry search with its actual permalink - if not primary and self.link_resolver: - entity = await self.link_resolver.resolve_link( - path, use_search=True, strict=False - ) - if entity: - logger.debug( - f"LinkResolver resolved '{path}' to permalink '{entity.permalink}'" - ) - normalized_path = entity.permalink + if has_wildcard: + parts = path.split("*") + normalized_parts = [ + generate_permalink(part, split_extension=False) if part else "" + for part in parts + ] + normalized_path = "*".join(normalized_parts) + logger.debug(f"Pattern search for '{normalized_path}'") primary = await self.search_repository.search( - permalink=entity.permalink, limit=fetch_limit, offset=offset + permalink_match=normalized_path, limit=fetch_limit, offset=offset ) - else: - logger.debug(f"Build context for '{types}'") - primary = await self.search_repository.search( - search_item_types=types, after_date=since, limit=fetch_limit, offset=offset + else: + normalized_path = generate_permalink(path, split_extension=False) + logger.debug(f"Direct lookup for '{normalized_path}'") + primary = await self.search_repository.search( + permalink=normalized_path, limit=fetch_limit, offset=offset + ) + + if not primary and self.link_resolver: + entity = await self.link_resolver.resolve_link( + path, use_search=True, strict=False + ) + if entity: + logger.debug( + f"LinkResolver resolved '{path}' to permalink '{entity.permalink}'" + ) + normalized_path = entity.permalink + primary = await self.search_repository.search( + permalink=entity.permalink, + limit=fetch_limit, + offset=offset, + ) + else: + logger.debug(f"Build context for '{types}'") + primary = await self.search_repository.search( + search_item_types=types, + after_date=since, + limit=fetch_limit, + offset=offset, + ) + + has_more = len(primary) > limit + if has_more: + primary = primary[:limit] + + type_id_pairs = [(r.type, r.id) for r in primary] if primary else [] + logger.debug(f"found primary type_id_pairs: {len(type_id_pairs)}") + + with telemetry.scope( + "memory.build_context.find_related", + domain="memory", + action="build_context", + phase="find_related", + ): + related = await self.find_related( + type_id_pairs, max_depth=depth, since=since, max_results=max_related + ) + logger.debug(f"Found {len(related)} related results") + + entity_ids = [] + for result in primary: + if result.type == SearchItemType.ENTITY.value: + entity_ids.append(result.id) + + for result in related: + if result.type == SearchItemType.ENTITY.value: + entity_ids.append(result.id) + + observations_by_entity = {} + if include_observations and entity_ids: + with telemetry.scope( + "memory.build_context.load_observations", + domain="memory", + action="build_context", + phase="load_observations", + result_count=len(entity_ids), + ): + observations_by_entity = await self.observation_repository.find_by_entities( + entity_ids + ) + logger.debug(f"Found observations for {len(observations_by_entity)} entities") + + metadata = ContextMetadata( + uri=normalized_path if memory_url else None, + types=types, + depth=depth, + timeframe=since.isoformat() if since else None, + primary_count=len(primary), + related_count=len(related), + total_observations=sum(len(obs) for obs in observations_by_entity.values()), + total_relations=sum(1 for r in related if r.type == SearchItemType.RELATION), + has_more=has_more, ) - # Trim to requested limit and set has_more flag - has_more = len(primary) > limit - if has_more: - primary = primary[:limit] + with telemetry.scope( + "memory.build_context.shape_results", + domain="memory", + action="build_context", + phase="shape_results", + result_count=len(primary), + ): + context_results = [] + for primary_item in primary: + related_to_primary = [r for r in related if r.root_id == primary_item.id] - # Get type_id pairs for traversal + item_observations = [] + if primary_item.type == SearchItemType.ENTITY.value and include_observations: + for obs in observations_by_entity.get(primary_item.id, []): + item_observations.append( + ContextResultRow( + type="observation", + id=obs.id, + title=f"{obs.category}: {obs.content[:50]}...", + permalink=generate_permalink( + f"{primary_item.permalink}/observations/{obs.category}/{obs.content}" + ), + file_path=primary_item.file_path, + content=obs.content, + category=obs.category, + entity_id=primary_item.id, + depth=0, + root_id=primary_item.id, + created_at=primary_item.created_at, + ) + ) - type_id_pairs = [(r.type, r.id) for r in primary] if primary else [] - logger.debug(f"found primary type_id_pairs: {len(type_id_pairs)}") - - # Find related content - related = await self.find_related( - type_id_pairs, max_depth=depth, since=since, max_results=max_related - ) - logger.debug(f"Found {len(related)} related results") - - # Collect entity IDs from primary and related results - entity_ids = [] - for result in primary: - if result.type == SearchItemType.ENTITY.value: - entity_ids.append(result.id) - - for result in related: - if result.type == SearchItemType.ENTITY.value: - entity_ids.append(result.id) - - # Fetch observations for all entities if requested - observations_by_entity = {} - if include_observations and entity_ids: - # Use our observation repository to get observations for all entities at once - observations_by_entity = await self.observation_repository.find_by_entities(entity_ids) - logger.debug(f"Found observations for {len(observations_by_entity)} entities") - - # Create metadata dataclass - metadata = ContextMetadata( - uri=normalized_path if memory_url else None, - types=types, - depth=depth, - timeframe=since.isoformat() if since else None, - primary_count=len(primary), - related_count=len(related), - total_observations=sum(len(obs) for obs in observations_by_entity.values()), - total_relations=sum(1 for r in related if r.type == SearchItemType.RELATION), - has_more=has_more, - ) - - # Build context results list directly with ContextResultItem objects - context_results = [] - - # For each primary result - for primary_item in primary: - # Find all related items with this primary item as root - related_to_primary = [r for r in related if r.root_id == primary_item.id] - - # Get observations for this item if it's an entity - item_observations = [] - if primary_item.type == SearchItemType.ENTITY.value and include_observations: - # Convert Observation models to ContextResultRows - for obs in observations_by_entity.get(primary_item.id, []): - item_observations.append( - ContextResultRow( - type="observation", - id=obs.id, - title=f"{obs.category}: {obs.content[:50]}...", - permalink=generate_permalink( - f"{primary_item.permalink}/observations/{obs.category}/{obs.content}" - ), - file_path=primary_item.file_path, - content=obs.content, - category=obs.category, - entity_id=primary_item.id, - depth=0, - root_id=primary_item.id, - created_at=primary_item.created_at, # created_at time from entity + context_results.append( + ContextResultItem( + primary_result=primary_item, + observations=item_observations, + related_results=related_to_primary, ) ) - # Create ContextResultItem directly - context_item = ContextResultItem( - primary_result=primary_item, - observations=item_observations, - related_results=related_to_primary, - ) - - context_results.append(context_item) - - # Return the structured ContextResult - return ContextResult(results=context_results, metadata=metadata) + return ContextResult(results=context_results, metadata=metadata) async def find_related( self, diff --git a/src/basic_memory/services/entity_service.py b/src/basic_memory/services/entity_service.py index e6a1aff9..61cbf133 100644 --- a/src/basic_memory/services/entity_service.py +++ b/src/basic_memory/services/entity_service.py @@ -10,7 +10,7 @@ import yaml from loguru import logger from sqlalchemy.exc import IntegrityError - +from basic_memory import telemetry from basic_memory.config import ProjectConfig, BasicMemoryConfig from basic_memory.file_utils import ( has_frontmatter, @@ -281,31 +281,54 @@ class EntityService(BaseService[EntityModel]): # Get unique permalink (prioritizing content frontmatter) unless disabled if self.app_config and self.app_config.disable_permalinks: - # Use empty string as sentinel to indicate permalinks are disabled - # The permalink property will return None when it sees empty string schema._permalink = "" else: - # Generate and set permalink - permalink = await self.resolve_permalink(file_path, content_markdown) + with telemetry.scope( + "entity_service.create.resolve_permalink", + domain="entity_service", + action="create", + phase="resolve_permalink", + ): + permalink = await self.resolve_permalink(file_path, content_markdown) schema._permalink = permalink post = await schema_to_markdown(schema) - # write file final_content = dump_frontmatter(post) - checksum = await self.file_service.write_file(file_path, final_content) + with telemetry.scope( + "entity_service.create.write_file", + domain="entity_service", + action="create", + phase="write_file", + ): + checksum = await self.file_service.write_file(file_path, final_content) - # parse entity from content we just wrote (avoids re-reading file for cloud compatibility) - entity_markdown = await self.entity_parser.parse_markdown_content( - file_path=file_path, - content=final_content, - ) + with telemetry.scope( + "entity_service.create.parse_markdown", + domain="entity_service", + action="create", + phase="parse_markdown", + ): + entity_markdown = await self.entity_parser.parse_markdown_content( + file_path=file_path, + content=final_content, + ) - # create entity and relations - entity = await self.upsert_entity_from_markdown(file_path, entity_markdown, is_new=True) + with telemetry.scope( + "entity_service.create.upsert_entity", + domain="entity_service", + action="create", + phase="upsert_entity", + ): + entity = await self.upsert_entity_from_markdown(file_path, entity_markdown, is_new=True) - # Set final checksum to mark complete - return await self.repository.update(entity.id, {"checksum": checksum}) + with telemetry.scope( + "entity_service.create.update_checksum", + domain="entity_service", + action="create", + phase="update_checksum", + ): + return await self.repository.update(entity.id, {"checksum": checksum}) async def update_entity(self, entity: EntityModel, schema: EntitySchema) -> EntityModel: """Update an entity's content and metadata.""" @@ -316,12 +339,23 @@ class EntityService(BaseService[EntityModel]): # Convert file path string to Path file_path = Path(entity.file_path) - # Read existing content via file_service (for cloud compatibility) - existing_content = await self.file_service.read_file_content(file_path) - existing_markdown = await self.entity_parser.parse_markdown_content( - file_path=file_path, - content=existing_content, - ) + with telemetry.scope( + "entity_service.update.read_file", + domain="entity_service", + action="update", + phase="read_file", + ): + existing_content = await self.file_service.read_file_content(file_path) + with telemetry.scope( + "entity_service.update.parse_markdown", + domain="entity_service", + action="update", + phase="parse_markdown", + ): + existing_markdown = await self.entity_parser.parse_markdown_content( + file_path=file_path, + content=existing_content, + ) # Parse content frontmatter to check for user-specified permalink and note_type content_markdown = None @@ -342,7 +376,13 @@ class EntityService(BaseService[EntityModel]): if self.app_config and not self.app_config.disable_permalinks: if content_markdown and content_markdown.frontmatter.permalink: # Resolve permalink with the new content frontmatter - resolved_permalink = await self.resolve_permalink(file_path, content_markdown) + with telemetry.scope( + "entity_service.update.resolve_permalink", + domain="entity_service", + action="update", + phase="resolve_permalink", + ): + resolved_permalink = await self.resolve_permalink(file_path, content_markdown) if resolved_permalink != entity.permalink: new_permalink = resolved_permalink # Update the schema to use the new permalink @@ -367,21 +407,41 @@ class EntityService(BaseService[EntityModel]): merged_post = frontmatter.Post(post.content) merged_post.metadata.update(existing_markdown.frontmatter.metadata) - # write file final_content = dump_frontmatter(merged_post) - checksum = await self.file_service.write_file(file_path, final_content) + with telemetry.scope( + "entity_service.update.write_file", + domain="entity_service", + action="update", + phase="write_file", + ): + checksum = await self.file_service.write_file(file_path, final_content) - # parse entity from content we just wrote (avoids re-reading file for cloud compatibility) - entity_markdown = await self.entity_parser.parse_markdown_content( - file_path=file_path, - content=final_content, - ) + with telemetry.scope( + "entity_service.update.parse_markdown", + domain="entity_service", + action="update", + phase="parse_markdown", + ): + entity_markdown = await self.entity_parser.parse_markdown_content( + file_path=file_path, + content=final_content, + ) - # update entity and relations - entity = await self.upsert_entity_from_markdown(file_path, entity_markdown, is_new=False) + with telemetry.scope( + "entity_service.update.upsert_entity", + domain="entity_service", + action="update", + phase="upsert_entity", + ): + entity = await self.upsert_entity_from_markdown(file_path, entity_markdown, is_new=False) - # Set final checksum to match file - entity = await self.repository.update(entity.id, {"checksum": checksum}) + with telemetry.scope( + "entity_service.update.update_checksum", + domain="entity_service", + action="update", + phase="update_checksum", + ): + entity = await self.repository.update(entity.id, {"checksum": checksum}) return entity @@ -399,7 +459,13 @@ class EntityService(BaseService[EntityModel]): ) # --- Identity & File Path --- - existing = await self.repository.get_by_external_id(external_id) if external_id else None + with telemetry.scope( + "entity_service.fast_write.resolve_entity", + domain="entity_service", + action="fast_write", + phase="resolve_entity", + ): + existing = await self.repository.get_by_external_id(external_id) if external_id else None # Trigger: external_id already exists # Why: avoid duplicate entities when title-derived paths change @@ -429,18 +495,35 @@ class EntityService(BaseService[EntityModel]): schema._permalink = "" else: if existing and not (content_markdown and content_markdown.frontmatter.permalink): - schema._permalink = existing.permalink or await self.resolve_permalink( - file_path, skip_conflict_check=True - ) + with telemetry.scope( + "entity_service.fast_write.resolve_permalink", + domain="entity_service", + action="fast_write", + phase="resolve_permalink", + ): + schema._permalink = existing.permalink or await self.resolve_permalink( + file_path, skip_conflict_check=True + ) else: - schema._permalink = await self.resolve_permalink( - file_path, content_markdown, skip_conflict_check=True - ) + with telemetry.scope( + "entity_service.fast_write.resolve_permalink", + domain="entity_service", + action="fast_write", + phase="resolve_permalink", + ): + schema._permalink = await self.resolve_permalink( + file_path, content_markdown, skip_conflict_check=True + ) - # --- File Write --- post = await schema_to_markdown(schema) final_content = dump_frontmatter(post) - checksum = await self.file_service.write_file(file_path, final_content) + with telemetry.scope( + "entity_service.fast_write.write_file", + domain="entity_service", + action="fast_write", + phase="write_file", + ): + checksum = await self.file_service.write_file(file_path, final_content) # --- Minimal DB Upsert --- metadata = normalize_frontmatter_metadata(post.metadata or {}) @@ -462,7 +545,13 @@ class EntityService(BaseService[EntityModel]): # Preserve existing created_by; only update last_updated_by if user_id is not None: update_data["last_updated_by"] = user_id - updated = await self.repository.update(existing.id, update_data) + with telemetry.scope( + "entity_service.fast_write.upsert_entity", + domain="entity_service", + action="fast_write", + phase="upsert_entity", + ): + updated = await self.repository.update(existing.id, update_data) if not updated: raise ValueError(f"Failed to update entity in database: {existing.id}") return updated @@ -473,7 +562,13 @@ class EntityService(BaseService[EntityModel]): if user_id is not None: create_data["created_by"] = user_id create_data["last_updated_by"] = user_id - return await self.repository.create(create_data) + with telemetry.scope( + "entity_service.fast_write.upsert_entity", + domain="entity_service", + action="fast_write", + phase="upsert_entity", + ): + return await self.repository.create(create_data) async def fast_edit_entity( self, @@ -487,13 +582,30 @@ class EntityService(BaseService[EntityModel]): """Edit an entity quickly and defer full indexing to background.""" logger.debug(f"Fast editing entity: {entity.external_id}, operation: {operation}") - # --- File Edit --- file_path = Path(entity.file_path) - current_content, _ = await self.file_service.read_file(file_path) - new_content = self.apply_edit_operation( - current_content, operation, content, section, find_text, expected_replacements - ) - checksum = await self.file_service.write_file(file_path, new_content) + with telemetry.scope( + "entity_service.fast_edit.read_file", + domain="entity_service", + action="fast_edit", + phase="read_file", + ): + current_content, _ = await self.file_service.read_file(file_path) + with telemetry.scope( + "entity_service.fast_edit.apply_operation", + domain="entity_service", + action="fast_edit", + phase="apply_operation", + ): + new_content = self.apply_edit_operation( + current_content, operation, content, section, find_text, expected_replacements + ) + with telemetry.scope( + "entity_service.fast_edit.write_file", + domain="entity_service", + action="fast_edit", + phase="write_file", + ): + checksum = await self.file_service.write_file(file_path, new_content) # --- Frontmatter Overrides --- update_data = { @@ -528,39 +640,84 @@ class EntityService(BaseService[EntityModel]): if self.app_config and self.app_config.disable_permalinks: update_data["permalink"] = None elif content_markdown and content_markdown.frontmatter.permalink: - update_data["permalink"] = await self.resolve_permalink( - file_path, content_markdown, skip_conflict_check=True - ) + with telemetry.scope( + "entity_service.fast_edit.resolve_permalink", + domain="entity_service", + action="fast_edit", + phase="resolve_permalink", + ): + update_data["permalink"] = await self.resolve_permalink( + file_path, content_markdown, skip_conflict_check=True + ) - updated = await self.repository.update(entity.id, update_data) + with telemetry.scope( + "entity_service.fast_edit.update_entity", + domain="entity_service", + action="fast_edit", + phase="update_entity", + ): + updated = await self.repository.update(entity.id, update_data) if not updated: raise ValueError(f"Failed to update entity in database: {entity.id}") return updated async def reindex_entity(self, entity_id: int) -> None: """Parse file content and rebuild observations/relations/search for an entity.""" - entity = await self.repository.find_by_id(entity_id) + with telemetry.scope( + "entity_service.reindex.load_entity", + domain="entity_service", + action="reindex", + phase="load_entity", + ): + entity = await self.repository.find_by_id(entity_id) if not entity: raise EntityNotFoundError(f"Entity not found: {entity_id}") - # --- Full Parse --- file_path = Path(entity.file_path) - content = await self.file_service.read_file_content(file_path) - entity_markdown = await self.entity_parser.parse_markdown_content( - file_path=file_path, - content=content, - ) + with telemetry.scope( + "entity_service.reindex.read_file", + domain="entity_service", + action="reindex", + phase="read_file", + ): + content = await self.file_service.read_file_content(file_path) + with telemetry.scope( + "entity_service.reindex.parse_markdown", + domain="entity_service", + action="reindex", + phase="parse_markdown", + ): + entity_markdown = await self.entity_parser.parse_markdown_content( + file_path=file_path, + content=content, + ) - # --- DB Reindex --- - updated = await self.upsert_entity_from_markdown(file_path, entity_markdown, is_new=False) - checksum = await self.file_service.compute_checksum(file_path) - updated = await self.repository.update(updated.id, {"checksum": checksum}) + with telemetry.scope( + "entity_service.reindex.upsert_entity", + domain="entity_service", + action="reindex", + phase="upsert_entity", + ): + updated = await self.upsert_entity_from_markdown(file_path, entity_markdown, is_new=False) + with telemetry.scope( + "entity_service.reindex.update_checksum", + domain="entity_service", + action="reindex", + phase="update_checksum", + ): + checksum = await self.file_service.compute_checksum(file_path) + updated = await self.repository.update(updated.id, {"checksum": checksum}) if not updated: raise ValueError(f"Failed to update entity in database: {entity.id}") - # --- Search Reindex --- if self.search_service: - await self.search_service.index_entity_data(updated, content=content) + with telemetry.scope( + "entity_service.reindex.search_index", + domain="entity_service", + action="reindex", + phase="search_index", + ): + await self.search_service.index_entity_data(updated, content=content) async def delete_entity(self, permalink_or_id: str | int) -> bool: """Delete entity and its file.""" @@ -808,34 +965,69 @@ class EntityService(BaseService[EntityModel]): """ logger.debug(f"Editing entity: {identifier}, operation: {operation}") - # Find the entity using the link resolver with strict mode for destructive operations - entity = await self.link_resolver.resolve_link(identifier, strict=True) + with telemetry.scope( + "entity_service.edit.resolve_entity", + domain="entity_service", + action="edit", + phase="resolve_entity", + ): + entity = await self.link_resolver.resolve_link(identifier, strict=True) if not entity: raise EntityNotFoundError(f"Entity not found: {identifier}") - # Read the current file content file_path = Path(entity.file_path) - current_content, _ = await self.file_service.read_file(file_path) + with telemetry.scope( + "entity_service.edit.read_file", + domain="entity_service", + action="edit", + phase="read_file", + ): + current_content, _ = await self.file_service.read_file(file_path) - # Apply the edit operation - new_content = self.apply_edit_operation( - current_content, operation, content, section, find_text, expected_replacements - ) + with telemetry.scope( + "entity_service.edit.apply_operation", + domain="entity_service", + action="edit", + phase="apply_operation", + ): + new_content = self.apply_edit_operation( + current_content, operation, content, section, find_text, expected_replacements + ) - # Write the updated content back to the file - checksum = await self.file_service.write_file(file_path, new_content) + with telemetry.scope( + "entity_service.edit.write_file", + domain="entity_service", + action="edit", + phase="write_file", + ): + checksum = await self.file_service.write_file(file_path, new_content) - # Parse the content we just wrote (avoids re-reading file for cloud compatibility) - entity_markdown = await self.entity_parser.parse_markdown_content( - file_path=file_path, - content=new_content, - ) + with telemetry.scope( + "entity_service.edit.parse_markdown", + domain="entity_service", + action="edit", + phase="parse_markdown", + ): + entity_markdown = await self.entity_parser.parse_markdown_content( + file_path=file_path, + content=new_content, + ) - # Update entity and its relationships - entity = await self.upsert_entity_from_markdown(file_path, entity_markdown, is_new=False) + with telemetry.scope( + "entity_service.edit.upsert_entity", + domain="entity_service", + action="edit", + phase="upsert_entity", + ): + entity = await self.upsert_entity_from_markdown(file_path, entity_markdown, is_new=False) - # Set final checksum to match file - entity = await self.repository.update(entity.id, {"checksum": checksum}) + with telemetry.scope( + "entity_service.edit.update_checksum", + domain="entity_service", + action="edit", + phase="update_checksum", + ): + entity = await self.repository.update(entity.id, {"checksum": checksum}) return entity diff --git a/src/basic_memory/services/file_service.py b/src/basic_memory/services/file_service.py index a1d0ea80..86cb3328 100644 --- a/src/basic_memory/services/file_service.py +++ b/src/basic_memory/services/file_service.py @@ -11,6 +11,7 @@ import aiofiles import yaml +from basic_memory import telemetry from basic_memory import file_utils if TYPE_CHECKING: # pragma: no cover @@ -79,13 +80,18 @@ class FileService: """ logger.debug(f"Reading entity content, entity_id={entity.id}, permalink={entity.permalink}") - # markdown_processor is required for entity content reads — fail fast if not configured - if self.markdown_processor is None: - raise ValueError("markdown_processor is required for read_entity_content") + with telemetry.scope( + "file_service.read_content", + domain="file_service", + action="read_content", + phase="read_content", + ): + if self.markdown_processor is None: + raise ValueError("markdown_processor is required for read_entity_content") - file_path = self.get_entity_path(entity) - markdown = await self.markdown_processor.read_file(file_path) - return markdown.content or "" + file_path = self.get_entity_path(entity) + markdown = await self.markdown_processor.read_file(file_path) + return markdown.content or "" async def delete_entity_file(self, entity: EntityModel) -> None: """Delete entity file from filesystem. @@ -176,32 +182,34 @@ class FileService: full_path = path_obj if path_obj.is_absolute() else self.base_path / path_obj try: - # Ensure parent directory exists - await self.ensure_directory(full_path.parent) + with telemetry.scope( + "file_service.write", + domain="file_service", + action="write", + phase="write", + ): + await self.ensure_directory(full_path.parent) - # Write content atomically - logger.info( - "Writing file: " - f"path={path_obj}, " - f"content_length={len(content)}, " - f"is_markdown={full_path.suffix.lower() == '.md'}" - ) - - await file_utils.write_file_atomic(full_path, content) - - # Format file if configured - final_content = content - if self.app_config: - formatted_content = await file_utils.format_file( - full_path, self.app_config, is_markdown=self.is_markdown(path) + logger.info( + "Writing file: " + f"path={path_obj}, " + f"content_length={len(content)}, " + f"is_markdown={full_path.suffix.lower() == '.md'}" ) - if formatted_content is not None: - final_content = formatted_content # pragma: no cover - # Compute and return checksum of final content - checksum = await file_utils.compute_checksum(final_content) - logger.debug(f"File write completed path={full_path}, {checksum=}") - return checksum + await file_utils.write_file_atomic(full_path, content) + + final_content = content + if self.app_config: + formatted_content = await file_utils.format_file( + full_path, self.app_config, is_markdown=self.is_markdown(path) + ) + if formatted_content is not None: + final_content = formatted_content # pragma: no cover + + checksum = await file_utils.compute_checksum(final_content) + logger.debug(f"File write completed path={full_path}, {checksum=}") + return checksum except Exception as e: logger.exception("File write error", path=str(full_path), error=str(e)) @@ -227,16 +235,24 @@ class FileService: full_path = path_obj if path_obj.is_absolute() else self.base_path / path_obj try: - logger.debug("Reading file content", operation="read_file_content", path=str(full_path)) - async with aiofiles.open(full_path, mode="r", encoding="utf-8") as f: - content = await f.read() + with telemetry.scope( + "file_service.read_content", + domain="file_service", + action="read_content", + phase="read_content", + ): + logger.debug( + "Reading file content", operation="read_file_content", path=str(full_path) + ) + async with aiofiles.open(full_path, mode="r", encoding="utf-8") as f: + content = await f.read() - logger.debug( - "File read completed", - path=str(full_path), - content_length=len(content), - ) - return content + logger.debug( + "File read completed", + path=str(full_path), + content_length=len(content), + ) + return content except FileNotFoundError: # Preserve FileNotFoundError so callers (e.g. sync) can treat it as deletion. @@ -266,16 +282,22 @@ class FileService: full_path = path_obj if path_obj.is_absolute() else self.base_path / path_obj try: - logger.debug("Reading file bytes", operation="read_file_bytes", path=str(full_path)) - async with aiofiles.open(full_path, mode="rb") as f: - content = await f.read() + with telemetry.scope( + "file_service.read_content", + domain="file_service", + action="read_content", + phase="read_content", + ): + logger.debug("Reading file bytes", operation="read_file_bytes", path=str(full_path)) + async with aiofiles.open(full_path, mode="rb") as f: + content = await f.read() - logger.debug( - "File read completed", - path=str(full_path), - content_length=len(content), - ) - return content + logger.debug( + "File read completed", + path=str(full_path), + content_length=len(content), + ) + return content except Exception as e: logger.exception("File read error", path=str(full_path), error=str(e)) @@ -303,21 +325,26 @@ class FileService: full_path = path_obj if path_obj.is_absolute() else self.base_path / path_obj try: - logger.debug("Reading file", operation="read_file", path=str(full_path)) + with telemetry.scope( + "file_service.read", + domain="file_service", + action="read", + phase="read", + ): + logger.debug("Reading file", operation="read_file", path=str(full_path)) - # Use aiofiles for non-blocking read - async with aiofiles.open(full_path, mode="r", encoding="utf-8") as f: - content = await f.read() + async with aiofiles.open(full_path, mode="r", encoding="utf-8") as f: + content = await f.read() - checksum = await file_utils.compute_checksum(content) + checksum = await file_utils.compute_checksum(content) - logger.debug( - "File read completed", - path=str(full_path), - checksum=checksum, - content_length=len(content), - ) - return content, checksum + logger.debug( + "File read completed", + path=str(full_path), + checksum=checksum, + content_length=len(content), + ) + return content, checksum except Exception as e: logger.exception("File read error", path=str(full_path), error=str(e)) diff --git a/src/basic_memory/services/search_service.py b/src/basic_memory/services/search_service.py index fe29e08e..6057482f 100644 --- a/src/basic_memory/services/search_service.py +++ b/src/basic_memory/services/search_service.py @@ -173,33 +173,52 @@ class SearchService: retrieval_mode = query.retrieval_mode or SearchRetrievalMode.FTS strict_search_text = query.text + has_query = bool( + strict_search_text + or query.title + or query.permalink + or query.permalink_match + ) + has_filters = bool( + metadata_filters + or query.note_types + or query.entity_types + or after_date + or query.tags + or query.status + ) with telemetry.scope( "search.execute", retrieval_mode=retrieval_mode.value, - has_text_query=bool(strict_search_text), - has_title_query=bool(query.title), - has_permalink_query=bool(query.permalink or query.permalink_match), - has_metadata_filters=bool(metadata_filters), + has_query=has_query, + has_filters=has_filters, limit=limit, offset=offset, ): logger.trace(f"Searching with query: {query}") - # First pass: preserve existing strict search behavior. - results = await self.repository.search( - search_text=strict_search_text, - permalink=query.permalink, - permalink_match=query.permalink_match, - title=query.title, - note_types=query.note_types, - search_item_types=query.entity_types, - after_date=after_date, - metadata_filters=metadata_filters, - retrieval_mode=retrieval_mode, - min_similarity=query.min_similarity, - limit=limit, - offset=offset, - ) + with telemetry.scope( + "search.repository_query", + retrieval_mode=retrieval_mode.value, + phase="repository_query", + has_query=has_query, + has_filters=has_filters, + ): + # First pass: preserve existing strict search behavior. + results = await self.repository.search( + search_text=strict_search_text, + permalink=query.permalink, + permalink_match=query.permalink_match, + title=query.title, + note_types=query.note_types, + search_item_types=query.entity_types, + after_date=after_date, + metadata_filters=metadata_filters, + retrieval_mode=retrieval_mode, + min_similarity=query.min_similarity, + limit=limit, + offset=offset, + ) # Trigger: strict FTS with plain multi-term text returned no results. # Why: natural-language queries often include stopwords that over-constrain implicit AND. @@ -225,20 +244,27 @@ class SearchService: limit=limit, offset=offset, ): - return await self.repository.search( - search_text=relaxed_search_text, - permalink=query.permalink, - permalink_match=query.permalink_match, - title=query.title, - note_types=query.note_types, - search_item_types=query.entity_types, - after_date=after_date, - metadata_filters=metadata_filters, - retrieval_mode=retrieval_mode, - min_similarity=query.min_similarity, - limit=limit, - offset=offset, - ) + with telemetry.scope( + "search.repository_query", + retrieval_mode=retrieval_mode.value, + phase="repository_query", + has_query=has_query, + has_filters=has_filters, + ): + return await self.repository.search( + search_text=relaxed_search_text, + permalink=query.permalink, + permalink_match=query.permalink_match, + title=query.title, + note_types=query.note_types, + search_item_types=query.entity_types, + after_date=after_date, + metadata_filters=metadata_filters, + retrieval_mode=retrieval_mode, + min_similarity=query.min_similarity, + limit=limit, + offset=offset, + ) @staticmethod def _tokenize_fts_text(search_text: str) -> list[str]: @@ -372,13 +398,22 @@ class SearchService: f"permalink={entity.permalink} project_id={entity.project_id}" ) try: - # delete all search index data associated with entity - await self.repository.delete_by_entity_id(entity_id=entity.id) + with telemetry.scope( + "search.index_entity_data", + phase="index_entity_data", + result_count=1, + ): + with telemetry.scope( + "search.index.delete_existing", + phase="delete_existing", + result_count=1, + ): + await self.repository.delete_by_entity_id(entity_id=entity.id) - # reindex - await self.index_entity_markdown( - entity, content - ) if entity.is_markdown else await self.index_entity_file(entity) + if entity.is_markdown: + await self.index_entity_markdown(entity, content) + else: + await self.index_entity_file(entity) logger.debug( f"[BackgroundTask] Completed search index for entity_id={entity.id} " @@ -490,23 +525,28 @@ class SearchService: self, entity: Entity, ) -> None: - # Index entity file with no content - await self.repository.index_item( - SearchIndexRow( - id=entity.id, - entity_id=entity.id, - type=SearchItemType.ENTITY.value, - title=_strip_nul(entity.title), - permalink=entity.permalink, # Required for Postgres NOT NULL constraint - file_path=entity.file_path, - metadata={ - "note_type": entity.note_type, - }, - created_at=entity.created_at, - updated_at=_mtime_to_datetime(entity), - project_id=entity.project_id, + with telemetry.scope( + "search.index_file", + phase="index_file", + result_count=1, + ): + # Index entity file with no content + await self.repository.index_item( + SearchIndexRow( + id=entity.id, + entity_id=entity.id, + type=SearchItemType.ENTITY.value, + title=_strip_nul(entity.title), + permalink=entity.permalink, # Required for Postgres NOT NULL constraint + file_path=entity.file_path, + metadata={ + "note_type": entity.note_type, + }, + created_at=entity.created_at, + updated_at=_mtime_to_datetime(entity), + project_id=entity.project_id, + ) ) - ) async def index_entity_markdown( self, @@ -539,129 +579,142 @@ class SearchService: The project_id is automatically added by the repository when indexing. """ - # Collect all search index rows to batch insert at the end - rows_to_index = [] + with telemetry.scope( + "search.index_markdown", + phase="index_markdown", + result_count=1, + ): + rows_to_index = [] - content_stems = [] - content_snippet = "" - title_variants = self._generate_variants(entity.title) - content_stems.extend(title_variants) + content_stems = [] + content_snippet = "" + title_variants = self._generate_variants(entity.title) + content_stems.extend(title_variants) - # Use provided content or read from file - if content is None: - content = await self.file_service.read_entity_content(entity) - if content: - content_stems.append(content) - # Store full content for vector embedding quality. - # The chunker in the vector pipeline splits this into - # appropriately-sized pieces for embedding. - content_snippet = _strip_nul(content) + if content is None: + with telemetry.scope( + "search.index.read_content", + phase="read_content", + result_count=1, + ): + content = await self.file_service.read_entity_content(entity) + if content: + content_stems.append(content) + content_snippet = _strip_nul(content) - if entity.permalink: - content_stems.extend(self._generate_variants(entity.permalink)) + with telemetry.scope( + "search.index.build_rows", + phase="build_rows", + result_count=1, + ): + if entity.permalink: + content_stems.extend(self._generate_variants(entity.permalink)) - content_stems.extend(self._generate_variants(entity.file_path)) + content_stems.extend(self._generate_variants(entity.file_path)) - # Add entity tags from frontmatter to search content - entity_tags = self._extract_entity_tags(entity) - if entity_tags: - content_stems.extend(entity_tags) + entity_tags = self._extract_entity_tags(entity) + if entity_tags: + content_stems.extend(entity_tags) - entity_content_stems = _strip_nul("\n".join(p for p in content_stems if p and p.strip())) - - # Truncate to stay under Postgres's 8KB index row limit - if len(entity_content_stems) > MAX_CONTENT_STEMS_SIZE: # pragma: no cover - entity_content_stems = entity_content_stems[:MAX_CONTENT_STEMS_SIZE] # pragma: no cover - - # Add entity row - rows_to_index.append( - SearchIndexRow( - id=entity.id, - type=SearchItemType.ENTITY.value, - title=_strip_nul(entity.title), - content_stems=entity_content_stems, - content_snippet=content_snippet, - permalink=entity.permalink, - file_path=entity.file_path, - entity_id=entity.id, - metadata={ - "note_type": entity.note_type, - }, - created_at=entity.created_at, - updated_at=_mtime_to_datetime(entity), - project_id=entity.project_id, - ) - ) - - # Add observation rows - dedupe by permalink to avoid unique constraint violations - # Two observations with same entity/category/content generate identical permalinks - seen_permalinks: set[str] = {entity.permalink} if entity.permalink else set() - for obs in entity.observations: - obs_permalink = obs.permalink - if obs_permalink in seen_permalinks: - logger.debug(f"Skipping duplicate observation permalink: {obs_permalink}") - continue - seen_permalinks.add(obs_permalink) - - # Index with parent entity's file path since that's where it's defined - obs_content_stems = _strip_nul( - "\n".join(p for p in self._generate_variants(obs.content) if p and p.strip()) - ) - # Truncate to stay under Postgres's 8KB index row limit - if len(obs_content_stems) > MAX_CONTENT_STEMS_SIZE: # pragma: no cover - obs_content_stems = obs_content_stems[:MAX_CONTENT_STEMS_SIZE] # pragma: no cover - rows_to_index.append( - SearchIndexRow( - id=obs.id, - type=SearchItemType.OBSERVATION.value, - title=_strip_nul(f"{obs.category}: {obs.content[:100]}..."), - content_stems=obs_content_stems, - content_snippet=_strip_nul(obs.content), - permalink=obs_permalink, - file_path=entity.file_path, - category=obs.category, - entity_id=entity.id, - metadata={ - "tags": obs.tags, - }, - created_at=entity.created_at, - updated_at=_mtime_to_datetime(entity), - project_id=entity.project_id, + entity_content_stems = _strip_nul( + "\n".join(p for p in content_stems if p and p.strip()) ) - ) - # Add relation rows (only outgoing relations defined in this file) - for rel in entity.outgoing_relations: - # Create descriptive title showing the relationship - relation_title = _strip_nul( - f"{rel.from_entity.title} → {rel.to_entity.title}" - if rel.to_entity - else f"{rel.from_entity.title}" - ) + if len(entity_content_stems) > MAX_CONTENT_STEMS_SIZE: # pragma: no cover + entity_content_stems = entity_content_stems[ + :MAX_CONTENT_STEMS_SIZE + ] # pragma: no cover - rel_content_stems = _strip_nul( - "\n".join(p for p in self._generate_variants(relation_title) if p and p.strip()) - ) - rows_to_index.append( - SearchIndexRow( - id=rel.id, - title=relation_title, - permalink=rel.permalink, - content_stems=rel_content_stems, - file_path=entity.file_path, - type=SearchItemType.RELATION.value, - entity_id=entity.id, - from_id=rel.from_id, - to_id=rel.to_id, - relation_type=rel.relation_type, - created_at=entity.created_at, - updated_at=_mtime_to_datetime(entity), - project_id=entity.project_id, + rows_to_index.append( + SearchIndexRow( + id=entity.id, + type=SearchItemType.ENTITY.value, + title=_strip_nul(entity.title), + content_stems=entity_content_stems, + content_snippet=content_snippet, + permalink=entity.permalink, + file_path=entity.file_path, + entity_id=entity.id, + metadata={ + "note_type": entity.note_type, + }, + created_at=entity.created_at, + updated_at=_mtime_to_datetime(entity), + project_id=entity.project_id, + ) ) - ) - # Batch insert all rows at once - await self.repository.bulk_index_items(rows_to_index) + seen_permalinks: set[str] = {entity.permalink} if entity.permalink else set() + for obs in entity.observations: + obs_permalink = obs.permalink + if obs_permalink in seen_permalinks: + logger.debug(f"Skipping duplicate observation permalink: {obs_permalink}") + continue + seen_permalinks.add(obs_permalink) + + obs_content_stems = _strip_nul( + "\n".join( + p for p in self._generate_variants(obs.content) if p and p.strip() + ) + ) + if len(obs_content_stems) > MAX_CONTENT_STEMS_SIZE: # pragma: no cover + obs_content_stems = obs_content_stems[:MAX_CONTENT_STEMS_SIZE] # pragma: no cover + rows_to_index.append( + SearchIndexRow( + id=obs.id, + type=SearchItemType.OBSERVATION.value, + title=_strip_nul(f"{obs.category}: {obs.content[:100]}..."), + content_stems=obs_content_stems, + content_snippet=_strip_nul(obs.content), + permalink=obs_permalink, + file_path=entity.file_path, + category=obs.category, + entity_id=entity.id, + metadata={ + "tags": obs.tags, + }, + created_at=entity.created_at, + updated_at=_mtime_to_datetime(entity), + project_id=entity.project_id, + ) + ) + + for rel in entity.outgoing_relations: + relation_title = _strip_nul( + f"{rel.from_entity.title} -> {rel.to_entity.title}" + if rel.to_entity + else f"{rel.from_entity.title}" + ) + + rel_content_stems = _strip_nul( + "\n".join( + p for p in self._generate_variants(relation_title) if p and p.strip() + ) + ) + rows_to_index.append( + SearchIndexRow( + id=rel.id, + title=relation_title, + permalink=rel.permalink, + content_stems=rel_content_stems, + file_path=entity.file_path, + type=SearchItemType.RELATION.value, + entity_id=entity.id, + from_id=rel.from_id, + to_id=rel.to_id, + relation_type=rel.relation_type, + created_at=entity.created_at, + updated_at=_mtime_to_datetime(entity), + project_id=entity.project_id, + ) + ) + + with telemetry.scope( + "search.index.bulk_upsert", + phase="bulk_upsert", + result_count=len(rows_to_index), + ): + await self.repository.bulk_index_items(rows_to_index) async def delete_by_permalink(self, permalink: str): """Delete an item from the search index.""" diff --git a/src/basic_memory/telemetry.py b/src/basic_memory/telemetry.py index 9f7aa3d6..adc48b35 100644 --- a/src/basic_memory/telemetry.py +++ b/src/basic_memory/telemetry.py @@ -8,7 +8,6 @@ helpers for manual spans and logger context binding. from __future__ import annotations from contextlib import contextmanager -from contextvars import ContextVar from dataclasses import dataclass, field from typing import Any, Iterator @@ -41,7 +40,6 @@ class TelemetryState: _STATE = TelemetryState() _LOGFIRE_HANDLER: dict[str, Any] | None = None -_ACTIVE_LOG_CONTEXT: ContextVar[dict[str, Any]] = ContextVar("basic_memory_log_context", default={}) def reset_telemetry_state() -> None: @@ -57,7 +55,6 @@ def reset_telemetry_state() -> None: _STATE.send_to_logfire = False _STATE.warnings.clear() _LOGFIRE_HANDLER = None - _ACTIVE_LOG_CONTEXT.set({}) def _filter_attributes(attrs: dict[str, Any]) -> dict[str, Any]: @@ -65,11 +62,6 @@ def _filter_attributes(attrs: dict[str, Any]) -> dict[str, Any]: return {key: value for key, value in attrs.items() if value is not None} -def _current_log_context() -> dict[str, Any]: - """Return the currently active telemetry context for this execution flow.""" - return dict(_ACTIVE_LOG_CONTEXT.get()) - - def configure_telemetry( service_name: str, *, @@ -144,26 +136,11 @@ def pop_telemetry_warnings() -> list[str]: return warnings -def bind_telemetry_context(**attrs: Any): - """Bind stable telemetry attributes onto the shared Loguru logger.""" - merged_attrs = _current_log_context() - merged_attrs.update(_filter_attributes(attrs)) - return logger.bind(**merged_attrs) - - @contextmanager def contextualize(**attrs: Any) -> Iterator[None]: - """Apply stable telemetry attributes to all Loguru calls in this scope.""" - filtered_attrs = _filter_attributes(attrs) - merged_attrs = _current_log_context() - merged_attrs.update(filtered_attrs) - context_token = _ACTIVE_LOG_CONTEXT.set(merged_attrs) - - try: - with logger.contextualize(**filtered_attrs): - yield - finally: - _ACTIVE_LOG_CONTEXT.reset(context_token) + """Apply filtered telemetry attributes to Loguru calls in this scope.""" + with logger.contextualize(**_filter_attributes(attrs)): + yield @contextmanager @@ -182,12 +159,8 @@ operation = scope @contextmanager def span(name: str, **attrs: Any) -> Iterator[None]: """Create a manual Logfire span when telemetry is enabled.""" - if not telemetry_enabled(): - yield - return - logfire = _load_logfire() - if logfire is None: # pragma: no cover + if logfire is None or not _STATE.configured: # pragma: no cover yield # pragma: no cover return # pragma: no cover @@ -196,7 +169,6 @@ def span(name: str, **attrs: Any) -> Iterator[None]: __all__ = [ - "bind_telemetry_context", "contextualize", "configure_telemetry", "get_logfire_handler", diff --git a/tests/api/v2/test_knowledge_router_telemetry.py b/tests/api/v2/test_knowledge_router_telemetry.py new file mode 100644 index 00000000..afadcc4b --- /dev/null +++ b/tests/api/v2/test_knowledge_router_telemetry.py @@ -0,0 +1,225 @@ +"""Telemetry coverage for the v2 knowledge router.""" + +from __future__ import annotations + +import importlib +from contextlib import contextmanager +from datetime import datetime, timezone +from types import SimpleNamespace + +import pytest +from fastapi import BackgroundTasks, Response + +from basic_memory.schemas.base import Entity +from basic_memory.schemas.request import EditEntityRequest + +knowledge_router_module = importlib.import_module("basic_memory.api.v2.routers.knowledge_router") + + +def _capture_spans(): + spans: list[tuple[str, dict]] = [] + + @contextmanager + def fake_span(name: str, **attrs): + spans.append((name, attrs)) + yield + + return spans, fake_span + + +def _fake_entity(*, external_id: str = "entity-123", file_path: str = "notes/test.md"): + now = datetime.now(timezone.utc) + return SimpleNamespace( + external_id=external_id, + id=1, + title="Telemetry Entity", + note_type="note", + content_type="text/markdown", + permalink="notes/test", + file_path=file_path, + entity_metadata=None, + observations=[], + relations=[], + created_at=now, + updated_at=now, + created_by=None, + last_updated_by=None, + ) + + +def _assert_names_in_order(names: list[str], expected: list[str]) -> None: + cursor = 0 + for expected_name in expected: + cursor = names.index(expected_name, cursor) + 1 + + +@pytest.mark.asyncio +async def test_create_entity_emits_root_and_nested_spans(monkeypatch) -> None: + spans, fake_span = _capture_spans() + monkeypatch.setattr(knowledge_router_module.telemetry, "span", fake_span) + + entity = _fake_entity() + + class FakeEntityService: + async def create_entity(self, data): + return entity + + class FakeSearchService: + async def index_entity(self, entity): + return None + + class FakeTaskScheduler: + def schedule(self, *args, **kwargs): + return None + + class FakeFileService: + async def read_file_content(self, path): + return "telemetry content" + + result = await knowledge_router_module.create_entity( + project_id="project-123", + data=Entity( + title="Telemetry Entity", + directory="notes", + note_type="note", + content_type="text/markdown", + content="telemetry content", + ), + background_tasks=BackgroundTasks(), + entity_service=FakeEntityService(), + search_service=FakeSearchService(), + task_scheduler=FakeTaskScheduler(), + file_service=FakeFileService(), + app_config=SimpleNamespace(semantic_search_enabled=False), + fast=False, + ) + + assert result.content == "telemetry content" + _assert_names_in_order( + [name for name, _ in spans], + [ + "api.request.knowledge.create_entity", + "api.knowledge.create_entity.write_entity", + "api.knowledge.create_entity.search_index", + "api.knowledge.create_entity.vector_sync", + "api.knowledge.create_entity.read_content", + ], + ) + + +@pytest.mark.asyncio +async def test_update_entity_emits_root_and_nested_spans(monkeypatch) -> None: + spans, fake_span = _capture_spans() + monkeypatch.setattr(knowledge_router_module.telemetry, "span", fake_span) + + entity = _fake_entity() + + class FakeEntityService: + async def update_entity(self, existing, data): + return entity + + class FakeSearchService: + async def index_entity(self, entity): + return None + + class FakeEntityRepository: + async def get_by_external_id(self, external_id): + return entity + + class FakeTaskScheduler: + def schedule(self, *args, **kwargs): + return None + + class FakeFileService: + async def read_file_content(self, path): + return "updated telemetry content" + + response = Response() + result = await knowledge_router_module.update_entity_by_id( + data=Entity( + title="Telemetry Entity", + directory="notes", + note_type="note", + content_type="text/markdown", + content="updated telemetry content", + ), + response=response, + background_tasks=BackgroundTasks(), + project_id="project-123", + entity_service=FakeEntityService(), + search_service=FakeSearchService(), + entity_repository=FakeEntityRepository(), + task_scheduler=FakeTaskScheduler(), + file_service=FakeFileService(), + app_config=SimpleNamespace(semantic_search_enabled=False), + entity_id=entity.external_id, + fast=False, + ) + + assert result.content == "updated telemetry content" + _assert_names_in_order( + [name for name, _ in spans], + [ + "api.request.knowledge.update_entity", + "api.knowledge.update_entity.load_entity", + "api.knowledge.update_entity.write_entity", + "api.knowledge.update_entity.search_index", + "api.knowledge.update_entity.vector_sync", + "api.knowledge.update_entity.read_content", + ], + ) + + +@pytest.mark.asyncio +async def test_edit_entity_emits_root_and_nested_spans(monkeypatch) -> None: + spans, fake_span = _capture_spans() + monkeypatch.setattr(knowledge_router_module.telemetry, "span", fake_span) + + entity = _fake_entity() + + class FakeEntityService: + async def edit_entity(self, **kwargs): + return entity + + class FakeSearchService: + async def index_entity(self, entity): + return None + + class FakeEntityRepository: + async def get_by_external_id(self, external_id): + return entity + + class FakeTaskScheduler: + def schedule(self, *args, **kwargs): + return None + + class FakeFileService: + async def read_file_content(self, path): + return "edited telemetry content" + + result = await knowledge_router_module.edit_entity_by_id( + data=EditEntityRequest(operation="append", content="edited telemetry content"), + background_tasks=BackgroundTasks(), + project_id="project-123", + entity_service=FakeEntityService(), + search_service=FakeSearchService(), + entity_repository=FakeEntityRepository(), + task_scheduler=FakeTaskScheduler(), + file_service=FakeFileService(), + app_config=SimpleNamespace(semantic_search_enabled=False), + entity_id=entity.external_id, + fast=False, + ) + + assert result.content == "edited telemetry content" + _assert_names_in_order( + [name for name, _ in spans], + [ + "api.request.knowledge.edit_entity", + "api.knowledge.edit_entity.load_entity", + "api.knowledge.edit_entity.write_entity", + "api.knowledge.edit_entity.search_index", + "api.knowledge.edit_entity.vector_sync", + "api.knowledge.edit_entity.read_content", + ], + ) diff --git a/tests/api/v2/test_search_router_telemetry.py b/tests/api/v2/test_search_router_telemetry.py index 809cf23f..bd850730 100644 --- a/tests/api/v2/test_search_router_telemetry.py +++ b/tests/api/v2/test_search_router_telemetry.py @@ -51,12 +51,13 @@ async def test_search_router_wraps_request_in_manual_operation() -> None: "api.request.search", { "entrypoint": "api", + "domain": "search", + "action": "search", "page": 2, "page_size": 5, "retrieval_mode": "fts", - "has_text_query": True, - "has_title_query": False, - "has_permalink_query": False, + "has_query": True, + "has_filters": False, }, ) ] diff --git a/tests/api/v2/test_utils_telemetry.py b/tests/api/v2/test_utils_telemetry.py new file mode 100644 index 00000000..6262c638 --- /dev/null +++ b/tests/api/v2/test_utils_telemetry.py @@ -0,0 +1,63 @@ +"""Telemetry coverage for API v2 hydration utilities.""" + +from __future__ import annotations + +import importlib +from contextlib import contextmanager +from datetime import datetime, timezone +from types import SimpleNamespace + +import pytest + +from basic_memory.repository.search_index_row import SearchIndexRow + +utils_module = importlib.import_module("basic_memory.api.v2.utils") + + +def _capture_spans(): + spans: list[tuple[str, dict]] = [] + + @contextmanager + def fake_span(name: str, **attrs): + spans.append((name, attrs)) + yield + + return spans, fake_span + + +@pytest.mark.asyncio +async def test_to_search_results_emits_hydration_spans(monkeypatch) -> None: + spans, fake_span = _capture_spans() + monkeypatch.setattr(utils_module.telemetry, "span", fake_span) + + class FakeEntityService: + async def get_entities_by_id(self, ids): + return [SimpleNamespace(permalink="notes/root"), SimpleNamespace(permalink="notes/child")] + + now = datetime.now(timezone.utc) + results = [ + SearchIndexRow( + project_id=1, + id=1, + type="relation", + file_path="notes/root.md", + created_at=now, + updated_at=now, + permalink="notes/root/relates_to/notes/child", + entity_id=1, + from_id=1, + to_id=2, + relation_type="relates_to", + title="Root relates to Child", + score=1.0, + ) + ] + + search_results = await utils_module.to_search_results(FakeEntityService(), results) + + assert search_results[0].relation_type == "relates_to" + assert [name for name, _ in spans] == [ + "search.hydrate_results", + "search.hydrate_results.fetch_entities", + "search.hydrate_results.shape_results", + ] diff --git a/tests/mcp/test_client_telemetry.py b/tests/mcp/test_client_telemetry.py new file mode 100644 index 00000000..8aa88211 --- /dev/null +++ b/tests/mcp/test_client_telemetry.py @@ -0,0 +1,91 @@ +"""Telemetry coverage for typed MCP clients and shared HTTP helpers.""" + +from __future__ import annotations + +import importlib +from contextlib import contextmanager + +import httpx +import pytest + +knowledge_client_module = importlib.import_module("basic_memory.mcp.clients.knowledge") +search_client_module = importlib.import_module("basic_memory.mcp.clients.search") + + +def _capture_spans(): + spans: list[tuple[str, dict]] = [] + + @contextmanager + def fake_span(name: str, **attrs): + spans.append((name, attrs)) + yield + + return spans, fake_span + + +@pytest.mark.asyncio +async def test_knowledge_client_resolve_entity_emits_client_and_http_spans(monkeypatch) -> None: + spans, fake_span = _capture_spans() + monkeypatch.setattr(knowledge_client_module.telemetry, "span", fake_span) + + async def handler(request: httpx.Request) -> httpx.Response: + assert request.method == "POST" + return httpx.Response(200, json={"external_id": "entity-123"}) + + transport = httpx.MockTransport(handler) + async with httpx.AsyncClient(transport=transport, base_url="https://example.test") as client: + knowledge_client = knowledge_client_module.KnowledgeClient(client, "project-123") + resolved = await knowledge_client.resolve_entity("notes/root", strict=True) + + assert resolved == "entity-123" + assert [name for name, _ in spans] == [ + "mcp.client.knowledge.resolve_entity", + "mcp.http.request", + ] + assert spans[1][1] == { + "method": "POST", + "client_name": "knowledge", + "operation": "resolve_entity", + "path_template": "/v2/projects/{project_id}/knowledge/resolve", + "phase": "request", + "has_query": False, + "has_body": True, + } + + +@pytest.mark.asyncio +async def test_search_client_emits_client_and_http_spans(monkeypatch) -> None: + spans, fake_span = _capture_spans() + monkeypatch.setattr(search_client_module.telemetry, "span", fake_span) + + async def handler(request: httpx.Request) -> httpx.Response: + assert request.method == "POST" + return httpx.Response( + 200, + json={ + "results": [], + "current_page": 2, + "page_size": 5, + "has_more": False, + }, + ) + + transport = httpx.MockTransport(handler) + async with httpx.AsyncClient(transport=transport, base_url="https://example.test") as client: + search_client = search_client_module.SearchClient(client, "project-123") + response = await search_client.search({"text": "telemetry"}, page=2, page_size=5) + + assert response.current_page == 2 + assert [name for name, _ in spans] == [ + "mcp.client.search.search", + "mcp.http.request", + ] + assert spans[1][1] == { + "method": "POST", + "client_name": "search", + "operation": "search", + "path_template": "/v2/projects/{project_id}/search/", + "phase": "request", + "has_query": True, + "has_body": True, + } diff --git a/tests/mcp/test_tool_telemetry.py b/tests/mcp/test_tool_telemetry.py index e07822c8..26e39aaa 100644 --- a/tests/mcp/test_tool_telemetry.py +++ b/tests/mcp/test_tool_telemetry.py @@ -51,20 +51,19 @@ async def test_write_note_emits_root_operation_and_project_context( output_format="json", ) - assert operations == [ - ( - "mcp.tool.write_note", - { - "entrypoint": "mcp", - "tool_name": "write_note", - "requested_project": test_project.name, - "workspace_id": None, - "note_type": "note", - "overwrite": False, - "output_format": "json", - }, - ) - ] + assert operations[0] == ( + "mcp.tool.write_note", + { + "entrypoint": "mcp", + "tool_name": "write_note", + "project_name": test_project.name, + "workspace_id": None, + "note_type": "note", + "overwrite": False, + "output_format": "json", + }, + ) + assert "api.request.knowledge.create_entity" in [name for name, _ in operations] assert _contains_context( contexts, { @@ -103,21 +102,23 @@ async def test_read_note_emits_root_operation_and_project_context( include_frontmatter=True, ) - assert operations == [ - ( - "mcp.tool.read_note", - { - "entrypoint": "mcp", - "tool_name": "read_note", - "requested_project": test_project.name, - "workspace_id": None, - "output_format": "json", - "page": 1, - "page_size": 10, - "include_frontmatter": True, - }, - ) - ] + assert operations[0] == ( + "mcp.tool.read_note", + { + "entrypoint": "mcp", + "tool_name": "read_note", + "project_name": test_project.name, + "workspace_id": None, + "output_format": "json", + "page": 1, + "page_size": 10, + "include_frontmatter": True, + }, + ) + operation_names = [name for name, _ in operations] + assert "api.request.knowledge.resolve_entity" in operation_names + assert "api.request.resource.get_content" in operation_names + assert "api.request.knowledge.get_entity" in operation_names assert _contains_context( contexts, { @@ -163,7 +164,7 @@ async def test_search_notes_emits_root_operation_and_project_context( { "entrypoint": "mcp", "tool_name": "search_notes", - "requested_project": test_project.name, + "project_name": test_project.name, "workspace_id": None, "search_type": "text", "output_format": "json", @@ -172,7 +173,7 @@ async def test_search_notes_emits_root_operation_and_project_context( "has_query": True, "note_type_filter_count": 0, "entity_type_filter_count": 0, - "has_metadata_filters": False, + "has_filters": True, "has_tags_filter": True, "has_status_filter": False, }, @@ -217,22 +218,23 @@ async def test_edit_note_emits_root_operation_and_project_context( output_format="json", ) - assert operations == [ - ( - "mcp.tool.edit_note", - { - "entrypoint": "mcp", - "tool_name": "edit_note", - "requested_project": test_project.name, - "workspace_id": None, - "edit_operation": "append", - "output_format": "json", - "has_section": False, - "has_find_text": False, - "expected_replacements": 1, - }, - ) - ] + assert operations[0] == ( + "mcp.tool.edit_note", + { + "entrypoint": "mcp", + "tool_name": "edit_note", + "project_name": test_project.name, + "workspace_id": None, + "edit_operation": "append", + "output_format": "json", + "has_section": False, + "has_find_text": False, + "expected_replacements": 1, + }, + ) + operation_names = [name for name, _ in operations] + assert "api.request.knowledge.resolve_entity" in operation_names + assert "api.request.knowledge.edit_entity" in operation_names assert _contains_context( contexts, { @@ -275,24 +277,23 @@ async def test_build_context_emits_root_operation_and_project_context( output_format="json", ) - assert operations == [ - ( - "mcp.tool.build_context", - { - "entrypoint": "mcp", - "tool_name": "build_context", - "requested_project": test_project.name, - "workspace_id": None, - "depth": 2, - "timeframe": "7d", - "page": 1, - "page_size": 5, - "max_related": 3, - "output_format": "json", - "is_memory_url": True, - }, - ) - ] + assert operations[0] == ( + "mcp.tool.build_context", + { + "entrypoint": "mcp", + "tool_name": "build_context", + "project_name": test_project.name, + "workspace_id": None, + "depth": 2, + "timeframe": "7d", + "page": 1, + "page_size": 5, + "max_related": 3, + "output_format": "json", + "is_memory_url": True, + }, + ) + assert "api.request.memory.build_context" in [name for name, _ in operations] assert _contains_context( contexts, { diff --git a/tests/services/test_entity_service_telemetry.py b/tests/services/test_entity_service_telemetry.py new file mode 100644 index 00000000..e13dfa44 --- /dev/null +++ b/tests/services/test_entity_service_telemetry.py @@ -0,0 +1,131 @@ +"""Telemetry coverage for entity service write/edit/reindex paths.""" + +from __future__ import annotations + +import importlib +from contextlib import contextmanager + +import pytest + +from basic_memory.schemas import Entity as EntitySchema + +entity_service_module = importlib.import_module("basic_memory.services.entity_service") + + +def _capture_spans(): + spans: list[tuple[str, dict]] = [] + + @contextmanager + def fake_span(name: str, **attrs): + spans.append((name, attrs)) + yield + + return spans, fake_span + + +def _assert_names_in_order(names: list[str], expected: list[str]) -> None: + cursor = 0 + for expected_name in expected: + cursor = names.index(expected_name, cursor) + 1 + + +@pytest.mark.asyncio +async def test_create_entity_emits_expected_phase_spans(entity_service, monkeypatch) -> None: + spans, fake_span = _capture_spans() + monkeypatch.setattr(entity_service_module.telemetry, "span", fake_span) + + schema = EntitySchema( + title="Telemetry Create", + directory="notes", + note_type="note", + content_type="text/markdown", + content="Create telemetry content", + ) + + entity = await entity_service.create_entity(schema) + + assert entity.title == "Telemetry Create" + span_names = [name for name, _ in spans] + _assert_names_in_order( + span_names, + [ + "entity_service.create.resolve_permalink", + "entity_service.create.write_file", + "file_service.write", + "entity_service.create.parse_markdown", + "entity_service.create.upsert_entity", + "entity_service.create.update_checksum", + ], + ) + + +@pytest.mark.asyncio +async def test_edit_entity_emits_expected_phase_spans(entity_service, monkeypatch) -> None: + created = await entity_service.create_entity( + EntitySchema( + title="Telemetry Edit", + directory="notes", + note_type="note", + content_type="text/markdown", + content="Before edit", + ) + ) + + spans, fake_span = _capture_spans() + monkeypatch.setattr(entity_service_module.telemetry, "span", fake_span) + + updated = await entity_service.edit_entity( + created.file_path, + operation="append", + content="\n\nAfter edit", + ) + + assert updated.id == created.id + span_names = [name for name, _ in spans] + _assert_names_in_order( + span_names, + [ + "entity_service.edit.resolve_entity", + "entity_service.edit.read_file", + "file_service.read", + "entity_service.edit.apply_operation", + "entity_service.edit.write_file", + "file_service.write", + "entity_service.edit.parse_markdown", + "entity_service.edit.upsert_entity", + "entity_service.edit.update_checksum", + ], + ) + + +@pytest.mark.asyncio +async def test_reindex_entity_emits_expected_phase_spans(entity_service, monkeypatch) -> None: + created = await entity_service.create_entity( + EntitySchema( + title="Telemetry Reindex", + directory="notes", + note_type="note", + content_type="text/markdown", + content="Reindex telemetry content", + ) + ) + + spans, fake_span = _capture_spans() + monkeypatch.setattr(entity_service_module.telemetry, "span", fake_span) + + await entity_service.reindex_entity(created.id) + + span_names = [name for name, _ in spans] + _assert_names_in_order( + span_names, + [ + "entity_service.reindex.load_entity", + "entity_service.reindex.read_file", + "file_service.read_content", + "entity_service.reindex.parse_markdown", + "entity_service.reindex.upsert_entity", + "entity_service.reindex.update_checksum", + ], + ) + if entity_service.search_service is not None: + assert "entity_service.reindex.search_index" in span_names diff --git a/tests/services/test_search_service_telemetry.py b/tests/services/test_search_service_telemetry.py index 6a47ddff..7150b5d5 100644 --- a/tests/services/test_search_service_telemetry.py +++ b/tests/services/test_search_service_telemetry.py @@ -34,14 +34,21 @@ async def test_search_service_wraps_repository_search(search_service, monkeypatc "search.execute", { "retrieval_mode": "fts", - "has_text_query": True, - "has_title_query": False, - "has_permalink_query": False, - "has_metadata_filters": False, + "has_query": True, + "has_filters": False, "limit": 10, "offset": 0, }, ) + assert spans[1] == ( + "search.repository_query", + { + "retrieval_mode": "fts", + "phase": "repository_query", + "has_query": True, + "has_filters": False, + }, + ) @pytest.mark.asyncio @@ -58,8 +65,13 @@ async def test_search_service_emits_relaxed_retry_span(search_service, monkeypat await search_service.search(SearchQuery(text="who are our main competitors and partners")) - assert [name for name, _ in spans] == ["search.execute", "search.relaxed_fts_retry"] - assert spans[1] == ( + assert [name for name, _ in spans] == [ + "search.execute", + "search.repository_query", + "search.relaxed_fts_retry", + "search.repository_query", + ] + assert spans[2] == ( "search.relaxed_fts_retry", { "retrieval_mode": "fts", diff --git a/tests/test_telemetry.py b/tests/test_telemetry.py index 1c12dab9..3bcf0f5a 100644 --- a/tests/test_telemetry.py +++ b/tests/test_telemetry.py @@ -91,24 +91,6 @@ def test_configure_telemetry_retries_without_send_to_logfire(monkeypatch) -> Non assert "send_to_logfire" not in fake_logfire.configure_calls[1] -def test_bind_telemetry_context_filters_nulls() -> None: - bound = telemetry.bind_telemetry_context(project_name="main", workspace_id=None) - extra = bound._options[-1] # type: ignore[attr-defined] - assert extra == {"project_name": "main"} - - -def test_bind_telemetry_context_merges_active_context() -> None: - with telemetry.contextualize(project_name="main", route_mode="local_asgi"): - bound = telemetry.bind_telemetry_context(tool_name="write_note", workspace_id=None) - - extra = bound._options[-1] # type: ignore[attr-defined] - assert extra == { - "project_name": "main", - "route_mode": "local_asgi", - "tool_name": "write_note", - } - - def test_contextualize_adds_filtered_loguru_context() -> None: records: list[dict] = [] sink_id = logger.add(lambda message: records.append(message.record["extra"].copy()))