api/app/agents/note_agent.py

"""Note agent — Markdown note management (list, get, create, update, delete)."""

from __future__ import annotations

import json
from typing import Any

from langchain_core.messages import HumanMessage, SystemMessage
from langchain_core.tools import tool

from app.core.agent_registry import ChatAgent, registry
from app.core.llm import embed, get_llm
from app.core.ws_context import execute_on_client

_SYSTEM_PROMPT = (
    "You are a note-taking assistant. You help users create, retrieve, update,\n"
    "and delete Markdown notes in their workspace.\n\n"
    "Rules:\n"
    "  - content is always Markdown; preserve formatting when updating\n"
    "  - project_id is optional; link a note to a project when mentioned\n"
    "  - When updating, call get_note first if you need to read existing content\n"
    "    before appending or replacing sections\n"
    "  - list_notes without project_id returns all notes; scope with project_id\n"
    "    when the user is working within a specific project\n"
    "  - Do not fabricate note content — reflect what the user provides or what\n"
    "    is already in the note (retrieved via get_note)."
)


@tool
async def list_notes(project_id: str = "") -> str:
    """List notes, optionally scoped to a project by project_id."""
    result = await execute_on_client(
        action="select",
        table="notes",
        filters={"projectId": project_id or None},
    )
    rows = result.get("rows", [])
    if not rows:
        return "No notes found."
    lines = [f"- {r['title']} (id: {r['id']})" for r in rows]
    return f"Found {len(rows)} note(s):\n" + "\n".join(lines)


@tool
async def get_note(note_id: str) -> str:
    """Fetch a single note by its UUID to read its full Markdown content."""
    result = await execute_on_client(action="get", table="notes", data={"id": note_id})
    row = result.get("row")
    if not row:
        return f"Note {note_id} not found."
    return f"Note '{row['title']}' (id: {row['id']}):\n\n{row['content']}"


@tool
async def create_note(
    title: str,
    content: str,
    project_id: str = "",
) -> str:
    """Create a new note.
    title: note heading (required)
    content: Markdown body text (required)
    project_id: optional UUID linking this note to a project
    """
    result = await execute_on_client(
        action="insert",
        table="notes",
        data={
            "title": title,
            "content": content,
            "projectId": project_id or None,
        },
    )
    row = result["row"]
    # Index the note content in the vector store.
    vector = await embed(content)
    await execute_on_client(
        action="vector_upsert",
        data={"id": row["id"], "projectId": row.get("projectId"), "content": content},
        vector=vector,
    )
    return f"Note created: '{row['title']}' (id: {row['id']})."


@tool
async def update_note(
    note_id: str,
    title: str = "",
    content: str = "",
) -> str:
    """Update an existing note. Only pass fields that should change.
    note_id: UUID of the note (required)
    If you need to preserve existing content, call get_note first.
    """
    updates: dict[str, Any] = {}
    if title:
        updates["title"] = title
    if content:
        updates["content"] = content
    result = await execute_on_client(
        action="update",
        table="notes",
        data={"id": note_id, "updates": updates},
    )
    row = result["row"]
    # Re-index if content changed.
    if content:
        vector = await embed(content)
        await execute_on_client(
            action="vector_upsert",
            data={"id": note_id, "projectId": row.get("projectId"), "content": content},
            vector=vector,
        )
    return f"Note updated: '{row['title']}' (id: {row['id']})."


@tool
async def delete_note(note_id: str) -> str:
    """Delete a note permanently by its UUID."""
    await execute_on_client(action="delete", table="notes", data={"id": note_id})
    return f"Note {note_id} deleted."


@registry.register
class NoteAgent(ChatAgent):
    def get_name(self) -> str:
        return "note_agent"

    def get_description(self) -> str:
        return "Manages notes: list, get, create, update, delete"

    def get_tools(self) -> list[Any]:
        return [list_notes, get_note, create_note, update_note, delete_note]

    async def handle(self, query: str, context: dict[str, Any]) -> str:
        llm = get_llm()
        messages = [
            SystemMessage(content=_SYSTEM_PROMPT),
            HumanMessage(
                content=f"User query: {query}\nContext: {json.dumps(context)[:1000]}"
            ),
        ]
        return await self._tool_loop(llm, messages, self.get_tools())