From d56197c2888afbfaca93b52e24de979ff6ddfa01 Mon Sep 17 00:00:00 2001 From: phernandez Date: Thu, 12 Dec 2024 20:29:11 -0600 Subject: [PATCH] default created_at for observation, relation --- .../20241213022126_fix-created-at-default.sql | 50 +++++++++++++++++++ db/schema.sql | 47 +++++++++-------- src/basic_memory/models.py | 15 +----- src/basic_memory/services/memory_service.py | 19 ++++--- 4 files changed, 84 insertions(+), 47 deletions(-) create mode 100644 db/migrations/20241213022126_fix-created-at-default.sql diff --git a/db/migrations/20241213022126_fix-created-at-default.sql b/db/migrations/20241213022126_fix-created-at-default.sql new file mode 100644 index 00000000..ccf384e6 --- /dev/null +++ b/db/migrations/20241213022126_fix-created-at-default.sql @@ -0,0 +1,50 @@ +-- migrate:up + +-- Create new observation table with correct default +CREATE TABLE observation_new ( + id INTEGER NOT NULL PRIMARY KEY AUTOINCREMENT, + entity_id VARCHAR NOT NULL, + content VARCHAR NOT NULL, + created_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP, + context VARCHAR, + FOREIGN KEY(entity_id) REFERENCES entity (id) ON DELETE CASCADE +); + +-- Copy data from old observation table +INSERT INTO observation_new +SELECT id, entity_id, content, COALESCE(created_at, CURRENT_TIMESTAMP), context +FROM observation; + +-- Drop old observation table and rename new one +DROP TABLE observation; +ALTER TABLE observation_new RENAME TO observation; + +-- Recreate observation index +CREATE INDEX ix_observation_entity_id ON observation (entity_id); + +-- Create new relation table with correct default +CREATE TABLE relation_new ( + id INTEGER NOT NULL PRIMARY KEY AUTOINCREMENT, + from_id VARCHAR NOT NULL, + to_id VARCHAR NOT NULL, + relation_type VARCHAR NOT NULL, + created_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP, + context VARCHAR, + FOREIGN KEY(from_id) REFERENCES entity (id) ON DELETE CASCADE, + FOREIGN KEY(to_id) REFERENCES entity (id) ON DELETE CASCADE +); + +-- Copy data from old relation table +INSERT INTO relation_new +SELECT id, from_id, to_id, relation_type, COALESCE(created_at, CURRENT_TIMESTAMP), context +FROM relation; + +-- Drop old relation table and rename new one +DROP TABLE relation; +ALTER TABLE relation_new RENAME TO relation; + +-- Recreate relation indexes +CREATE INDEX ix_relation_from_id ON relation (from_id); +CREATE INDEX ix_relation_to_id ON relation (to_id); + +-- migrate:down diff --git a/db/schema.sql b/db/schema.sql index da9a6c1f..b980064b 100644 --- a/db/schema.sql +++ b/db/schema.sql @@ -1,26 +1,3 @@ -CREATE TABLE observation ( - id INTEGER NOT NULL, - entity_id VARCHAR NOT NULL, - content VARCHAR NOT NULL, - created_at DATETIME NOT NULL, - context VARCHAR, - PRIMARY KEY (id), - FOREIGN KEY(entity_id) REFERENCES entity (id) ON DELETE CASCADE -); -CREATE INDEX ix_observation_entity_id ON observation (entity_id); -CREATE TABLE relation ( - id INTEGER NOT NULL, - from_id VARCHAR NOT NULL, - to_id VARCHAR NOT NULL, - relation_type VARCHAR NOT NULL, - created_at DATETIME NOT NULL, - context VARCHAR, - PRIMARY KEY (id), - FOREIGN KEY(from_id) REFERENCES entity (id) ON DELETE CASCADE, - FOREIGN KEY(to_id) REFERENCES entity (id) ON DELETE CASCADE -); -CREATE INDEX ix_relation_to_id ON relation (to_id); -CREATE INDEX ix_relation_from_id ON relation (from_id); CREATE TABLE IF NOT EXISTS "schema_migrations" (version varchar(128) primary key); CREATE TABLE IF NOT EXISTS "entity" ( id TEXT PRIMARY KEY, @@ -31,10 +8,32 @@ CREATE TABLE IF NOT EXISTS "entity" ( updated_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP ); CREATE UNIQUE INDEX idx_entity_type_name ON entity(entity_type, name); +CREATE TABLE IF NOT EXISTS "observation" ( + id INTEGER NOT NULL PRIMARY KEY AUTOINCREMENT, + entity_id VARCHAR NOT NULL, + content VARCHAR NOT NULL, + created_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP, + context VARCHAR, + FOREIGN KEY(entity_id) REFERENCES entity (id) ON DELETE CASCADE +); +CREATE INDEX ix_observation_entity_id ON observation (entity_id); +CREATE TABLE IF NOT EXISTS "relation" ( + id INTEGER NOT NULL PRIMARY KEY AUTOINCREMENT, + from_id VARCHAR NOT NULL, + to_id VARCHAR NOT NULL, + relation_type VARCHAR NOT NULL, + created_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP, + context VARCHAR, + FOREIGN KEY(from_id) REFERENCES entity (id) ON DELETE CASCADE, + FOREIGN KEY(to_id) REFERENCES entity (id) ON DELETE CASCADE +); +CREATE INDEX ix_relation_from_id ON relation (from_id); +CREATE INDEX ix_relation_to_id ON relation (to_id); -- Dbmate schema migrations INSERT INTO "schema_migrations" (version) VALUES ('20240101000000'), ('20241210213454'), ('20241211034719'), ('20241211052101'), - ('20241211190000'); + ('20241211190000'), + ('20241213022126'); diff --git a/src/basic_memory/models.py b/src/basic_memory/models.py index 3f4d84c3..59b956d1 100644 --- a/src/basic_memory/models.py +++ b/src/basic_memory/models.py @@ -4,7 +4,6 @@ from typing import List, Optional from sqlalchemy import String, DateTime, ForeignKey, Text, TypeDecorator, Integer, text, UniqueConstraint from sqlalchemy.orm import Mapped, mapped_column, relationship, DeclarativeBase from sqlalchemy.ext.asyncio import AsyncAttrs -from sqlalchemy import orm class UTCDateTime(TypeDecorator): @@ -29,20 +28,10 @@ def utc_now() -> datetime: """Helper to get current UTC time""" return datetime.now(UTC) -def lenient_constructor(self, **kwargs): - cls_ = type(self) - for k in kwargs: - if not hasattr(cls_, k): - print(f'Skipping invalid attr {k!r}') - continue - setattr(self, k, kwargs[k]) - -registry = orm.registry(constructor=lenient_constructor) - class Base(AsyncAttrs, DeclarativeBase): """Base class for all models""" - registry = registry + pass class Entity(Base): @@ -127,7 +116,7 @@ class Observation(Base): ) content: Mapped[str] = mapped_column(String) created_at: Mapped[datetime] = mapped_column( - UTCDateTime, + DateTime, server_default=text('CURRENT_TIMESTAMP') ) context: Mapped[str | None] = mapped_column(String, nullable=True) diff --git a/src/basic_memory/services/memory_service.py b/src/basic_memory/services/memory_service.py index a876de7e..c89e0d5b 100644 --- a/src/basic_memory/services/memory_service.py +++ b/src/basic_memory/services/memory_service.py @@ -22,8 +22,11 @@ class MemoryService: relation_service: RelationService, observation_service: ObservationService ): - self.project_path = project_path - self.entities_path = project_path / "entities" if project_path else None + if project_path: + assert project_path.is_dir(), "Path does not exist or is not a directory: {project_path}" + self.project_path = project_path + self.entities_path = project_path / "entities" + self.entity_service = entity_service self.relation_service = relation_service self.observation_service = observation_service @@ -64,13 +67,9 @@ class MemoryService: created_entity = await self.entity_service.create_entity(entity_in) logger.debug(f"Created base entity: {created_entity.id}") - # Convert ObservationIn to Observation instances - if entity_in.observations: - created_observations = await self.observation_service.add_observations( - created_entity.id, - [ObservationIn(**obs.model_dump()) for obs in entity_in.observations] - ) - logger.debug(f"Added {len(created_observations)} observations to {created_entity.id}") + # Add observations + await self.observation_service.add_observations(created_entity.id, entity_in.observations) + logger.debug(f"Added {len(entity_in.observations)} observations to {created_entity.id}") # Add relations for relation in entity_in.relations: @@ -104,7 +103,7 @@ class MemoryService: if path.exists(): path.unlink() except Exception as cleanup_error: - logger.error(f"Failed to clean up file for {entity_id}: {cleanup_error}") + logger.error(f"Failed to clean up file for {entity.id}: {cleanup_error}") raise async def create_relations(self, relations_data: List[RelationIn]) -> List[Relation]: