basic-memory 0.0.0__py3-none-any.whl

basic_memory/repository/repository.py

@@ -0,0 +1,303 @@
+"""Base repository implementation."""
+
+from datetime import datetime
+from typing import Type, Optional, Any, Sequence, TypeVar, List
+
+from loguru import logger
+from sqlalchemy import (
+    select,
+    func,
+    Select,
+    Executable,
+    inspect,
+    Result,
+    Column,
+    and_,
+    delete,
+)
+from sqlalchemy.exc import NoResultFound
+from sqlalchemy.ext.asyncio import async_sessionmaker, AsyncSession
+from sqlalchemy.orm.interfaces import LoaderOption
+
+from basic_memory import db
+from basic_memory.models import Base
+
+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]):
+        self.session_maker = session_maker
+        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]
+
+    def get_model_data(self, entity_data):
+        model_data = {
+            k: v for k, v in entity_data.items() if k in self.valid_columns and v is not None
+        }
+        return model_data
+
+    async def select_by_id(self, session: AsyncSession, entity_id: int) -> Optional[T]:
+        """Select an entity by ID using an existing session."""
+        query = (
+            select(self.Model)
+            .filter(self.primary_key == entity_id)
+            .options(*self.get_load_options())
+        )
+        result = await session.execute(query)
+        return result.scalars().one_or_none()
+
+    async def select_by_ids(self, session: AsyncSession, ids: List[int]) -> Sequence[T]:
+        """Select multiple entities by IDs using an existing session."""
+        query = (
+            select(self.Model).where(self.primary_key.in_(ids)).options(*self.get_load_options())
+        )
+        result = await session.execute(query)
+        return result.scalars().all()
+
+    async def add(self, model: T) -> T:
+        """
+        Add a model to the repository. This will also add related objects
+        :param model: the model to add
+        :return: the added model instance
+        """
+        async with db.scoped_session(self.session_maker) as session:
+            session.add(model)
+            await session.flush()
+
+            # Query within same session
+            found = await self.select_by_id(session, model.id)  # pyright: ignore [reportAttributeAccessIssue]
+            assert found is not None, "can't find model after session.add"
+            return found
+
+    async def add_all(self, models: List[T]) -> Sequence[T]:
+        """
+        Add a list of models to the repository. This will also add related objects
+        :param models: the models to add
+        :return: the added models instances
+        """
+        async with db.scoped_session(self.session_maker) as session:
+            session.add_all(models)
+            await session.flush()
+
+            # Query within same session
+            return await self.select_by_ids(session, [m.id for m in models])  # pyright: ignore [reportAttributeAccessIssue]
+
+    def select(self, *entities: Any) -> Select:
+        """Create a new SELECT statement.
+
+        Returns:
+            A SQLAlchemy Select object configured with the provided entities
+            or this repository's model if no entities provided.
+        """
+        if not entities:
+            entities = (self.Model,)
+        return select(*entities)
+
+    async def refresh(self, instance: T, relationships: list[str] | None = None) -> None:
+        """Refresh instance and optionally specified relationships."""
+        logger.debug(f"Refreshing {self.Model.__name__} instance: {getattr(instance, 'id', None)}")
+        async with db.scoped_session(self.session_maker) as session:
+            await session.refresh(instance, relationships or [])
+            logger.debug(f"Refreshed relationships: {relationships}")
+
+    async def find_all(self, skip: int = 0, limit: Optional[int] = 0) -> Sequence[T]:
+        """Fetch records from the database with pagination."""
+        logger.debug(f"Finding all {self.Model.__name__} (skip={skip}, limit={limit})")
+
+        async with db.scoped_session(self.session_maker) as session:
+            query = select(self.Model).offset(skip).options(*self.get_load_options())
+            if limit:
+                query = query.limit(limit)
+
+            result = await session.execute(query)
+
+            items = result.scalars().all()
+            logger.debug(f"Found {len(items)} {self.Model.__name__} records")
+            return items
+
+    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}")
+
+        async with db.scoped_session(self.session_maker) as session:
+            return await self.select_by_id(session, entity_id)
+
+    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}")
+
+        async with db.scoped_session(self.session_maker) as session:
+            return await self.select_by_ids(session, ids)
+
+    async def find_one(self, query: Select[tuple[T]]) -> Optional[T]:
+        """Execute a query and retrieve a single record."""
+        logger.debug(f"Finding one {self.Model.__name__} with query: {query}")
+
+        # add in load options
+        query = query.options(*self.get_load_options())
+        result = await self.execute_query(query)
+        entity = result.scalars().one_or_none()
+
+        if entity:
+            logger.debug(f"Found {self.Model.__name__}: {getattr(entity, 'id', None)}")
+        else:
+            logger.debug(f"No {self.Model.__name__} found")
+        return entity
+
+    async def find_modified_since(self, since: datetime) -> Sequence[T]:
+        """Find all records modified since the given timestamp.
+
+        This method assumes the model has an updated_at column. Override
+        in subclasses if a different column should be used.
+
+        Args:
+            since: Datetime to search from
+
+        Returns:
+            Sequence of records modified since the timestamp
+        """
+        logger.debug(f"Finding {self.Model.__name__} modified since: {since}")
+
+        if not hasattr(self.Model, "updated_at"):
+            raise AttributeError(f"{self.Model.__name__} does not have updated_at column")
+
+        query = (
+            select(self.Model)
+            .filter(self.Model.updated_at >= since)
+            .options(*self.get_load_options())
+        )
+
+        async with db.scoped_session(self.session_maker) as session:
+            result = await session.execute(query)
+            items = result.scalars().all()
+            logger.debug(f"Found {len(items)} modified {self.Model.__name__} records")
+            return items
+
+    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}")
+        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)
+            model = self.Model(**model_data)
+            session.add(model)
+            await session.flush()
+
+            return_instance = await self.select_by_id(session, model.id)  # pyright: ignore [reportAttributeAccessIssue]
+            assert return_instance is not None, "can't find model after session.add"
+            return return_instance
+
+    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")
+
+        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
+            ]
+            session.add_all(model_list)
+            await session.flush()
+
+            return await self.select_by_ids(session, [model.id for model in model_list])  # pyright: ignore [reportAttributeAccessIssue]
+
+    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}")
+        async with db.scoped_session(self.session_maker) as session:
+            try:
+                result = await session.execute(
+                    select(self.Model).filter(self.primary_key == entity_id)
+                )
+                entity = result.scalars().one()
+
+                if isinstance(entity_data, dict):
+                    for key, value in entity_data.items():
+                        if key in self.valid_columns:
+                            setattr(entity, key, value)
+                            
+                elif isinstance(entity_data, self.Model):
+                    for column in self.Model.__table__.columns.keys():
+                        setattr(entity, column, getattr(entity_data, column))
+                        
+                await session.flush()  # Make sure changes are flushed
+                await session.refresh(entity)  # Refresh
+
+                logger.debug(f"Updated {self.Model.__name__}: {entity_id}")
+                return await self.select_by_id(session, entity.id)  # pyright: ignore [reportAttributeAccessIssue]
+
+            except NoResultFound:
+                logger.debug(f"No {self.Model.__name__} found to update: {entity_id}")
+                return None
+
+    async def delete(self, entity_id: int) -> bool:
+        """Delete an entity from the database."""
+        logger.debug(f"Deleting {self.Model.__name__}: {entity_id}")
+        async with db.scoped_session(self.session_maker) as session:
+            try:
+                result = await session.execute(
+                    select(self.Model).filter(self.primary_key == entity_id)
+                )
+                entity = result.scalars().one()
+                await session.delete(entity)
+
+                logger.debug(f"Deleted {self.Model.__name__}: {entity_id}")
+                return True
+            except NoResultFound:
+                logger.debug(f"No {self.Model.__name__} found to delete: {entity_id}")
+                return False
+
+    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}")
+        async with db.scoped_session(self.session_maker) as session:
+            query = delete(self.Model).where(self.primary_key.in_(ids))
+            result = await session.execute(query)
+            logger.debug(f"Deleted {result.rowcount} records")
+            return result.rowcount
+
+    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}")
+        async with db.scoped_session(self.session_maker) as session:
+            conditions = [getattr(self.Model, field) == value for field, value in filters.items()]
+            query = delete(self.Model).where(and_(*conditions))
+            result = await session.execute(query)
+            deleted = result.rowcount > 0
+            logger.debug(f"Deleted {result.rowcount} records")
+            return deleted
+
+    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:
+            if query is None:
+                query = select(func.count()).select_from(self.Model)
+            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]:
+        """Execute a query asynchronously."""
+
+        query = query.options(*self.get_load_options()) if use_query_options else query
+
+        logger.debug(f"Executing query: {query}")
+        async with db.scoped_session(self.session_maker) as session:
+            result = await session.execute(query)
+            logger.debug("Query executed successfully")
+            return result
+
+    def get_load_options(self) -> List[LoaderOption]:
+        """Get list of loader options for eager loading relationships.
+        Override in subclasses to specify what to load."""
+        return []

basic_memory/repository/search_repository.py

@@ -0,0 +1,259 @@
+"""Repository for search operations."""
+
+import json
+import time
+from dataclasses import dataclass
+from datetime import datetime
+from typing import List, Optional, Any, Dict
+
+from loguru import logger
+from sqlalchemy import text, Executable, Result
+from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker
+
+from basic_memory import db
+from basic_memory.models.search import CREATE_SEARCH_INDEX
+from basic_memory.schemas.search import SearchItemType
+
+
+@dataclass
+class SearchIndexRow:
+    """Search result with score and metadata."""
+
+    id: int
+    type: str
+    metadata: Optional[dict] = None
+
+    # date values
+    created_at: Optional[datetime] = None
+    updated_at: Optional[datetime] = None
+
+    # assigned in result
+    score: Optional[float] = None
+
+    # Common fields
+    permalink: Optional[str] = None
+    file_path: Optional[str] = None
+
+    # Type-specific fields
+    title: Optional[int] = None  # entity
+    content: Optional[int] = None  # entity, observation
+    entity_id: Optional[int] = None  # observations
+    category: Optional[str] = None  # observations
+    from_id: Optional[int] = None  # relations
+    to_id: Optional[int] = None  # relations
+    relation_type: Optional[str] = None  # relations
+
+    def to_insert(self):
+        return {
+            "id": self.id,
+            "title": self.title,
+            "content": self.content,
+            "permalink": self.permalink,
+            "file_path": self.file_path,
+            "type": self.type,
+            "metadata": json.dumps(self.metadata),
+            "from_id": self.from_id,
+            "to_id": self.to_id,
+            "relation_type": self.relation_type,
+            "entity_id": self.entity_id,
+            "category": self.category,
+            "created_at": self.created_at if self.created_at else None,
+            "updated_at": self.updated_at if self.updated_at else None,
+        }
+
+
+class SearchRepository:
+    """Repository for search index operations."""
+
+    def __init__(self, session_maker: async_sessionmaker[AsyncSession]):
+        self.session_maker = session_maker
+
+    async def init_search_index(self):
+        """Create or recreate the search index."""
+        async with db.scoped_session(self.session_maker) as session:
+            await session.execute(CREATE_SEARCH_INDEX)
+            await session.commit()
+
+    def _quote_search_term(self, term: str) -> str:
+        """Add quotes if term contains special characters.
+        For FTS5, special characters and phrases need to be quoted to be treated as a single token.
+        """
+        # List of special characters that need quoting
+        special_chars = ['/', '*', '-', '.', ' ', '(', ')', '[', ']', '"', "'"]
+        
+        # Check if term contains any special characters
+        if any(c in term for c in special_chars):
+            # If the term already contains quotes, escape them
+            term = term.replace('"', '""')
+            return f'"{term}"'
+        return term
+
+    async def search(
+        self,
+        search_text: Optional[str] = None,
+        permalink: Optional[str] = None,
+        permalink_match: Optional[str] = None,
+        title: Optional[str] = None,
+        types: List[SearchItemType] = None,
+        after_date: datetime = None,
+        entity_types: List[str] = None,
+        limit: int = 10,
+    ) -> List[SearchIndexRow]:
+        """Search across all indexed content with fuzzy matching."""
+        conditions = []
+        params = {}
+
+        # Handle text search for title and content
+        if search_text:
+            search_text = self._quote_search_term(search_text.lower().strip())
+            params["text"] = f"{search_text}*"
+            conditions.append("(title MATCH :text OR content MATCH :text)")
+
+        # Handle title match search
+        if title:
+            title_text = self._quote_search_term(title.lower().strip())
+            params["text"] = f"{title_text}*"
+            conditions.append("title MATCH :text")
+
+        # Handle permalink exact search
+        if permalink:
+            params["permalink"] = permalink
+            conditions.append("permalink = :permalink")
+
+        # Handle permalink match search, supports *
+        if permalink_match:
+            params["permalink"] = self._quote_search_term(permalink_match)
+            conditions.append("permalink MATCH :permalink")
+            
+        # Handle type filter
+        if types:
+            type_list = ", ".join(f"'{t.value}'" for t in types)
+            conditions.append(f"type IN ({type_list})")
+
+        # Handle entity type filter
+        if entity_types:
+            entity_type_list = ", ".join(f"'{t}'" for t in entity_types)
+            conditions.append(f"json_extract(metadata, '$.entity_type') IN ({entity_type_list})")
+
+        # Handle date filter using datetime() for proper comparison
+        if after_date:
+            params["after_date"] = after_date
+            conditions.append("datetime(created_at) > datetime(:after_date)")
+
+        # set limit on search query
+        params["limit"] = limit
+        
+        # Build WHERE clause
+        where_clause = " AND ".join(conditions) if conditions else "1=1"
+
+        sql = f"""
+            SELECT 
+                id, 
+                title, 
+                permalink,
+                file_path,
+                type,
+                metadata,
+                from_id,
+                to_id,
+                relation_type,
+                entity_id,
+                content,
+                category,
+                created_at,
+                updated_at,
+                bm25(search_index) as score
+            FROM search_index 
+            WHERE {where_clause}
+            ORDER BY score ASC
+            LIMIT :limit
+        """
+
+        #logger.debug(f"Search {sql} params: {params}")
+        async with db.scoped_session(self.session_maker) as session:
+            result = await session.execute(text(sql), params)
+            rows = result.fetchall()
+
+        results = [
+            SearchIndexRow(
+                id=row.id,
+                title=row.title,
+                permalink=row.permalink,
+                file_path=row.file_path,
+                type=row.type,
+                score=row.score,
+                metadata=json.loads(row.metadata),
+                from_id=row.from_id,
+                to_id=row.to_id,
+                relation_type=row.relation_type,
+                entity_id=row.entity_id,
+                content=row.content,
+                category=row.category,
+                created_at=row.created_at,
+                updated_at=row.updated_at,
+            )
+            for row in rows
+        ]
+
+        #for r in results:
+        #    logger.debug(f"Search result: type:{r.type} title: {r.title} permalink: {r.permalink} score: {r.score}")
+        return results
+
+    async def index_item(
+        self,
+        search_index_row: SearchIndexRow,
+    ):
+        """Index or update a single item."""
+        async with db.scoped_session(self.session_maker) as session:
+            # Delete existing record if any
+            await session.execute(
+                text("DELETE FROM search_index WHERE permalink = :permalink"),
+                {"permalink": search_index_row.permalink},
+            )
+
+            # Insert new record
+            await session.execute(
+                text("""
+                    INSERT INTO search_index (
+                        id, title, content, permalink, file_path, type, metadata,
+                        from_id, to_id, relation_type,
+                        entity_id, category,
+                        created_at, updated_at
+                    ) VALUES (
+                        :id, :title, :content, :permalink, :file_path, :type, :metadata,
+                        :from_id, :to_id, :relation_type,
+                        :entity_id, :category,
+                        :created_at, :updated_at
+                    )
+                """),
+                search_index_row.to_insert(),
+            )
+            logger.debug(f"indexed permalink {search_index_row.permalink}")
+            await session.commit()
+
+    async def delete_by_permalink(self, permalink: str):
+        """Delete an item from the search index."""
+        async with db.scoped_session(self.session_maker) as session:
+            await session.execute(
+                text("DELETE FROM search_index WHERE permalink = :permalink"),
+                {"permalink": permalink},
+            )
+            await session.commit()
+
+    async def execute_query(
+        self,
+        query: Executable,
+        params: Optional[Dict[str, Any]] = None,
+    ) -> Result[Any]:
+        """Execute a query asynchronously."""
+        #logger.debug(f"Executing query: {query}")
+        async with db.scoped_session(self.session_maker) as session:
+            start_time = time.perf_counter()
+            if params:
+                result = await session.execute(query, params)
+            else:
+                result = await session.execute(query)
+            end_time = time.perf_counter()
+            elapsed_time = end_time - start_time
+            logger.debug(f"Query executed successfully in {elapsed_time:.2f}s.")
+            return result

basic_memory/schemas/__init__.py

@@ -0,0 +1,73 @@
+"""Knowledge graph schema exports.
+
+This module exports all schema classes to simplify imports.
+Rather than importing from individual schema files, you can
+import everything from basic_memory.schemas.
+"""
+
+# Base types and models
+from basic_memory.schemas.base import (
+    Observation,
+    EntityType,
+    RelationType,
+    Relation,
+    Entity,
+)
+
+# Delete operation models
+from basic_memory.schemas.delete import (
+    DeleteEntitiesRequest,
+)
+
+# Request models
+from basic_memory.schemas.request import (
+    SearchNodesRequest,
+    GetEntitiesRequest,
+    CreateRelationsRequest, UpdateEntityRequest,
+)
+
+# Response models
+from basic_memory.schemas.response import (
+    SQLAlchemyModel,
+    ObservationResponse,
+    RelationResponse,
+    EntityResponse,
+    EntityListResponse,
+    SearchNodesResponse,
+    DeleteEntitiesResponse,
+)
+
+# Discovery and analytics models
+from basic_memory.schemas.discovery import (
+    EntityTypeList,
+    ObservationCategoryList, TypedEntityList,
+)
+
+# For convenient imports, export all models
+__all__ = [
+    # Base
+    "Observation",
+    "EntityType",
+    "RelationType",
+    "Relation",
+    "Entity",
+    # Requests
+    "SearchNodesRequest",
+    "GetEntitiesRequest",
+    "CreateRelationsRequest",
+    "UpdateEntityRequest",
+    # Responses
+    "SQLAlchemyModel",
+    "ObservationResponse",
+    "RelationResponse",
+    "EntityResponse",
+    "EntityListResponse",
+    "SearchNodesResponse",
+    "DeleteEntitiesResponse",
+    # Delete Operations
+    "DeleteEntitiesRequest",
+    # Discovery and Analytics
+    "EntityTypeList",
+    "ObservationCategoryList",
+    "TypedEntityList"
+]