From 3199ab0217bf0258b8c8b2973fe70a3cd04f1aea Mon Sep 17 00:00:00 2001 From: phernandez Date: Sat, 1 Feb 2025 13:42:48 -0600 Subject: [PATCH] simplify MemoryUrl to be validated string --- src/basic_memory/api/app.py | 2 +- src/basic_memory/api/routers/memory_router.py | 7 +- src/basic_memory/mcp/tools/memory.py | 6 +- src/basic_memory/schemas/memory.py | 87 ++++++++++--------- src/basic_memory/services/context_service.py | 17 ++-- tests/mcp/test_tool_memory.py | 3 +- tests/mcp/test_tool_notes.py | 1 + tests/schemas/test_memory_url.py | 35 ++++---- tests/services/test_context_service.py | 18 ++-- 9 files changed, 89 insertions(+), 87 deletions(-) diff --git a/src/basic_memory/api/app.py b/src/basic_memory/api/app.py index a3808b13..e52c0b8c 100644 --- a/src/basic_memory/api/app.py +++ b/src/basic_memory/api/app.py @@ -39,4 +39,4 @@ async def exception_handler(request, exc): logger.exception( f"An unhandled exception occurred for request '{request.url}', exception: {exc}" ) - return await http_exception_handler(request, HTTPException(status_code=500, detail=exc.args[0])) + return await http_exception_handler(request, HTTPException(status_code=500, detail=str(exc))) diff --git a/src/basic_memory/api/routers/memory_router.py b/src/basic_memory/api/routers/memory_router.py index 0caf0e43..07dd3156 100644 --- a/src/basic_memory/api/routers/memory_router.py +++ b/src/basic_memory/api/routers/memory_router.py @@ -13,12 +13,11 @@ from basic_memory.repository import EntityRepository from basic_memory.repository.search_repository import SearchIndexRow from basic_memory.schemas.base import TimeFrame from basic_memory.schemas.memory import ( - MemoryUrl, GraphContext, RelationSummary, EntitySummary, ObservationSummary, - MemoryMetadata, + MemoryMetadata, normalize_memory_url, ) from basic_memory.schemas.search import SearchItemType from basic_memory.services.context_service import ContextResultRow @@ -50,7 +49,7 @@ async def to_graph_context(context, entity_repository: EntityRepository): permalink=item.permalink, type=item.type, from_id=from_entity.permalink, - to_id=to_entity.permalink, + to_id=to_entity.permalink if to_entity else None, ) primary_results = [await to_summary(r) for r in context["primary_results"]] @@ -109,7 +108,7 @@ async def get_memory_context( logger.debug( f"Getting context for URI: `{uri}` depth: `{depth}` timeframe: `{timeframe}` max_results: `{max_results}`" ) - memory_url = MemoryUrl(f"memory://{uri}") + memory_url = normalize_memory_url(uri) # Parse timeframe since = parse(timeframe) diff --git a/src/basic_memory/mcp/tools/memory.py b/src/basic_memory/mcp/tools/memory.py index c40f3f19..84b56c7e 100644 --- a/src/basic_memory/mcp/tools/memory.py +++ b/src/basic_memory/mcp/tools/memory.py @@ -7,7 +7,7 @@ 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_get -from basic_memory.schemas.memory import GraphContext, MemoryUrl +from basic_memory.schemas.memory import GraphContext, MemoryUrl, memory_url, memory_url_path, normalize_memory_url from basic_memory.schemas.base import TimeFrame @@ -61,10 +61,10 @@ async def build_context( build_context("memory://features/knowledge-graph", timeframe="3 months ago") """ logger.info(f"Building context from {url}") - memory_url = MemoryUrl.validate(url) + url = normalize_memory_url(url) response = await call_get( client, - f"/memory/{memory_url.relative_path()}", + f"/memory/{memory_url_path(url)}", params={"depth": depth, "timeframe": timeframe, "max_results": max_results}, ) return GraphContext.model_validate(response.json()) diff --git a/src/basic_memory/schemas/memory.py b/src/basic_memory/schemas/memory.py index 7fd0488d..2865f849 100644 --- a/src/basic_memory/schemas/memory.py +++ b/src/basic_memory/schemas/memory.py @@ -1,54 +1,59 @@ """Schemas for memory context.""" from datetime import datetime -from typing import Dict, List, Any, Optional +from typing import Dict, List, Any, Optional, Annotated -from pydantic import BaseModel, field_validator, Field +import pydantic +from annotated_types import MinLen, MaxLen +from pydantic import BaseModel, field_validator, Field, BeforeValidator, TypeAdapter, AnyUrl from basic_memory.schemas.search import SearchItemType -class MemoryUrl(BaseModel): - """memory:// URL scheme for knowledge addressing. - - Example URLs: - memory://specs/search # Direct reference - memory://specs/search/* # Pattern matching - memory://related/xyz # Special lookup +def normalize_memory_url(url: str) -> str: + """Normalize a MemoryUrl string. + + Args: + url: A path like "specs/search" or "memory://specs/search" + + Returns: + Normalized URL starting with memory:// + + Examples: + >>> normalize_memory_url("specs/search") + 'memory://specs/search' + >>> normalize_memory_url("memory://specs/search") + 'memory://specs/search' """ - url: str - path: str = "" - - @field_validator("url") - @classmethod - def validate_url(cls, v: str) -> str: - """Validate the URL starts with memory://.""" - if isinstance(v, MemoryUrl): - return v.url - if not isinstance(v, str): - raise ValueError(f"URL must be a string, got {type(v)}") - if not v.startswith("memory://"): - raise ValueError(f"Invalid memory URL: {v}. Must start with memory://") - return v - - def __init__(self, url: str, **data): - if isinstance(url, MemoryUrl): - url = url.url - super().__init__(url=url, **data) - self.path = url.removeprefix("memory://") - - @classmethod - def validate(cls, url: str) -> "MemoryUrl": - """Validate and construct a MemoryUrl.""" - return cls(url=url) + clean_path = url.removeprefix("memory://") + return f"memory://{clean_path}" - def relative_path(self) -> str: - """Get the path.""" - return self.path - def __str__(self) -> str: - """Convert back to URL string.""" - return self.url +MemoryUrl = Annotated[ + str, + BeforeValidator(str.strip), # Clean whitespace + MinLen(1), + MaxLen(2028), +] + +memory_url = TypeAdapter(MemoryUrl) + +def memory_url_path(url: memory_url) -> str: + """ + Returns the uri for a url value by removing the prefix "memory://" from a given MemoryUrl. + + This function processes a given MemoryUrl by removing the "memory://" + prefix and returns the resulting string. If the provided url does not + begin with "memory://", the function will simply return the input url + unchanged. + + :param url: A MemoryUrl object representing the URL with a "memory://" prefix. + :type url: MemoryUrl + :return: A string representing the URL with the "memory://" prefix removed. + :rtype: str + """ + return url.removeprefix("memory://") + class EntitySummary(BaseModel): @@ -66,7 +71,7 @@ class RelationSummary(BaseModel): permalink: str type: str from_id: str - to_id: str + to_id: Optional[str] = None class ObservationSummary(BaseModel): diff --git a/src/basic_memory/services/context_service.py b/src/basic_memory/services/context_service.py index 5bae8766..a683bcd5 100644 --- a/src/basic_memory/services/context_service.py +++ b/src/basic_memory/services/context_service.py @@ -9,7 +9,7 @@ from sqlalchemy import text from basic_memory.repository.entity_repository import EntityRepository from basic_memory.repository.search_repository import SearchRepository -from basic_memory.schemas.memory import MemoryUrl +from basic_memory.schemas.memory import MemoryUrl, memory_url_path from basic_memory.schemas.search import SearchItemType @@ -62,17 +62,18 @@ class ContextService: ) if memory_url: + path = memory_url_path(memory_url) # Pattern matching - use search - if "*" in memory_url.relative_path(): - logger.debug(f"Pattern search for '{memory_url.relative_path()}'") + if "*" in path: + logger.debug(f"Pattern search for '{path}'") primary = await self.search_repository.search( - permalink_match=memory_url.relative_path() + permalink_match=path ) - + # Direct lookup for exact path else: - logger.debug(f"Direct lookup for '{memory_url.relative_path()}'") - primary = await self.search_repository.search(permalink=memory_url.relative_path()) + logger.debug(f"Direct lookup for '{path}'") + primary = await self.search_repository.search(permalink=path) else: logger.debug(f"Build context for '{types}'") primary = await self.search_repository.search(types=types) @@ -95,7 +96,7 @@ class ContextService: "primary_results": primary, "related_results": related, "metadata": { - "uri": memory_url.relative_path() if memory_url else None, + "uri": memory_url_path(memory_url) if memory_url else None, "types": types if types else None, "depth": depth, "timeframe": since.isoformat() if since else None, diff --git a/tests/mcp/test_tool_memory.py b/tests/mcp/test_tool_memory.py index ddcd0ec2..496503b1 100644 --- a/tests/mcp/test_tool_memory.py +++ b/tests/mcp/test_tool_memory.py @@ -8,7 +8,6 @@ from mcp.server.fastmcp.exceptions import ToolError from basic_memory.mcp.tools.memory import build_context, recent_activity from basic_memory.schemas.memory import ( GraphContext, - MemoryUrl, EntitySummary, ObservationSummary, RelationSummary, @@ -132,7 +131,7 @@ async def test_recent_activity_type_filters(client, test_graph): @pytest.mark.asyncio async def test_build_context_timeframe_formats(client, test_graph): """Test that build_context accepts various timeframe formats.""" - test_url = MemoryUrl.validate("memory://specs/test") + test_url = "memory://specs/test" # Test each valid timeframe for timeframe in valid_timeframes: diff --git a/tests/mcp/test_tool_notes.py b/tests/mcp/test_tool_notes.py index 4bc7f056..5d9ccbb2 100644 --- a/tests/mcp/test_tool_notes.py +++ b/tests/mcp/test_tool_notes.py @@ -7,6 +7,7 @@ from basic_memory.mcp.tools import notes from basic_memory.schemas import EntityResponse + @pytest.mark.asyncio async def test_write_note(app): """Test creating a new note. diff --git a/tests/schemas/test_memory_url.py b/tests/schemas/test_memory_url.py index 086c2b96..eee33234 100644 --- a/tests/schemas/test_memory_url.py +++ b/tests/schemas/test_memory_url.py @@ -1,53 +1,48 @@ """Tests for MemoryUrl parsing.""" import pytest -from basic_memory.schemas.memory import MemoryUrl +from basic_memory.schemas.memory import MemoryUrl, memory_url, memory_url_path def test_basic_permalink(): """Test basic permalink parsing.""" - url = MemoryUrl.validate("memory://specs/search") + url = memory_url.validate_strings("memory://specs/search") assert str(url) == "memory://specs/search" - assert url.path == "specs/search" + assert memory_url_path(url) == "specs/search" def test_glob_pattern(): """Test pattern matching.""" - url = MemoryUrl.validate("memory://specs/search/*") - assert url.path == "specs/search/*" + url = memory_url.validate_python("memory://specs/search/*") + assert memory_url_path(url) == "specs/search/*" def test_related_prefix(): """Test related content prefix.""" - url = MemoryUrl.validate("memory://related/specs/search") - assert url.path == "related/specs/search" + url = memory_url.validate_python("memory://related/specs/search") + assert memory_url_path(url) == "related/specs/search" def test_context_prefix(): """Test context prefix.""" - url = MemoryUrl.validate("memory://context/current") - assert url.path == "context/current" + url = memory_url.validate_python("memory://context/current") + assert memory_url_path(url) == "context/current" def test_complex_pattern(): """Test multiple glob patterns.""" - url = MemoryUrl.validate("memory://specs/*/search/*") - assert url.path == "specs/*/search/*" + url = memory_url.validate_python("memory://specs/*/search/*") + assert memory_url_path(url) == "specs/*/search/*" def test_path_with_dashes(): """Test path with dashes and other chars.""" - url = MemoryUrl.validate("memory://file-sync-and-note-updates-implementation") - assert url.path == "file-sync-and-note-updates-implementation" + url = memory_url.validate_python("memory://file-sync-and-note-updates-implementation") + assert memory_url_path(url) == "file-sync-and-note-updates-implementation" - -def test_invalid_url(): - """Test URL must start with memory://.""" - with pytest.raises(ValueError, match="Invalid memory URL"): - MemoryUrl.validate("http://specs/search") def test_str_representation(): """Test converting back to string.""" - url = MemoryUrl.validate("memory://specs/search") - assert str(url) == "memory://specs/search" \ No newline at end of file + url = memory_url.validate_python("memory://specs/search") + assert url == "memory://specs/search" \ No newline at end of file diff --git a/tests/services/test_context_service.py b/tests/services/test_context_service.py index 1441f924..4cd06fec 100644 --- a/tests/services/test_context_service.py +++ b/tests/services/test_context_service.py @@ -5,9 +5,8 @@ from datetime import datetime, timedelta, UTC import pytest import pytest_asyncio -from basic_memory.config import config from basic_memory.repository.search_repository import SearchIndexRow -from basic_memory.schemas.memory import MemoryUrl +from basic_memory.schemas.memory import memory_url, memory_url_path from basic_memory.schemas.search import SearchItemType from basic_memory.services.context_service import ContextService @@ -39,7 +38,10 @@ async def test_find_connected_basic(context_service, test_graph, search_service) assert test_graph["connected1"].id in entity_ids # Verify we found relation - assert any(r.type == "relation" and "Connected Entity 1 → Connected Entity 2" in r.title for r in results) + assert any( + r.type == "relation" and "Connected Entity 1 → Connected Entity 2" in r.title + for r in results + ) @pytest.mark.asyncio @@ -137,14 +139,14 @@ async def test_find_connected_timeframe(context_service, test_graph, search_repo @pytest.mark.asyncio async def test_build_context(context_service, test_graph): """Test exact permalink lookup.""" - url = MemoryUrl(f"memory://test/root") + url = memory_url.validate_strings("memory://test/root") results = await context_service.build_context(url) matched_results = results["metadata"]["matched_results"] primary_results = results["primary_results"] related_results = results["related_results"] total_results = results["metadata"]["total_results"] - assert results["metadata"]["uri"] == url.relative_path() + assert results["metadata"]["uri"] == memory_url_path(url) assert results["metadata"]["depth"] == 1 assert matched_results == 1 assert len(primary_results) == 1 @@ -155,7 +157,7 @@ async def test_build_context(context_service, test_graph): @pytest.mark.asyncio async def test_build_context_pattern(context_service, test_graph): """Test exact permalink lookup.""" - url = MemoryUrl("memory://test/connected*") + url = memory_url.validate_strings("memory://test/connected*") results = await context_service.build_context(url) matched_results = results["metadata"]["matched_results"] primary_results = results["primary_results"] @@ -168,7 +170,7 @@ async def test_build_context_pattern(context_service, test_graph): @pytest.mark.asyncio async def test_build_context_not_found(context_service): """Test handling non-existent permalinks.""" - context = await context_service.build_context(MemoryUrl("memory://does/not/exist")) + context = await context_service.build_context("memory://does/not/exist") assert len(context["primary_results"]) == 0 assert len(context["related_results"]) == 0 @@ -176,7 +178,7 @@ async def test_build_context_not_found(context_service): @pytest.mark.asyncio async def test_context_metadata(context_service, test_graph): """Test metadata is correctly populated.""" - context = await context_service.build_context(MemoryUrl("memory://test/root"), depth=2) + context = await context_service.build_context("memory://test/root", depth=2) metadata = context["metadata"] assert metadata["uri"] == "test/root" assert metadata["depth"] == 2