mirror of
https://github.com/basicmachines-co/basic-memory
synced 2026-06-21 13:47:35 +00:00
cleaned up schemas for api/mcp server
This commit is contained in:
phernandez
@@ -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."""
|
||||
|
||||
@@ -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]
|
||||
|
||||
@@ -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")
|
||||
|
||||
+13
-18
@@ -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."""
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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())
|
||||
|
||||
Reference in New Issue
Block a user