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

cleaned up schemas for api/mcp server

Browse Source
This commit is contained in:
phernandez
2024-12-14 19:09:56 -06:00
parent 9b359b884c
commit 2b4bd19da1
10 changed files with 81 additions and 89 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/basic_memory/api/routers/knowledge.py
+15 -15
View File
@@ -4,13 +4,13 @@ from fastapi import APIRouter
from basic_memory.deps import MemoryServiceDep
from basic_memory.schemas import (
CreateEntitiesInput, CreateEntitiesResponse,
SearchNodesInput, SearchNodesResponse,
CreateRelationsInput, CreateRelationsResponse,
EntityOut, RelationOut, ObservationsIn, ObservationsOut, ObservationOut,
OpenNodesInput, OpenNodesResponse,
CreateEntitiesRequest, CreateEntitiesResponse,
SearchNodesRequest, SearchNodesResponse,
CreateRelationsRequest, CreateRelationsResponse,
EntityOut, RelationOut, AddObservationsRequest, ObservationOut,
OpenNodesRequest, OpenNodesResponse,
DeleteEntitiesResponse,
DeleteObservationsInput, DeleteObservationsResponse
DeleteObservationsRequest, DeleteObservationsResponse, AddObservationsResponse
)
router = APIRouter(prefix="/knowledge", tags=["knowledge"])
@@ -18,7 +18,7 @@ router = APIRouter(prefix="/knowledge", tags=["knowledge"])
@router.post("/entities", response_model=CreateEntitiesResponse)
async def create_entities(
data: CreateEntitiesInput,
data: CreateEntitiesRequest,
memory_service: MemoryServiceDep
) -> CreateEntitiesResponse:
"""Create new entities in the knowledge graph."""
@@ -48,7 +48,7 @@ async def delete_entity(
@router.post("/nodes", response_model=OpenNodesResponse)
async def open_nodes(
data: OpenNodesInput,
data: OpenNodesRequest,
memory_service: MemoryServiceDep
) -> OpenNodesResponse:
"""Open specific nodes by their names."""
@@ -58,7 +58,7 @@ async def open_nodes(
@router.post("/relations", response_model=CreateRelationsResponse)
async def create_relations(
data: CreateRelationsInput,
data: CreateRelationsRequest,
memory_service: MemoryServiceDep
) -> CreateRelationsResponse:
"""Create relations between entities."""
@@ -76,19 +76,19 @@ async def delete_relation(
raise NotImplementedError("Delete relation not implemented yet")
@router.post("/observations", response_model=ObservationsOut)
@router.post("/observations", response_model=AddObservationsResponse)
async def add_observations(
data: ObservationsIn,
data: AddObservationsRequest,
memory_service: MemoryServiceDep
) -> ObservationsOut:
) -> AddObservationsResponse:
"""Add observations to an entity."""
observations = await memory_service.add_observations(data)
return ObservationsOut(entity_id=data.entity_id, observations=[ObservationOut.model_validate(observation) for observation in observations])
return AddObservationsResponse(entity_id=data.entity_id, observations=[ObservationOut.model_validate(observation) for observation in observations])
@router.delete("/observations", response_model=DeleteObservationsResponse)
async def delete_observations(
data: DeleteObservationsInput,
data: DeleteObservationsRequest,
memory_service: MemoryServiceDep
) -> DeleteObservationsResponse:
"""Delete observations from an entity."""
@@ -98,7 +98,7 @@ async def delete_observations(
@router.post("/search", response_model=SearchNodesResponse)
async def search_nodes(
data: SearchNodesInput,
data: SearchNodesRequest,
memory_service: MemoryServiceDep
) -> SearchNodesResponse:
"""Search for entities in the knowledge graph."""
src/basic_memory/fileio.py
+2 -2
View File
@@ -6,7 +6,7 @@ from pathlib import Path
from loguru import logger
from basic_memory.schemas import EntityIn, Relation
from basic_memory.schemas import EntityIn, RelationIn
class FileOperationError(Exception):
@@ -174,7 +174,7 @@ async def read_entity_file(project_entities_path: Path, entity_id: str) -> Entit
relation_type = parts[0]
context = parts[1] if len(parts) > 1 else None
relations.append(Relation( # pyright: ignore [reportCallIssue]
relations.append(RelationIn( # pyright: ignore [reportCallIssue]
from_id=entity_id, # pyright: ignore [reportCallIssue]
to_id=target_id, # pyright: ignore [reportCallIssue]
relation_type=relation_type, # pyright: ignore [reportCallIssue]
src/basic_memory/mcp/server.py
+18 -20
View File
@@ -21,15 +21,13 @@ from basic_memory.repository.observation_repository import ObservationRepository
from basic_memory.repository.relation_repository import RelationRepository
from basic_memory.schemas import (
# Tool inputs
CreateEntitiesInput, SearchNodesInput, OpenNodesInput,
AddObservationsInput, CreateRelationsInput, DeleteEntitiesInput,
DeleteObservationsInput,
CreateEntitiesRequest, SearchNodesRequest, OpenNodesRequest,
CreateRelationsRequest, DeleteEntitiesRequest,
DeleteObservationsRequest,
# Tool responses
CreateEntitiesResponse, SearchNodesResponse, OpenNodesResponse,
AddObservationsResponse, CreateRelationsResponse, DeleteEntitiesResponse,
DeleteObservationsResponse,
# Base models
EntityOut, ObservationOut, RelationOut, ObservationsIn
EntityOut, ObservationOut, RelationOut, AddObservationsRequest
)
from basic_memory.services import EntityService, ObservationService, RelationService
from basic_memory.services.memory_service import MemoryService
@@ -105,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 = CreateEntitiesInput.model_validate(args)
input_args = CreateEntitiesRequest.model_validate(args)
logger.debug(f"Validated input: {len(input_args.entities)} entities")
# Call service with validated data
@@ -124,7 +122,7 @@ async def handle_search_nodes(
) -> EmbeddedResource:
"""Handle search_nodes tool call."""
logger.debug(f"Searching nodes with query: {args.get('query')}")
input_args = SearchNodesInput.model_validate(args)
input_args = SearchNodesRequest.model_validate(args)
results = await service.search_nodes(input_args.query)
logger.debug(f"Found {len(results)} matches for query '{input_args.query}'")
response = SearchNodesResponse(
@@ -140,7 +138,7 @@ async def handle_open_nodes(
) -> EmbeddedResource:
"""Handle open_nodes tool call."""
logger.debug(f"Opening nodes: {args.get('names')}")
input_args = OpenNodesInput.model_validate(args)
input_args = OpenNodesRequest.model_validate(args)
entities = await service.open_nodes(input_args.names)
logger.debug(f"Opened {len(entities)} entities")
response = OpenNodesResponse(entities=[EntityOut.model_validate(entity) for entity in entities])
@@ -154,7 +152,7 @@ async def handle_add_observations(
"""Handle add_observations tool call."""
# Validate input
logger.debug(f"Adding observations: {args}")
input_args = ObservationsIn.model_validate(args)
input_args = AddObservationsRequest.model_validate(args)
logger.debug(f"Adding {len(input_args.observations)} observations to entity {input_args.entity_id}")
# Call service with validated data
@@ -164,7 +162,7 @@ async def handle_add_observations(
# Format response
response = AddObservationsResponse(
entity_id=input_args.entity_id,
added_observations=[ObservationOut.model_validate(obs) for obs in observations]
observations=[ObservationOut.model_validate(obs) for obs in observations]
)
return create_response(response)
@@ -176,7 +174,7 @@ async def handle_create_relations(
"""Handle create_relations tool call."""
# Validate input
logger.debug(f"Creating relations: {args}")
input_args = CreateRelationsInput.model_validate(args)
input_args = CreateRelationsRequest.model_validate(args)
logger.debug(f"Creating {len(input_args.relations)} relations")
# Call service with validated data
@@ -194,7 +192,7 @@ async def handle_delete_entities(
) -> EmbeddedResource:
"""Handle delete_entities tool call."""
logger.debug(f"Deleting entities: {args}")
input_args = DeleteEntitiesInput.model_validate(args)
input_args = DeleteEntitiesRequest.model_validate(args)
deleted = await service.delete_entities(input_args.names)
logger.debug(f"Deleted entities: {deleted}")
response = DeleteEntitiesResponse(deleted=deleted)
@@ -242,37 +240,37 @@ class MemoryServer(Server):
Tool(
name="create_entities",
description="Create multiple new entities in the knowledge graph",
inputSchema=CreateEntitiesInput.model_json_schema()
inputSchema=CreateEntitiesRequest.model_json_schema()
),
Tool(
name="search_nodes",
description="Search for nodes in the knowledge graph",
inputSchema=SearchNodesInput.model_json_schema()
inputSchema=SearchNodesRequest.model_json_schema()
),
Tool(
name="open_nodes",
description="Open specific nodes by their names",
inputSchema=OpenNodesInput.model_json_schema()
inputSchema=OpenNodesRequest.model_json_schema()
),
Tool(
name="add_observations",
description="Add observations to existing entities",
inputSchema=AddObservationsInput.model_json_schema()
inputSchema=AddObservationsRequest.model_json_schema()
),
Tool(
name="create_relations",
description="Create relations between entities",
inputSchema=CreateRelationsInput.model_json_schema()
inputSchema=CreateRelationsRequest.model_json_schema()
),
Tool(
name="delete_entities",
description="Delete entities from the knowledge graph",
inputSchema=DeleteEntitiesInput.model_json_schema()
inputSchema=DeleteEntitiesRequest.model_json_schema()
),
Tool(
name="delete_observations",
description="Delete observations from entities",
inputSchema=DeleteObservationsInput.model_json_schema()
inputSchema=DeleteObservationsRequest.model_json_schema()
)
]
logger.debug(f"Returning {len(tools)} available tools")
src/basic_memory/schemas.py
+13 -18
View File
@@ -1,5 +1,5 @@
"""Core pydantic models for basic-memory entities, observations, and relations."""
from typing import List, Optional, Dict, Any, Annotated
from typing import List, Optional, Annotated
from annotated_types import Len
from pydantic import BaseModel, ConfigDict
@@ -9,7 +9,7 @@ class SQLAlchemyOut(BaseModel):
model_config = ConfigDict(from_attributes=True)
# Base Models
class ObservationsIn(BaseModel):
class AddObservationsRequest(BaseModel):
"""Schema for adding observations to an entity."""
entity_id: str
context: Optional[str] = None
@@ -27,7 +27,7 @@ class ObservationsOut(SQLAlchemyOut):
observations: List[ObservationOut]
model_config = ConfigDict(populate_by_name=True)
class Relation(BaseModel):
class RelationIn(BaseModel):
"""
Represents a directed edge between entities in the knowledge graph.
Relations are always stored in active voice (e.g. "created", "teaches", etc.)
@@ -67,7 +67,7 @@ class EntityIn(EntityBase):
associated observations.
"""
observations: List[str] = []
relations: List[Relation] = []
relations: List[RelationIn] = []
model_config = ConfigDict(populate_by_name=True)
class EntityOut(EntityBase, SQLAlchemyOut):
@@ -77,36 +77,31 @@ class EntityOut(EntityBase, SQLAlchemyOut):
model_config = ConfigDict(populate_by_name=True)
# Tool Input Schemas
class CreateEntitiesInput(BaseModel):
class CreateEntitiesRequest(BaseModel):
"""Input schema for create_entities tool."""
entities: Annotated[List[EntityIn], Len(min_length=1)]
class SearchNodesInput(BaseModel):
class SearchNodesRequest(BaseModel):
"""Input schema for search_nodes tool."""
query: str
class OpenNodesInput(BaseModel):
class OpenNodesRequest(BaseModel):
"""Input schema for open_nodes tool."""
names: Annotated[List[str], Len(min_length=1)]
class AddObservationsInput(ObservationsIn):
"""Input schema for add_observations tool."""
pass
class CreateRelationsInput(BaseModel):
class CreateRelationsRequest(BaseModel):
"""Input schema for create_relations tool."""
relations: List[Relation]
relations: List[RelationIn]
class DeleteEntitiesInput(BaseModel):
class DeleteEntitiesRequest(BaseModel):
"""Input schema for delete_entities tool."""
names: List[str]
class DeleteObservationsInput(BaseModel):
class DeleteObservationsRequest(BaseModel):
"""Input schema for delete_observations tool."""
entity_id: str
deletions: List[Dict[str, Any]] # TODO: Make this more specific
deletions: List[str] # TODO: Make this more specific
# Tool Response Schemas
class CreateEntitiesResponse(SQLAlchemyOut):
"""Response for create_entities tool."""
entities: List[EntityOut]
@@ -123,7 +118,7 @@ class OpenNodesResponse(SQLAlchemyOut):
class AddObservationsResponse(SQLAlchemyOut):
"""Response for add_observations tool."""
entity_id: str
added_observations: List[ObservationOut]
observations: List[ObservationOut]
class CreateRelationsResponse(SQLAlchemyOut):
"""Response for create_relations tool."""
src/basic_memory/services/memory_service.py
+4 -4
View File
@@ -3,9 +3,9 @@ import asyncio
from typing import List, Dict, Any, Optional, Sequence
from pathlib import Path
from basic_memory.models import Entity, Observation, Relation
from basic_memory.models import Entity, Observation
from basic_memory.schemas import (
ObservationsIn, EntityIn, Relation
AddObservationsRequest, EntityIn, RelationIn
)
from basic_memory.fileio import write_entity_file, read_entity_file, EntityNotFoundError
from basic_memory.services import EntityService, RelationService, ObservationService
@@ -112,7 +112,7 @@ class MemoryService:
logger.debug(f"Found entity {entity}")
return entity
async def create_relations(self, relations_data: List[Relation]) -> List[Relation]:
async def create_relations(self, relations_data: List[RelationIn]) -> List[RelationIn]:
"""Create multiple relations between entities."""
logger.debug(f"Creating {len(relations_data)} relations")
@@ -153,7 +153,7 @@ class MemoryService:
logger.debug(f"Successfully created {len(relations)} relations")
return relations
async def add_observations(self, observations_in: ObservationsIn) -> List[Observation]:
async def add_observations(self, observations_in: AddObservationsRequest) -> List[Observation]:
"""Add observations to an existing entity."""
logger.debug(f"Adding observations to entity: {observations_in.entity_id}")
try:
src/basic_memory/services/relation_service.py
+2 -3
View File
@@ -1,9 +1,8 @@
"""Service for managing relations in the database."""
from pathlib import Path
from basic_memory.models import Relation
from basic_memory.repository.relation_repository import RelationRepository
from basic_memory.schemas import EntityIn, Relation
from basic_memory.schemas import EntityIn, RelationIn
from . import DatabaseSyncError
@@ -17,7 +16,7 @@ class RelationService:
self.project_path = project_path
self.relation_repo = relation_repo
async def create_relation(self, relation: Relation) -> Relation:
async def create_relation(self, relation: RelationIn) -> RelationIn:
"""Create a new relation in the database."""
try:
return await self.relation_repo.create(relation.model_dump())
tests/test_mcp_server.py
+2 -2
View File
@@ -186,8 +186,8 @@ async def test_add_observations(test_entity_data, memory_service, test_config):
response = AddObservationsResponse.model_validate_json(result[0].resource.text) # pyright: ignore [reportAttributeAccessIssue]
assert response.entity_id == entity_id
assert len(response.added_observations) == 1
assert response.added_observations[0].content == "A new observation"
assert len(response.observations) == 1
assert response.observations[0].content == "A new observation"
@pytest.mark.anyio
async def test_invalid_tool_name(test_config):
tests/test_memory_service.py
+9 -9
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 CreateEntitiesInput, CreateRelationsInput, ObservationsIn, Relation
from basic_memory.schemas import CreateEntitiesRequest, CreateRelationsRequest, AddObservationsRequest, RelationIn
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 = CreateEntitiesInput.model_validate({"entities": test_entities_data})
entity_input = CreateEntitiesRequest.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 = CreateEntitiesInput.model_validate({"entities": test_entities_data})
entity_input = CreateEntitiesRequest.model_validate({"entities": test_entities_data})
entities = await memory_service.create_entities([entity_input.entities[0]])
entity = entities[0]
@@ -65,7 +65,7 @@ async def test_add_observations(memory_service: MemoryService):
}
# Add observations - returns List[models.Observation]
observation_input = ObservationsIn.model_validate(observations_data)
observation_input = AddObservationsRequest.model_validate(observations_data)
added_observations = await memory_service.add_observations(observation_input)
# Check the SQLAlchemy model results
@@ -94,13 +94,13 @@ async def test_add_observations_nonexistent_entity(memory_service: MemoryService
}
with pytest.raises(Exception) as exc: # We might want to define a specific error type
observation_input = ObservationsIn.model_validate(observations_data)
observation_input = AddObservationsRequest.model_validate(observations_data)
await memory_service.add_observations(observation_input)
@pytest.mark.asyncio
async def test_create_relations(memory_service: MemoryService):
"""Should create relations between entities and update both filesystem and database."""
entity_input = CreateEntitiesInput.model_validate({"entities": test_entities_data})
entity_input = CreateEntitiesRequest.model_validate({"entities": test_entities_data})
entities = await memory_service.create_entities(entity_input.entities)
entity1, entity2 = entities
@@ -120,7 +120,7 @@ async def test_create_relations(memory_service: MemoryService):
]
# Create relations - returns List[models.Relation]
input_args = CreateRelationsInput.model_validate({"relations": test_relations_data})
input_args = CreateRelationsRequest.model_validate({"relations": test_relations_data})
relations = await memory_service.create_relations(input_args.relations)
# Verify SQLAlchemy Relation models were created
@@ -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 = CreateEntitiesInput.model_validate({"entities": test_entities_data})
entity_input = CreateEntitiesRequest.model_validate({"entities": test_entities_data})
entities = await memory_service.create_entities([entity_input.entities[0]])
entity1 = entities[0]
@@ -195,4 +195,4 @@ async def test_create_relations_with_invalid_entity_id(memory_service: MemorySer
}
with pytest.raises(Exception) as exc: # We might want to define a specific error type
await memory_service.create_relations([Relation.model_validate(bad_relation)])
await memory_service.create_relations([RelationIn.model_validate(bad_relation)])
tests/test_relation_service.py
+3 -3
View File
@@ -2,7 +2,7 @@
import pytest
import pytest_asyncio
from basic_memory.schemas import EntityIn, Relation
from basic_memory.schemas import EntityIn, RelationIn
pytestmark = pytest.mark.asyncio
@@ -36,7 +36,7 @@ async def test_create_relation(relation_service, sample_entities):
"""Test creating a basic relation between two entities"""
entity1, entity2 = sample_entities
relation_data = Relation(
relation_data = RelationIn(
from_id=entity1.id,
to_id=entity2.id,
relation_type="test_relation"
@@ -62,7 +62,7 @@ async def test_create_relation_with_context(relation_service, sample_entities):
"""Test creating a relation with context information"""
entity1, entity2 = sample_entities
relation_data = Relation(
relation_data = RelationIn(
from_id=entity1.id,
to_id=entity2.id,
relation_type="test_relation",
tests/test_schemas.py
+13 -13
View File
@@ -4,10 +4,10 @@ from pydantic import ValidationError
from basic_memory.schemas import (
EntityIn,
EntityOut,
Relation,
CreateEntitiesInput,
SearchNodesInput,
OpenNodesInput,
RelationIn,
CreateEntitiesRequest,
SearchNodesRequest,
OpenNodesRequest,
)
def test_entity_in_minimal():
@@ -67,7 +67,7 @@ def test_relation_in_validation():
"to_id": "456",
"relation_type": "test"
}
relation = Relation.model_validate(data)
relation = RelationIn.model_validate(data)
assert relation.from_id == "123"
assert relation.to_id == "456"
assert relation.relation_type == "test"
@@ -75,12 +75,12 @@ def test_relation_in_validation():
# With context
data["context"] = "test context"
relation = Relation.model_validate(data)
relation = RelationIn.model_validate(data)
assert relation.context == "test context"
# Missing required fields
with pytest.raises(ValidationError):
Relation.model_validate({"from_id": "123", "to_id": "456"}) # Missing relationType
RelationIn.model_validate({"from_id": "123", "to_id": "456"}) # Missing relationType
def test_create_entities_input():
"""Test CreateEntitiesInput validation."""
@@ -97,13 +97,13 @@ def test_create_entities_input():
}
]
}
create_input = CreateEntitiesInput.model_validate(data)
create_input = CreateEntitiesRequest.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):
CreateEntitiesInput.model_validate({"entities": []})
CreateEntitiesRequest.model_validate({"entities": []})
def test_entity_out_from_attributes():
"""Test EntityOut creation from database model attributes."""
@@ -167,17 +167,17 @@ def test_optional_fields():
def test_search_nodes_input():
"""Test SearchNodesInput validation."""
search = SearchNodesInput.model_validate({"query": "test query"})
search = SearchNodesRequest.model_validate({"query": "test query"})
assert search.query == "test query"
with pytest.raises(ValidationError):
SearchNodesInput.model_validate({}) # Missing required query
SearchNodesRequest.model_validate({}) # Missing required query
def test_open_nodes_input():
"""Test OpenNodesInput validation."""
open_input = OpenNodesInput.model_validate({"names": ["entity1", "entity2"]})
open_input = OpenNodesRequest.model_validate({"names": ["entity1", "entity2"]})
assert len(open_input.names) == 2
# Empty names list should fail
with pytest.raises(ValidationError):
OpenNodesInput.model_validate({"names": []})
OpenNodesRequest.model_validate({"names": []})