From 353342a5e2096639b193ebb2f8a5b6e2e64d32f0 Mon Sep 17 00:00:00 2001 From: phernandez Date: Wed, 25 Dec 2024 10:18:12 -0600 Subject: [PATCH] fix all endpoints for tools --- src/basic_memory/api/routers/documents.py | 4 +- src/basic_memory/mcp/__init__.py | 3 +- src/basic_memory/mcp/async_client.py | 8 +++ src/basic_memory/mcp/server.py | 78 ++++++++++++----------- tests/api/test_documents_router.py | 16 ++--- uv.lock | 7 +- 6 files changed, 64 insertions(+), 52 deletions(-) create mode 100644 src/basic_memory/mcp/async_client.py diff --git a/src/basic_memory/api/routers/documents.py b/src/basic_memory/api/routers/documents.py index 1c35f1fd..69b66b65 100644 --- a/src/basic_memory/api/routers/documents.py +++ b/src/basic_memory/api/routers/documents.py @@ -17,7 +17,7 @@ from basic_memory.services.document_service import ( router = APIRouter(prefix="/documents", tags=["documents"]) -@router.post("/", response_model=DocumentCreateResponse, status_code=201) +@router.post("/create", response_model=DocumentCreateResponse, status_code=201) async def create_document( doc: DocumentRequest, service: DocumentServiceDep, @@ -41,7 +41,7 @@ async def create_document( raise HTTPException(status_code=400, detail=str(e)) -@router.get("/", response_model=List[DocumentCreateResponse]) +@router.get("/list", response_model=List[DocumentCreateResponse]) async def list_documents( service: DocumentServiceDep, ) -> List[DocumentCreateResponse]: diff --git a/src/basic_memory/mcp/__init__.py b/src/basic_memory/mcp/__init__.py index 0d9332e8..a37f2d34 100644 --- a/src/basic_memory/mcp/__init__.py +++ b/src/basic_memory/mcp/__init__.py @@ -1,2 +1 @@ -"""MCP server for basic-memory.""" -from .server import server \ No newline at end of file +"""MCP server for basic-memory.""" \ No newline at end of file diff --git a/src/basic_memory/mcp/async_client.py b/src/basic_memory/mcp/async_client.py new file mode 100644 index 00000000..42c98248 --- /dev/null +++ b/src/basic_memory/mcp/async_client.py @@ -0,0 +1,8 @@ +from httpx import ASGITransport, AsyncClient + +from basic_memory.api.app import app as fastapi_app + +BASE_URL = "http://test" + +# Create shared async client +client = AsyncClient(transport=ASGITransport(app=fastapi_app), base_url=BASE_URL) diff --git a/src/basic_memory/mcp/server.py b/src/basic_memory/mcp/server.py index e0707d5c..4cfd6dba 100644 --- a/src/basic_memory/mcp/server.py +++ b/src/basic_memory/mcp/server.py @@ -4,17 +4,13 @@ import sys from typing import Any, List from fastmcp import FastMCP -from httpx import AsyncClient, ASGITransport from loguru import logger -from basic_memory.api.app import app as fastapi_app +from basic_memory.mcp.async_client import client # Create FastMCP server mcp = FastMCP("Basic Memory") -# Create shared async client -client = AsyncClient(transport=ASGITransport(app=fastapi_app), base_url="http://test") - def setup_logging(log_file: str = "basic-memory-mcp.log"): """Configure logging for the application.""" @@ -48,7 +44,8 @@ async def log_api_call(method: str, url: str, data: Any, response: Any): logger.debug(f"API Request: {method} {url}") logger.debug(f"Request Data: {data}") logger.debug(f"Response Status: {response.status_code}") - logger.debug(f"Response Data: {response.json()}") + if response.status_code != 204: # Only try to log response data if it's not No Content + logger.debug(f"Response Data: {response.json()}") # Knowledge Graph Tools @@ -59,25 +56,28 @@ async def log_api_call(method: str, url: str, data: Any, response: Any): @mcp.tool() async def create_entities(entities: list[dict]) -> dict: """Create new entities in the knowledge graph.""" - response = await client.post("/knowledge/entities", json={"entities": entities}) - await log_api_call("POST", "/knowledge/entities", entities, response) + url = "/knowledge/entities" + data = {"entities": entities} + response = await client.post(url, json=data) return response.json() @mcp.tool() async def create_relations(relations: list[dict]) -> dict: """Create relations between entities.""" - response = await client.post("/knowledge/relations", json={"relations": relations}) - await log_api_call("POST", "/knowledge/relations", relations, response) + url = "/knowledge/relations" + data = {"relations": relations} + response = await client.post(url, json=data) return response.json() @mcp.tool() async def add_observations(path_id: str, observations: list[str]) -> dict: """Add observations to an entity.""" + url = "/knowledge/observations" data = {"path_id": path_id, "observations": observations} - response = await client.post("/knowledge/observations", json=data) - await log_api_call("POST", "/knowledge/observations", data, response) + response = await client.post(url, json=data) + await log_api_call("POST", url, data, response) return response.json() @@ -87,24 +87,26 @@ async def add_observations(path_id: str, observations: list[str]) -> dict: @mcp.tool() async def get_entity(path_id: str) -> dict: """Get a specific entity by path_id.""" - response = await client.get(f"/knowledge/entities/{path_id}") - await log_api_call("GET", f"/knowledge/entities/{path_id}", None, response) + url = f"/knowledge/entities/{path_id}" + response = await client.get(url) return response.json() @mcp.tool() async def search_nodes(query: str) -> dict: """Search for entities in the knowledge graph.""" - response = await client.post("/knowledge/search", json={"query": query}) - await log_api_call("POST", "/knowledge/search", {"query": query}, response) + url = "/knowledge/search" + data = {"query": query} + response = await client.post(url, json=data) return response.json() @mcp.tool() async def open_nodes(path_ids: List[str]) -> dict: """Search for entities in the knowledge graph.""" - response = await client.post("/knowledge/nodes", json={"path_ids": path_ids}) - await log_api_call("POST", "/knowledge/nodes", {"path_ids": path_ids}, response) + url = "/knowledge/nodes" + data = {"path_ids": path_ids} + response = await client.post(url, json=data) return response.json() @@ -114,32 +116,34 @@ async def open_nodes(path_ids: List[str]) -> dict: @mcp.tool() async def delete_entities(path_ids: List[str]) -> dict: """Search for entities in the knowledge graph.""" - response = await client.post("/knowledge/entities/delete", json={"path_ids": path_ids}) - await log_api_call("POST", "/knowledge/entities/delete", {"path_ids": path_ids}, response) + url = "/knowledge/entities/delete" + data = {"path_ids": path_ids} + response = await client.post(url, json=data) return response.json() @mcp.tool() async def delete_observations(path_id: str, observations: list[str]) -> dict: """Delete observations from an entity.""" + url = "/knowledge/observations/delete" data = { "path_id": path_id, "observations": observations, - } # Match the parameter name with what we're using + } response = await client.post( - "/knowledge/observations/delete", json=data - ) # Change to observations endpoint - await log_api_call("POST", "/knowledge/observations/delete", data, response) + url, json=data + ) return response.json() @mcp.tool() async def delete_relations(relations: list[dict]) -> dict: """Delete relations between entities.""" + url = "/knowledge/relations/delete" + data = {"relations": relations} response = await client.post( - "/knowledge/relations/delete", json={"relations": relations} - ) # Change to relations endpoint - await log_api_call("POST", "/knowledge/relations/delete", {"relations": relations}, response) + url, json=data + ) return response.json() @@ -149,42 +153,42 @@ async def delete_relations(relations: list[dict]) -> dict: @mcp.tool() async def create_document(path: str, content: str, metadata: dict = None) -> dict: """Create a new document.""" + url = "/documents/create" data = {"path": path, "content": content, "metadata": metadata} - response = await client.post("/documents/", json=data) - await log_api_call("POST", "/documents/", data, response) + response = await client.post(url, json=data) return response.json() @mcp.tool() async def get_document(path: str) -> dict: """Get a document by path_id.""" - response = await client.get(f"/documents/{path}/") - await log_api_call("GET", f"/documents/{path}/", None, response) + url = f"/documents/{path}" + response = await client.get(url) return response.json() @mcp.tool() async def update_document(path: str, content: str, metadata: dict = None) -> dict: """Update an existing document.""" + url = f"/documents/{path}" data = {"path": path, "content": content, "metadata": metadata} - response = await client.put(f"/documents/{path}/", json=data) - await log_api_call("PUT", f"/documents/{path}/", data, response) + response = await client.put(url, json=data) return response.json() @mcp.tool() async def list_documents() -> list: """List all documents.""" - response = await client.get("/documents/") - await log_api_call("GET", "/documents/", None, response) + url = "/documents/list" + response = await client.get(url) return response.json() @mcp.tool() async def delete_document(path: str) -> dict: """Delete an existing document.""" - response = await client.delete(f"/documents/{path}/") - await log_api_call("DELETE", f"/documents/{path}/", None, response) + url = f"/documents/{path}" + response = await client.delete(url) if response.status_code == 204: return {"deleted": True} diff --git a/tests/api/test_documents_router.py b/tests/api/test_documents_router.py index a2590035..25bab399 100644 --- a/tests/api/test_documents_router.py +++ b/tests/api/test_documents_router.py @@ -17,7 +17,7 @@ async def test_create_document(client: AsyncClient, test_config): "doc_metadata": {"type": "test", "tags": ["documentation", "test"]}, } - response = await client.post("/documents/", json=test_doc) + response = await client.post("/documents/create", json=test_doc) assert response.status_code == 201 data = response.json() @@ -44,7 +44,7 @@ async def test_create_document_should_create_path(client: AsyncClient): "doc_metadata": {"type": "test"}, } - response = await client.post("/documents/", json=test_doc) + response = await client.post("/documents/create", json=test_doc) assert response.status_code == 201 @@ -57,7 +57,7 @@ async def test_create_document_absolute_path(client: AsyncClient, tmp_path: Path "doc_metadata": {"type": "test"}, } - response = await client.post("/documents/", json=test_doc) + response = await client.post("/documents/create", json=test_doc) assert response.status_code == 400 @@ -71,7 +71,7 @@ async def test_get_document(client: AsyncClient): } # Create document - create_response = await client.post("/documents/", json=test_doc) + create_response = await client.post("/documents/create", json=test_doc) assert create_response.status_code == 201 created = create_response.json() @@ -108,7 +108,7 @@ async def test_update_document(client: AsyncClient, test_config: ProjectConfig): "content": "# Original\nOriginal content.", "doc_metadata": {"type": "test", "status": "draft"}, } - create_response = await client.post("/documents/", json=test_doc) + create_response = await client.post("/documents/create", json=test_doc) assert create_response.status_code == 201 created = create_response.json() @@ -157,7 +157,7 @@ async def test_delete_document(client: AsyncClient, test_config: ProjectConfig): } # Create document - create_response = await client.post("/documents/", json=test_doc) + create_response = await client.post("/documents/create", json=test_doc) assert create_response.status_code == 201 created = create_response.json() @@ -202,12 +202,12 @@ async def test_list_documents(client: AsyncClient, tmp_path: Path): # Create all documents created_docs = [] for doc in docs: - response = await client.post("/documents/", json=doc) + response = await client.post("/documents/create", json=doc) assert response.status_code == 201 created_docs.append(response.json()) # List all documents - response = await client.get("/documents/") + response = await client.get("/documents/list") assert response.status_code == 200 data = response.json() diff --git a/uv.lock b/uv.lock index bf60b480..86c3a417 100644 --- a/uv.lock +++ b/uv.lock @@ -755,17 +755,18 @@ wheels = [ [[package]] name = "httpx" -version = "0.28.1" +version = "0.27.2" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "anyio" }, { name = "certifi" }, { name = "httpcore" }, { name = "idna" }, + { name = "sniffio" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/b1/df/48c586a5fe32a0f01324ee087459e112ebb7224f646c0b5023f5e79e9956/httpx-0.28.1.tar.gz", hash = "sha256:75e98c5f16b0f35b567856f597f06ff2270a374470a5c2392242528e3e3e42fc", size = 141406 } +sdist = { url = "https://files.pythonhosted.org/packages/78/82/08f8c936781f67d9e6b9eeb8a0c8b4e406136ea4c3d1f89a5db71d42e0e6/httpx-0.27.2.tar.gz", hash = "sha256:f7c2be1d2f3c3c3160d441802406b206c2b76f5947b11115e6df10c6c65e66c2", size = 144189 } wheels = [ - { url = "https://files.pythonhosted.org/packages/2a/39/e50c7c3a983047577ee07d2a9e53faf5a69493943ec3f6a384bdc792deb2/httpx-0.28.1-py3-none-any.whl", hash = "sha256:d909fcccc110f8c7faf814ca82a9a4d816bc5a6dbfea25d6591d6985b8ba59ad", size = 73517 }, + { url = "https://files.pythonhosted.org/packages/56/95/9377bcb415797e44274b51d46e3249eba641711cf3348050f76ee7b15ffc/httpx-0.27.2-py3-none-any.whl", hash = "sha256:7bb2708e112d8fdd7829cd4243970f0c223274051cb35ee80c03301ee29a3df0", size = 76395 }, ] [[package]]