tcsenpai/pensieve
1
0
Fork 0
mirror of https://github.com/tcsenpai/pensieve.git synced 2025-06-06 03:05:25 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat(entity): update tags and metadata

Browse Source
This commit is contained in:
arkohut 2024-06-10 00:29:27 +08:00
parent 0db275fba0
commit 8ec89219f3
5 changed files with 253 additions and 60 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

43
memos/crud.py
View File

@ -11,6 +11,7 @@ from .schemas import (
NewPluginParam,
UpdateEntityParam,
NewFolderParam,
MetadataSource,
)
from .models import (
LibraryModel,
@ -19,6 +20,9 @@ from .models import (
EntityModel,
PluginModel,
LibraryPluginModel,
TagModel,
EntityMetadataModel,
EntityTagModel,
)
@ -149,8 +153,45 @@ def update_entity(
entity_id: int, updated_entity: UpdateEntityParam, db: Session
) -> Entity:
db_entity = get_entity_by_id(entity_id, db)
# Update the main fields of the entity
for key, value in updated_entity.model_dump().items():
setattr(db_entity, key, value)
if key not in ["tags", "attrs"] and value is not None:
setattr(db_entity, key, value)
# Handle tags separately
if updated_entity.tags is not None:
db_entity.tags = []
for tag_name in updated_entity.tags:
tag = db.query(TagModel).filter(TagModel.name == tag_name).first()
if not tag:
tag = TagModel(name=tag_name)
db.add(tag)
db.commit()
db.refresh(tag)
entity_tag = EntityTagModel(
entity_id=db_entity.id,
tag_id=tag.id,
source=MetadataSource.PLUGIN_GENERATED,
)
db.add(entity_tag)
db.commit()
# Handle attrs separately
if updated_entity.attrs is not None:
db_entity.attrs = []
for attr in updated_entity.attrs:
entity_metadata = EntityMetadataModel(
entity_id=db_entity.id,
key=attr.key,
value=attr.value,
source=attr.source if attr.source is not None else None,
source_type=MetadataSource.PLUGIN_GENERATED if attr.source is not None else None,
data_type=attr.data_type,
)
db.add(entity_metadata)
db_entity.attrs.append(entity_metadata)
db.commit()
db.refresh(db_entity)
return Entity(**db_entity.__dict__)

10
memos/models.py
View File

@ -69,9 +69,10 @@ class EntityModel(Base):
"FolderModel", back_populates="entities"
)
metadata_entries: Mapped[List["EntityMetadataModel"]] = relationship(
"EntityMetadataModel"
"EntityMetadataModel", lazy="joined"
)
tags: Mapped[List["TagModel"]] = relationship("EntityTagModel")
tags: Mapped[List["TagModel"]] = relationship("TagModel", secondary="entity_tags", lazy="joined")
class TagModel(Base):
@ -79,7 +80,7 @@ class TagModel(Base):
name: Mapped[str] = mapped_column(String, nullable=False)
description: Mapped[str | None] = mapped_column(Text, nullable=True)
color: Mapped[str | None] = mapped_column(String, nullable=True)
source: Mapped[str | None] = mapped_column(String, nullable=True)
# source: Mapped[str | None] = mapped_column(String, nullable=True)
class EntityTagModel(Base):
@ -102,7 +103,7 @@ class EntityMetadataModel(Base):
Enum(MetadataSource), nullable=False
)
source: Mapped[str | None] = mapped_column(String, nullable=True)
date_type: Mapped[MetadataType] = mapped_column(Enum(MetadataType), nullable=False)
data_type: Mapped[MetadataType] = mapped_column(Enum(MetadataType), nullable=False)
entity = relationship("EntityModel", back_populates="metadata_entries")
@ -121,7 +122,6 @@ class LibraryPluginModel(Base):
plugin_id: Mapped[int] = mapped_column(
Integer, ForeignKey("plugins.id"), nullable=False
)
# Create the database engine with the path from config

69
memos/schemas.py
View File

@ -34,11 +34,20 @@ class NewEntityParam(BaseModel):
folder_id: int
class EntityMetadataParam(BaseModel):
key: str
value: str
source: str
data_type: MetadataType
class UpdateEntityParam(BaseModel):
size: int
file_created_at: datetime
file_last_modified_at: datetime
file_type: str
size: int | None = None
file_created_at: datetime | None = None
file_last_modified_at: datetime | None = None
file_type: str | None = None
tags: List[str] = []
attrs: List[EntityMetadataParam] = []
class UpdateTagParam(BaseModel):
@ -50,13 +59,6 @@ class UpdateEntityTagsParam(BaseModel):
tags: List[str] = []
class EntityMetadataParam(BaseModel):
key: str
value: str
source: MetadataSource
data_type: MetadataType
class UpdateEntityMetadataParam(BaseModel):
metadata_entries: List[EntityMetadataParam]
@ -96,6 +98,28 @@ class Library(BaseModel):
model_config = ConfigDict(from_attributes=True)
class Tag(BaseModel):
id: int
name: str
description: str | None
color: str | None
created_at: datetime
# source: str
model_config = ConfigDict(from_attributes=True)
class EntityMetadata(BaseModel):
id: int
entity_id: int
key: str
value: str
source: str
data_type: MetadataType
model_config = ConfigDict(from_attributes=True)
class Entity(BaseModel):
id: int
filepath: str
@ -107,27 +131,8 @@ class Entity(BaseModel):
last_scan_at: datetime | None
folder_id: int
library_id: int
tags: List[Tag] = []
metadata_entries: List[EntityMetadata] = []
model_config = ConfigDict(from_attributes=True)
class Tag(BaseModel):
id: int
name: str
description: str | None
color: str | None
created_at: datetime
source: str
model_config = ConfigDict(from_attributes=True)
class EntityMetadata(BaseModel):
id: int
entity_id: int
key: str
value: str
source: str
date_type: MetadataType
model_config = ConfigDict(from_attributes=True)

4
memos/server.py
View File

@ -197,6 +197,7 @@ async def update_entity(
updated_entity: UpdateEntityParam,
request: Request,
db: Session = Depends(get_db),
trigger_webhooks_flag: bool = False,
):
entity = crud.get_entity_by_id(entity_id, db)
if entity is None or entity.library_id != library_id:
@ -212,7 +213,8 @@ async def update_entity(
)
entity = crud.update_entity(entity_id, updated_entity, db)
await trigger_webhooks(library, entity, request)
if trigger_webhooks_flag:
await trigger_webhooks(library, entity, request)
return entity

187
memos/test_server.py
View File

@ -15,6 +15,8 @@ from memos.schemas import (
NewEntityParam,
UpdateEntityParam,
NewFolderParam,
EntityMetadataParam,
MetadataType,
)
from memos.models import Base
@ -72,7 +74,9 @@ def test_new_library(client):
duplicate_response = client.post("/libraries", json=library_param.model_dump())
# Check that the response indicates a failure due to duplicate name
assert duplicate_response.status_code == 400
assert duplicate_response.json() == {"detail": "Library with this name already exists"}
assert duplicate_response.json() == {
"detail": "Library with this name already exists"
}
def test_list_libraries(client):
@ -265,9 +269,7 @@ def test_get_entity_by_filepath(client):
def test_list_entities_in_folder(client):
# Setup data: Create a new library and folder
new_library = NewLibraryParam(
name="Library for List Entities Test", folders=[]
)
new_library = NewLibraryParam(name="Library for List Entities Test", folders=[])
library_response = client.post(
"/libraries", json=new_library.model_dump(mode="json")
)
@ -325,12 +327,16 @@ def test_list_entities_in_folder(client):
def test_remove_entity(client):
# Create a new library
new_library = NewLibraryParam(name="Test Library")
library_response = client.post("/libraries", json=new_library.model_dump(mode="json"))
library_response = client.post(
"/libraries", json=new_library.model_dump(mode="json")
)
library_id = library_response.json()["id"]
# Create a new folder in the library
new_folder = NewFolderParam(path="/tmp")
folder_response = client.post(f"/libraries/{library_id}/folders", json=new_folder.model_dump(mode="json"))
folder_response = client.post(
f"/libraries/{library_id}/folders", json=new_folder.model_dump(mode="json")
)
folder_id = folder_response.json()["id"]
# Create a new entity to be deleted
@ -360,7 +366,9 @@ def test_remove_entity(client):
# Test for entity not found in the specified library
invalid_delete_response = client.delete(f"/libraries/{library_id}/entities/9999")
assert invalid_delete_response.status_code == 404
assert invalid_delete_response.json() == {"detail": "Entity not found in the specified library"}
assert invalid_delete_response.json() == {
"detail": "Entity not found in the specified library"
}
def test_add_folder_to_library(client):
@ -374,12 +382,16 @@ def test_add_folder_to_library(client):
# Create a new library
new_library = NewLibraryParam(name="Test Library")
library_response = client.post("/libraries", json=new_library.model_dump(mode="json"))
library_response = client.post(
"/libraries", json=new_library.model_dump(mode="json")
)
library_id = library_response.json()["id"]
# Add a new folder to the library
new_folder = NewFolderParam(path="/tmp/new_folder")
folder_response = client.post(f"/libraries/{library_id}/folders", json=new_folder.model_dump(mode="json"))
folder_response = client.post(
f"/libraries/{library_id}/folders", json=new_folder.model_dump(mode="json")
)
assert folder_response.status_code == 200
assert folder_response.json()["path"] == "/tmp/new_folder"
@ -391,41 +403,174 @@ def test_add_folder_to_library(client):
assert "/tmp/new_folder" in folder_paths
# Test for adding a folder that already exists
duplicate_folder_response = client.post(f"/libraries/{library_id}/folders", json=new_folder.model_dump(mode="json"))
duplicate_folder_response = client.post(
f"/libraries/{library_id}/folders", json=new_folder.model_dump(mode="json")
)
assert duplicate_folder_response.status_code == 400
assert duplicate_folder_response.json() == {"detail": "Folder already exists in the library"}
assert duplicate_folder_response.json() == {
"detail": "Folder already exists in the library"
}
# Test for adding a folder to a non-existent library
invalid_folder_response = client.post(f"/libraries/9999/folders", json=new_folder.model_dump(mode="json"))
invalid_folder_response = client.post(
f"/libraries/9999/folders", json=new_folder.model_dump(mode="json")
)
assert invalid_folder_response.status_code == 404
assert invalid_folder_response.json() == {"detail": "Library not found"}
def test_new_plugin(client):
new_plugin = NewPluginParam(name="Test Plugin", description="A test plugin", webhook_url="http://example.com/webhook")
new_plugin = NewPluginParam(
name="Test Plugin",
description="A test plugin",
webhook_url="http://example.com/webhook",
)
# Make a POST request to the /plugins endpoint
response = client.post("/plugins", json=new_plugin.model_dump(mode="json"))
# Check that the response is successful
assert response.status_code == 200
# Check the response data
plugin_data = response.json()
assert plugin_data["name"] == "Test Plugin"
assert plugin_data["description"] == "A test plugin"
assert plugin_data["webhook_url"] == "http://example.com/webhook"
# Test for duplicate plugin name
duplicate_response = client.post("/plugins", json=new_plugin.model_dump(mode="json"))
duplicate_response = client.post(
"/plugins", json=new_plugin.model_dump(mode="json")
)
# Check that the response indicates a failure due to duplicate name
assert duplicate_response.status_code == 400
assert duplicate_response.json() == {"detail": "Plugin with this name already exists"}
assert duplicate_response.json() == {
"detail": "Plugin with this name already exists"
}
# Test for another duplicate plugin name
another_duplicate_response = client.post("/plugins", json=new_plugin.model_dump(mode="json"))
another_duplicate_response = client.post(
"/plugins", json=new_plugin.model_dump(mode="json")
)
# Check that the response indicates a failure due to duplicate name
assert another_duplicate_response.status_code == 400
assert another_duplicate_response.json() == {"detail": "Plugin with this name already exists"}
assert another_duplicate_response.json() == {
"detail": "Plugin with this name already exists"
}
def test_update_entity_with_tags(client):
# Create a new library
new_library = NewLibraryParam(name="Test Library")
library_response = client.post(
"/libraries", json=new_library.model_dump(mode="json")
)
assert library_response.status_code == 200
library_id = library_response.json()["id"]
# Create a new folder in the library
new_folder = NewFolderParam(path="/tmp/new_folder")
folder_response = client.post(
f"/libraries/{library_id}/folders", json=new_folder.model_dump(mode="json")
)
assert folder_response.status_code == 200
folder_id = folder_response.json()["id"]
# Create a new entity in the folder
new_entity = NewEntityParam(
filename="test_file.txt",
filepath="/tmp/new_folder/test_file.txt",
size=1234,
file_created_at="2023-01-01T00:00:00",
file_last_modified_at="2023-01-01T00:00:00",
file_type="text/plain",
folder_id=folder_id,
)
entity_response = client.post(
f"/libraries/{library_id}/entities", json=new_entity.model_dump(mode="json")
)
assert entity_response.status_code == 200
entity_id = entity_response.json()["id"]
# Update the entity with tags
update_entity_param = UpdateEntityParam(tags=["tag1", "tag2"])
# Make a PUT request to the /libraries/{library_id}/entities/{entity_id} endpoint
update_response = client.put(
f"/libraries/{library_id}/entities/{entity_id}",
json=update_entity_param.model_dump(mode="json"),
)
# Check that the response is successful
assert update_response.status_code == 200
# Check the response data
updated_entity_data = update_response.json()
assert "tags" in updated_entity_data
assert len(updated_entity_data["tags"]) == 2
assert "tag1" in [tag["name"] for tag in updated_entity_data["tags"]]
assert "tag2" in [tag["name"] for tag in updated_entity_data["tags"]]
def test_add_metadata_entry_to_entity_success(client):
# Create a new library
new_library = NewLibraryParam(name="Test Library for Metadata")
library_response = client.post(
"/libraries", json=new_library.model_dump(mode="json")
)
assert library_response.status_code == 200
library_id = library_response.json()["id"]
# Create a new folder in the library
new_folder = NewFolderParam(path="/tmp")
folder_response = client.post(
f"/libraries/{library_id}/folders", json=new_folder.model_dump(mode="json")
)
assert folder_response.status_code == 200
folder_id = folder_response.json()["id"]
# Create a new entity in the folder
new_entity = NewEntityParam(
filename="metadata_test_file.txt",
filepath="/tmp/metadata_folder/metadata_test_file.txt",
size=5678,
file_created_at="2023-01-01T00:00:00",
file_last_modified_at="2023-01-01T00:00:00",
file_type="text/plain",
folder_id=folder_id,
)
entity_response = client.post(
f"/libraries/{library_id}/entities", json=new_entity.model_dump(mode="json")
)
assert entity_response.status_code == 200
entity_id = entity_response.json()["id"]
# Add metadata entry to the entity
metadata_entry = EntityMetadataParam(
key="author",
value="John Doe",
source="plugin_generated",
data_type=MetadataType.ATTRIBUTE,
)
update_entity_param = UpdateEntityParam(attrs=[metadata_entry])
# Make a PUT request to the /libraries/{library_id}/entities/{entity_id} endpoint
update_response = client.put(
f"/libraries/{library_id}/entities/{entity_id}",
json=update_entity_param.model_dump(mode="json"),
)
# Check that the response is successful
assert update_response.status_code == 200
# Check the response data
updated_entity_data = update_response.json()
assert "metadata_entries" in updated_entity_data
assert len(updated_entity_data["metadata_entries"]) == 1
assert updated_entity_data["metadata_entries"][0]["key"] == "author"
assert updated_entity_data["metadata_entries"][0]["value"] == "John Doe"
assert updated_entity_data["metadata_entries"][0]["source"] == "plugin_generated"
assert (
updated_entity_data["metadata_entries"][0]["data_type"]
== MetadataType.ATTRIBUTE.value
)