add logfire instrumentation to services and repository code

Signed-off-by: phernandez <paul@basicmachines.co>

src/basic_memory/services/context_service.py

@@ -4,6 +4,7 @@ from dataclasses import dataclass, field
 from datetime import datetime, timezone
 from typing import List, Optional, Tuple
 
+import logfire
 from loguru import logger
 from sqlalchemy import text
 
@@ -85,6 +86,7 @@ class ContextService:
         self.entity_repository = entity_repository
         self.observation_repository = observation_repository
 
+    @logfire.instrument(record_return=True)
     async def build_context(
         self,
         memory_url: Optional[MemoryUrl] = None,
@@ -215,6 +217,7 @@ class ContextService:
         # Return the structured ContextResult
         return ContextResult(results=context_results, metadata=metadata)
 
+    @logfire.instrument(record_return=True)
     async def find_related(
         self,
         type_id_pairs: List[Tuple[str, int]],

src/basic_memory/services/directory_service.py

@@ -4,6 +4,7 @@ import fnmatch
 import logging
 import os
 from typing import Dict, List, Optional, Sequence
+import logfire
 
 from basic_memory.models import Entity
 from basic_memory.repository import EntityRepository
@@ -23,6 +24,7 @@ class DirectoryService:
         """
         self.entity_repository = entity_repository
 
+    @logfire.instrument(record_return=True)
     async def get_directory_tree(self) -> DirectoryNode:
         """Build a hierarchical directory tree from indexed files."""
 
@@ -90,6 +92,7 @@ class DirectoryService:
         # Return the root node with its children
         return root_node
 
+    @logfire.instrument(record_return=True)
     async def get_directory_structure(self) -> DirectoryNode:
         """Build a hierarchical directory structure without file details.
 
@@ -133,6 +136,7 @@ class DirectoryService:
 
         return root_node
 
+    @logfire.instrument(record_return=True)
     async def list_directory(
         self,
         dir_name: str = "/",

src/basic_memory/services/entity_service.py

@@ -7,6 +7,7 @@ import frontmatter
 import yaml
 from loguru import logger
 from sqlalchemy.exc import IntegrityError
+import logfire
 
 from basic_memory.config import ProjectConfig, BasicMemoryConfig
 from basic_memory.file_utils import (
@@ -52,6 +53,7 @@ class EntityService(BaseService[EntityModel]):
         self.link_resolver = link_resolver
         self.app_config = app_config
 
+    @logfire.instrument(record_return=True)
     async def detect_file_path_conflicts(
         self, file_path: str, skip_check: bool = False
     ) -> List[Entity]:
@@ -91,6 +93,7 @@ class EntityService(BaseService[EntityModel]):
 
         return conflicts
 
+    @logfire.instrument(record_return=True)
     async def resolve_permalink(
         self,
         file_path: Permalink | Path,
@@ -149,6 +152,7 @@ class EntityService(BaseService[EntityModel]):
 
         return permalink
 
+    @logfire.instrument(record_return=True)
     async def create_or_update_entity(self, schema: EntitySchema) -> Tuple[EntityModel, bool]:
         """Create new entity or update existing one.
         Returns: (entity, is_new) where is_new is True if a new entity was created
@@ -170,6 +174,7 @@ class EntityService(BaseService[EntityModel]):
             # Create new entity
             return await self.create_entity(schema), True
 
+    @logfire.instrument(record_return=True)
     async def create_entity(self, schema: EntitySchema) -> EntityModel:
         """Create a new entity and write to filesystem."""
         logger.debug(f"Creating entity: {schema.title}")
@@ -236,6 +241,7 @@ class EntityService(BaseService[EntityModel]):
         # Set final checksum to mark complete
         return await self.repository.update(entity.id, {"checksum": checksum})
 
+    @logfire.instrument(record_return=True)
     async def update_entity(self, entity: EntityModel, schema: EntitySchema) -> EntityModel:
         """Update an entity's content and metadata."""
         logger.debug(
@@ -316,6 +322,7 @@ class EntityService(BaseService[EntityModel]):
 
         return entity
 
+    @logfire.instrument(record_return=True)
     async def delete_entity(self, permalink_or_id: str | int) -> bool:
         """Delete entity and its file."""
         logger.debug(f"Deleting entity: {permalink_or_id}")
@@ -345,6 +352,7 @@ class EntityService(BaseService[EntityModel]):
             logger.info(f"Entity not found: {permalink_or_id}")
             return True  # Already deleted
 
+    @logfire.instrument(record_return=True)
     async def get_by_permalink(self, permalink: str) -> EntityModel:
         """Get entity by type and name combination."""
         logger.debug(f"Getting entity by permalink: {permalink}")
@@ -353,20 +361,24 @@ class EntityService(BaseService[EntityModel]):
             raise EntityNotFoundError(f"Entity not found: {permalink}")
         return db_entity
 
+    @logfire.instrument(record_return=True)
     async def get_entities_by_id(self, ids: List[int]) -> Sequence[EntityModel]:
         """Get specific entities and their relationships."""
         logger.debug(f"Getting entities: {ids}")
         return await self.repository.find_by_ids(ids)
 
+    @logfire.instrument(record_return=True)
     async def get_entities_by_permalinks(self, permalinks: List[str]) -> Sequence[EntityModel]:
         """Get specific nodes and their relationships."""
         logger.debug(f"Getting entities permalinks: {permalinks}")
         return await self.repository.find_by_permalinks(permalinks)
 
+    @logfire.instrument(record_return=True)
     async def delete_entity_by_file_path(self, file_path: Union[str, Path]) -> None:
         """Delete entity by file path."""
         await self.repository.delete_by_file_path(str(file_path))
 
+    @logfire.instrument(record_return=True)
     async def create_entity_from_markdown(
         self, file_path: Path, markdown: EntityMarkdown
     ) -> EntityModel:
@@ -390,6 +402,7 @@ class EntityService(BaseService[EntityModel]):
             logger.error(f"Failed to upsert entity for {file_path}: {e}")
             raise EntityCreationError(f"Failed to create entity: {str(e)}") from e
 
+    @logfire.instrument(record_return=True)
     async def update_entity_and_observations(
         self, file_path: Path, markdown: EntityMarkdown
     ) -> EntityModel:
@@ -430,6 +443,7 @@ class EntityService(BaseService[EntityModel]):
             db_entity,
         )
 
+    @logfire.instrument(record_return=True)
     async def update_entity_relations(
         self,
         path: str,
@@ -501,6 +515,7 @@ class EntityService(BaseService[EntityModel]):
 
         return await self.repository.get_by_file_path(path)
 
+    @logfire.instrument(record_return=True)
     async def edit_entity(
         self,
         identifier: str,
@@ -558,6 +573,7 @@ class EntityService(BaseService[EntityModel]):
 
         return entity
 
+    @logfire.instrument(record_return=True)
     def apply_edit_operation(
         self,
         current_content: str,
@@ -610,6 +626,7 @@ class EntityService(BaseService[EntityModel]):
         else:
             raise ValueError(f"Unsupported operation: {operation}")
 
+    @logfire.instrument(record_return=True)
     def replace_section_content(
         self, current_content: str, section_header: str, new_content: str
     ) -> str:
@@ -729,6 +746,7 @@ class EntityService(BaseService[EntityModel]):
             return content + "\n" + current_content  # pragma: no cover
         return content + current_content  # pragma: no cover
 
+    @logfire.instrument(record_return=True)
     async def move_entity(
         self,
         identifier: str,

src/basic_memory/services/file_service.py

@@ -8,6 +8,7 @@ from pathlib import Path
 from typing import Any, Dict, Tuple, Union
 
 import aiofiles
+import logfire
 import yaml
 
 from basic_memory import file_utils
@@ -59,6 +60,7 @@ class FileService:
         """
         return self.base_path / entity.file_path
 
+    @logfire.instrument(record_return=True)
     async def read_entity_content(self, entity: EntityModel) -> str:
         """Get entity's content without frontmatter or structured sections.
 
@@ -77,6 +79,7 @@ class FileService:
         markdown = await self.markdown_processor.read_file(file_path)
         return markdown.content or ""
 
+    @logfire.instrument(record_return=True)
     async def delete_entity_file(self, entity: EntityModel) -> None:
         """Delete entity file from filesystem.
 
@@ -89,6 +92,7 @@ class FileService:
         path = self.get_entity_path(entity)
         await self.delete_file(path)
 
+    @logfire.instrument(record_return=True)
     async def exists(self, path: FilePath) -> bool:
         """Check if file exists at the provided path.
 
@@ -115,6 +119,7 @@ class FileService:
             logger.error("Failed to check file existence", path=str(path), error=str(e))
             raise FileOperationError(f"Failed to check file existence: {e}")
 
+    @logfire.instrument(record_return=True)
     async def ensure_directory(self, path: FilePath) -> None:
         """Ensure directory exists, creating if necessary.
 
@@ -142,6 +147,7 @@ class FileService:
             logger.error("Failed to create directory", path=str(path), error=str(e))
             raise FileOperationError(f"Failed to create directory {path}: {e}")
 
+    @logfire.instrument(record_return=True)
     async def write_file(self, path: FilePath, content: str) -> str:
         """Write content to file and return checksum.
 
@@ -185,6 +191,7 @@ class FileService:
             logger.exception("File write error", path=str(full_path), error=str(e))
             raise FileOperationError(f"Failed to write file: {e}")
 
+    @logfire.instrument(record_return=True)
     async def read_file_content(self, path: FilePath) -> str:
         """Read file content using true async I/O with aiofiles.
 
@@ -220,6 +227,7 @@ class FileService:
             logger.exception("File read error", path=str(full_path), error=str(e))
             raise FileOperationError(f"Failed to read file: {e}")
 
+    @logfire.instrument(record_return=True)
     async def read_file(self, path: FilePath) -> Tuple[str, str]:
         """Read file and compute checksum using true async I/O.
 
@@ -262,6 +270,7 @@ class FileService:
             logger.exception("File read error", path=str(full_path), error=str(e))
             raise FileOperationError(f"Failed to read file: {e}")
 
+    @logfire.instrument(record_return=True)
     async def delete_file(self, path: FilePath) -> None:
         """Delete file if it exists.
 
@@ -276,6 +285,7 @@ class FileService:
         full_path = path_obj if path_obj.is_absolute() else self.base_path / path_obj
         full_path.unlink(missing_ok=True)
 
+    @logfire.instrument(record_return=True)
     async def update_frontmatter(self, path: FilePath, updates: Dict[str, Any]) -> str:
         """Update frontmatter fields in a file while preserving all content.
 
@@ -344,6 +354,7 @@ class FileService:
                 )
             raise FileOperationError(f"Failed to update frontmatter: {e}")
 
+    @logfire.instrument(record_return=True)
     async def compute_checksum(self, path: FilePath) -> str:
         """Compute checksum for a file using true async I/O.
 

src/basic_memory/services/initialization.py

@@ -7,6 +7,7 @@ to ensure consistent application startup across all entry points.
 import asyncio
 from pathlib import Path
 
+import logfire
 from loguru import logger
 
 from basic_memory import db
@@ -17,6 +18,7 @@ from basic_memory.repository import (
 )
 
 
+@logfire.instrument(record_return=True)
 async def initialize_database(app_config: BasicMemoryConfig) -> None:
     """Initialize database with migrations handled automatically by get_or_create_db.
 
@@ -38,6 +40,7 @@ async def initialize_database(app_config: BasicMemoryConfig) -> None:
         # more specific error if the database is actually unusable
 
 
+@logfire.instrument(record_return=True)
 async def reconcile_projects_with_config(app_config: BasicMemoryConfig):
     """Ensure all projects in config.json exist in the projects table and vice versa.
 
@@ -71,6 +74,7 @@ async def reconcile_projects_with_config(app_config: BasicMemoryConfig):
         logger.info("Continuing with initialization despite synchronization error")
 
 
+@logfire.instrument(record_return=True)
 async def initialize_file_sync(
     app_config: BasicMemoryConfig,
 ):
@@ -141,6 +145,7 @@ async def initialize_file_sync(
     return None
 
 
+@logfire.instrument(record_return=True)
 async def initialize_app(
     app_config: BasicMemoryConfig,
 ):

src/basic_memory/services/link_resolver.py

@@ -2,6 +2,7 @@
 
 from typing import Optional, Tuple
 
+import logfire
 from loguru import logger
 
 from basic_memory.models import Entity
@@ -26,6 +27,7 @@ class LinkResolver:
         self.entity_repository = entity_repository
         self.search_service = search_service
 
+    @logfire.instrument(record_return=True)
     async def resolve_link(
         self, link_text: str, use_search: bool = True, strict: bool = False
     ) -> Optional[Entity]:

src/basic_memory/services/project_service.py

@@ -8,6 +8,7 @@ from datetime import datetime
 from pathlib import Path
 from typing import Dict, Optional, Sequence
 
+import logfire
 from loguru import logger
 from sqlalchemy import text
 
@@ -81,6 +82,7 @@ class ProjectService:
         """
         return os.environ.get("BASIC_MEMORY_PROJECT", self.config_manager.default_project)
 
+    @logfire.instrument(record_return=True)
     async def list_projects(self) -> Sequence[Project]:
         """List all projects without loading entity relationships.
 
@@ -90,6 +92,7 @@ class ProjectService:
         """
         return await self.repository.find_all(use_load_options=False)
 
+    @logfire.instrument(record_return=True)
     async def get_project(self, name: str) -> Optional[Project]:
         """Get the file path for a project by name or permalink."""
         return await self.repository.get_by_name(name) or await self.repository.get_by_permalink(
@@ -130,6 +133,7 @@ class ProjectService:
                 # Not nested in either direction
                 return False
 
+    @logfire.instrument(record_return=True)
     async def add_project(self, name: str, path: str, set_default: bool = False) -> None:
         """Add a new project to the configuration and database.
 
@@ -221,6 +225,7 @@ class ProjectService:
 
         logger.info(f"Project '{name}' added at {resolved_path}")
 
+    @logfire.instrument(record_return=True)
     async def remove_project(self, name: str, delete_notes: bool = False) -> None:
         """Remove a project from configuration and database.
 
@@ -271,6 +276,7 @@ class ProjectService:
             except Exception as e:
                 logger.warning(f"Failed to delete project directory {project_path}: {e}")
 
+    @logfire.instrument(record_return=True)
     async def set_default_project(self, name: str) -> None:
         """Set the default project in configuration and database.
 
@@ -295,6 +301,7 @@ class ProjectService:
 
         logger.info(f"Project '{name}' set as default in configuration and database")
 
+    @logfire.instrument(record_return=True)
     async def _ensure_single_default_project(self) -> None:
         """Ensure only one project has is_default=True.
 
@@ -336,6 +343,7 @@ class ProjectService:
                     f"Set '{config_default}' as default project (was missing)"
                 )  # pragma: no cover
 
+    @logfire.instrument(record_return=True)
     async def synchronize_projects(self) -> None:  # pragma: no cover
         """Synchronize projects between database and configuration.
 
@@ -420,6 +428,7 @@ class ProjectService:
 
         logger.info("Project synchronization complete")
 
+    @logfire.instrument(record_return=True)
     async def move_project(self, name: str, new_path: str) -> None:
         """Move a project to a new location.
 
@@ -461,6 +470,7 @@ class ProjectService:
             self.config_manager.save_config(config)
             raise ValueError(f"Project '{name}' not found in database")
 
+    @logfire.instrument(record_return=True)
     async def update_project(  # pragma: no cover
         self, name: str, updated_path: Optional[str] = None, is_active: Optional[bool] = None
     ) -> None:
@@ -520,6 +530,7 @@ class ProjectService:
                     f"Changed default project to '{new_default.name}' as '{name}' was deactivated"
                 )
 
+    @logfire.instrument(record_return=True)
     async def get_project_info(self, project_name: Optional[str] = None) -> ProjectInfoResponse:
         """Get comprehensive information about the specified Basic Memory project.
 
@@ -587,6 +598,7 @@ class ProjectService:
             system=system,
         )
 
+    @logfire.instrument(record_return=True)
     async def get_statistics(self, project_id: int) -> ProjectStatistics:
         """Get statistics about the specified project.
 
@@ -703,6 +715,7 @@ class ProjectService:
             isolated_entities=isolated_count,
         )
 
+    @logfire.instrument(record_return=True)
     async def get_activity_metrics(self, project_id: int) -> ActivityMetrics:
         """Get activity metrics for the specified project.
 

src/basic_memory/services/search_service.py

@@ -4,6 +4,7 @@ import ast
 from datetime import datetime
 from typing import List, Optional, Set
 
+import logfire
 from dateparser import parse
 from fastapi import BackgroundTasks
 from loguru import logger
@@ -35,10 +36,12 @@ class SearchService:
         self.entity_repository = entity_repository
         self.file_service = file_service
 
+    @logfire.instrument(record_return=True)
     async def init_search_index(self):
         """Create FTS5 virtual table if it doesn't exist."""
         await self.repository.init_search_index()
 
+    @logfire.instrument(record_return=True)
     async def reindex_all(self, background_tasks: Optional[BackgroundTasks] = None) -> None:
         """Reindex all content from database."""
 
@@ -55,6 +58,7 @@ class SearchService:
 
         logger.info("Reindex complete")
 
+    @logfire.instrument(record_return=True)
     async def search(self, query: SearchQuery, limit=10, offset=0) -> List[SearchIndexRow]:
         """Search across all indexed content.
 
@@ -152,6 +156,7 @@ class SearchService:
 
         return []  # pragma: no cover
 
+    @logfire.instrument(record_return=True)
     async def index_entity(
         self,
         entity: Entity,
@@ -163,6 +168,7 @@ class SearchService:
         else:
             await self.index_entity_data(entity, content)
 
+    @logfire.instrument(record_return=True)
     async def index_entity_data(
         self,
         entity: Entity,
@@ -176,6 +182,7 @@ class SearchService:
             entity, content
         ) if entity.is_markdown else await self.index_entity_file(entity)
 
+    @logfire.instrument(record_return=True)
     async def index_entity_file(
         self,
         entity: Entity,
@@ -198,6 +205,7 @@ class SearchService:
             )
         )
 
+    @logfire.instrument(record_return=True)
     async def index_entity_markdown(
         self,
         entity: Entity,
@@ -335,14 +343,17 @@ class SearchService:
         # Batch insert all rows at once
         await self.repository.bulk_index_items(rows_to_index)
 
+    @logfire.instrument(record_return=True)
     async def delete_by_permalink(self, permalink: str):
         """Delete an item from the search index."""
         await self.repository.delete_by_permalink(permalink)
 
+    @logfire.instrument(record_return=True)
     async def delete_by_entity_id(self, entity_id: int):
         """Delete an item from the search index."""
         await self.repository.delete_by_entity_id(entity_id)
 
+    @logfire.instrument(record_return=True)
     async def handle_delete(self, entity: Entity):
         """Handle complete entity deletion from search index including observations and relations.