tcsenpai/pensieve
1
0
Fork 0
mirror of https://github.com/tcsenpai/pensieve.git synced 2025-06-10 13:07:15 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat: support scan and trigger specific plugins

Browse Source
This commit is contained in:
arkohut 2024-07-10 18:35:44 +08:00
parent d1b83571e6
commit b5d524f7b0
3 changed files with 25 additions and 26 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
memos/commands.py
View File

@ -98,7 +98,7 @@ def show(library_id: int):
@lib_app.command("scan")
def scan(library_id: int, force: bool = False):
def scan(library_id: int, force: bool = False, plugins: List[int] = typer.Option(None, "--plugin", "-p")):
response = httpx.get(f"{BASE_URL}/libraries/{library_id}")
if response.status_code != 200:
@ -187,7 +187,7 @@ def scan(library_id: int, force: bool = False):
update_response = httpx.put(
f"{BASE_URL}/entities/{existing_entity['id']}",
json=new_entity,
params={"trigger_webhooks_flag": "true"},
params={"trigger_webhooks_flag": "true", **({"plugins": plugins} if plugins else {})},
timeout=30
)
if 200 <= update_response.status_code < 300:
@ -207,6 +207,7 @@ def scan(library_id: int, force: bool = False):
post_response = httpx.post(
f"{BASE_URL}/libraries/{library_id}/entities",
json=new_entity,
params={"plugins": plugins} if plugins else {},
timeout=30
)
if 200 <= post_response.status_code < 300:

40
memos/server.py
View File

@ -152,29 +152,31 @@ def new_folder(
return crud.add_folder(library_id=library.id, folder=folder, db=db)
async def trigger_webhooks(library: Library, entity: Entity, request: Request):
async def trigger_webhooks(library: Library, entity: Entity, request: Request, plugins: List[int] = None):
async with httpx.AsyncClient() as client:
tasks = []
for plugin in library.plugins:
if plugin.webhook_url:
location = str(request.url_for("get_entity_by_id", entity_id=entity.id))
task = client.post(
plugin.webhook_url,
json=entity.model_dump(mode="json"),
headers={"Location": location},
timeout=10.0, # Adding a timeout of 10 seconds
)
tasks.append(task)
if plugins is None or plugin.id in plugins:
if plugin.webhook_url:
location = str(request.url_for("get_entity_by_id", entity_id=entity.id))
task = client.post(
plugin.webhook_url,
json=entity.model_dump(mode="json"),
headers={"Location": location},
timeout=10.0, # Adding a timeout of 10 seconds
)
tasks.append(task)
responses = await asyncio.gather(*tasks, return_exceptions=True)
for plugin, response in zip(library.plugins, responses):
if isinstance(response, Exception):
print(f"Error triggering webhook for plugin {plugin.id}: {response}")
elif response.status_code >= 400:
print(
f"Error triggering webhook for plugin {plugin.id}: {response.status_code} - {response.text}"
)
if plugins is None or plugin.id in plugins:
if isinstance(response, Exception):
print(f"Error triggering webhook for plugin {plugin.id}: {response}")
elif response.status_code >= 400:
print(
f"Error triggering webhook for plugin {plugin.id}: {response.status_code} - {response.text}"
)
@app.post("/libraries/{library_id}/entities", response_model=Entity, tags=["entity"])
@ -183,6 +185,7 @@ async def new_entity(
library_id: int,
request: Request,
db: Session = Depends(get_db),
plugins: Annotated[List[int] | None, Query()] = None,
):
library = crud.get_library_by_id(library_id, db)
if library is None:
@ -191,7 +194,7 @@ async def new_entity(
)
entity = crud.create_entity(library_id, new_entity, db)
await trigger_webhooks(library, entity, request)
await trigger_webhooks(library, entity, request, plugins)
return entity
@ -271,6 +274,7 @@ async def update_entity(
request: Request,
db: Session = Depends(get_db),
trigger_webhooks_flag: bool = False,
plugins: Annotated[List[int] | None, Query()] = None,
):
entity = crud.get_entity_by_id(entity_id, db)
if entity is None:
@ -286,7 +290,7 @@ async def update_entity(
raise HTTPException(
status_code=status.HTTP_404_NOT_FOUND, detail="Library not found"
)
await trigger_webhooks(library, entity, request)
await trigger_webhooks(library, entity, request, plugins)
return entity

6
memos/test_server.py
View File

@ -93,12 +93,6 @@ def client():
Base.metadata.drop_all(bind=engine)
def test_read_main(client):
response = client.get("/")
assert response.status_code == 200
assert response.json() == {"healthy": True}
# Test the new_library endpoint
def test_new_library(client):
library_param = NewLibraryParam(name="Test Library")