basic-memory 0.12.3__py3-none-any.whl → 0.13.0__py3-none-any.whl

basic_memory/repository/observation_repository.py

@@ -1,6 +1,6 @@
 """Repository for managing Observation objects."""
 
-from typing import Sequence
+from typing import Dict, List, Sequence
 
 from sqlalchemy import select
 from sqlalchemy.ext.asyncio import async_sessionmaker
@@ -12,8 +12,14 @@ from basic_memory.repository.repository import Repository
 class ObservationRepository(Repository[Observation]):
     """Repository for Observation model with memory-specific operations."""
 
-    def __init__(self, session_maker: async_sessionmaker):
-        super().__init__(session_maker, Observation)
+    def __init__(self, session_maker: async_sessionmaker, project_id: int):
+        """Initialize with session maker and project_id filter.
+
+        Args:
+            session_maker: SQLAlchemy session maker
+            project_id: Project ID to filter all operations by
+        """
+        super().__init__(session_maker, Observation, project_id=project_id)
 
     async def find_by_entity(self, entity_id: int) -> Sequence[Observation]:
         """Find all observations for a specific entity."""
@@ -38,3 +44,29 @@ class ObservationRepository(Repository[Observation]):
         query = select(Observation.category).distinct()
         result = await self.execute_query(query, use_query_options=False)
         return result.scalars().all()
+
+    async def find_by_entities(self, entity_ids: List[int]) -> Dict[int, List[Observation]]:
+        """Find all observations for multiple entities in a single query.
+
+        Args:
+            entity_ids: List of entity IDs to fetch observations for
+
+        Returns:
+            Dictionary mapping entity_id to list of observations
+        """
+        if not entity_ids:  # pragma: no cover
+            return {}
+
+        # Query observations for all entities in the list
+        query = select(Observation).filter(Observation.entity_id.in_(entity_ids))
+        result = await self.execute_query(query)
+        observations = result.scalars().all()
+
+        # Group observations by entity_id
+        observations_by_entity = {}
+        for obs in observations:
+            if obs.entity_id not in observations_by_entity:
+                observations_by_entity[obs.entity_id] = []
+            observations_by_entity[obs.entity_id].append(obs)
+
+        return observations_by_entity

basic_memory/repository/project_info_repository.py

@@ -1,9 +1,10 @@
 from basic_memory.repository.repository import Repository
+from basic_memory.models.project import Project
 
 
 class ProjectInfoRepository(Repository):
     """Repository for statistics queries."""
 
     def __init__(self, session_maker):
-        # Initialize with a dummy model since we're just using the execute_query method
-        super().__init__(session_maker, None)  # type: ignore
+        # Initialize with Project model as a reference
+        super().__init__(session_maker, Project)

basic_memory/repository/project_repository.py

@@ -0,0 +1,85 @@
+"""Repository for managing projects in Basic Memory."""
+
+from pathlib import Path
+from typing import Optional, Sequence, Union
+
+from sqlalchemy import text
+from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker
+
+from basic_memory import db
+from basic_memory.models.project import Project
+from basic_memory.repository.repository import Repository
+
+
+class ProjectRepository(Repository[Project]):
+    """Repository for Project model.
+
+    Projects represent collections of knowledge entities grouped together.
+    Each entity, observation, and relation belongs to a specific project.
+    """
+
+    def __init__(self, session_maker: async_sessionmaker[AsyncSession]):
+        """Initialize with session maker."""
+        super().__init__(session_maker, Project)
+
+    async def get_by_name(self, name: str) -> Optional[Project]:
+        """Get project by name.
+
+        Args:
+            name: Unique name of the project
+        """
+        query = self.select().where(Project.name == name)
+        return await self.find_one(query)
+
+    async def get_by_permalink(self, permalink: str) -> Optional[Project]:
+        """Get project by permalink.
+
+        Args:
+            permalink: URL-friendly identifier for the project
+        """
+        query = self.select().where(Project.permalink == permalink)
+        return await self.find_one(query)
+
+    async def get_by_path(self, path: Union[Path, str]) -> Optional[Project]:
+        """Get project by filesystem path.
+
+        Args:
+            path: Path to the project directory (will be converted to string internally)
+        """
+        query = self.select().where(Project.path == str(path))
+        return await self.find_one(query)
+
+    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)
+
+    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())
+
+    async def set_as_default(self, project_id: int) -> Optional[Project]:
+        """Set a project as the default and unset previous default.
+
+        Args:
+            project_id: ID of the project to set as default
+
+        Returns:
+            The updated project if found, None otherwise
+        """
+        async with db.scoped_session(self.session_maker) as session:
+            # First, clear the default flag for all projects using direct SQL
+            await session.execute(
+                text("UPDATE project SET is_default = NULL WHERE is_default IS NOT NULL")
+            )
+            await session.flush()
+
+            # Set the new default project
+            target_project = await self.select_by_id(session, project_id)
+            if target_project:
+                target_project.is_default = True
+                await session.flush()
+                return target_project
+            return None  # pragma: no cover

basic_memory/repository/relation_repository.py

@@ -16,8 +16,14 @@ from basic_memory.repository.repository import Repository
 class RelationRepository(Repository[Relation]):
     """Repository for Relation model with memory-specific operations."""
 
-    def __init__(self, session_maker: async_sessionmaker):
-        super().__init__(session_maker, Relation)
+    def __init__(self, session_maker: async_sessionmaker, project_id: int):
+        """Initialize with session maker and project_id filter.
+
+        Args:
+            session_maker: SQLAlchemy session maker
+            project_id: Project ID to filter all operations by
+        """
+        super().__init__(session_maker, Relation, project_id=project_id)
 
     async def find_relation(
         self, from_permalink: str, to_permalink: str, relation_type: str

basic_memory/repository/repository.py

@@ -1,6 +1,6 @@
 """Base repository implementation."""
 
-from typing import Type, Optional, Any, Sequence, TypeVar, List
+from typing import Type, Optional, Any, Sequence, TypeVar, List, Dict
 
 from loguru import logger
 from sqlalchemy import (
@@ -27,13 +27,30 @@ T = TypeVar("T", bound=Base)
 class Repository[T: Base]:
     """Base repository implementation with generic CRUD operations."""
 
-    def __init__(self, session_maker: async_sessionmaker[AsyncSession], Model: Type[T]):
+    def __init__(
+        self,
+        session_maker: async_sessionmaker[AsyncSession],
+        Model: Type[T],
+        project_id: Optional[int] = None,
+    ):
         self.session_maker = session_maker
+        self.project_id = project_id
         if Model:
             self.Model = Model
             self.mapper = inspect(self.Model).mapper
             self.primary_key: Column[Any] = self.mapper.primary_key[0]
             self.valid_columns = [column.key for column in self.mapper.columns]
+            # Check if this model has a project_id column
+            self.has_project_id = "project_id" in self.valid_columns
+
+    def _set_project_id_if_needed(self, model: T) -> None:
+        """Set project_id on model if needed and available."""
+        if (
+            self.has_project_id
+            and self.project_id is not None
+            and getattr(model, "project_id", None) is None
+        ):
+            setattr(model, "project_id", self.project_id)
 
     def get_model_data(self, entity_data):
         model_data = {
@@ -41,6 +58,19 @@ class Repository[T: Base]:
         }
         return model_data
 
+    def _add_project_filter(self, query: Select) -> Select:
+        """Add project_id filter to query if applicable.
+
+        Args:
+            query: The SQLAlchemy query to modify
+
+        Returns:
+            Updated query with project filter if applicable
+        """
+        if self.has_project_id and self.project_id is not None:
+            query = query.filter(getattr(self.Model, "project_id") == self.project_id)
+        return query
+
     async def select_by_id(self, session: AsyncSession, entity_id: int) -> Optional[T]:
         """Select an entity by ID using an existing session."""
         query = (
@@ -48,6 +78,9 @@ class Repository[T: Base]:
             .filter(self.primary_key == entity_id)
             .options(*self.get_load_options())
         )
+        # Add project filter if applicable
+        query = self._add_project_filter(query)
+
         result = await session.execute(query)
         return result.scalars().one_or_none()
 
@@ -56,6 +89,9 @@ class Repository[T: Base]:
         query = (
             select(self.Model).where(self.primary_key.in_(ids)).options(*self.get_load_options())
         )
+        # Add project filter if applicable
+        query = self._add_project_filter(query)
+
         result = await session.execute(query)
         return result.scalars().all()
 
@@ -66,6 +102,9 @@ class Repository[T: Base]:
         :return: the added model instance
         """
         async with db.scoped_session(self.session_maker) as session:
+            # Set project_id if applicable and not already set
+            self._set_project_id_if_needed(model)
+
             session.add(model)
             await session.flush()
 
@@ -89,6 +128,10 @@ class Repository[T: Base]:
         :return: the added models instances
         """
         async with db.scoped_session(self.session_maker) as session:
+            # set the project id if not present in models
+            for model in models:
+                self._set_project_id_if_needed(model)
+
             session.add_all(models)
             await session.flush()
 
@@ -104,7 +147,10 @@ class Repository[T: Base]:
         """
         if not entities:
             entities = (self.Model,)
-        return select(*entities)
+        query = select(*entities)
+
+        # Add project filter if applicable
+        return self._add_project_filter(query)
 
     async def find_all(self, skip: int = 0, limit: Optional[int] = None) -> Sequence[T]:
         """Fetch records from the database with pagination."""
@@ -112,6 +158,9 @@ class Repository[T: Base]:
 
         async with db.scoped_session(self.session_maker) as session:
             query = select(self.Model).offset(skip).options(*self.get_load_options())
+            # Add project filter if applicable
+            query = self._add_project_filter(query)
+
             if limit:
                 query = query.limit(limit)
 
@@ -143,9 +192,9 @@ class Repository[T: Base]:
         entity = result.scalars().one_or_none()
 
         if entity:
-            logger.debug(f"Found {self.Model.__name__}: {getattr(entity, 'id', None)}")
+            logger.trace(f"Found {self.Model.__name__}: {getattr(entity, 'id', None)}")
         else:
-            logger.debug(f"No {self.Model.__name__} found")
+            logger.trace(f"No {self.Model.__name__} found")
         return entity
 
     async def create(self, data: dict) -> T:
@@ -154,6 +203,15 @@ class Repository[T: Base]:
         async with db.scoped_session(self.session_maker) as session:
             # Only include valid columns that are provided in entity_data
             model_data = self.get_model_data(data)
+
+            # Add project_id if applicable and not already provided
+            if (
+                self.has_project_id
+                and self.project_id is not None
+                and "project_id" not in model_data
+            ):
+                model_data["project_id"] = self.project_id
+
             model = self.Model(**model_data)
             session.add(model)
             await session.flush()
@@ -176,12 +234,20 @@ class Repository[T: Base]:
 
         async with db.scoped_session(self.session_maker) as session:
             # Only include valid columns that are provided in entity_data
-            model_list = [
-                self.Model(
-                    **self.get_model_data(d),
-                )
-                for d in data_list
-            ]
+            model_list = []
+            for d in data_list:
+                model_data = self.get_model_data(d)
+
+                # Add project_id if applicable and not already provided
+                if (
+                    self.has_project_id
+                    and self.project_id is not None
+                    and "project_id" not in model_data
+                ):
+                    model_data["project_id"] = self.project_id  # pragma: no cover
+
+                model_list.append(self.Model(**model_data))
+
             session.add_all(model_list)
             await session.flush()
 
@@ -237,7 +303,13 @@ class Repository[T: Base]:
         """Delete records matching given IDs."""
         logger.debug(f"Deleting {self.Model.__name__} by ids: {ids}")
         async with db.scoped_session(self.session_maker) as session:
-            query = delete(self.Model).where(self.primary_key.in_(ids))
+            conditions = [self.primary_key.in_(ids)]
+
+            # Add project_id filter if applicable
+            if self.has_project_id and self.project_id is not None:  # pragma: no cover
+                conditions.append(getattr(self.Model, "project_id") == self.project_id)
+
+            query = delete(self.Model).where(and_(*conditions))
             result = await session.execute(query)
             logger.debug(f"Deleted {result.rowcount} records")
             return result.rowcount
@@ -247,6 +319,11 @@ class Repository[T: Base]:
         logger.debug(f"Deleting {self.Model.__name__} by fields: {filters}")
         async with db.scoped_session(self.session_maker) as session:
             conditions = [getattr(self.Model, field) == value for field, value in filters.items()]
+
+            # Add project_id filter if applicable
+            if self.has_project_id and self.project_id is not None:
+                conditions.append(getattr(self.Model, "project_id") == self.project_id)
+
             query = delete(self.Model).where(and_(*conditions))
             result = await session.execute(query)
             deleted = result.rowcount > 0
@@ -258,19 +335,34 @@ class Repository[T: Base]:
         async with db.scoped_session(self.session_maker) as session:
             if query is None:
                 query = select(func.count()).select_from(self.Model)
+                # Add project filter if applicable
+                if (
+                    isinstance(query, Select)
+                    and self.has_project_id
+                    and self.project_id is not None
+                ):
+                    query = query.where(
+                        getattr(self.Model, "project_id") == self.project_id
+                    )  # pragma: no cover
+
             result = await session.execute(query)
             scalar = result.scalar()
             count = scalar if scalar is not None else 0
             logger.debug(f"Counted {count} {self.Model.__name__} records")
             return count
 
-    async def execute_query(self, query: Executable, use_query_options: bool = True) -> Result[Any]:
+    async def execute_query(
+        self,
+        query: Executable,
+        params: Optional[Dict[str, Any]] = None,
+        use_query_options: bool = True,
+    ) -> Result[Any]:
         """Execute a query asynchronously."""
 
         query = query.options(*self.get_load_options()) if use_query_options else query
-        logger.debug(f"Executing query: {query}")
+        logger.trace(f"Executing query: {query}, params: {params}")
         async with db.scoped_session(self.session_maker) as session:
-            result = await session.execute(query)
+            result = await session.execute(query, params)
             return result
 
     def get_load_options(self) -> List[LoaderOption]: