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

feat(plugin): allow unbind plugin

Browse Source
This commit is contained in:
arkohut 2024-08-25 00:57:44 +08:00
parent cead5d9755
commit ec7ba1f989
3 changed files with 102 additions and 12 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

63
memos/commands.py
View File

@ -198,19 +198,31 @@ async def loop_files(library_id, folder, folder_path, force, plugins):
if file_type_group == "image":
metadata = read_metadata(absolute_file_path)
if metadata:
if "active_window" in metadata and "active_app" not in metadata:
metadata["active_app"] = metadata["active_window"].split(" - ")[0]
if (
"active_window" in metadata
and "active_app" not in metadata
):
metadata["active_app"] = metadata[
"active_window"
].split(" - ")[0]
new_entity["metadata_entries"] = [
{
"key": key,
"value": str(value),
"source": MetadataSource.SYSTEM_GENERATED.value,
"data_type": "number" if isinstance(value, (int, float)) else "text",
"data_type": (
"number"
if isinstance(value, (int, float))
else "text"
),
}
for key, value in metadata.items() if key != IS_THUMBNAIL
for key, value in metadata.items()
if key != IS_THUMBNAIL
]
if "active_app" in metadata:
new_entity.setdefault("tags", []).append(metadata["active_app"])
new_entity.setdefault("tags", []).append(
metadata["active_app"]
)
is_thumbnail = metadata.get(IS_THUMBNAIL, False)
existing_entity = existing_entities_dict.get(
@ -232,18 +244,31 @@ async def loop_files(library_id, folder, folder_path, force, plugins):
# Ignore file changes for thumbnails
if is_thumbnail:
new_entity["file_created_at"] = existing_entity["file_created_at"]
new_entity["file_last_modified_at"] = existing_entity["file_last_modified_at"]
new_entity["file_created_at"] = existing_entity[
"file_created_at"
]
new_entity["file_last_modified_at"] = existing_entity[
"file_last_modified_at"
]
new_entity["file_type"] = existing_entity["file_type"]
new_entity["file_type_group"] = existing_entity["file_type_group"]
new_entity["file_type_group"] = existing_entity[
"file_type_group"
]
new_entity["size"] = existing_entity["size"]
# Merge existing metadata with new metadata
if new_entity.get("metadata_entries"):
new_metadata_keys = {entry["key"] for entry in new_entity["metadata_entries"]}
for existing_entry in existing_entity["metadata_entries"]:
new_metadata_keys = {
entry["key"]
for entry in new_entity["metadata_entries"]
}
for existing_entry in existing_entity[
"metadata_entries"
]:
if existing_entry["key"] not in new_metadata_keys:
new_entity["metadata_entries"].append(existing_entry)
new_entity["metadata_entries"].append(
existing_entry
)
if (
force
@ -650,5 +675,21 @@ def bind(
)
@plugin_app.command("unbind")
def unbind(
library_id: int = typer.Option(..., "--lib", help="ID of the library"),
plugin_id: int = typer.Option(..., "--plugin", help="ID of the plugin"),
):
response = httpx.delete(
f"{BASE_URL}/libraries/{library_id}/plugins/{plugin_id}",
)
if response.status_code == 204:
print("Plugin unbound from library successfully")
else:
print(
f"Failed to unbind plugin from library: {response.status_code} - {response.text}"
)
if __name__ == "__main__":
app()

16
memos/crud.py
View File

@ -366,3 +366,19 @@ def update_entity_metadata_entries(
db.commit()
db.refresh(db_entity)
return Entity(**db_entity.__dict__)
def get_plugin_by_id(plugin_id: int, db: Session) -> Plugin | None:
return db.query(PluginModel).filter(PluginModel.id == plugin_id).first()
def remove_plugin_from_library(library_id: int, plugin_id: int, db: Session):
library_plugin = db.query(LibraryPluginModel).filter(
LibraryPluginModel.library_id == library_id,
LibraryPluginModel.plugin_id == plugin_id
).first()
if library_plugin:
db.delete(library_plugin)
db.commit()
else:
raise ValueError(f"Plugin {plugin_id} not found in library {library_id}")

33
memos/server.py
View File

@ -21,6 +21,7 @@ import typesense
from .config import get_database_path, settings
from .plugins.vlm import main as vlm_main
from .plugins.ocr import main as ocr_main # Add this import
from . import crud
from . import indexing
from .schemas import (
@ -582,6 +583,31 @@ def add_library_plugin(
crud.add_plugin_to_library(library_id, new_plugin.plugin_id, db)
@app.delete(
"/libraries/{library_id}/plugins/{plugin_id}",
status_code=status.HTTP_204_NO_CONTENT,
tags=["plugin"]
)
def delete_library_plugin(
library_id: int, plugin_id: int, db: Session = Depends(get_db)
):
library = crud.get_library_by_id(library_id, db)
if library is None:
raise HTTPException(
status_code=status.HTTP_404_NOT_FOUND,
detail="Library not found"
)
plugin = crud.get_plugin_by_id(plugin_id, db)
if plugin is None:
raise HTTPException(
status_code=status.HTTP_404_NOT_FOUND,
detail="Plugin not found"
)
crud.remove_plugin_from_library(library_id, plugin_id, db)
def is_image(file_path: Path) -> bool:
return file_path.suffix.lower() in [".png", ".jpg", ".jpeg"]
@ -663,6 +689,12 @@ if settings.vlm.enabled:
vlm_main.init_plugin(settings.vlm)
app.include_router(vlm_main.router, prefix=f"/plugins/{vlm_main.PLUGIN_NAME}")
# Add OCR plugin router
if settings.ocr.enabled:
print("OCR plugin is enabled")
ocr_main.init_plugin(settings.ocr)
app.include_router(ocr_main.router, prefix=f"/plugins/{ocr_main.PLUGIN_NAME}")
def run_server():
print("Database path:", get_database_path())
@ -670,6 +702,7 @@ def run_server():
f"Typesense connection info: Host: {settings.typesense_host}, Port: {settings.typesense_port}, Protocol: {settings.typesense_protocol}"
)
print(f"VLM plugin enabled: {settings.vlm}")
print(f"OCR plugin enabled: {settings.ocr}") # Add this line
uvicorn.run(
"memos.server:app",