fix all endpoints for tools

This commit is contained in:
phernandez
2024-12-25 10:18:12 -06:00
parent 59fa354c77
commit 353342a5e2
6 changed files with 64 additions and 52 deletions
+2 -2
View File
@@ -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]:
+1 -2
View File
@@ -1,2 +1 @@
"""MCP server for basic-memory."""
from .server import server
"""MCP server for basic-memory."""
+8
View File
@@ -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)
+41 -37
View File
@@ -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}
+8 -8
View File
@@ -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()
Generated
+4 -3
View File
@@ -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]]