305 lines
9.3 KiB
Python
305 lines
9.3 KiB
Python
from datetime import datetime
|
|
|
|
from pydantic import BaseModel, ConfigDict, Field
|
|
|
|
from app.agent.tools import ToolExecutionContext, ToolRegistry
|
|
from app.contracts import SearchMode, SearchRequest, TaskStatus, ToolDefinition
|
|
from app.retrieval.engine import engine
|
|
from app.services import note_service
|
|
from app.services import attachment_service, task_service, transcription_service
|
|
|
|
|
|
class ToolArguments(BaseModel):
|
|
model_config = ConfigDict(extra="forbid")
|
|
|
|
|
|
class EchoArguments(ToolArguments):
|
|
text: str
|
|
|
|
|
|
class AddArguments(ToolArguments):
|
|
left: float
|
|
right: float
|
|
|
|
|
|
class NoteSearchArguments(ToolArguments):
|
|
query: str = Field(min_length=1)
|
|
mode: SearchMode = SearchMode.hybrid
|
|
folders: list[str] = Field(default_factory=list)
|
|
note_ids: list[str] = Field(default_factory=list)
|
|
tags: list[str] = Field(default_factory=list)
|
|
limit: int = Field(default=10, ge=1, le=100)
|
|
offset: int = Field(default=0, ge=0)
|
|
|
|
|
|
class NoteReadArguments(ToolArguments):
|
|
note_id: str = Field(min_length=1)
|
|
|
|
|
|
class NoteCreateArguments(ToolArguments):
|
|
title: str = Field(min_length=1)
|
|
markdown: str = ""
|
|
folder: str | None = None
|
|
tags: list[str] = Field(default_factory=list)
|
|
|
|
|
|
class NoteUpdateArguments(ToolArguments):
|
|
note_id: str = Field(min_length=1)
|
|
title: str | None = None
|
|
markdown: str | None = None
|
|
tags: list[str] | None = None
|
|
|
|
|
|
class NoteListArguments(ToolArguments):
|
|
limit: int = Field(default=50, ge=1, le=100)
|
|
offset: int = Field(default=0, ge=0)
|
|
folder: str | None = None
|
|
tag: str | None = None
|
|
|
|
|
|
class NoteMoveArguments(ToolArguments):
|
|
note_id: str = Field(min_length=1)
|
|
folder: str
|
|
|
|
|
|
class TaskCreateArguments(ToolArguments):
|
|
title: str = Field(min_length=1)
|
|
description: str = ""
|
|
note_id: str | None = None
|
|
due_at: datetime | None = None
|
|
|
|
|
|
class TaskUpdateArguments(ToolArguments):
|
|
task_id: str = Field(min_length=1)
|
|
title: str | None = None
|
|
description: str | None = None
|
|
status: TaskStatus | None = None
|
|
note_id: str | None = None
|
|
due_at: datetime | None = None
|
|
|
|
|
|
class TaskListArguments(ToolArguments):
|
|
limit: int = Field(default=50, ge=1, le=100)
|
|
offset: int = Field(default=0, ge=0)
|
|
|
|
|
|
class AttachmentReadArguments(ToolArguments):
|
|
attachment_id: str = Field(min_length=1)
|
|
max_chars: int = Field(default=100_000, ge=1, le=1_000_000)
|
|
|
|
|
|
class AudioTranscribeArguments(ToolArguments):
|
|
attachment_id: str = Field(min_length=1)
|
|
language: str | None = None
|
|
|
|
|
|
async def echo(arguments: EchoArguments, _: ToolExecutionContext) -> dict[str, str]:
|
|
return {"text": arguments.text}
|
|
|
|
|
|
async def add(arguments: AddArguments, _: ToolExecutionContext) -> dict[str, float]:
|
|
return {"value": arguments.left + arguments.right}
|
|
|
|
|
|
async def search_notes(arguments: NoteSearchArguments, _: ToolExecutionContext) -> dict:
|
|
request = SearchRequest(**arguments.model_dump(), include_snippet=True)
|
|
return (await engine.search(request)).model_dump(mode="json")
|
|
|
|
|
|
async def read_note(arguments: NoteReadArguments, _: ToolExecutionContext) -> dict:
|
|
note = await note_service.get_note(arguments.note_id)
|
|
if note is None:
|
|
raise LookupError(f"Note does not exist: {arguments.note_id}")
|
|
import hashlib
|
|
return {**note.model_dump(mode="json"), "content_hash": hashlib.sha256(note.markdown.encode()).hexdigest()}
|
|
|
|
|
|
async def create_note(arguments: NoteCreateArguments, _: ToolExecutionContext) -> dict:
|
|
note = await note_service.create_note(**arguments.model_dump())
|
|
return note.model_dump(mode="json")
|
|
|
|
|
|
async def update_note(arguments: NoteUpdateArguments, _: ToolExecutionContext) -> dict:
|
|
values = arguments.model_dump()
|
|
note_id = values.pop("note_id")
|
|
note = await note_service.update_note(note_id, **values)
|
|
return note.model_dump(mode="json")
|
|
|
|
|
|
def list_notes(arguments: NoteListArguments, _: ToolExecutionContext) -> dict:
|
|
items, total = note_service.list_notes(**arguments.model_dump())
|
|
return {
|
|
"items": [item.model_dump(mode="json") for item in items],
|
|
"page": {"total": total, "limit": arguments.limit, "offset": arguments.offset},
|
|
}
|
|
|
|
|
|
async def move_note(arguments: NoteMoveArguments, _: ToolExecutionContext) -> dict:
|
|
note = await note_service.move_note(arguments.note_id, folder=arguments.folder)
|
|
return note.model_dump(mode="json")
|
|
|
|
|
|
def create_task(arguments: TaskCreateArguments, _: ToolExecutionContext) -> dict:
|
|
return task_service.create_task(**arguments.model_dump()).model_dump(mode="json")
|
|
|
|
|
|
def update_task(arguments: TaskUpdateArguments, _: ToolExecutionContext) -> dict:
|
|
values = arguments.model_dump(exclude_unset=True)
|
|
task_id = values.pop("task_id")
|
|
return task_service.update_task(task_id, values).model_dump(mode="json")
|
|
|
|
|
|
def list_tasks(arguments: TaskListArguments, _: ToolExecutionContext) -> dict:
|
|
items, total = task_service.list_tasks(**arguments.model_dump())
|
|
return {
|
|
"items": [item.model_dump(mode="json") for item in items],
|
|
"page": {"total": total, "limit": arguments.limit, "offset": arguments.offset},
|
|
}
|
|
|
|
|
|
def read_attachment(arguments: AttachmentReadArguments, _: ToolExecutionContext) -> dict:
|
|
return attachment_service.read_attachment(**arguments.model_dump())
|
|
|
|
|
|
async def transcribe_audio(arguments: AudioTranscribeArguments, _: ToolExecutionContext) -> dict:
|
|
job = await transcription_service.create_transcription(
|
|
arguments.attachment_id, arguments.language
|
|
)
|
|
return job.model_dump(mode="json")
|
|
|
|
|
|
def _register(
|
|
registry: ToolRegistry,
|
|
*,
|
|
name: str,
|
|
description: str,
|
|
arguments_model: type[BaseModel],
|
|
executor,
|
|
permission: str | None = None,
|
|
) -> None:
|
|
registry.register(
|
|
ToolDefinition(
|
|
name=name,
|
|
description=description,
|
|
parameters=arguments_model.model_json_schema(),
|
|
permission=permission,
|
|
),
|
|
arguments_model,
|
|
executor,
|
|
)
|
|
|
|
|
|
def register_builtin_tools(registry: ToolRegistry) -> None:
|
|
from app.agent.markdown_tools import register
|
|
register(registry)
|
|
_register(
|
|
registry,
|
|
name="system.echo",
|
|
description="Echo text for local Agent integration testing.",
|
|
arguments_model=EchoArguments,
|
|
executor=echo,
|
|
)
|
|
_register(
|
|
registry,
|
|
name="math.add",
|
|
description="Add two numbers without external side effects.",
|
|
arguments_model=AddArguments,
|
|
executor=add,
|
|
)
|
|
_register(
|
|
registry,
|
|
name="notes.search",
|
|
description="Search indexed notes and return snippets with citations.",
|
|
arguments_model=NoteSearchArguments,
|
|
executor=search_notes,
|
|
permission="notes.search",
|
|
)
|
|
_register(
|
|
registry,
|
|
name="rag.search",
|
|
description="Retrieve relevant note blocks for Agent context with citations.",
|
|
arguments_model=NoteSearchArguments,
|
|
executor=search_notes,
|
|
permission="notes.search",
|
|
)
|
|
_register(
|
|
registry,
|
|
name="notes.read",
|
|
description="Read a note and its parsed blocks by note_id.",
|
|
arguments_model=NoteReadArguments,
|
|
executor=read_note,
|
|
permission="notes.read",
|
|
)
|
|
_register(
|
|
registry,
|
|
name="notes.create",
|
|
description="Create a Markdown note in the current Vault.",
|
|
arguments_model=NoteCreateArguments,
|
|
executor=create_note,
|
|
permission="notes.write",
|
|
)
|
|
_register(
|
|
registry,
|
|
name="notes.update",
|
|
description="Update an existing Markdown note.",
|
|
arguments_model=NoteUpdateArguments,
|
|
executor=update_note,
|
|
permission="notes.write",
|
|
)
|
|
_register(
|
|
registry,
|
|
name="notes.list",
|
|
description="List note summaries with folder and tag filters.",
|
|
arguments_model=NoteListArguments,
|
|
executor=list_notes,
|
|
permission="notes.read",
|
|
)
|
|
_register(
|
|
registry,
|
|
name="notes.move",
|
|
description="Move a note to another folder while preserving note_id.",
|
|
arguments_model=NoteMoveArguments,
|
|
executor=move_note,
|
|
permission="notes.write",
|
|
)
|
|
_register(
|
|
registry,
|
|
name="tasks.create",
|
|
description="Create a persistent task.",
|
|
arguments_model=TaskCreateArguments,
|
|
executor=create_task,
|
|
permission="tasks.write",
|
|
)
|
|
_register(
|
|
registry,
|
|
name="tasks.update",
|
|
description="Update a persistent task.",
|
|
arguments_model=TaskUpdateArguments,
|
|
executor=update_task,
|
|
permission="tasks.write",
|
|
)
|
|
_register(
|
|
registry,
|
|
name="tasks.list",
|
|
description="List persistent tasks.",
|
|
arguments_model=TaskListArguments,
|
|
executor=list_tasks,
|
|
permission="tasks.read",
|
|
)
|
|
_register(
|
|
registry,
|
|
name="attachments.read",
|
|
description="Read a UTF-8 attachment from host-managed attachment storage.",
|
|
arguments_model=AttachmentReadArguments,
|
|
executor=read_attachment,
|
|
permission="attachments.read",
|
|
)
|
|
_register(
|
|
registry,
|
|
name="audio.transcribe",
|
|
description="Read a host-generated transcript for an audio attachment.",
|
|
arguments_model=AudioTranscribeArguments,
|
|
executor=transcribe_audio,
|
|
permission="attachments.read",
|
|
)
|