admin/basicmachines-co-basic-memory
1
0
Fork 0
mirror of https://github.com/basicmachines-co/basic-memory synced 2026-06-21 13:47:35 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix schema plural namings

Browse Source
This commit is contained in:
phernandez
2024-12-14 19:30:01 -06:00
parent 635c2851a7
commit b471e1f8d7
6 changed files with 44 additions and 44 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/basic_memory/api/routers/knowledge.py
+11 -11
View File
@@ -4,26 +4,26 @@ from fastapi import APIRouter
from basic_memory.deps import MemoryServiceDep
from basic_memory.schemas import (
CreateEntitiesRequest, CreateEntitiesResponse,
CreateEntityRequest, CreateEntityResponse,
SearchNodesRequest, SearchNodesResponse,
CreateRelationsRequest, CreateRelationsResponse,
EntityResponse, RelationResponse, AddObservationsRequest, ObservationResponse,
OpenNodesRequest, OpenNodesResponse,
DeleteEntitiesResponse,
DeleteEntityResponse,
DeleteObservationsRequest, DeleteObservationsResponse, AddObservationsResponse
)
router = APIRouter(prefix="/knowledge", tags=["knowledge"])
@router.post("/entities", response_model=CreateEntitiesResponse)
@router.post("/entities", response_model=CreateEntityResponse)
async def create_entities(
data: CreateEntitiesRequest,
data: CreateEntityRequest,
memory_service: MemoryServiceDep
) -> CreateEntitiesResponse:
) -> CreateEntityResponse:
"""Create new entities in the knowledge graph."""
entities = await memory_service.create_entities(data.entities)
return CreateEntitiesResponse(entities=[EntityResponse.model_validate(entity) for entity in entities])
return CreateEntityResponse(entities=[EntityResponse.model_validate(entity) for entity in entities])
@router.get("/entities/{entity_id:path}", response_model=EntityResponse)
@@ -36,14 +36,14 @@ async def get_entity(
return EntityResponse.model_validate(entity)
@router.delete("/entities/{entity_id}", response_model=DeleteEntitiesResponse)
@router.delete("/entities/{entity_id}", response_model=DeleteEntityResponse)
async def delete_entity(
entity_id: str,
memory_service: MemoryServiceDep
) -> DeleteEntitiesResponse:
) -> DeleteEntityResponse:
"""Delete a specific entity by ID."""
deleted = await memory_service.delete_entities([entity_id])
return DeleteEntitiesResponse(deleted=deleted) # pyright: ignore [reportArgumentType]
return DeleteEntityResponse(deleted=deleted) # pyright: ignore [reportArgumentType]
@router.post("/nodes", response_model=OpenNodesResponse)
@@ -66,11 +66,11 @@ async def create_relations(
return CreateRelationsResponse(relations=[RelationResponse.model_validate(relation) for relation in relations])
@router.delete("/relations/{relation_id}", response_model=DeleteEntitiesResponse)
@router.delete("/relations/{relation_id}", response_model=DeleteEntityResponse)
async def delete_relation(
relation_id: int,
memory_service: MemoryServiceDep
) -> DeleteEntitiesResponse:
) -> DeleteEntityResponse:
"""Delete a specific relation by ID."""
# TODO: Implement delete_relation in memory service
raise NotImplementedError("Delete relation not implemented yet")
src/basic_memory/mcp/server.py
+10 -10
View File
@@ -21,12 +21,12 @@ from basic_memory.repository.observation_repository import ObservationRepository
from basic_memory.repository.relation_repository import RelationRepository
from basic_memory.schemas import (
# Tool inputs
CreateEntitiesRequest, SearchNodesRequest, OpenNodesRequest,
CreateRelationsRequest, DeleteEntitiesRequest,
CreateEntityRequest, SearchNodesRequest, OpenNodesRequest,
CreateRelationsRequest, DeleteEntityRequest,
DeleteObservationsRequest,
# Tool responses
CreateEntitiesResponse, SearchNodesResponse, OpenNodesResponse,
AddObservationsResponse, CreateRelationsResponse, DeleteEntitiesResponse,
CreateEntityResponse, SearchNodesResponse, OpenNodesResponse,
AddObservationsResponse, CreateRelationsResponse, DeleteEntityResponse,
EntityResponse, ObservationResponse, RelationResponse, AddObservationsRequest
)
from basic_memory.services import EntityService, ObservationService, RelationService
@@ -103,7 +103,7 @@ async def handle_create_entities(
"""Handle create_entities tool call."""
# Validate input
logger.debug(f"Creating entities with args: {args}")
input_args = CreateEntitiesRequest.model_validate(args)
input_args = CreateEntityRequest.model_validate(args)
logger.debug(f"Validated input: {len(input_args.entities)} entities")
# Call service with validated data
@@ -111,7 +111,7 @@ async def handle_create_entities(
logger.debug(f"Created {len(entities)} entities")
# Format response
response = CreateEntitiesResponse(entities=[EntityResponse.model_validate(entity) for entity in entities])
response = CreateEntityResponse(entities=[EntityResponse.model_validate(entity) for entity in entities])
logger.debug("Formatted create_entities response")
return create_response(response)
@@ -192,10 +192,10 @@ async def handle_delete_entities(
) -> EmbeddedResource:
"""Handle delete_entities tool call."""
logger.debug(f"Deleting entities: {args}")
input_args = DeleteEntitiesRequest.model_validate(args)
input_args = DeleteEntityRequest.model_validate(args)
deleted = await service.delete_entities(input_args.names)
logger.debug(f"Deleted entities: {deleted}")
response = DeleteEntitiesResponse(deleted=deleted)
response = DeleteEntityResponse(deleted=deleted)
return create_response(response)
@@ -240,7 +240,7 @@ class MemoryServer(Server):
Tool(
name="create_entities",
description="Create multiple new entities in the knowledge graph",
inputSchema=CreateEntitiesRequest.model_json_schema()
inputSchema=CreateEntityRequest.model_json_schema()
),
Tool(
name="search_nodes",
@@ -265,7 +265,7 @@ class MemoryServer(Server):
Tool(
name="delete_entities",
description="Delete entities from the knowledge graph",
inputSchema=DeleteEntitiesRequest.model_json_schema()
inputSchema=DeleteEntityRequest.model_json_schema()
),
Tool(
name="delete_observations",
src/basic_memory/schemas.py
+11 -11
View File
@@ -76,33 +76,33 @@ class EntityResponse(EntityBase, SQLAlchemyModel):
relations: List[RelationResponse] = []
model_config = ConfigDict(populate_by_name=True)
# Tool Input Schemas
class CreateEntitiesRequest(BaseModel):
"""Input schema for create_entities tool."""
# Tool Request schemas
class CreateEntityRequest(BaseModel):
"""Request schema for create_entities tool."""
entities: Annotated[List[EntityRequest], Len(min_length=1)]
class SearchNodesRequest(BaseModel):
"""Input schema for search_nodes tool."""
"""Request schema for search_nodes tool."""
query: str
class OpenNodesRequest(BaseModel):
"""Input schema for open_nodes tool."""
"""Request schema for open_nodes tool."""
names: Annotated[List[str], Len(min_length=1)]
class CreateRelationsRequest(BaseModel):
"""Input schema for create_relations tool."""
"""Request schema for create_relations tool."""
relations: List[RelationRequest]
class DeleteEntitiesRequest(BaseModel):
"""Input schema for delete_entities tool."""
class DeleteEntityRequest(BaseModel):
"""Request schema for delete_entities tool."""
names: List[str]
class DeleteObservationsRequest(BaseModel):
"""Input schema for delete_observations tool."""
"""Request schema for delete_observations tool."""
entity_id: str
deletions: List[str] # TODO: Make this more specific
class CreateEntitiesResponse(SQLAlchemyModel):
class CreateEntityResponse(SQLAlchemyModel):
"""Response for create_entities tool."""
entities: List[EntityResponse]
@@ -124,7 +124,7 @@ class CreateRelationsResponse(SQLAlchemyModel):
"""Response for create_relations tool."""
relations: List[RelationResponse]
class DeleteEntitiesResponse(SQLAlchemyModel):
class DeleteEntityResponse(SQLAlchemyModel):
"""Response for delete_entities tool."""
deleted: List[str]
tests/test_mcp_server.py
+4 -4
View File
@@ -5,7 +5,7 @@ from mcp.types import EmbeddedResource
from mcp.shared.exceptions import McpError
from basic_memory.mcp.server import MemoryServer, MIME_TYPE, BASIC_MEMORY_URI
from basic_memory.schemas import (
CreateEntitiesResponse, SearchNodesResponse, AddObservationsResponse,
CreateEntityResponse, SearchNodesResponse, AddObservationsResponse,
)
@pytest.fixture
@@ -90,7 +90,7 @@ async def test_create_directory_entity(test_directory_entity_data, memory_servic
assert result[0].type == "resource"
# Verify entity creation
response = CreateEntitiesResponse.model_validate_json(result[0].resource.text) # pyright: ignore [reportAttributeAccessIssue]
response = CreateEntityResponse.model_validate_json(result[0].resource.text) # pyright: ignore [reportAttributeAccessIssue]
assert len(response.entities) == 1
assert response.entities[0].name == "Directory Organization"
assert response.entities[0].entity_type == "memory"
@@ -115,7 +115,7 @@ async def test_create_entities_snake_case(test_entity_snake_case, memory_service
assert str(result[0].resource.uri) == str(BASIC_MEMORY_URI)
assert result[0].resource.mimeType == MIME_TYPE
response = CreateEntitiesResponse.model_validate_json(result[0].resource.text) # pyright: ignore [reportAttributeAccessIssue]
response = CreateEntityResponse.model_validate_json(result[0].resource.text) # pyright: ignore [reportAttributeAccessIssue]
assert len(response.entities) == 1
assert response.entities[0].name == "Test Entity"
assert response.entities[0].entity_type == "test"
@@ -164,7 +164,7 @@ async def test_add_observations(test_entity_data, memory_service, test_config):
memory_service=memory_service
)
create_response = CreateEntitiesResponse.model_validate_json(create_result[0].resource.text) # pyright: ignore [reportAttributeAccessIssue]
create_response = CreateEntityResponse.model_validate_json(create_result[0].resource.text) # pyright: ignore [reportAttributeAccessIssue]
entity_id = create_response.entities[0].id
# Add new observations using camelCase
tests/test_memory_service.py
+5 -5
View File
@@ -2,7 +2,7 @@
import pytest
from basic_memory.services import MemoryService
from basic_memory.fileio import read_entity_file
from basic_memory.schemas import CreateEntitiesRequest, CreateRelationsRequest, AddObservationsRequest, RelationRequest
from basic_memory.schemas import CreateEntityRequest, CreateRelationsRequest, AddObservationsRequest, RelationRequest
test_entities_data = [
{
@@ -21,7 +21,7 @@ test_entities_data = [
async def test_create_entities(memory_service: MemoryService):
"""Should create multiple entities in parallel with their observations."""
entity_input = CreateEntitiesRequest.model_validate({"entities": test_entities_data})
entity_input = CreateEntityRequest.model_validate({"entities": test_entities_data})
entities = await memory_service.create_entities(entity_input.entities)
# Verify the SQLAlchemy models were created
@@ -51,7 +51,7 @@ async def test_create_entities(memory_service: MemoryService):
@pytest.mark.asyncio
async def test_add_observations(memory_service: MemoryService):
"""Should add observations to an existing entity."""
entity_input = CreateEntitiesRequest.model_validate({"entities": test_entities_data})
entity_input = CreateEntityRequest.model_validate({"entities": test_entities_data})
entities = await memory_service.create_entities([entity_input.entities[0]])
entity = entities[0]
@@ -100,7 +100,7 @@ async def test_add_observations_nonexistent_entity(memory_service: MemoryService
@pytest.mark.asyncio
async def test_create_relations(memory_service: MemoryService):
"""Should create relations between entities and update both filesystem and database."""
entity_input = CreateEntitiesRequest.model_validate({"entities": test_entities_data})
entity_input = CreateEntityRequest.model_validate({"entities": test_entities_data})
entities = await memory_service.create_entities(entity_input.entities)
entity1, entity2 = entities
@@ -183,7 +183,7 @@ async def test_create_relations(memory_service: MemoryService):
async def test_create_relations_with_invalid_entity_id(memory_service: MemoryService):
"""Should raise an appropriate error when trying to create relations with non-existent entity IDs."""
# Create one entity - returns SQLAlchemy Entity
entity_input = CreateEntitiesRequest.model_validate({"entities": test_entities_data})
entity_input = CreateEntityRequest.model_validate({"entities": test_entities_data})
entities = await memory_service.create_entities([entity_input.entities[0]])
entity1 = entities[0]
tests/test_schemas.py
+3 -3
View File
@@ -5,7 +5,7 @@ from basic_memory.schemas import (
EntityRequest,
EntityResponse,
RelationRequest,
CreateEntitiesRequest,
CreateEntityRequest,
SearchNodesRequest,
OpenNodesRequest,
)
@@ -97,13 +97,13 @@ def test_create_entities_input():
}
]
}
create_input = CreateEntitiesRequest.model_validate(data)
create_input = CreateEntityRequest.model_validate(data)
assert len(create_input.entities) == 2
assert create_input.entities[1].description == "test description"
# Empty entities list should fail
with pytest.raises(ValidationError):
CreateEntitiesRequest.model_validate({"entities": []})
CreateEntityRequest.model_validate({"entities": []})
def test_entity_out_from_attributes():
"""Test EntityOut creation from database model attributes."""