diff --git a/pyproject.toml b/pyproject.toml index 6e60505e..594ac814 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -29,7 +29,7 @@ dependencies = [ "fastapi[standard]>=0.115.8", "alembic>=1.14.1", "qasync>=0.27.1", - "logfire[fastapi,sqlalchemy,sqlite3]>=3.6.0", + "logfire[fastapi,httpx,sqlalchemy,sqlite3]>=3.6.0", ] diff --git a/src/basic_memory/cli/commands/mcp.py b/src/basic_memory/cli/commands/mcp.py index 79cc3d7d..8bd53150 100644 --- a/src/basic_memory/cli/commands/mcp.py +++ b/src/basic_memory/cli/commands/mcp.py @@ -17,4 +17,4 @@ def mcp(): # pragma: no cover home_dir = config.home logger.info(f"Starting Basic Memory MCP server {basic_memory.__version__}") logger.info(f"Home directory: {home_dir}") - mcp_server.run() + mcp_server.run() \ No newline at end of file diff --git a/src/basic_memory/config.py b/src/basic_memory/config.py index e583510e..6e74b56a 100644 --- a/src/basic_memory/config.py +++ b/src/basic_memory/config.py @@ -9,7 +9,7 @@ from pydantic_settings import BaseSettings, SettingsConfigDict DATABASE_NAME = "memory.db" DATA_DIR_NAME = ".basic-memory" -Environment = Literal["test", "dev", "prod"] +Environment = Literal["test", "dev", "user"] class ProjectConfig(BaseSettings): diff --git a/src/basic_memory/db.py b/src/basic_memory/db.py index edb70889..bdb988b1 100644 --- a/src/basic_memory/db.py +++ b/src/basic_memory/db.py @@ -4,6 +4,8 @@ from enum import Enum, auto from pathlib import Path from typing import AsyncGenerator, Optional +import logfire + from basic_memory.config import ProjectConfig from alembic import command from alembic.config import Config diff --git a/src/basic_memory/deps.py b/src/basic_memory/deps.py index 11e0e004..87fd21bb 100644 --- a/src/basic_memory/deps.py +++ b/src/basic_memory/deps.py @@ -2,6 +2,7 @@ from typing import Annotated +import logfire from fastapi import Depends from sqlalchemy.ext.asyncio import ( AsyncSession, @@ -43,7 +44,10 @@ async def get_engine_factory( project_config: ProjectConfigDep, ) -> tuple[AsyncEngine, async_sessionmaker[AsyncSession]]: # pragma: no cover """Get engine and session maker.""" - return await db.get_or_create_db(project_config.database_path) + engine, session_maker = await db.get_or_create_db(project_config.database_path) + if project_config.env != "test": + logfire.instrument_sqlalchemy(engine=engine) + return engine, session_maker EngineFactoryDep = Annotated[ diff --git a/src/basic_memory/mcp/tools/knowledge.py b/src/basic_memory/mcp/tools/knowledge.py index 4efa29c1..7b223ebb 100644 --- a/src/basic_memory/mcp/tools/knowledge.py +++ b/src/basic_memory/mcp/tools/knowledge.py @@ -1,8 +1,11 @@ """Knowledge graph management tools for Basic Memory MCP server.""" +import logfire + from basic_memory.mcp.server import mcp from basic_memory.mcp.tools.utils import call_get, call_post from basic_memory.schemas.base import Permalink +from basic_memory.schemas.memory import memory_url_path from basic_memory.schemas.request import ( GetEntitiesRequest, ) @@ -16,15 +19,17 @@ from basic_memory.mcp.async_client import client @mcp.tool( description="Get complete information about a specific entity including observations and relations", ) -async def get_entity(permalink: Permalink) -> EntityResponse: +async def get_entity(identifier: str) -> EntityResponse: """Get a specific entity info by its permalink. Args: - permalink: Path identifier for the entity + identifier: Path identifier for the entity """ - url = f"/knowledge/entities/{permalink}" - response = await call_get(client, url) - return EntityResponse.model_validate(response.json()) + with logfire.span("Getting entity", permalink=identifier) as s: + permalink = memory_url_path(identifier) + url = f"/knowledge/entities/{permalink}" + response = await call_get(client, url) + return EntityResponse.model_validate(response.json()) @mcp.tool( @@ -39,11 +44,12 @@ async def get_entities(request: GetEntitiesRequest) -> EntityListResponse: Returns: EntityListResponse containing complete details for each requested entity """ - url = "/knowledge/entities" - response = await call_get( - client, url, params=[("permalink", permalink) for permalink in request.permalinks] - ) - return EntityListResponse.model_validate(response.json()) + with logfire.span("Getting multiple entities", permalink_count=len(request.permalinks)) as s: + url = "/knowledge/entities" + response = await call_get( + client, url, params=[("permalink", memory_url_path(identifier)) for identifier in request.permalinks] + ) + return EntityListResponse.model_validate(response.json()) @mcp.tool( @@ -51,6 +57,9 @@ async def get_entities(request: GetEntitiesRequest) -> EntityListResponse: ) async def delete_entities(request: DeleteEntitiesRequest) -> DeleteEntitiesResponse: """Delete entities from the knowledge graph.""" - url = "/knowledge/entities/delete" - response = await call_post(client, url, json=request.model_dump()) - return DeleteEntitiesResponse.model_validate(response.json()) + with logfire.span("Deleting entities", permalink_count=len(request.permalinks)) as s: + url = "/knowledge/entities/delete" + + request.permalinks = [memory_url_path(permlink) for permlink in request.permalinks] + response = await call_post(client, url, json=request.model_dump()) + return DeleteEntitiesResponse.model_validate(response.json()) \ No newline at end of file diff --git a/src/basic_memory/mcp/tools/memory.py b/src/basic_memory/mcp/tools/memory.py index 9daffb91..ab3f4cff 100644 --- a/src/basic_memory/mcp/tools/memory.py +++ b/src/basic_memory/mcp/tools/memory.py @@ -3,6 +3,7 @@ from typing import Optional, Literal, List from loguru import logger +import logfire from basic_memory.mcp.async_client import client from basic_memory.mcp.server import mcp @@ -65,14 +66,15 @@ async def build_context( # Research the history of a feature build_context("memory://features/knowledge-graph", timeframe="3 months ago") """ - logger.info(f"Building context from {url}") - url = normalize_memory_url(url) - response = await call_get( - client, - f"/memory/{memory_url_path(url)}", - params={"depth": depth, "timeframe": timeframe, "max_results": max_results}, - ) - return GraphContext.model_validate(response.json()) + with logfire.span("Building context", url=url, depth=depth, timeframe=timeframe) as s: + logger.info(f"Building context from {url}") + url = normalize_memory_url(url) + response = await call_get( + client, + f"/memory/{memory_url_path(url)}", + params={"depth": depth, "timeframe": timeframe, "max_results": max_results}, + ) + return GraphContext.model_validate(response.json()) @mcp.tool( @@ -132,20 +134,21 @@ async def recent_activity( - For focused queries, consider using build_context with a specific URI - Max timeframe is 1 year in the past """ - 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, - } - if type: - params["type"] = type + with logfire.span("Getting recent activity", type=type, depth=depth, timeframe=timeframe) as s: + 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, + } + if type: + params["type"] = type - response = await call_get( - client, - "/memory/recent", - params=params, - ) - return GraphContext.model_validate(response.json()) + response = await call_get( + client, + "/memory/recent", + params=params, + ) + return GraphContext.model_validate(response.json()) \ No newline at end of file diff --git a/src/basic_memory/mcp/tools/notes.py b/src/basic_memory/mcp/tools/notes.py index 21affe0b..9bf832ef 100644 --- a/src/basic_memory/mcp/tools/notes.py +++ b/src/basic_memory/mcp/tools/notes.py @@ -7,6 +7,7 @@ while leveraging the underlying knowledge graph structure. from typing import Optional, List from loguru import logger +import logfire from basic_memory.mcp.server import mcp from basic_memory.mcp.async_client import client @@ -60,71 +61,54 @@ async def write_note( - Observation counts by category - Relation counts (resolved/unresolved) - Tags if present - - Examples: - write_note( - title="Search Implementation", - content="# Search Component\\n\\n" - "Implementation of the search feature, building on [[Core Search]].\\n\\n" - "## Observations\\n" - "- [tech] Using FTS5 for full-text search #implementation\\n" - "- [design] Need pagination support #todo\\n\\n" - "## Relations\\n" - "- implements [[Search Spec]]\\n" - "- depends_on [[Database Schema]]", - folder="docs/components" - ) """ - logger.info(f"Writing note folder:'{folder}' title: '{title}'") + with logfire.span("Writing note", title=title, folder=folder) as s: + logger.info(f"Writing note folder:'{folder}' title: '{title}'") - # Create the entity request - metadata = {"tags": [f"#{tag}" for tag in tags]} if tags else None - entity = Entity( - title=title, - folder=folder, - entity_type="note", - content_type="text/markdown", - content=content, - entity_metadata=metadata, - ) + # Create the entity request + metadata = {"tags": [f"#{tag}" for tag in tags]} if tags else None + entity = Entity( + title=title, + folder=folder, + entity_type="note", + content_type="text/markdown", + content=content, + entity_metadata=metadata, + ) - # Create or update via knowledge API - logger.info(f"Creating {entity.permalink}") - url = f"/knowledge/entities/{entity.permalink}" - response = await call_put(client, url, json=entity.model_dump()) - result = EntityResponse.model_validate(response.json()) + # Create or update via knowledge API + logger.info(f"Creating {entity.permalink}") + url = f"/knowledge/entities/{entity.permalink}" + response = await call_put(client, url, json=entity.model_dump()) + result = EntityResponse.model_validate(response.json()) + + # Format semantic summary based on status code + action = "Created" if response.status_code == 201 else "Updated" + summary = [f"# {action} {result.file_path} ({result.checksum[:8]})", f"permalink: {result.permalink}"] - # Format semantic summary based on status code - action = "Created" if response.status_code == 201 else "Updated" - assert result.checksum is not None - summary = [ - f"# {action} {result.file_path} ({result.checksum[:8]})", - f"permalink: {result.permalink}", - ] + if result.observations: + categories = {} + for obs in result.observations: + categories[obs.category] = categories.get(obs.category, 0) + 1 + + summary.append("\n## Observations") + for category, count in sorted(categories.items()): + summary.append(f"- {category}: {count}") + + if result.relations: + unresolved = sum(1 for r in result.relations if not r.to_id) + resolved = len(result.relations) - unresolved + + summary.append("\n## Relations") + summary.append(f"- Resolved: {resolved}") + if unresolved: + summary.append(f"- Unresolved: {unresolved}") + summary.append("\nUnresolved relations will be retried on next sync.") - if result.observations: - categories = {} - for obs in result.observations: - categories[obs.category] = categories.get(obs.category, 0) + 1 + if tags: + summary.append(f"\n## Tags\n- {', '.join(tags)}") - summary.append("\n## Observations") - for category, count in sorted(categories.items()): - summary.append(f"- {category}: {count}") - - if result.relations: - unresolved = sum(1 for r in result.relations if not r.to_id) - resolved = len(result.relations) - unresolved - - summary.append("\n## Relations") - summary.append(f"- Resolved: {resolved}") - if unresolved: - summary.append(f"- Unresolved: {unresolved}") - summary.append("\nUnresolved relations will be retried on next sync.") - - if tags: - summary.append(f"\n## Tags\n- {', '.join(tags)}") - - return "\n".join(summary) + return "\n".join(summary) @mcp.tool(description="Read note content by title, permalink, relation, or pattern") @@ -180,10 +164,11 @@ async def read_note(identifier: str) -> str: - Last modified timestamp - Content checksum """ - logger.info(f"Reading note {identifier}") - url = memory_url_path(identifier) - response = await call_get(client, f"/resource/{url}") - return response.text + with logfire.span("Reading note", identifier=identifier) as s: + logger.info(f"Reading note {identifier}") + url = memory_url_path(identifier) + response = await call_get(client, f"/resource/{url}") + return response.text @mcp.tool(description="Delete a note by title or permalink") @@ -203,6 +188,7 @@ async def delete_note(identifier: str) -> bool: # Delete by permalink delete_note("notes/project-planning") """ - response = await call_delete(client, f"/knowledge/entities/{identifier}") - result = DeleteEntitiesResponse.model_validate(response.json()) - return result.deleted + with logfire.span("Deleting note", identifier=identifier) as s: + response = await call_delete(client, f"/knowledge/entities/{identifier}") + result = DeleteEntitiesResponse.model_validate(response.json()) + return result.deleted \ No newline at end of file diff --git a/src/basic_memory/mcp/tools/search.py b/src/basic_memory/mcp/tools/search.py index a894f774..eaa00b68 100644 --- a/src/basic_memory/mcp/tools/search.py +++ b/src/basic_memory/mcp/tools/search.py @@ -1,5 +1,5 @@ """Search tools for Basic Memory MCP server.""" - +import logfire from loguru import logger from basic_memory.mcp.server import mcp @@ -24,6 +24,7 @@ async def search(query: SearchQuery) -> SearchResponse: Returns: SearchResponse with search results and metadata """ - logger.info(f"Searching for {query.text}") - response = await call_post(client, "/search/", json=query.model_dump()) - return SearchResponse.model_validate(response.json()) + with logfire.span("Searching for {query}", qurey=query): + logger.info(f"Searching for {query}") + response = await call_post(client, "/search/", json=query.model_dump()) + return SearchResponse.model_validate(response.json()) diff --git a/src/basic_memory/schemas/delete.py b/src/basic_memory/schemas/delete.py index 86d82fe3..812ccc8b 100644 --- a/src/basic_memory/schemas/delete.py +++ b/src/basic_memory/schemas/delete.py @@ -34,4 +34,4 @@ class DeleteEntitiesRequest(BaseModel): 4. Deletes the corresponding markdown file """ - permalinks: Annotated[List[Permalink], MinLen(1)] + permalinks: Annotated[List[Permalink], MinLen(1)] \ No newline at end of file diff --git a/src/basic_memory/schemas/request.py b/src/basic_memory/schemas/request.py index d8803d85..77fa966d 100644 --- a/src/basic_memory/schemas/request.py +++ b/src/basic_memory/schemas/request.py @@ -55,4 +55,4 @@ class GetEntitiesRequest(BaseModel): class CreateRelationsRequest(BaseModel): - relations: List[Relation] + relations: List[Relation] \ No newline at end of file diff --git a/src/basic_memory/utils.py b/src/basic_memory/utils.py index 30282b66..9a075cdd 100644 --- a/src/basic_memory/utils.py +++ b/src/basic_memory/utils.py @@ -87,17 +87,13 @@ def setup_logging( root_path="/src/basic_memory", ), environment=config.env, + console=False, ) logger.configure(handlers=[logfire.loguru_handler()]) # instrument code spans logfire.instrument_sqlite3() - logfire.instrument_pydantic() - - from basic_memory.db import _engine as engine - - if engine: - logfire.instrument_sqlalchemy(engine=engine) + logfire.instrument_httpx() # setup logger log_path = home_dir / log_file diff --git a/tests/mcp/test_tool_knowledge.py b/tests/mcp/test_tool_knowledge.py index 085bba6f..dd0f254e 100644 --- a/tests/mcp/test_tool_knowledge.py +++ b/tests/mcp/test_tool_knowledge.py @@ -32,6 +32,29 @@ async def test_get_single_entity(client): assert entity.permalink == "test/test-note" assert len(entity.observations) == 1 +@pytest.mark.asyncio +async def test_get_single_entity_memory_url(client): + """Test retrieving a single entity.""" + # First create an entity + result = await notes.write_note( + title="Test Note", + folder="test", + content=""" +# Test\nThis is a test note +- [note] First observation +""", + tags=["test", "documentation"], + ) + assert result + + # Get the entity + entity = await get_entity("memory://test/test-note") + + # Verify entity details + assert entity.title == "Test Note" + assert entity.permalink == "test/test-note" + assert len(entity.observations) == 1 + @pytest.mark.asyncio async def test_get_multiple_entities(client): @@ -58,6 +81,31 @@ async def test_get_multiple_entities(client): assert "test/test-note-1" in permalinks assert "test/test-note-2" in permalinks +@pytest.mark.asyncio +async def test_get_multiple_entities_memory_ur(client): + """Test retrieving multiple entities.""" + # Create two test entities + await notes.write_note( + title="Test Note 1", + folder="test", + content="# Test 1", + ) + await notes.write_note( + title="Test Note 2", + folder="test", + content="# Test 2", + ) + + # Get both entities + request = GetEntitiesRequest(permalinks=["memory://test/test-note-1", "memory://test/test-note-2"]) + response = await get_entities(request) + + # Verify we got both entities + assert len(response.entities) == 2 + permalinks = {e.permalink for e in response.entities} + assert "test/test-note-1" in permalinks + assert "test/test-note-2" in permalinks + @pytest.mark.asyncio async def test_delete_entities(client): @@ -80,6 +128,27 @@ async def test_delete_entities(client): with pytest.raises(ToolError): await get_entity("test/test-note") +@pytest.mark.asyncio +async def test_delete_entities_memory_url(client): + """Test deleting entities.""" + # Create a test entity + await notes.write_note( + title="Test Note", + folder="test", + content="# Test Note to Delete", + ) + + # Delete the entity + request = DeleteEntitiesRequest(permalinks=["memory://test/test-note"]) + response = await delete_entities(request) + + # Verify deletion + assert response.deleted is True + + # Verify entity no longer exists + with pytest.raises(ToolError): + await get_entity("test/test-note") + @pytest.mark.asyncio async def test_get_nonexistent_entity(client): diff --git a/uv.lock b/uv.lock index dc292830..93e27bfa 100644 --- a/uv.lock +++ b/uv.lock @@ -79,7 +79,7 @@ wheels = [ [[package]] name = "basic-memory" -version = "0.5.0" +version = "0.6.0" source = { editable = "." } dependencies = [ { name = "aiosqlite" }, @@ -88,7 +88,7 @@ dependencies = [ { name = "fastapi", extra = ["standard"] }, { name = "greenlet" }, { name = "icecream" }, - { name = "logfire", extra = ["fastapi", "sqlalchemy", "sqlite3"] }, + { name = "logfire", extra = ["fastapi", "httpx", "sqlalchemy", "sqlite3"] }, { name = "loguru" }, { name = "markdown-it-py" }, { name = "mcp" }, @@ -126,7 +126,7 @@ requires-dist = [ { name = "fastapi", extras = ["standard"], specifier = ">=0.115.8" }, { name = "greenlet", specifier = ">=3.1.1" }, { name = "icecream", specifier = ">=2.1.3" }, - { name = "logfire", extras = ["fastapi", "sqlalchemy", "sqlite3"], specifier = ">=3.6.0" }, + { name = "logfire", extras = ["fastapi", "httpx", "sqlalchemy", "sqlite3"], specifier = ">=3.6.0" }, { name = "loguru", specifier = ">=0.7.3" }, { name = "markdown-it-py", specifier = ">=3.0.0" }, { name = "mcp", specifier = ">=1.2.0" }, @@ -680,6 +680,9 @@ wheels = [ fastapi = [ { name = "opentelemetry-instrumentation-fastapi" }, ] +httpx = [ + { name = "opentelemetry-instrumentation-httpx" }, +] sqlalchemy = [ { name = "opentelemetry-instrumentation-sqlalchemy" }, ] @@ -913,6 +916,22 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/55/1c/ec2d816b78edf2404d7b3df6d09eefb690b70bfd191b7da06f76634f1bdc/opentelemetry_instrumentation_fastapi-0.51b0-py3-none-any.whl", hash = "sha256:10513bbc11a1188adb9c1d2c520695f7a8f2b5f4de14e8162098035901cd6493", size = 12117 }, ] +[[package]] +name = "opentelemetry-instrumentation-httpx" +version = "0.51b0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "opentelemetry-api" }, + { name = "opentelemetry-instrumentation" }, + { name = "opentelemetry-semantic-conventions" }, + { name = "opentelemetry-util-http" }, + { name = "wrapt" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/b7/d5/4a3990c461ae7e55212115e0f8f3aa412b5ce6493579e85c292245ac69ea/opentelemetry_instrumentation_httpx-0.51b0.tar.gz", hash = "sha256:061d426a04bf5215a859fea46662e5074f920e5cbde7e6ad6825a0a1b595802c", size = 17700 } +wheels = [ + { url = "https://files.pythonhosted.org/packages/c3/ba/23d4ab6402408c01f1c3f32e0c04ea6dae575bf19bcb9a0049c9e768c983/opentelemetry_instrumentation_httpx-0.51b0-py3-none-any.whl", hash = "sha256:2e3fdf755ba6ead6ab43031497c3d55d4c796d0368eccc0ce48d304b7ec6486a", size = 14109 }, +] + [[package]] name = "opentelemetry-instrumentation-sqlalchemy" version = "0.51b0"