mirror of
https://github.com/basicmachines-co/basic-memory
synced 2026-06-21 13:47:35 +00:00
Compare commits
1 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| 6e5fbf46eb |
@@ -4,7 +4,7 @@ from datetime import datetime
|
||||
from typing import List, Optional, Annotated, Sequence, Literal, Union
|
||||
|
||||
from annotated_types import MinLen, MaxLen
|
||||
from pydantic import BaseModel, Field, BeforeValidator, TypeAdapter, ConfigDict
|
||||
from pydantic import BaseModel, Field, BeforeValidator, TypeAdapter, model_serializer
|
||||
|
||||
from basic_memory.schemas.search import SearchItemType
|
||||
|
||||
@@ -118,8 +118,6 @@ def memory_url_path(url: memory_url) -> str: # pyright: ignore
|
||||
class EntitySummary(BaseModel):
|
||||
"""Simplified entity representation."""
|
||||
|
||||
model_config = ConfigDict(json_encoders={datetime: lambda dt: dt.isoformat()})
|
||||
|
||||
type: Literal["entity"] = "entity"
|
||||
permalink: Optional[str]
|
||||
title: str
|
||||
@@ -127,12 +125,17 @@ class EntitySummary(BaseModel):
|
||||
file_path: str
|
||||
created_at: datetime
|
||||
|
||||
@model_serializer
|
||||
def serialize_model(self) -> dict:
|
||||
"""Ensure datetime fields are serialized as ISO strings for MCP compliance."""
|
||||
data = self.__dict__.copy()
|
||||
data["created_at"] = self.created_at.isoformat()
|
||||
return data
|
||||
|
||||
|
||||
class RelationSummary(BaseModel):
|
||||
"""Simplified relation representation."""
|
||||
|
||||
model_config = ConfigDict(json_encoders={datetime: lambda dt: dt.isoformat()})
|
||||
|
||||
type: Literal["relation"] = "relation"
|
||||
title: str
|
||||
file_path: str
|
||||
@@ -142,12 +145,17 @@ class RelationSummary(BaseModel):
|
||||
to_entity: Optional[str] = None
|
||||
created_at: datetime
|
||||
|
||||
@model_serializer
|
||||
def serialize_model(self) -> dict:
|
||||
"""Ensure datetime fields are serialized as ISO strings for MCP compliance."""
|
||||
data = self.__dict__.copy()
|
||||
data["created_at"] = self.created_at.isoformat()
|
||||
return data
|
||||
|
||||
|
||||
class ObservationSummary(BaseModel):
|
||||
"""Simplified observation representation."""
|
||||
|
||||
model_config = ConfigDict(json_encoders={datetime: lambda dt: dt.isoformat()})
|
||||
|
||||
type: Literal["observation"] = "observation"
|
||||
title: str
|
||||
file_path: str
|
||||
@@ -156,12 +164,17 @@ class ObservationSummary(BaseModel):
|
||||
content: str
|
||||
created_at: datetime
|
||||
|
||||
@model_serializer
|
||||
def serialize_model(self) -> dict:
|
||||
"""Ensure datetime fields are serialized as ISO strings for MCP compliance."""
|
||||
data = self.__dict__.copy()
|
||||
data["created_at"] = self.created_at.isoformat()
|
||||
return data
|
||||
|
||||
|
||||
class MemoryMetadata(BaseModel):
|
||||
"""Simplified response metadata."""
|
||||
|
||||
model_config = ConfigDict(json_encoders={datetime: lambda dt: dt.isoformat()})
|
||||
|
||||
uri: Optional[str] = None
|
||||
types: Optional[List[SearchItemType]] = None
|
||||
depth: int
|
||||
@@ -173,6 +186,13 @@ class MemoryMetadata(BaseModel):
|
||||
total_relations: Optional[int] = None
|
||||
total_observations: Optional[int] = None
|
||||
|
||||
@model_serializer
|
||||
def serialize_model(self) -> dict:
|
||||
"""Ensure datetime fields are serialized as ISO strings for MCP compliance."""
|
||||
data = self.__dict__.copy()
|
||||
data["generated_at"] = self.generated_at.isoformat()
|
||||
return data
|
||||
|
||||
|
||||
class ContextResult(BaseModel):
|
||||
"""Context result containing a primary item with its observations and related items."""
|
||||
|
||||
@@ -197,18 +197,40 @@ class TestDateTimeSerialization:
|
||||
assert "T" in datetime_str # Contains date-time separator
|
||||
assert len(datetime_str) >= 19 # At least YYYY-MM-DDTHH:MM:SS format
|
||||
|
||||
def test_all_models_have_json_encoders_configured(self):
|
||||
"""Test that all memory schema models have datetime json_encoders configured."""
|
||||
def test_all_models_have_model_serializers_configured(self):
|
||||
"""Test that all memory schema models have model serializers for datetime fields."""
|
||||
models_to_test = [EntitySummary, RelationSummary, ObservationSummary, MemoryMetadata]
|
||||
|
||||
for model_class in models_to_test:
|
||||
# Check that ConfigDict with json_encoders is configured
|
||||
assert hasattr(model_class, "model_config")
|
||||
assert "json_encoders" in model_class.model_config
|
||||
assert datetime in model_class.model_config["json_encoders"]
|
||||
# Check that model serializer method is defined
|
||||
assert hasattr(model_class, "serialize_model")
|
||||
assert callable(getattr(model_class, "serialize_model"))
|
||||
|
||||
# Verify the encoder function produces ISO format
|
||||
encoder = model_class.model_config["json_encoders"][datetime]
|
||||
# Test that serializer produces proper datetime format
|
||||
test_datetime = datetime(2023, 12, 8, 10, 30, 0)
|
||||
result = encoder(test_datetime)
|
||||
assert result == "2023-12-08T10:30:00"
|
||||
|
||||
# Create instance with minimal required fields
|
||||
if model_class == EntitySummary:
|
||||
instance = model_class(
|
||||
title="Test", file_path="test.md", created_at=test_datetime
|
||||
)
|
||||
json_data = instance.model_dump()
|
||||
assert json_data["created_at"] == "2023-12-08T10:30:00"
|
||||
elif model_class == RelationSummary:
|
||||
instance = model_class(
|
||||
title="Test", file_path="test.md", permalink="test",
|
||||
relation_type="relates", created_at=test_datetime
|
||||
)
|
||||
json_data = instance.model_dump()
|
||||
assert json_data["created_at"] == "2023-12-08T10:30:00"
|
||||
elif model_class == ObservationSummary:
|
||||
instance = model_class(
|
||||
title="Test", file_path="test.md", permalink="test",
|
||||
category="note", content="test", created_at=test_datetime
|
||||
)
|
||||
json_data = instance.model_dump()
|
||||
assert json_data["created_at"] == "2023-12-08T10:30:00"
|
||||
elif model_class == MemoryMetadata:
|
||||
instance = model_class(depth=1, generated_at=test_datetime)
|
||||
json_data = instance.model_dump()
|
||||
assert json_data["generated_at"] == "2023-12-08T10:30:00"
|
||||
|
||||
Reference in New Issue
Block a user