add logfire instrumentation to services and repository code

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

src/basic_memory/repository/entity_repository.py

@@ -3,6 +3,7 @@
 from pathlib import Path
 from typing import List, Optional, Sequence, Union, Any
 
+import logfire
 from loguru import logger
 from sqlalchemy import select
 from sqlalchemy.exc import IntegrityError
@@ -32,6 +33,7 @@ class EntityRepository(Repository[Entity]):
         """
         super().__init__(session_maker, Entity, project_id=project_id)
 
+    @logfire.instrument(record_return=True)
     async def get_by_id(self, entity_id: int) -> Optional[Entity]:
         """Get entity by numeric ID.
 
@@ -44,6 +46,7 @@ class EntityRepository(Repository[Entity]):
         async with db.scoped_session(self.session_maker) as session:
             return await self.select_by_id(session, entity_id)
 
+    @logfire.instrument(record_return=True)
     async def get_by_permalink(self, permalink: str) -> Optional[Entity]:
         """Get entity by permalink.
 
@@ -53,6 +56,7 @@ class EntityRepository(Repository[Entity]):
         query = self.select().where(Entity.permalink == permalink).options(*self.get_load_options())
         return await self.find_one(query)
 
+    @logfire.instrument(record_return=True)
     async def get_by_title(self, title: str) -> Sequence[Entity]:
         """Get entity by title.
 
@@ -63,6 +67,7 @@ class EntityRepository(Repository[Entity]):
         result = await self.execute_query(query)
         return list(result.scalars().all())
 
+    @logfire.instrument(record_return=True)
     async def get_by_file_path(self, file_path: Union[Path, str]) -> Optional[Entity]:
         """Get entity by file_path.
 
@@ -76,6 +81,7 @@ class EntityRepository(Repository[Entity]):
         )
         return await self.find_one(query)
 
+    @logfire.instrument(record_return=True)
     async def get_by_file_paths(
         self, session: AsyncSession, file_paths: Sequence[Union[Path, str]]
     ) -> List[Row[Any]]:
@@ -104,6 +110,7 @@ class EntityRepository(Repository[Entity]):
         result = await session.execute(query)
         return list(result.all())
 
+    @logfire.instrument(record_return=True)
     async def find_by_checksum(self, checksum: str) -> Sequence[Entity]:
         """Find entities with the given checksum.
 
@@ -121,6 +128,7 @@ class EntityRepository(Repository[Entity]):
         result = await self.execute_query(query, use_query_options=False)
         return list(result.scalars().all())
 
+    @logfire.instrument(record_return=True)
     async def find_by_checksums(self, checksums: Sequence[str]) -> Sequence[Entity]:
         """Find entities with any of the given checksums (batch query for move detection).
 
@@ -149,6 +157,7 @@ class EntityRepository(Repository[Entity]):
         result = await self.execute_query(query, use_query_options=False)
         return list(result.scalars().all())
 
+    @logfire.instrument(record_return=True)
     async def delete_by_file_path(self, file_path: Union[Path, str]) -> bool:
         """Delete entity with the provided file_path.
 
@@ -169,6 +178,7 @@ class EntityRepository(Repository[Entity]):
             selectinload(Entity.incoming_relations).selectinload(Relation.to_entity),
         ]
 
+    @logfire.instrument(record_return=True)
     async def find_by_permalinks(self, permalinks: List[str]) -> Sequence[Entity]:
         """Find multiple entities by their permalink.
 
@@ -187,6 +197,7 @@ class EntityRepository(Repository[Entity]):
         result = await self.execute_query(query)
         return list(result.scalars().all())
 
+    @logfire.instrument(record_return=True)
     async def upsert_entity(self, entity: Entity) -> Entity:
         """Insert or update entity using simple try/catch with database-level conflict resolution.
 
@@ -288,6 +299,7 @@ class EntityRepository(Repository[Entity]):
                     entity = await self._handle_permalink_conflict(entity, session)
                     return entity
 
+    @logfire.instrument(record_return=True)
     async def get_all_file_paths(self) -> List[str]:
         """Get all file paths for this project - optimized for deletion detection.
 
@@ -303,6 +315,7 @@ class EntityRepository(Repository[Entity]):
         result = await self.execute_query(query, use_query_options=False)
         return list(result.scalars().all())
 
+    @logfire.instrument(record_return=True)
     async def get_distinct_directories(self) -> List[str]:
         """Extract unique directory paths from file_path column.
 
@@ -331,6 +344,7 @@ class EntityRepository(Repository[Entity]):
 
         return sorted(directories)
 
+    @logfire.instrument(record_return=True)
     async def find_by_directory_prefix(self, directory_prefix: str) -> Sequence[Entity]:
         """Find entities whose file_path starts with the given directory prefix.
 
@@ -363,6 +377,7 @@ class EntityRepository(Repository[Entity]):
         result = await self.execute_query(query, use_query_options=False)
         return list(result.scalars().all())
 
+    @logfire.instrument(record_return=True)
     async def _handle_permalink_conflict(self, entity: Entity, session: AsyncSession) -> Entity:
         """Handle permalink conflicts by generating a unique permalink."""
         base_permalink = entity.permalink

src/basic_memory/repository/observation_repository.py

@@ -2,6 +2,7 @@
 
 from typing import Dict, List, Sequence
 
+import logfire
 from sqlalchemy import select
 from sqlalchemy.ext.asyncio import async_sessionmaker
 
@@ -21,30 +22,35 @@ class ObservationRepository(Repository[Observation]):
         """
         super().__init__(session_maker, Observation, project_id=project_id)
 
+    @logfire.instrument(record_return=True)
     async def find_by_entity(self, entity_id: int) -> Sequence[Observation]:
         """Find all observations for a specific entity."""
         query = select(Observation).filter(Observation.entity_id == entity_id)
         result = await self.execute_query(query)
         return result.scalars().all()
 
+    @logfire.instrument(record_return=True)
     async def find_by_context(self, context: str) -> Sequence[Observation]:
         """Find observations with a specific context."""
         query = select(Observation).filter(Observation.context == context)
         result = await self.execute_query(query)
         return result.scalars().all()
 
+    @logfire.instrument(record_return=True)
     async def find_by_category(self, category: str) -> Sequence[Observation]:
         """Find observations with a specific context."""
         query = select(Observation).filter(Observation.category == category)
         result = await self.execute_query(query)
         return result.scalars().all()
 
+    @logfire.instrument(record_return=True)
     async def observation_categories(self) -> Sequence[str]:
         """Return a list of all observation categories."""
         query = select(Observation.category).distinct()
         result = await self.execute_query(query, use_query_options=False)
         return result.scalars().all()
 
+    @logfire.instrument(record_return=True)
     async def find_by_entities(self, entity_ids: List[int]) -> Dict[int, List[Observation]]:
         """Find all observations for multiple entities in a single query.
 

src/basic_memory/repository/postgres_search_repository.py

@@ -5,6 +5,7 @@ import re
 from datetime import datetime
 from typing import List, Optional
 
+import logfire
 from loguru import logger
 from sqlalchemy import text
 
@@ -25,6 +26,7 @@ class PostgresSearchRepository(SearchRepositoryBase):
     - JSONB containment operators for metadata search
     """
 
+    @logfire.instrument(record_return=True)
     async def init_search_index(self):
         """Create Postgres table with tsvector column and GIN indexes.
 
@@ -145,6 +147,7 @@ class PostgresSearchRepository(SearchRepositoryBase):
         else:
             return cleaned_term
 
+    @logfire.instrument(record_return=True)
     async def search(
         self,
         search_text: Optional[str] = None,
@@ -312,6 +315,7 @@ class PostgresSearchRepository(SearchRepositoryBase):
 
         return results
 
+    @logfire.instrument(record_return=True)
     async def bulk_index_items(self, search_index_rows: List[SearchIndexRow]) -> None:
         """Index multiple items in a single batch operation using UPSERT.
 

src/basic_memory/repository/project_repository.py

@@ -3,6 +3,7 @@
 from pathlib import Path
 from typing import Optional, Sequence, Union
 
+import logfire
 from sqlalchemy import text
 from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker
 
@@ -22,6 +23,7 @@ class ProjectRepository(Repository[Project]):
         """Initialize with session maker."""
         super().__init__(session_maker, Project)
 
+    @logfire.instrument(record_return=True)
     async def get_by_name(self, name: str) -> Optional[Project]:
         """Get project by name.
 
@@ -31,6 +33,7 @@ class ProjectRepository(Repository[Project]):
         query = self.select().where(Project.name == name)
         return await self.find_one(query)
 
+    @logfire.instrument(record_return=True)
     async def get_by_permalink(self, permalink: str) -> Optional[Project]:
         """Get project by permalink.
 
@@ -40,6 +43,7 @@ class ProjectRepository(Repository[Project]):
         query = self.select().where(Project.permalink == permalink)
         return await self.find_one(query)
 
+    @logfire.instrument(record_return=True)
     async def get_by_path(self, path: Union[Path, str]) -> Optional[Project]:
         """Get project by filesystem path.
 
@@ -49,6 +53,7 @@ class ProjectRepository(Repository[Project]):
         query = self.select().where(Project.path == Path(path).as_posix())
         return await self.find_one(query)
 
+    @logfire.instrument(record_return=True)
     async def get_by_id(self, project_id: int) -> Optional[Project]:
         """Get project by numeric ID.
 
@@ -61,17 +66,20 @@ class ProjectRepository(Repository[Project]):
         async with db.scoped_session(self.session_maker) as session:
             return await self.select_by_id(session, project_id)
 
+    @logfire.instrument(record_return=True)
     async def get_default_project(self) -> Optional[Project]:
         """Get the default project (the one marked as is_default=True)."""
         query = self.select().where(Project.is_default.is_not(None))
         return await self.find_one(query)
 
+    @logfire.instrument(record_return=True)
     async def get_active_projects(self) -> Sequence[Project]:
         """Get all active projects."""
         query = self.select().where(Project.is_active == True)  # noqa: E712
         result = await self.execute_query(query)
         return list(result.scalars().all())
 
+    @logfire.instrument(record_return=True)
     async def set_as_default(self, project_id: int) -> Optional[Project]:
         """Set a project as the default and unset previous default.
 
@@ -96,6 +104,7 @@ class ProjectRepository(Repository[Project]):
                 return target_project
             return None  # pragma: no cover
 
+    @logfire.instrument(record_return=True)
     async def update_path(self, project_id: int, new_path: str) -> Optional[Project]:
         """Update project path.
 

src/basic_memory/repository/relation_repository.py

@@ -1,9 +1,9 @@
 """Repository for managing Relation objects."""
 
-from sqlalchemy import and_, delete
 from typing import Sequence, List, Optional
 
-from sqlalchemy import select
+import logfire
+from sqlalchemy import and_, delete, select
 from sqlalchemy.ext.asyncio import async_sessionmaker
 from sqlalchemy.orm import selectinload, aliased
 from sqlalchemy.orm.interfaces import LoaderOption
@@ -25,6 +25,7 @@ class RelationRepository(Repository[Relation]):
         """
         super().__init__(session_maker, Relation, project_id=project_id)
 
+    @logfire.instrument(record_return=True)
     async def find_relation(
         self, from_permalink: str, to_permalink: str, relation_type: str
     ) -> Optional[Relation]:
@@ -46,18 +47,21 @@ class RelationRepository(Repository[Relation]):
         )
         return await self.find_one(query)
 
+    @logfire.instrument(record_return=True)
     async def find_by_entities(self, from_id: int, to_id: int) -> Sequence[Relation]:
         """Find all relations between two entities."""
         query = select(Relation).where((Relation.from_id == from_id) & (Relation.to_id == to_id))
         result = await self.execute_query(query)
         return result.scalars().all()
 
+    @logfire.instrument(record_return=True)
     async def find_by_type(self, relation_type: str) -> Sequence[Relation]:
         """Find all relations of a specific type."""
         query = select(Relation).filter(Relation.relation_type == relation_type)
         result = await self.execute_query(query)
         return result.scalars().all()
 
+    @logfire.instrument(record_return=True)
     async def delete_outgoing_relations_from_entity(self, entity_id: int) -> None:
         """Delete outgoing relations for an entity.
 
@@ -67,12 +71,14 @@ class RelationRepository(Repository[Relation]):
         async with db.scoped_session(self.session_maker) as session:
             await session.execute(delete(Relation).where(Relation.from_id == entity_id))
 
+    @logfire.instrument(record_return=True)
     async def find_unresolved_relations(self) -> Sequence[Relation]:
         """Find all unresolved relations, where to_id is null."""
         query = select(Relation).filter(Relation.to_id.is_(None))
         result = await self.execute_query(query)
         return result.scalars().all()
 
+    @logfire.instrument(record_return=True)
     async def find_unresolved_relations_for_entity(self, entity_id: int) -> Sequence[Relation]:
         """Find unresolved relations for a specific entity.
 

src/basic_memory/repository/repository.py

@@ -2,6 +2,7 @@
 
 from typing import Type, Optional, Any, Sequence, TypeVar, List, Dict
 
+import logfire
 from loguru import logger
 from sqlalchemy import (
     select,
@@ -84,6 +85,7 @@ class Repository[T: Base]:
         result = await session.execute(query)
         return result.scalars().one_or_none()
 
+    @logfire.instrument(record_return=True)
     async def select_by_ids(self, session: AsyncSession, ids: List[int]) -> Sequence[T]:
         """Select multiple entities by IDs using an existing session."""
         query = (
@@ -95,6 +97,7 @@ class Repository[T: Base]:
         result = await session.execute(query)
         return result.scalars().all()
 
+    @logfire.instrument(record_return=True)
     async def add(self, model: T) -> T:
         """
         Add a model to the repository. This will also add related objects
@@ -121,6 +124,7 @@ class Repository[T: Base]:
                 )
             return found
 
+    @logfire.instrument(record_return=True)
     async def add_all(self, models: List[T]) -> Sequence[T]:
         """
         Add a list of models to the repository. This will also add related objects
@@ -152,6 +156,7 @@ class Repository[T: Base]:
         # Add project filter if applicable
         return self._add_project_filter(query)
 
+    @logfire.instrument(record_return=True)
     async def find_all(
         self, skip: int = 0, limit: Optional[int] = None, use_load_options: bool = True
     ) -> Sequence[T]:
@@ -183,6 +188,7 @@ class Repository[T: Base]:
             logger.debug(f"Found {len(items)} {self.Model.__name__} records")
             return items
 
+    @logfire.instrument(record_return=True)
     async def find_by_id(self, entity_id: int) -> Optional[T]:
         """Fetch an entity by its unique identifier."""
         logger.debug(f"Finding {self.Model.__name__} by ID: {entity_id}")
@@ -190,6 +196,7 @@ class Repository[T: Base]:
         async with db.scoped_session(self.session_maker) as session:
             return await self.select_by_id(session, entity_id)
 
+    @logfire.instrument(record_return=True)
     async def find_by_ids(self, ids: List[int]) -> Sequence[T]:
         """Fetch multiple entities by their identifiers in a single query."""
         logger.debug(f"Finding {self.Model.__name__} by IDs: {ids}")
@@ -197,6 +204,7 @@ class Repository[T: Base]:
         async with db.scoped_session(self.session_maker) as session:
             return await self.select_by_ids(session, ids)
 
+    @logfire.instrument(record_return=True)
     async def find_one(self, query: Select[tuple[T]]) -> Optional[T]:
         """Execute a query and retrieve a single record."""
         # add in load options
@@ -210,6 +218,7 @@ class Repository[T: Base]:
             logger.trace(f"No {self.Model.__name__} found")
         return entity
 
+    @logfire.instrument(record_return=True)
     async def create(self, data: dict) -> T:
         """Create a new record from a model instance."""
         logger.debug(f"Creating {self.Model.__name__} from entity_data: {data}")
@@ -241,6 +250,7 @@ class Repository[T: Base]:
                 )
             return return_instance
 
+    @logfire.instrument(record_return=True)
     async def create_all(self, data_list: List[dict]) -> Sequence[T]:
         """Create multiple records in a single transaction."""
         logger.debug(f"Bulk creating {len(data_list)} {self.Model.__name__} instances")
@@ -266,6 +276,7 @@ class Repository[T: Base]:
 
             return await self.select_by_ids(session, [model.id for model in model_list])  # pyright: ignore [reportAttributeAccessIssue]
 
+    @logfire.instrument(record_return=True)
     async def update(self, entity_id: int, entity_data: dict | T) -> Optional[T]:
         """Update an entity with the given data."""
         logger.debug(f"Updating {self.Model.__name__} {entity_id} with data: {entity_data}")
@@ -295,6 +306,7 @@ class Repository[T: Base]:
                 logger.debug(f"No {self.Model.__name__} found to update: {entity_id}")
                 return None
 
+    @logfire.instrument(record_return=True)
     async def delete(self, entity_id: int) -> bool:
         """Delete an entity from the database."""
         logger.debug(f"Deleting {self.Model.__name__}: {entity_id}")
@@ -312,6 +324,7 @@ class Repository[T: Base]:
                 logger.debug(f"No {self.Model.__name__} found to delete: {entity_id}")
                 return False
 
+    @logfire.instrument(record_return=True)
     async def delete_by_ids(self, ids: List[int]) -> int:
         """Delete records matching given IDs."""
         logger.debug(f"Deleting {self.Model.__name__} by ids: {ids}")
@@ -327,6 +340,7 @@ class Repository[T: Base]:
             logger.debug(f"Deleted {result.rowcount} records")
             return result.rowcount
 
+    @logfire.instrument(record_return=True)
     async def delete_by_fields(self, **filters: Any) -> bool:
         """Delete records matching given field values."""
         logger.debug(f"Deleting {self.Model.__name__} by fields: {filters}")
@@ -343,6 +357,7 @@ class Repository[T: Base]:
             logger.debug(f"Deleted {result.rowcount} records")
             return deleted
 
+    @logfire.instrument(record_return=True)
     async def count(self, query: Executable | None = None) -> int:
         """Count entities in the database table."""
         async with db.scoped_session(self.session_maker) as session:
@@ -364,6 +379,7 @@ class Repository[T: Base]:
             logger.debug(f"Counted {count} {self.Model.__name__} records")
             return count
 
+    @logfire.instrument(record_return=True)
     async def execute_query(
         self,
         query: Executable,

src/basic_memory/repository/search_repository_base.py

@@ -4,6 +4,7 @@ from abc import ABC, abstractmethod
 from datetime import datetime
 from typing import Any, Dict, List, Optional
 
+import logfire
 from loguru import logger
 from sqlalchemy import Executable, Result, text
 from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker
@@ -102,6 +103,7 @@ class SearchRepositoryBase(ABC):
         """
         pass
 
+    @logfire.instrument(record_return=True)
     async def index_item(self, search_index_row: SearchIndexRow) -> None:
         """Index or update a single item.
 
@@ -145,6 +147,7 @@ class SearchRepositoryBase(ABC):
             logger.debug(f"indexed row {search_index_row}")
             await session.commit()
 
+    @logfire.instrument(record_return=True)
     async def bulk_index_items(self, search_index_rows: List[SearchIndexRow]) -> None:
         """Index multiple items in a single batch operation.
 
@@ -192,6 +195,7 @@ class SearchRepositoryBase(ABC):
             logger.debug(f"Bulk indexed {len(search_index_rows)} rows")
             await session.commit()
 
+    @logfire.instrument(record_return=True)
     async def delete_by_entity_id(self, entity_id: int) -> None:
         """Delete all search index entries for an entity.
 
@@ -206,6 +210,7 @@ class SearchRepositoryBase(ABC):
             )
             await session.commit()
 
+    @logfire.instrument(record_return=True)
     async def delete_by_permalink(self, permalink: str) -> None:
         """Delete a search index entry by permalink.
 
@@ -220,6 +225,7 @@ class SearchRepositoryBase(ABC):
             )
             await session.commit()
 
+    @logfire.instrument(record_return=True)
     async def execute_query(
         self,
         query: Executable,

src/basic_memory/repository/sqlite_search_repository.py

@@ -5,6 +5,7 @@ import re
 from datetime import datetime
 from typing import List, Optional
 
+import logfire
 from loguru import logger
 from sqlalchemy import text
 
@@ -25,6 +26,7 @@ class SQLiteSearchRepository(SearchRepositoryBase):
     - Prefix wildcard matching with *
     """
 
+    @logfire.instrument(record_return=True)
     async def init_search_index(self):
         """Create FTS5 virtual table for search.
 
@@ -279,6 +281,7 @@ class SQLiteSearchRepository(SearchRepositoryBase):
         # For non-Boolean queries, use the single term preparation logic
         return self._prepare_single_term(term, is_prefix)
 
+    @logfire.instrument(record_return=True)
     async def search(
         self,
         search_text: Optional[str] = None,