diff --git a/src/basic_memory/mcp/tools/memory.py b/src/basic_memory/mcp/tools/memory.py index 7a63dcde..c40f3f19 100644 --- a/src/basic_memory/mcp/tools/memory.py +++ b/src/basic_memory/mcp/tools/memory.py @@ -1,6 +1,6 @@ """Discussion context tools for Basic Memory MCP server.""" -from typing import Optional, List +from typing import Optional, Literal from loguru import logger @@ -8,7 +8,6 @@ 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.search import SearchItemType from basic_memory.schemas.base import TimeFrame @@ -84,7 +83,7 @@ async def build_context( """, ) async def recent_activity( - type: List[SearchItemType] = None, + type: Literal["entity", "observation", "relation"] = None, depth: Optional[int] = 1, timeframe: Optional[TimeFrame] = "7d", max_results: int = 10, @@ -128,9 +127,16 @@ async def recent_activity( logger.info( f"Getting recent activity from {type}, depth={depth}, timeframe={timeframe}, max_results={max_results}" ) + params = { + "depth": depth, + "timeframe": timeframe, + "max_results": max_results, + "type": type if type else None, + } + response = await call_get( client, "/memory/recent", - params={"depth": depth, "timeframe": timeframe, "max_results": max_results, "type": type}, + params=params, ) return GraphContext.model_validate(response.json()) diff --git a/tests/mcp/test_tool_memory.py b/tests/mcp/test_tool_memory.py index 92c8a1ba..ddcd0ec2 100644 --- a/tests/mcp/test_tool_memory.py +++ b/tests/mcp/test_tool_memory.py @@ -1,13 +1,18 @@ """Tests for discussion context MCP tool.""" + import pytest from datetime import datetime -from httpx import HTTPStatusError from mcp.server.fastmcp.exceptions import ToolError from basic_memory.mcp.tools.memory import build_context, recent_activity -from basic_memory.schemas.base import TimeFrame -from basic_memory.schemas.memory import GraphContext, MemoryUrl +from basic_memory.schemas.memory import ( + GraphContext, + MemoryUrl, + EntitySummary, + ObservationSummary, + RelationSummary, +) @pytest.mark.asyncio @@ -85,11 +90,7 @@ async def test_recent_activity_timeframe_formats(client, test_graph): # Test each valid timeframe for timeframe in valid_timeframes: try: - result = await recent_activity( - type=["entity"], - timeframe=timeframe, - max_results=1 - ) + result = await recent_activity(type=["entity"], timeframe=timeframe, max_results=1) assert result is not None except Exception as e: pytest.fail(f"Failed with valid timeframe '{timeframe}': {str(e)}") @@ -100,6 +101,34 @@ async def test_recent_activity_timeframe_formats(client, test_graph): await recent_activity(timeframe=timeframe) +@pytest.mark.asyncio +async def test_recent_activity_type_filters(client, test_graph): + """Test that recent_activity correctly filters by types.""" + # Test single type + result = await recent_activity(type=["entity"]) + assert result is not None + assert all(isinstance(r, EntitySummary) for r in result.primary_results) + + # Test multiple types + result = await recent_activity(type=["entity", "observation"]) + assert result is not None + assert all( + isinstance(r, EntitySummary) or isinstance(r, ObservationSummary) + for r in result.primary_results + ) + + # Test all types + result = await recent_activity(type=["entity", "observation", "relation"]) + assert result is not None + # Results can be any type + assert all( + isinstance(r, EntitySummary) + or isinstance(r, ObservationSummary) + or isinstance(r, RelationSummary) + for r in result.primary_results + ) + + @pytest.mark.asyncio async def test_build_context_timeframe_formats(client, test_graph): """Test that build_context accepts various timeframe formats.""" @@ -108,11 +137,7 @@ async def test_build_context_timeframe_formats(client, test_graph): # Test each valid timeframe for timeframe in valid_timeframes: try: - result = await build_context( - url=test_url, - timeframe=timeframe, - max_results=1 - ) + result = await build_context(url=test_url, timeframe=timeframe, max_results=1) assert result is not None except Exception as e: pytest.fail(f"Failed with valid timeframe '{timeframe}': {str(e)}") @@ -121,5 +146,3 @@ async def test_build_context_timeframe_formats(client, test_graph): for timeframe in invalid_timeframes: with pytest.raises(ToolError): await build_context(url=test_url, timeframe=timeframe) - -