add relations and observations to search index

This commit is contained in:
phernandez
2025-01-14 19:06:00 -06:00
parent 9dff24746d
commit 73e8c5bee7
8 changed files with 117 additions and 111 deletions
+1 -10
View File
@@ -29,15 +29,6 @@ CREATE VIRTUAL TABLE IF NOT EXISTS search_index USING fts5(
-- Configuration
tokenize='porter unicode61',
prefix='2,3',
content='title content' -- Default searchable columns
prefix='2,3'
);
-- Add index for relation traversal
CREATE INDEX IF NOT EXISTS idx_search_relations
ON search_index(from_id, to_id) WHERE type = 'relation';
-- Add index for type filtering
CREATE INDEX IF NOT EXISTS idx_search_type_perm
ON search_index(type);
""")
@@ -26,7 +26,7 @@ class SearchRepository:
def _quote_search_term(self, term: str) -> str:
"""Add quotes if term contains special characters."""
if any(c in term for c in '/-'):
if any(c in term for c in "/-"):
return f'"{term}"'
return term
@@ -112,12 +112,18 @@ class SearchRepository:
# Insert new record
await session.execute(
text("""
INSERT INTO search_index (
id, title, content, permalink, file_path, type, metadata
) VALUES (
:id, :title, :content, :permalink, :file_path, :type, :metadata
)
"""),
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
)
"""),
{
"id": id,
"title": title,
@@ -126,6 +132,14 @@ class SearchRepository:
"file_path": file_path,
"type": type.value,
"metadata": json.dumps(metadata),
# Optional fields based on type
"from_id": metadata.get("from_id"),
"to_id": metadata.get("to_id"),
"relation_type": metadata.get("relation_type"),
"entity_id": metadata.get("entity_id"),
"category": metadata.get("category"),
"created_at": metadata.get("created_at"),
"updated_at": metadata.get("updated_at")
},
)
logger.debug(f"indexed {permalink}")
+5 -6
View File
@@ -8,11 +8,11 @@ from pydantic import BaseModel, field_validator
class SearchItemType(str, Enum):
"""Types of searchable items."""
ENTITY = "entity"
OBSERVATION = "observation"
RELATION = "relation"
class SearchQuery(BaseModel):
"""Search query parameters."""
text: str
@@ -33,15 +33,15 @@ class SearchQuery(BaseModel):
class SearchResult(BaseModel):
"""Search result item."""
id: int
type: SearchItemType
score: float
metadata: dict
# Entity Observation-specific fields
# File-based fields (optional since observations/relations
# don't have their own files)
permalink: Optional[str] = None
file_path: Optional[str] = None
metadata: dict
# Observation-specific fields
entity_id: Optional[int] = None
@@ -55,5 +55,4 @@ class SearchResult(BaseModel):
class SearchResponse(BaseModel):
"""Wrapper for search results list."""
results: List[SearchResult]
results: List[SearchResult]
+36 -12
View File
@@ -73,11 +73,22 @@ class SearchService:
async def index_entity(
self, entity: Entity, background_tasks: Optional[BackgroundTasks] = None
) -> None:
"""Index an entity's content for search.
"""Index an entity and all its observations and relations.
Indexing structure:
1. Entities
- permalink: direct from entity (e.g., "specs/search")
- file_path: physical file location
2. Observations
- permalink: entity permalink + /observations/id (e.g., "specs/search/observations/123")
- file_path: parent entity's file (where observation is defined)
3. Relations (only index outgoing relations defined in this file)
- permalink: from_entity/relation_type/to_entity (e.g., "specs/search/implements/features/search-ui")
- file_path: source entity's file (where relation is defined)
Each type gets its own row in the search index with appropriate metadata
and type-specific fields populated.
- Content with context preservation
Each type gets its own row in the search index with appropriate metadata.
"""
content_parts = []
title_variants = self._generate_variants(entity.title)
@@ -122,15 +133,21 @@ class SearchService:
}
)
# Index each observation
# Index each observation with synthetic permalink
for obs in entity.observations:
# Create synthetic permalink for the observation
# We can construct these because observations are always
# defined in and owned by a single entity
observation_permalink = f"{entity.permalink}/observations/{obs.id}"
# Index with parent entity's file path since that's where it's defined
if background_tasks:
background_tasks.add_task(
self._do_index,
id=obs.id,
title=f"{obs.category}: {obs.content[:50]}...",
content=obs.content,
permalink=f"{entity.permalink}/observations/{obs.id}",
permalink=observation_permalink,
file_path=entity.file_path,
type=SearchItemType.OBSERVATION,
metadata={
@@ -146,7 +163,7 @@ class SearchService:
id=obs.id,
title=f"{obs.category}: {obs.content[:50]}...",
content=obs.content,
permalink=f"{entity.permalink}/observations/{obs.id}",
permalink=observation_permalink,
file_path=entity.file_path,
type=SearchItemType.OBSERVATION,
metadata={
@@ -158,15 +175,21 @@ class SearchService:
}
)
# Index each relation
for rel in entity.relations:
# Only index outgoing relations (ones defined in this file)
# Relations are indexed from the source entity's perspective since
# that's where they are defined in the markdown
for rel in entity.outgoing_relations:
# Create relation permalink showing the semantic connection:
# source/relation_type/target
# e.g., "specs/search/implements/features/search-ui"
relation_permalink = f"{rel.from_entity.permalink}/{rel.relation_type}/{rel.to_entity.permalink}"
if background_tasks:
background_tasks.add_task(
self._do_index,
id=rel.id,
title=f"{rel.relation_type}",
content=rel.context or "",
permalink=f"{entity.permalink}/relations/{rel.id}",
permalink=relation_permalink,
file_path=entity.file_path,
type=SearchItemType.RELATION,
metadata={
@@ -181,7 +204,7 @@ class SearchService:
id=rel.id,
title=f"{rel.relation_type}",
content=rel.context or "",
permalink=f"{entity.permalink}/relations/{rel.id}",
permalink=relation_permalink,
file_path=entity.file_path,
type=SearchItemType.RELATION,
metadata={
@@ -191,6 +214,7 @@ class SearchService:
"updated_at": rel.updated_at.isoformat()
}
)
async def _do_index(
self,
id: int,
@@ -214,4 +238,4 @@ class SearchService:
async def delete_by_permalink(self, path_id: str):
"""Delete an item from the search index."""
await self.repository.delete_by_permalink(path_id)
await self.repository.delete_by_permalink(path_id)
+2 -2
View File
@@ -400,10 +400,10 @@ async def test_full_knowledge_flow(client: AsyncClient):
assert len(main_entity["observations"]) == 2
assert len(main_entity["relations"]) == 2
# 6. Search should find all related entities
# 6. Search should find all related entities/relations/observations
search = await client.post("/search/", json={"text": "Related"})
matches = search.json()["results"]
assert len(matches) == 2
assert len(matches) == 4
# 7. Delete main entity
response = await client.post(
+9 -27
View File
@@ -11,38 +11,18 @@ from basic_memory.schemas import Entity as EntitySchema
from basic_memory.schemas.search import SearchItemType, SearchResponse
@pytest.fixture
def test_entity():
"""Create a test entity."""
class Entity:
id = 1
title = "TestComponent"
entity_type = "test"
entity_metadata = {"test": "test"}
permalink = "component/test_component"
file_path = "entities/component/test_component.md"
summary = "A test component for search testing"
content_type = "text/markdown"
created_at = datetime.now(timezone.utc)
updated_at = datetime.now(timezone.utc)
observations = []
relations = []
return Entity()
@pytest_asyncio.fixture
async def indexed_entity(init_search_index, test_entity, search_service):
async def indexed_entity(init_search_index, full_entity, search_service):
"""Create an entity and index it."""
await search_service.index_entity(test_entity)
return test_entity
await search_service.index_entity(full_entity)
return full_entity
@pytest.mark.asyncio
async def test_search_basic(client, indexed_entity):
"""Test basic text search."""
response = await client.post("/search/", json={"text": "test component"})
response = await client.post("/search/", json={"text": "searchable"})
assert response.status_code == 200
search_results = SearchResponse.model_validate(response.json())
assert len(search_results.results) == 1
@@ -62,7 +42,7 @@ async def test_search_with_type_filter(client, indexed_entity):
# Should not find with wrong type
response = await client.post(
"/search/", json={"text": "test", "types": [SearchItemType.DOCUMENT.value]}
"/search/", json={"text": "test", "types": [SearchItemType.RELATION.value]}
)
assert response.status_code == 200
search_results = SearchResponse.model_validate(response.json())
@@ -138,7 +118,7 @@ async def test_search_empty(search_service, client):
@pytest.mark.asyncio
async def test_reindex(client, search_service, entity_service, test_entity, session_maker):
async def test_reindex(client, search_service, entity_service, session_maker):
"""Test reindex endpoint."""
# Create test entity and document
await entity_service.create_entity(
@@ -168,7 +148,9 @@ async def test_reindex(client, search_service, entity_service, test_entity, sess
# Verify content is searchable again
search_response = await client.post("/search/", json={"text": "test"})
search_results = SearchResponse.model_validate(search_response.json())
assert len(search_results.results) == 1
assert len(search_results.results) == 2
@pytest.mark.asyncio
+27 -1
View File
@@ -13,7 +13,7 @@ from basic_memory.db import DatabaseType
from basic_memory.markdown import EntityParser
from basic_memory.markdown.knowledge_writer import KnowledgeWriter
from basic_memory.models import Base
from basic_memory.models.knowledge import Entity
from basic_memory.models.knowledge import Entity, Observation, ObservationCategory, Relation
from basic_memory.repository.entity_repository import EntityRepository
from basic_memory.repository.observation_repository import ObservationRepository
from basic_memory.repository.relation_repository import RelationRepository
@@ -226,3 +226,29 @@ async def sample_entity(entity_repository: EntityRepository) -> Entity:
"content_type": "text/markdown",
}
return await entity_repository.create(entity_data)
@pytest_asyncio.fixture
async def full_entity(sample_entity, entity_repository):
"""Create a search test entity."""
search_entity = await entity_repository.create({
"title": "Search Entity",
"entity_type": "test",
"summary": "A searchable entity",
"permalink": "test/search_entity",
"file_path": "test/search_entity.md",
"content_type": "text/markdown",
})
observations = [
Observation(content="Tech note", category=ObservationCategory.TECH),
Observation(content="Design note", category=ObservationCategory.DESIGN),
]
relations = [
Relation(from_id=search_entity.id, to_id=sample_entity.id, relation_type="out1"),
Relation(from_id=search_entity.id, to_id=sample_entity.id, relation_type="out2"),
]
search_entity.observations = observations
search_entity.outgoing_relations = relations
return await entity_repository.add(search_entity)
+16 -46
View File
@@ -1,14 +1,11 @@
"""Tests for search service."""
from datetime import datetime, timezone
import pytest
import pytest_asyncio
from sqlalchemy import text
from basic_memory import db
from basic_memory.models import Entity
from basic_memory.schemas import Entity as EntitySchema
from basic_memory.models import Entity, Observation, ObservationCategory, Relation
from basic_memory.schemas.search import SearchQuery, SearchItemType
@@ -55,9 +52,9 @@ async def test_entities(entity_repository):
summary="API documentation and examples",
file_path="docs/api/documentation.md",
content_type="text/markdown",
)
),
]
return await entity_repository.add_all(entities)
@@ -96,7 +93,7 @@ async def test_case_insensitive_search(indexed_search):
for search_text in test_cases:
results = await indexed_search.search(SearchQuery(text=search_text))
assert len(results) == 2, f"Failed for '{search_text}'"
file_paths = [r.file_path for r in results]
assert "components/core-service.md" in file_paths
assert "specs/features/core.md" in file_paths
@@ -106,9 +103,9 @@ async def test_case_insensitive_search(indexed_search):
async def test_whitespace_handling(indexed_search):
"""Test that whitespace is handled correctly."""
test_cases = [
" API ", # Extra spaces
"API Documentation", # Normal spacing
"API Documentation", # Multiple spaces
" API ", # Extra spaces
"API Documentation", # Normal spacing
"API Documentation", # Multiple spaces
]
for search_text in test_cases:
results = await indexed_search.search(SearchQuery(text=search_text))
@@ -133,11 +130,7 @@ async def test_search_filters(indexed_search):
"""Test search filtering."""
# Search with correct type filter
results = await indexed_search.search(
SearchQuery(
text="service",
types=[SearchItemType.ENTITY],
entity_types=["component"]
)
SearchQuery(text="service", types=[SearchItemType.ENTITY], entity_types=["component"])
)
assert len(results) == 2
@@ -147,34 +140,11 @@ async def test_search_filters(indexed_search):
# Search with non-matching type (should return empty)
results = await indexed_search.search(
SearchQuery(
text="service",
types=[SearchItemType.DOCUMENT]
)
SearchQuery(text="service", types=[SearchItemType.RELATION])
)
assert len(results) == 0
# Basic operation tests
@pytest_asyncio.fixture
def test_entity():
"""Create a basic test entity."""
class Entity:
id = 1
title = "TestComponent"
entity_type = "knowledge"
entity_metadata = {"test": "test"}
permalink = "component/test_component"
file_path = "entities/component/test_component.md"
summary = "A test component for search"
content_type = "text/markdown"
created_at = datetime.now(timezone.utc)
updated_at = datetime.now(timezone.utc)
observations = []
relations = []
return Entity()
@pytest.mark.asyncio
async def test_init_search_index(search_service, session_maker):
"""Test search index initialization."""
@@ -186,14 +156,14 @@ async def test_init_search_index(search_service, session_maker):
@pytest.mark.asyncio
async def test_update_index(search_service, test_entity):
async def test_update_index(search_service, full_entity):
"""Test updating indexed content."""
await search_service.index_entity(test_entity)
await search_service.index_entity(full_entity)
# Update entity
test_entity.summary = "Updated description with new terms"
await search_service.index_entity(test_entity)
full_entity.summary = "Updated description with new terms"
await search_service.index_entity(full_entity)
# Search for new terms
results = await search_service.search(SearchQuery(text="new terms"))
assert len(results) == 1
assert len(results) == 1