diff --git a/src/basic_memory/services/context_service.py b/src/basic_memory/services/context_service.py index 705f0183..d553ad0a 100644 --- a/src/basic_memory/services/context_service.py +++ b/src/basic_memory/services/context_service.py @@ -1,6 +1,6 @@ """Service for building rich context from the knowledge graph.""" -from datetime import datetime +from datetime import datetime, UTC, timezone from typing import List, Optional, Tuple from loguru import logger from sqlalchemy import text @@ -8,106 +8,112 @@ from sqlalchemy import text from basic_memory.repository.search_repository import SearchRepository from basic_memory.repository.entity_repository import EntityRepository from basic_memory.schemas.memory_url import MemoryUrl -from basic_memory.schemas.search import SearchItemType, SearchQuery +from basic_memory.schemas.search import SearchQuery, SearchItemType class ContextService: - """Service for building rich context from memory:// URIs.""" - + """Service for building rich context from memory:// URIs. + + Handles three types of context building: + 1. Direct permalink lookup - exact match on path + 2. Pattern matching - using * wildcards + 3. Special modes via params (e.g., 'related') + """ + def __init__( - self, - search_repository: SearchRepository, - entity_repository: EntityRepository, + self, + search_repository: SearchRepository, + entity_repository: EntityRepository, ): self.search_repository = search_repository self.entity_repository = entity_repository async def build_context( - self, - uri: str, - depth: int = 2, - since: Optional[datetime] = None, + self, + uri: str, + depth: int = 2, + since: Optional[datetime] = None, ): """Build rich context from a memory:// URI.""" logger.debug(f"Building context for URI {uri}") - + # Parse the URI memory_url = MemoryUrl.parse(uri) - - # Handle different URL types - if memory_url.pattern: - # Pattern matching (*) + + # Find primary entities based on URL type + if memory_url.params.get("type") == "related": + # Special mode for finding related content + target = memory_url.params["target"] + primary = await self.find_related(target) + elif memory_url.pattern: + # Pattern matching with * primary = await self.find_by_pattern(memory_url.pattern) - elif memory_url.fuzzy: - # Fuzzy search (~) - primary = await self.find_by_fuzzy(memory_url.fuzzy) - elif memory_url.params.get("type") == "related": - # Related content - primary = await self.find_related(memory_url.params["target"]) else: # Direct permalink lookup primary = await self.find_by_permalink(memory_url.relative_path()) - # Get type_id pairs from primary results + # Get type_id pairs for traversal type_id_pairs = [(r.type, r.id) for r in primary] if primary else [] - # Find connected entities through relations + # Find connected content related = await self.find_connected( type_id_pairs, max_depth=depth, since=since ) - + + # Build response return { - "primary": primary, - "related": related + "primary_entities": primary, + "related_entities": related, + "metadata": { + "uri": uri, + "depth": depth, + "timeframe": since.isoformat() if since else None, + "generated_at": datetime.now(timezone.utc).isoformat(), + "matched_entities": len(primary), + "total_entities": len(primary) + len(related), + "total_relations": sum(1 for r in related if r.type == SearchItemType.RELATION) + } } async def find_by_pattern(self, pattern: str): - """Find entities matching a glob pattern.""" - # Convert glob pattern to SQL LIKE pattern - # Use search with permalink pattern - query = SearchQuery( - permalink=pattern, - ) - + """Find entities matching a pattern with * wildcards.""" + query = SearchQuery(permalink_pattern=pattern) return await self.search_repository.search(query) - async def find_by_fuzzy(self, search_terms: str): - """Find entities using fuzzy text search.""" - query = SearchQuery( - text=search_terms, - ) - return await self.search_repository.search(query) - - async def find_related(self, permalink: str): - """Find entities related to a given permalink.""" - # First find the target entity - query = SearchQuery(permalink=permalink) - results = await self.search_repository.search(query) - - if not results: - return [] - - # Use find_connected to get related items - entity = results[0] - return await self.find_connected( - [(entity.type, entity.id)], - max_depth=1 # Only immediate relations - ) - async def find_by_permalink(self, permalink: str): """Find an entity by exact permalink.""" query = SearchQuery(permalink=permalink) return await self.search_repository.search(query) + async def find_related(self, permalink: str): + """Find entities related to a given permalink.""" + # First find the target entity + target = await self.find_by_permalink(permalink) + if not target: + return [] + + # Use find_connected to get related items + type_id_pairs = [(r.type.value, r.id) for r in target] + return await self.find_connected( + type_id_pairs, + max_depth=1 # Only immediate relations + ) + async def find_connected( - self, - type_id_pairs: List[Tuple[str, int]], - max_depth: int = 2, - since: Optional[datetime] = None, + self, + type_id_pairs: List[Tuple[str, int]], + max_depth: int = 2, + since: Optional[datetime] = None, ): - """Find items connected through relations.""" + """Find items connected through relations. + + Uses recursive CTE to find: + - Connected entities + - Their observations + - Relations that connect them + """ if not type_id_pairs: return [] diff --git a/tests/services/test_context_service.py b/tests/services/test_context_service.py index 8f9fcd6b..fb4b336a 100644 --- a/tests/services/test_context_service.py +++ b/tests/services/test_context_service.py @@ -1,12 +1,13 @@ """Tests for context service.""" +from datetime import datetime, timedelta, UTC + import pytest import pytest_asyncio -from datetime import datetime, timedelta, UTC from basic_memory.models import Entity, Relation, Observation, ObservationCategory -from basic_memory.services.context_service import ContextService from basic_memory.schemas.search import SearchItemType +from basic_memory.services.context_service import ContextService @pytest_asyncio.fixture @@ -29,7 +30,7 @@ async def test_graph(entity_repository, search_service): ), Entity( title="Connected Entity 1", - entity_type="test", + entity_type="test", permalink="test/connected1", file_path="test/connected1.md", content_type="text/markdown", @@ -37,7 +38,7 @@ async def test_graph(entity_repository, search_service): Entity( title="Connected Entity 2", entity_type="test", - permalink="test/connected2", + permalink="test/connected2", file_path="test/connected2.md", content_type="text/markdown", ), @@ -47,7 +48,7 @@ async def test_graph(entity_repository, search_service): permalink="test/deep", file_path="test/deep.md", content_type="text/markdown", - ) + ), ] entities = await entity_repository.add_all(entities) root, conn1, conn2, deep = entities @@ -55,7 +56,7 @@ async def test_graph(entity_repository, search_service): # Add some observations root.observations = [ Observation(content="Root note 1", category=ObservationCategory.NOTE), - Observation(content="Root tech note", category=ObservationCategory.TECH) + Observation(content="Root tech note", category=ObservationCategory.TECH), ] conn1.observations = [ @@ -67,9 +68,8 @@ async def test_graph(entity_repository, search_service): # Direct connections to root Relation(from_id=root.id, to_id=conn1.id, relation_type="connects_to"), Relation(from_id=conn2.id, to_id=root.id, relation_type="connected_from"), - # Deep connection - Relation(from_id=conn1.id, to_id=deep.id, relation_type="deep_connection") + Relation(from_id=conn1.id, to_id=deep.id, relation_type="deep_connection"), ] root.outgoing_relations = [relations[0]] @@ -86,42 +86,64 @@ async def test_graph(entity_repository, search_service): await search_service.index_entity(entity) return { - 'root': root, - 'connected1': conn1, - 'connected2': conn2, - 'deep': deep, - 'observations': root.observations + conn1.observations, - 'relations': relations + "root": root, + "connected1": conn1, + "connected2": conn2, + "deep": deep, + "observations": root.observations + conn1.observations, + "relations": relations, } +@pytest.mark.asyncio +async def test_find_by_pattern(context_service, test_graph): + """Test pattern matching.""" + results = await context_service.find_by_pattern("test/*") + assert len(results) > 0 + assert all("test/" in r.permalink for r in results) + + +@pytest.mark.asyncio +async def test_find_by_permalink(context_service, test_graph): + """Test exact permalink lookup.""" + results = await context_service.find_by_permalink("test/root") + assert len(results) == 1 + assert results[0].permalink == "test/root" + + +@pytest.mark.asyncio +async def test_find_related(context_service, test_graph): + """Test finding related content.""" + results = await context_service.find_related("test/root") + # Should get immediate connections + assert any("connected1" in r.permalink for r in results) + assert any("connected2" in r.permalink for r in results) + + @pytest.mark.asyncio async def test_find_connected_basic(context_service, test_graph, search_service): """Test basic connectivity traversal.""" # Start with root entity and one of its observations type_id_pairs = [ - ('entity', test_graph['root'].id), - ('observation', test_graph['observations'][0].id) + ("entity", test_graph["root"].id), + ("observation", test_graph["observations"][0].id), ] - + results = await context_service.find_connected(type_id_pairs) # Verify types types_found = {r.type for r in results} - assert 'entity' in types_found - assert 'relation' in types_found - assert 'observation' in types_found + assert "entity" in types_found + assert "relation" in types_found + assert "observation" in types_found # Verify we found directly connected entities - entity_ids = {r.id for r in results if r.type == 'entity'} - assert test_graph['connected1'].id in entity_ids - assert test_graph['connected2'].id in entity_ids + entity_ids = {r.id for r in results if r.type == "entity"} + assert test_graph["connected1"].id in entity_ids + assert test_graph["connected2"].id in entity_ids # Verify we found observations - assert any( - r.type == 'observation' and "Root note 1" in r.content - for r in results - ) + assert any(r.type == "observation" and "Root note 1" in r.content for r in results) @pytest.mark.asyncio @@ -132,34 +154,22 @@ async def test_find_connected_depth_limit(context_service, test_graph): - Depth 1: Relations + directly connected entities (Connected1, Connected2) - Depth 2: Relations + next level entities (Deep) """ - type_id_pairs = [('entity', test_graph['root'].id)] + type_id_pairs = [("entity", test_graph["root"].id)] # With depth=1, we get direct connections - shallow_results = await context_service.find_connected( - type_id_pairs, - max_depth=1 - ) - shallow_entities = { - (r.id, r.type) for r in shallow_results - if r.type == 'entity' - } + shallow_results = await context_service.find_connected(type_id_pairs, max_depth=1) + shallow_entities = {(r.id, r.type) for r in shallow_results if r.type == "entity"} # Should find Connected1 and Connected2 - assert (test_graph['connected1'].id, 'entity') in shallow_entities - assert (test_graph['connected2'].id, 'entity') in shallow_entities + assert (test_graph["connected1"].id, "entity") in shallow_entities + assert (test_graph["connected2"].id, "entity") in shallow_entities # But not Deep entity - assert (test_graph['deep'].id, 'entity') not in shallow_entities + assert (test_graph["deep"].id, "entity") not in shallow_entities # With depth=2, we get the next level - deep_results = await context_service.find_connected( - type_id_pairs, - max_depth=2 - ) - deep_entities = { - (r.id, r.type) for r in deep_results - if r.type == 'entity' - } + deep_results = await context_service.find_connected(type_id_pairs, max_depth=2) + deep_entities = {(r.id, r.type) for r in deep_results if r.type == "entity"} # Should now include Deep entity - assert (test_graph['deep'].id, 'entity') in deep_entities + assert (test_graph["deep"].id, "entity") in deep_entities @pytest.mark.asyncio @@ -176,48 +186,86 @@ async def test_find_connected_timeframe(context_service, test_graph, search_repo # Index root and its relation as old await search_repository.index_item( - id=test_graph['root'].id, - title=test_graph['root'].title, + id=test_graph["root"].id, + title=test_graph["root"].title, content="Root content", - permalink=test_graph['root'].permalink, - file_path=test_graph['root'].file_path, + permalink=test_graph["root"].permalink, + file_path=test_graph["root"].file_path, type=SearchItemType.ENTITY, metadata={"created_at": old_date.isoformat()}, ) await search_repository.index_item( - id=test_graph['relations'][0].id, + id=test_graph["relations"][0].id, title="Root Entity → Connected Entity 1", content="", permalink=f"{test_graph['root'].permalink}/connects_to/{test_graph['connected1'].permalink}", - file_path=test_graph['root'].file_path, + file_path=test_graph["root"].file_path, type=SearchItemType.RELATION, - from_id=test_graph['root'].id, - to_id=test_graph['connected1'].id, - relation_type='connects_to', + from_id=test_graph["root"].id, + to_id=test_graph["connected1"].id, + relation_type="connects_to", metadata={"created_at": old_date.isoformat()}, ) - + # Index connected1 as recent await search_repository.index_item( - id=test_graph['connected1'].id, - title=test_graph['connected1'].title, + id=test_graph["connected1"].id, + title=test_graph["connected1"].title, content="Connected 1 content", - permalink=test_graph['connected1'].permalink, - file_path=test_graph['connected1'].file_path, + permalink=test_graph["connected1"].permalink, + file_path=test_graph["connected1"].file_path, type=SearchItemType.ENTITY, metadata={"created_at": recent_date.isoformat()}, ) - type_id_pairs = [('entity', test_graph['root'].id)] + type_id_pairs = [("entity", test_graph["root"].id)] # Search with a 7-day cutoff since_date = now - timedelta(days=7) - results = await context_service.find_connected( - type_id_pairs, - since=since_date - ) + results = await context_service.find_connected(type_id_pairs, since=since_date) - # Only connected1 is recent, but we can't get to it + # Only connected1 is recent, but we can't get to it # because its connecting relation is too old - entity_ids = {r.id for r in results if r.type == 'entity'} - assert len(entity_ids) == 0 # No accessible entities within timeframe \ No newline at end of file + entity_ids = {r.id for r in results if r.type == "entity"} + assert len(entity_ids) == 0 # No accessible entities within timeframe + +@pytest.mark.asyncio +async def test_build_context_pattern(context_service, test_graph): + """Test building context from pattern.""" + context = await context_service.build_context("memory://project/test/*") + assert len(context["primary_entities"]) > 0 + assert "uri" in context["metadata"] + assert "total_entities" in context["metadata"] + +@pytest.mark.asyncio +async def test_build_context_related(context_service, test_graph): + """Test building context from related mode.""" + context = await context_service.build_context( + "memory://project/related/test/root" + ) + assert len(context["primary_entities"]) > 0 + assert len(context["related_entities"]) > 0 + + +@pytest.mark.asyncio +async def test_build_context_not_found(context_service): + """Test handling non-existent permalinks.""" + context = await context_service.build_context( + "memory://project/does/not/exist" + ) + assert len(context["primary_entities"]) == 0 + assert len(context["related_entities"]) == 0 + + +@pytest.mark.asyncio +async def test_context_metadata(context_service, test_graph): + """Test metadata is correctly populated.""" + context = await context_service.build_context( + "memory://project/test/root", + depth=2 + ) + metadata = context["metadata"] + assert metadata["uri"] == "memory://project/test/root" + assert metadata["depth"] == 2 + assert metadata["generated_at"] is not None + assert metadata["matched_entities"] > 0 \ No newline at end of file