Refactor search API to consistently handle types and entity_types

Updated the naming conventions and handling of `types` and `entity_types` across the codebase for consistency. Modified function definitions, queries, and filters to ensure clear separation and proper usage. Adjusted and added tests to align with the new structure.

Author: phernandez

src/basic_memory/cli/commands/tool.py

@@ -2,31 +2,29 @@
 
 import asyncio
 import sys
-from typing import Optional, List, Annotated
+from typing import Annotated, List, Optional
 
 import typer
 from loguru import logger
 from rich import print as rprint
 
 from basic_memory.cli.app import app
-from basic_memory.mcp.tools import build_context as mcp_build_context
-from basic_memory.mcp.tools import read_note as mcp_read_note
-from basic_memory.mcp.tools import recent_activity as mcp_recent_activity
-from basic_memory.mcp.tools import search_notes as mcp_search
-from basic_memory.mcp.tools import write_note as mcp_write_note
 
 # Import prompts
 from basic_memory.mcp.prompts.continue_conversation import (
     continue_conversation as mcp_continue_conversation,
 )
-
 from basic_memory.mcp.prompts.recent_activity import (
     recent_activity_prompt as recent_activity_prompt,
 )
-
+from basic_memory.mcp.tools import build_context as mcp_build_context
+from basic_memory.mcp.tools import read_note as mcp_read_note
+from basic_memory.mcp.tools import recent_activity as mcp_recent_activity
+from basic_memory.mcp.tools import search_notes as mcp_search
+from basic_memory.mcp.tools import write_note as mcp_write_note
 from basic_memory.schemas.base import TimeFrame
 from basic_memory.schemas.memory import MemoryUrl
-from basic_memory.schemas.search import SearchQuery, SearchItemType
+from basic_memory.schemas.search import SearchItemType
 
 tool_app = typer.Typer()
 app.add_typer(tool_app, name="tool", help="Access to MCP tools via CLI")
@@ -198,13 +196,28 @@ def search_notes(
         raise typer.Abort()
 
     try:
-        search_query = SearchQuery(
-            permalink_match=query if permalink else None,
-            text=query if not (permalink or title) else None,
-            title=query if title else None,
-            after_date=after_date,
+        if permalink and title:  # pragma: no cover
+            typer.echo(
+                "Use either --permalink or --title, not both. Exiting.",
+                err=True,
+            )
+            raise typer.Exit(1)
+
+        # set search type
+        search_type = ("permalink" if permalink else None,)
+        search_type = ("permalink_match" if permalink and "*" in query else None,)
+        search_type = ("title" if title else None,)
+        search_type = "text" if search_type is None else search_type
+
+        results = asyncio.run(
+            mcp_search(
+                query,
+                search_type=search_type,
+                page=page,
+                after_date=after_date,
+                page_size=page_size,
+            )
         )
-        results = asyncio.run(mcp_search(query=search_query, page=page, page_size=page_size))
         # Use json module for more controlled serialization
         import json
 

src/basic_memory/config.py

@@ -234,6 +234,7 @@ def get_process_name():  # pragma: no cover
 # Global flag to track if logging has been set up
 _LOGGING_SETUP = False
 
+
 def setup_basic_memory_logging():  # pragma: no cover
     """Set up logging for basic-memory, ensuring it only happens once."""
     global _LOGGING_SETUP

src/basic_memory/mcp/prompts/continue_conversation.py

@@ -5,19 +5,19 @@
 """
 
 from textwrap import dedent
-from typing import Optional, Annotated
+from typing import Annotated, Optional
 
 from loguru import logger
 from pydantic import Field
 
-from basic_memory.mcp.prompts.utils import format_prompt_context, PromptContext, PromptContextItem
+from basic_memory.mcp.prompts.utils import PromptContext, PromptContextItem, format_prompt_context
 from basic_memory.mcp.server import mcp
 from basic_memory.mcp.tools.build_context import build_context
 from basic_memory.mcp.tools.recent_activity import recent_activity
 from basic_memory.mcp.tools.search import search_notes
 from basic_memory.schemas.base import TimeFrame
 from basic_memory.schemas.memory import GraphContext
-from basic_memory.schemas.search import SearchQuery, SearchItemType
+from basic_memory.schemas.search import SearchItemType
 
 
 @mcp.prompt(
@@ -48,7 +48,7 @@ async def continue_conversation(
     # If topic provided, search for it
     if topic:
         search_results = await search_notes(
-            SearchQuery(text=topic, after_date=timeframe, types=[SearchItemType.ENTITY])
+            query=topic, after_date=timeframe, entity_types=[SearchItemType.ENTITY]
         )
 
         # Build context from results

src/basic_memory/mcp/prompts/search.py

@@ -12,7 +12,7 @@
 from basic_memory.mcp.server import mcp
 from basic_memory.mcp.tools.search import search_notes as search_tool
 from basic_memory.schemas.base import TimeFrame
-from basic_memory.schemas.search import SearchQuery, SearchResponse
+from basic_memory.schemas.search import SearchResponse
 
 
 @mcp.prompt(

src/basic_memory/mcp/tools/read_note.py

@@ -9,7 +9,6 @@
 from basic_memory.mcp.tools.search import search_notes
 from basic_memory.mcp.tools.utils import call_get
 from basic_memory.schemas.memory import memory_url_path
-from basic_memory.schemas.search import SearchQuery
 
 
 @mcp.tool(

src/basic_memory/mcp/tools/search.py

@@ -1,20 +1,21 @@
 """Search tools for Basic Memory MCP server."""
 
-from typing import Optional, List
+from typing import List, Optional
+
 from loguru import logger
 
+from basic_memory.mcp.async_client import client
 from basic_memory.mcp.server import mcp
 from basic_memory.mcp.tools.utils import call_post
-from basic_memory.schemas.search import SearchQuery, SearchResponse, SearchItemType
-from basic_memory.mcp.async_client import client
+from basic_memory.schemas.search import SearchItemType, SearchQuery, SearchResponse
 
 
 @mcp.tool(
     description="Search across all content in the knowledge base.",
 )
 async def search_notes(
-    query: str, 
-    page: int = 1, 
+    query: str,
+    page: int = 1,
     page_size: int = 10,
     search_type: str = "text",
     types: Optional[List[str]] = None,
@@ -31,9 +32,9 @@ async def search_notes(
         query: The search query string
         page: The page number of results to return (default 1)
         page_size: The number of results to return per page (default 10)
-        search_type: Type of search to perform, one of: "text", "title", "permalink", "permalink_match" (default: "text")
-        types: Optional list of content types to search (e.g., ["entity", "observation"])
-        entity_types: Optional list of entity types to filter by (e.g., ["note", "person"])
+        search_type: Type of search to perform, one of: "text", "title", "permalink" (default: "text")
+        types: Optional list of note types to search (e.g., ["note", "person"])
+        entity_types: Optional list of entity types to filter by (e.g., ["entity", "observation"])
         after_date: Optional date filter for recent content (e.g., "1 week", "2d")
 
     Returns:
@@ -61,6 +62,12 @@ async def search_notes(
             types=["entity"],
         )
 
+        # Search with entity type filter, e.g., note vs
+        results = await search_notes(
+            query="meeting notes",
+            types=["entity"],
+        )
+
         # Search for recent content
         results = await search_notes(
             query="bug report",
@@ -70,32 +77,32 @@ async def search_notes(
         # Pattern matching on permalinks
         results = await search_notes(
             query="docs/meeting-*",
-            search_type="permalink_match"
+            search_type="permalink"
         )
     """
     # Create a SearchQuery object based on the parameters
     search_query = SearchQuery()
-    
+
     # Set the appropriate search field based on search_type
     if search_type == "text":
         search_query.text = query
     elif search_type == "title":
         search_query.title = query
+    elif search_type == "permalink" and "*" in query:
+        search_query.permalink_match = query
     elif search_type == "permalink":
         search_query.permalink = query
-    elif search_type == "permalink_match":
-        search_query.permalink_match = query
     else:
         search_query.text = query  # Default to text search
-    
+
     # Add optional filters if provided
-    if types:
-        search_query.types = [SearchItemType(t) for t in types]
     if entity_types:
-        search_query.entity_types = entity_types
+        search_query.entity_types = [SearchItemType(t) for t in entity_types]
+    if types:
+        search_query.types = types
     if after_date:
         search_query.after_date = after_date
-    
+
     logger.info(f"Searching for {search_query}")
     response = await call_post(
         client,

src/basic_memory/repository/search_repository.py

@@ -4,10 +4,10 @@
 import time
 from dataclasses import dataclass
 from datetime import datetime
-from typing import List, Optional, Any, Dict
+from typing import Any, Dict, List, Optional
 
 from loguru import logger
-from sqlalchemy import text, Executable, Result
+from sqlalchemy import Executable, Result, text
 from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker
 
 from basic_memory import db
@@ -123,9 +123,9 @@ async def search(
         permalink: Optional[str] = None,
         permalink_match: Optional[str] = None,
         title: Optional[str] = None,
-        types: Optional[List[SearchItemType]] = None,
+        types: Optional[List[str]] = None,
         after_date: Optional[datetime] = None,
-        entity_types: Optional[List[str]] = None,
+        entity_types: Optional[List[SearchItemType]] = None,
         limit: int = 10,
         offset: int = 0,
     ) -> List[SearchIndexRow]:
@@ -174,15 +174,15 @@ async def search(
             else:
                 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})")
+            type_list = ", ".join(f"'{t.value}'" for t in entity_types)
+            conditions.append(f"type IN ({type_list})")
+
+        # Handle type filter
+        if types:
+            type_list = ", ".join(f"'{t}'" for t in types)
+            conditions.append(f"json_extract(metadata, '$.entity_type') IN ({type_list})")
 
         # Handle date filter using datetime() for proper comparison
         if after_date:

src/basic_memory/schemas/search.py

@@ -49,8 +49,8 @@ class SearchQuery(BaseModel):
     title: Optional[str] = None  # title only search
 
     # Optional filters
-    types: Optional[List[SearchItemType]] = None  # Filter by item type
-    entity_types: Optional[List[str]] = None  # Filter by entity type
+    types: Optional[List[str]] = None  # Filter by type
+    entity_types: Optional[List[SearchItemType]] = None  # Filter by entity type
     after_date: Optional[Union[datetime, str]] = None  # Time-based filter
 
     @field_validator("after_date")

src/basic_memory/services/context_service.py

@@ -81,7 +81,7 @@ async def build_context(
         else:
             logger.debug(f"Build context for '{types}'")
             primary = await self.search_repository.search(
-                types=types, after_date=since, limit=limit, offset=offset
+                entity_types=types, after_date=since, limit=limit, offset=offset
             )
 
         # Get type_id pairs for traversal

src/basic_memory/services/link_resolver.py

@@ -48,15 +48,15 @@ async def resolve_link(self, link_text: str, use_search: bool = True) -> Optiona
 
         # 3. Try file path
         found_path = await self.entity_repository.get_by_file_path(clean_text)
-        if found_path :
+        if found_path:
             logger.debug(f"Found entity with path: {found_path.file_path}")
             return found_path
-        
+
         # search if indicated
         if use_search and "*" not in clean_text:
             # 3. Fall back to search for fuzzy matching on title
             results = await self.search_service.search(
-                query=SearchQuery(title=clean_text, types=[SearchItemType.ENTITY]),
+                query=SearchQuery(title=clean_text, entity_types=[SearchItemType.ENTITY]),
             )
 
             if results: