diff --git a/app.py b/app.py index 7b3ba5a..1fa4fa5 100644 --- a/app.py +++ b/app.py @@ -16,6 +16,7 @@ from quart import ( from quart_jwt_extended import JWTManager, get_jwt_identity from tortoise import Tortoise +import blueprints.artifacts import blueprints.conversation import blueprints.conversation.logic import blueprints.email @@ -71,6 +72,7 @@ app.register_blueprint(blueprints.imessage.imessage_blueprint) app.register_blueprint(blueprints.scheduled_messages.scheduled_messages_blueprint) app.register_blueprint(blueprints.cron.cron_blueprint) app.register_blueprint(blueprints.beli.beli_blueprint) +app.register_blueprint(blueprints.artifacts.artifacts_blueprint) async def _obsidian_sync_loop(): diff --git a/blueprints/artifacts/__init__.py b/blueprints/artifacts/__init__.py new file mode 100644 index 0000000..cab7643 --- /dev/null +++ b/blueprints/artifacts/__init__.py @@ -0,0 +1,56 @@ +import uuid + +from quart import Blueprint, jsonify, request +from quart_jwt_extended import get_jwt_identity + +from blueprints.conversation.models import Conversation +from blueprints.users.decorators import jwt_access_required + +from .logic import artifact_full, artifact_summary +from .models import Artifact + +artifacts_blueprint = Blueprint("artifacts_api", __name__, url_prefix="/api/artifacts") + + +def _parse_uuid(value: str) -> uuid.UUID | None: + """Validate ids before querying: asyncpg raises OperationalError on bad UUIDs.""" + try: + return uuid.UUID(value) + except (ValueError, AttributeError, TypeError): + return None + + +@artifacts_blueprint.get("") +@jwt_access_required +async def list_artifacts(): + user_uuid = get_jwt_identity() + raw_conversation_id = request.args.get("conversation_id", "").strip() + if not raw_conversation_id: + return jsonify({"error": "conversation_id is required"}), 400 + conversation_id = _parse_uuid(raw_conversation_id) + if conversation_id is None: + return jsonify({"error": "conversation_id must be a UUID"}), 400 + + conversation = await Conversation.get_or_none(id=conversation_id, user_id=user_uuid) + if conversation is None: + return jsonify({"error": "conversation not found"}), 404 + + artifacts = await Artifact.filter( + conversation_id=conversation.id, user_id=user_uuid + ).order_by("created_at") + return jsonify([artifact_summary(a) for a in artifacts]) + + +@artifacts_blueprint.get("/") +@jwt_access_required +async def get_artifact(artifact_id: str): + user_uuid = get_jwt_identity() + parsed_id = _parse_uuid(artifact_id) + artifact = ( + await Artifact.get_or_none(id=parsed_id, user_id=user_uuid) + if parsed_id is not None + else None + ) + if artifact is None: + return jsonify({"error": "artifact not found"}), 404 + return jsonify(artifact_full(artifact)) diff --git a/blueprints/artifacts/logic.py b/blueprints/artifacts/logic.py new file mode 100644 index 0000000..d76a76d --- /dev/null +++ b/blueprints/artifacts/logic.py @@ -0,0 +1,78 @@ +import json +import logging + +from .models import Artifact + +logger = logging.getLogger(__name__) + +# Artifacts whose serialised `data` exceeds this are dropped (the tool text still persists). +ARTIFACT_MAX_BYTES = 256 * 1024 + + +def artifact_ref(artifact: Artifact) -> dict: + """The lightweight reference sent in SSE events and the conversation payload.""" + return {"id": str(artifact.id), "kind": artifact.kind, "title": artifact.title} + + +def artifact_summary(artifact: Artifact) -> dict: + """ArtifactRef plus listing metadata, without `data`.""" + return { + **artifact_ref(artifact), + "tool": artifact.tool, + "message_id": str(artifact.message_id) if artifact.message_id else None, + "created_at": artifact.created_at.isoformat(), + } + + +def artifact_full(artifact: Artifact) -> dict: + return { + **artifact_summary(artifact), + "data": artifact.data, + "source_ref": artifact.source_ref, + "conversation_id": str(artifact.conversation_id), + } + + +async def create_artifact_from_output( + *, output, tool: str, user, conversation, message=None +) -> Artifact | None: + """Persist `output.artifact` if it is a ToolArtifact dict; otherwise return None. + + Tools opt in via `@tool(response_format="content_and_artifact")`. Artifacts + without a `kind` (e.g. simba_search's raw Documents) are ignored, as are + ones whose `data` is not JSON-serialisable or over ARTIFACT_MAX_BYTES. + """ + raw = getattr(output, "artifact", None) + if not isinstance(raw, dict) or not raw.get("kind"): + return None + + data = raw.get("data") or {} + source_ref = raw.get("source_ref") + try: + size = len(json.dumps(data).encode("utf-8")) + json.dumps(source_ref) + except (TypeError, ValueError) as e: + logger.warning( + "Dropping %s artifact from %s: not serialisable (%s)", raw["kind"], tool, e + ) + return None + if size > ARTIFACT_MAX_BYTES: + logger.warning( + "Dropping %s artifact from %s: %d bytes exceeds %d", + raw["kind"], + tool, + size, + ARTIFACT_MAX_BYTES, + ) + return None + + return await Artifact.create( + user=user, + conversation=conversation, + message=message, + tool=tool, + kind=str(raw["kind"])[:64], + title=str(raw.get("title") or "")[:255], + data=data, + source_ref=source_ref, + ) diff --git a/blueprints/artifacts/models.py b/blueprints/artifacts/models.py new file mode 100644 index 0000000..ef61659 --- /dev/null +++ b/blueprints/artifacts/models.py @@ -0,0 +1,29 @@ +from tortoise import fields +from tortoise.models import Model + + +class Artifact(Model): + """A structured tool result (recipe, budget, note…) rendered in the artifact panel.""" + + id = fields.UUIDField(primary_key=True) + user = fields.ForeignKeyField("models.User", related_name="artifacts") + conversation = fields.ForeignKeyField( + "models.Conversation", related_name="artifacts", on_delete=fields.CASCADE + ) + message = fields.ForeignKeyField( + "models.ConversationMessage", + related_name="artifacts", + null=True, + on_delete=fields.SET_NULL, + ) + tool = fields.CharField(max_length=255) + kind = fields.CharField(max_length=64) + title = fields.CharField(max_length=255) + data = fields.JSONField() + source_ref = fields.JSONField(null=True) + version = fields.IntField(default=1) + created_at = fields.DatetimeField(auto_now_add=True) + updated_at = fields.DatetimeField(auto_now=True) + + class Meta: + table = "artifacts" diff --git a/blueprints/conversation/__init__.py b/blueprints/conversation/__init__.py index 1df322e..a9ea890 100644 --- a/blueprints/conversation/__init__.py +++ b/blueprints/conversation/__init__.py @@ -10,6 +10,8 @@ from quart import Blueprint, jsonify, make_response, request from quart_jwt_extended import get_jwt_identity import blueprints.users.models +from blueprints.artifacts.logic import artifact_ref, create_artifact_from_output +from blueprints.artifacts.models import Artifact from blueprints.users.decorators import jwt_access_required from metrics import count_result, observe_response, time_response from utils.image_process import image_content_block @@ -32,6 +34,8 @@ from .models import ( PydListConversation, Speaker, ) +from .tool_events import tool_end_payload, tool_name_from_message, tool_start_payload +from .tool_labels import tool_label TTFT_TIMEOUT = float(os.getenv("TTFT_TIMEOUT_SECONDS", "30")) @@ -340,28 +344,47 @@ async def stream_query(): if event_type == "on_tool_start": tool_timers[event["run_id"]] = time.monotonic() - yield f"data: {json.dumps({'type': 'tool_start', 'tool': event['name']})}\n\n" + yield f"data: {json.dumps(tool_start_payload(event))}\n\n" - elif event_type == "on_tool_end": + elif event_type in ("on_tool_end", "on_tool_error"): + now = time.monotonic() start = tool_timers.pop(event["run_id"], None) tool_spans.append( { "tool": event["name"], - "ms": round((time.monotonic() - start) * 1000) - if start - else None, + "ms": round((now - start) * 1000) if start else None, } ) - yield f"data: {json.dumps({'type': 'tool_end', 'tool': event['name']})}\n\n" + ref = None # Persist tool output so later turns keep the fetched context output = event.get("data", {}).get("output") - if output is not None: - await add_message_to_conversation( + if event_type == "on_tool_end" and output is not None: + tool_message = await add_message_to_conversation( conversation=conversation, message=tool_result_text(event["name"], output), speaker="tool", user=user, ) + try: + artifact = await create_artifact_from_output( + output=output, + tool=event["name"], + user=user, + conversation=conversation, + message=tool_message, + ) + except Exception as e: + logging.error( + f"Failed to save artifact from {event['name']}: {e}" + ) + artifact = None + if artifact is not None: + ref = artifact_ref(artifact) + status = "error" if event_type == "on_tool_error" else "ok" + end_payload = tool_end_payload( + event, start, ref, now=now, status=status + ) + yield f"data: {json.dumps(end_payload)}\n\n" elif event_type == "on_chat_model_start": turn_has_content = False @@ -453,22 +476,34 @@ async def get_conversation(conversation_id: str): if conversation is None: return jsonify({"error": "conversation not found"}), 404 await conversation.fetch_related("messages") + # Keyed by str so the match doesn't depend on the driver's UUID types. + artifacts_by_message = { + str(a.message_id): a + for a in await Artifact.filter( + conversation_id=conversation.id, message_id__isnull=False + ).order_by("created_at") + } # Manually serialize the conversation with messages messages = [] for msg in conversation.messages: - # Tool messages hold raw fetched context for the LLM; only surface - # their header line (e.g. "[web_search result]") in the UI. - text = msg.text.split("\n", 1)[0] if msg.speaker == Speaker.TOOL else msg.text - messages.append( - { - "id": str(msg.id), - "text": text, - "speaker": msg.speaker.value, - "created_at": msg.created_at.isoformat(), - "image_key": msg.image_key, - } - ) + item = { + "id": str(msg.id), + "text": msg.text, + "speaker": msg.speaker.value, + "created_at": msg.created_at.isoformat(), + "image_key": msg.image_key, + } + if msg.speaker == Speaker.TOOL: + # Tool messages hold raw fetched context for the LLM; only surface + # their header line (e.g. "[web_search result]") in the UI. + item["text"] = msg.text.split("\n", 1)[0] + tool = tool_name_from_message(msg.text) + artifact = artifacts_by_message.get(str(msg.id)) + item["tool"] = tool + item["label"] = tool_label(tool) + item["artifact"] = artifact_ref(artifact) if artifact else None + messages.append(item) return jsonify( { "id": str(conversation.id), diff --git a/blueprints/conversation/tool_events.py b/blueprints/conversation/tool_events.py new file mode 100644 index 0000000..1d43458 --- /dev/null +++ b/blueprints/conversation/tool_events.py @@ -0,0 +1,77 @@ +"""Pure shaping of LangChain tool events into SSE `tool_start` / `tool_end` payloads.""" + +import json +import re + +from .tool_labels import tool_label + +INPUT_VALUE_MAX_CHARS = 200 +# Injected by LangChain rather than chosen by the model; never shown to the client. +_HIDDEN_INPUT_KEYS = frozenset({"config", "callbacks", "run_manager"}) + +_TOOL_HEADER_RE = re.compile(r"^\[(\S+) result\]") + + +def _truncate(text: str) -> str: + if len(text) > INPUT_VALUE_MAX_CHARS: + return text[: INPUT_VALUE_MAX_CHARS - 1] + "…" + return text + + +def summarize_tool_input(raw) -> dict[str, str | int | float | bool]: + """Reduce a tool's input to flat, display-safe scalars.""" + if raw is None: + return {} + if not isinstance(raw, dict): + return {"input": _truncate(str(raw))} + summary: dict[str, str | int | float | bool] = {} + for key, value in raw.items(): + if key in _HIDDEN_INPUT_KEYS or value is None: + continue + if isinstance(value, (bool, int, float)): + summary[key] = value + elif isinstance(value, str): + summary[key] = _truncate(value) + else: + try: + summary[key] = _truncate(json.dumps(value, default=str)) + except (TypeError, ValueError): + summary[key] = _truncate(str(value)) + return summary + + +def tool_start_payload(event: dict) -> dict: + name = event["name"] + return { + "type": "tool_start", + "run_id": event["run_id"], + "tool": name, + "label": tool_label(name), + "input": summarize_tool_input(event.get("data", {}).get("input")), + } + + +def tool_end_payload( + event: dict, + started_at: float | None, + artifact_ref: dict | None, + *, + now: float, + status: str = "ok", +) -> dict: + name = event["name"] + return { + "type": "tool_end", + "run_id": event["run_id"], + "tool": name, + "label": tool_label(name), + "status": status, + "ms": round((now - started_at) * 1000) if started_at is not None else 0, + "artifact": artifact_ref, + } + + +def tool_name_from_message(text: str) -> str: + """Recover the tool name from a persisted `[tool_name result]` header.""" + match = _TOOL_HEADER_RE.match(text) + return match.group(1) if match else "tool" diff --git a/blueprints/conversation/tool_labels.py b/blueprints/conversation/tool_labels.py new file mode 100644 index 0000000..4623f53 --- /dev/null +++ b/blueprints/conversation/tool_labels.py @@ -0,0 +1,64 @@ +"""Friendly, user-facing labels for agent tool calls (shown in the activity timeline).""" + +TOOL_LABELS: dict[str, str] = { + "get_current_date": "Checking the date…", + "web_search": "Searching the web…", + "simba_search": "Searching household documents…", + "save_user_memory": "Saving to memory…", + # YNAB (current names) + "ynab_budget_summary": "Checking budget…", + "ynab_search_transactions": "Searching transactions…", + "ynab_category_spending": "Reviewing category spending…", + "ynab_category_allocations": "Reviewing category budgets…", + "ynab_insights": "Analyzing spending trends…", + # YNAB (names after the #55 mirror) + "ynab_account_balances": "Checking account balances…", + "ynab_spending_report": "Building spending report…", + "ynab_category_budget": "Reviewing category budgets…", + "ynab_recent_changes": "Checking recent budget changes…", + "ynab_sql": "Querying budget data…", + # Mealie + "mealie_todays_meals": "Checking today's meals…", + "mealie_meal_plans": "Checking the meal plan…", + "mealie_get_recipe": "Fetching recipe…", + "mealie_shopping_lists": "Checking shopping lists…", + "mealie_create_meal_plan": "Adding to the meal plan…", + # Obsidian / journal + "obsidian_search_notes": "Searching notes…", + "obsidian_read_note": "Reading note…", + "obsidian_create_note": "Creating note…", + "obsidian_create_task": "Creating task…", + "journal_get_today": "Reading today's journal…", + "journal_get_tasks": "Checking tasks…", + "journal_add_task": "Adding task…", + "journal_complete_task": "Completing task…", + # Google + "get_calendar_events": "Checking calendar…", + "create_calendar_event": "Creating calendar event…", + "gmail_search_emails": "Searching email…", + "gmail_get_email": "Reading email…", + # SimbaKit + "simbakit_device_status": "Checking Simba's devices…", + "simbakit_weight_history": "Checking Simba's weight…", + "simbakit_litter_activity": "Checking litter box activity…", + "simbakit_drinking_activity": "Checking water fountain activity…", + # Beli + "beli_search_restaurants": "Searching restaurants…", +} + +# MCP tools are discovered at runtime, so they match by name prefix. +TOOL_PREFIX_LABELS: tuple[tuple[str, str], ...] = ( + ("homebox_", "Checking home inventory…"), + ("bookstack_", "Checking the wiki…"), +) + + +def tool_label(tool_name: str) -> str: + """Return the friendly label for a tool, falling back to "Using …".""" + label = TOOL_LABELS.get(tool_name) + if label: + return label + for prefix, prefix_label in TOOL_PREFIX_LABELS: + if tool_name.startswith(prefix): + return prefix_label + return f"Using {tool_name}…" diff --git a/config/db.py b/config/db.py index 73c2be0..745ef1e 100644 --- a/config/db.py +++ b/config/db.py @@ -19,6 +19,7 @@ TORTOISE_CONFIG = { "blueprints.scheduled_messages.models", "blueprints.cron.models", "blueprints.beli.models", + "blueprints.artifacts.models", "aerich.models", ], "default_connection": "default", diff --git a/migrations/models/17_20261012175338_add_artifacts.py b/migrations/models/17_20261012175338_add_artifacts.py new file mode 100644 index 0000000..491a04c --- /dev/null +++ b/migrations/models/17_20261012175338_add_artifacts.py @@ -0,0 +1,108 @@ +from tortoise import BaseDBAsyncClient + +RUN_IN_TRANSACTION = True + + +async def upgrade(db: BaseDBAsyncClient) -> str: + return """ + CREATE TABLE IF NOT EXISTS "artifacts" ( + "id" UUID NOT NULL PRIMARY KEY, + "tool" VARCHAR(255) NOT NULL, + "kind" VARCHAR(64) NOT NULL, + "title" VARCHAR(255) NOT NULL, + "data" JSONB NOT NULL, + "source_ref" JSONB, + "version" INT NOT NULL DEFAULT 1, + "created_at" TIMESTAMPTZ NOT NULL DEFAULT CURRENT_TIMESTAMP, + "updated_at" TIMESTAMPTZ NOT NULL DEFAULT CURRENT_TIMESTAMP, + "conversation_id" UUID NOT NULL REFERENCES "conversations" ("id") ON DELETE CASCADE, + "message_id" UUID REFERENCES "conversation_messages" ("id") ON DELETE SET NULL, + "user_id" UUID NOT NULL REFERENCES "users" ("id") ON DELETE CASCADE +); +COMMENT ON TABLE "artifacts" IS 'A structured tool result (recipe, budget, note…) rendered in the artifact panel.';""" + + +async def downgrade(db: BaseDBAsyncClient) -> str: + return """ + DROP TABLE IF EXISTS "artifacts";""" + + +MODELS_STATE = ( + "eJztXW1zo7YW/isaf9qdSXPjZJNuM3fujJN423QTJxN7b9/coRjkWDdYUASb9bT971cSYI" + "MQGGyDwasPbdZCR8Cjg3TOc46kvzpz24QWOb628WfoEt1DNu5cgr86WJ9D+g/p9SPQ0R1n" + "dZUVePrE4gJGrCa/ok+I5+qGRy9OdYtAWmRCYrjICW+GfctihbZBKyL8vCryMfrTh5pnP0" + "NvBl164bffaTHCJvwCSfTTedGmCFpm4rmRye7NyzVv4fCyT59ubz7wmux2E82wLX+OV7Wd" + "hTez8bK67yPzmMmwa88QQ1f3oBl7DfaU4WtHRcET0wLP9eHyUc1VgQmnum8xMDr/nvrYYB" + "gAfif2v3f/6ZSAh0LNoEXYY1j89U/wVqt35qUddqvrH3pPb84u3vK3tIn37PKLHJHOP1xQ" + "9/RAlOO6ApL/TUF5PdNdOZRRfQFM+qCbwBgVrHBc6VAEZATQZqh15voXzYL42ZvRn6fn5z" + "kw/rf3xJGktTiUNtXrQOsH4aXT4BqDdAWh4UL2yprupYG8oVc8NIdyMJOSAqRmKHoc/aOh" + "ANN3MB+wtQg/ghx8R7f3/eGod//I3mROyJ8Wh6g36rMrp7x0IZS+uRC6YtkI+Ol29ANgP8" + "GvD4O+qPvLeqNfO+yZdN+zNWy/aroZ+16j0giYRMf6jrlhxyYlVcfutWPDh499sDMdY2iV" + "GfZiIhuNfOFTleq+ziucbDFfCCPfSZGB7yR73DsRhz2fQFcrNw/HRLaYjDfBcrtPYfOplx" + "kw0xfpzMvASIP3wXYhesYf4YJDeEufQ8eGbMoNTbdPYTNNBW1VunoKV39d2nRxpaBvR98J" + "esHH2Bte9276HY7hRDdeXnXX1DLAnENC9GdI0oBehZIfPj5Ba2ngyrGMm8H3QYvNHJuzsE" + "18oLrroSm1zLcEpRc20zIkmNbYp3ZMWxJ6lL40P52LJTqmGmCG92Z3ytGTNV5VTJ2KOVda" + "XKWVk9VqJ8uDXyQm5IiWyqGM6rfFycozDfs/jxJWYWRSvLnv/fw2YRnePQy+j6rHTJDru4" + "cr5XEdomGe9riIA/UXmVnELPM+9ucpoyjRuzHxPX86nU/D/tMlYLbNGA9v7696l4Cg+UQf" + "49HDwx1Fw7a5N1HWju8WseO72XZ8N2XHozmdZbQXuCjjDyWEavOIdkkFnXdPCyBJa2VCya" + "8lsQznbI16i6ZVildLS7YS1WoItrhZVM4okYju0kJpqetpCHT/li6oGD1oLIprfVGJumzu" + "k9bvfzXHya/U/WKUxz2c2+6iI/G6YleP8pwtzjvMWUWknKwDcLLoTTyIS/lZMRHlailXS2" + "LcHK6rpYJbB9GxKdZ3b2GZZg2GbYnLNNgY3iIwU6X193B7cz30dE9Kua8u5tp+NjINjbB6" + "xSy/ziPEJsUUsPaBZT8jfAnmukMANffAH7ylP4Cju/oceDZAHgGPH6/7gDoTiL6WeyxSPV" + "s3OMZjPPSoUpoAYfAYajh4g20POK5tQEIAN24XbwGxgR7cg7aue4C27tIGqT1nYzjGE8uH" + "/6LSEAM6Jtgu0LEJpggjMoO8Fnskm9m1VBJZFvisW4gN74S/1ha2L4l6qihLshTYjbm23v" + "bdJTfSPX1fhLs7fZ9N3rFr+Vaw82JALdKSMtCmBNtiEVcG8gpU+MVB9OvawFZKStZjK+1Y" + "iVtiGkU45Bq9yps5CKM36NiUDbcfe4Qbehk81HoGSjFP7WeeWD+WzaOOy7Rzpq0k1OPohL" + "zarkQzrxDW3UWG7RKTEsCcLEIbv6nUuQy8K8bDxQdeCtnV7aD39Iucyovqx/X26pdRvyfa" + "MXMdlUp7XQq00t4+KRYrzwuWp6Llr9SBInRI17A/n5SzsCWiO4nx1ovqWRFQz7IxPZMkII" + "Sx7/KQSkQVpNF3q0HMXk82ltq2BXWc883HZcURNcxcqWJ+Kmv9lBhSHx7uEkPq1a0Y/fh0" + "f9Wn4wGHmVZCAdt1OxhJwZ3NdYOaXS9QEkFfM6AKsi3U2O5FkZFV9AdiA+uFqLEBI+dP/g" + "cNiW+WDago10IwKzGkqLc00xzX/ozMckNqSrA+67Rj2YZubWHeCxleRYbV8+xh9Tw1rFqm" + "7mjPru07kpyOH4cPAzmkgphIHiDDA38DC5HKUuw7v/1ejdPE3jk/5CxGlwUrlTWgQs4HTN" + "KokPPBdewyv6sg+SbPO9zhUq0mu/f5S7Ti6WebY5HMeGue5heCIrCMdcOwfbxtymSftdUL" + "mmoxJMSYQdOnXpi2m+WNw6i99q9tnEALMSt1iqxtUbmiTT0GLbUYELXYs6r4Dh9MOpIATz" + "8iR7MjPHxQK5hiwpsD4ZdO/3o6CzjwZIwwTzWdRVJAZpkoQlg6CYEm4A8FeMVX5M3A2ck3" + "pr4ALmT16cMcgzuEXwjLO7meufZcv7ka46ntgs/Ux6Z/CNRdYwYM3dEnyEIe2j4RRIWidh" + "SKivhQGaDrFx9JgW0F118Jh7IBHbUtE7X3qN55oaDJeU7Q5DwdNJnSUYS5E1SPJfNTNpqi" + "XDshrUQ3qTeWCWjOwnK7jXDWvebBlGYl5jMUkUw7uYmWcBGFcrwmtrnQyu64kBBqycrbuj" + "8LjtHMm0syCNYAGwkpYOVrrLiVTe9IH62k5SYRbQnIamNGxTdvGkhoVz62mqyzezLke0su" + "WEtKqTVrESBpDEsvW2svdX4kLF9LKkmTVrAlMM4iFmN9sIZfTARNivKMoQyjCafo2Xc5Lw" + "wYy3d733tkxRhyuotksY4lWohxkEspF5qMbqTPGLCQEBvuwqGKBqJs3qNYG4BAz6O3J0dj" + "zNjNqBm2IMwnx2DoO47NFrXNqYohx4JRDQIc6PJ9mBRH2RCOMlTa8nRQSlAxlRFjgYhj6Q" + "ut7CoEUU75DrFd0RxtZpNS7G9CqC3sWl1gsvE5DeYtziAOEjIClqg62+S77862gJLvA/vN" + "affdt+/en128e0+r8EdZlnybA24605hjsMnqopSgUkZBGTPXGfUjOySb00q10BZ06ya2EH" + "W0qT3yWaK7uUsQEnI1Lj9YTvwNXn1g6cTToOvakqzubJVNSrVkmq+dh1V84WHyhSrx+CA6" + "Vu11pfa6at5eVwKxsYvU5XZhmcyZWuBgZy1/F0AMaWvDZWMtgqRy+jiGTBaDnARvHYksdN" + "x6HnlEL78QwOTYjmNs0CE885Rv0sVpVxgnitM8cukWGI98z2Cm/xHAbNpAmM29VGbuHAWF" + "4QZo0FymxtKRMGSOp7RB34WAvd0L24HNswEJ2GOAqO8H55C2bwUNMwFuNQO+Lzutrljkhr" + "DI3KPhSrtJulZaege2XbM8pAaZcoVCwbxLokxkX/adZHKFMtH6KMOTLUzzHROGEQYZ0eBM" + "AFNyXyV6rHlo+Ix40sJpQmLEZIKYJf5VYsk/yRCFzUdosQE1SO95kFZMymEyKSoP68DysK" + "pOamh/FlZqNbTEj5atmM52pOXrtdXmnq32M11oIAdJD5bJjr8nhNoSHa4jIz/v5PX15ztu" + "ewr77kDtsBl32Pu+fwmiPQjHuH/fu727BMtNM/d/Rrs6E6mKOPEe1gUfZI7dar4s706Isu" + "10KFriQBRyDbOCPgWP7l1K17jfYGjVBUGS48f+4OZ28H3aAumEVy6BE5zaMsbD/mB0CQgd" + "KMf4Ax31+zeXPLIAzTG+7g2u+3e8yGDvaoX7qO5/OqBmie+6MMR/k25KtlBjVz0tb3w8CB" + "VY6CRWTIHlR9rc0C755RKYtEcWY/xTv/+R/XyF8IX9vn8YjH5gBXM6182sRTM6xydQo8oo" + "m61zk+UScmqvXmGvXha8i9yxMmZQSrAl07bKmVOEoMqZUx2byfRGn91kUfYodVFQ8b0xTH" + "ZA+bY/iy6lIk3ie69dG/9oTzoSmje6dJTH7hq0kvY/e6JI3faTumXXU6llVOklp7EnK2FX" + "C2LKqpZa1Y5rz51StO1Koi06WjemKpCjAjmtDeRseJTTXg5xas0qStfHG7i3gqhKSWtC3n" + "Be8CEvA3y7qMO+o3IVnHCkVhcrplQRaoop/fo6tuyxRlVyVfEzSSR8lXBkSTZnlTomZbe8" + "1W/LBcT8RoG9/LtiszaZc7Zgs2LglzCAklLKtxXOFtpmg6BUC21Bt26biAOVuYvDGt3N3s" + "qhqSZnXcrrwikdPWZZxxMX1eBUMy2BuW4tDhiBAKyN7EBpA4pYaAKxoHzhKr4YRWfuls5U" + "3MJBuKCKWzjQjlU7l6mdy5q8c5nLtk/yXX3rk5cZM/W0bKxdyFa6aZcATAapl4RuDa8ndN" + "quqT1nRTDy2018gjAkJFS2WMTNRRCbmuiT898G7flndjC5YgX3wgoKnVaKXxFk28Jh1UCy" + "ZGh8UXQzxGtcTbWFstYHb9kkTUGsJa53XZzgciQuOwzEBduoot0iSRLd7CSJbipJ4uvKHq" + "7kRGgDeaU0MaqvPuqlElJHZzax3Zltl5p8RDmF6FInfcQMnjSYPw4fBhlquRIRGRZkeOBv" + "YCFSmRcUs0MnPrI8hMkxu19FpihDIZ8IFjlfwWZlDaQT3ZEBNerDyFbnZiuxINZKHd79vN" + "S+HRjkKdlX/f7gEkwgxGP8U28w0kYP2uiJLeCnvi71UjXPbciafaqBLxp/wTTqmfuEJoU2" + "2h10D7q74/1BiWG7sj1BoYHmuiUHbikjjrWB0HEo3Ej8csC56V/f3vfu3pwfnQqBnkgn36" + "VNUNuT7U2bHZZcCrRkrKw7IunpzxI8s2f+qL6a9jeZ9qfQMzZLlUhKqrBWs8JaKg59EP2a" + "jkOHwYCSIcuklIpaxmMqWwYuhRTwxuK3Nn6Z1JEmbRjRcz00DQJqqWDd8tpRXphOD2sVPF" + "GnB2gN3/B8F5rAs20LUMWifQHe8GXIVGLim8/QOwLMmhv7pyenF29pHdqHTAJh4M0giG4K" + "HB1DK33qTmV3URtb1DO+HOUE/bwwEa4okRLVbyc1XQkN+EKfvgyEUf12QnjxrgCCF+8yAW" + "SXBE8KebIZLkcHI4F2IljN7ir0Ay/jjkb163dHK7OyKnE8ie27jDmG0zLoJqW2wrhRTEol" + "EH+GLpFuCpRJiMYk6jsrqbvF97/rc6eU13yYXrPK3j6Ijk3TXDZmYxbPhS27i2Za9GvhRW" + "QHFZYDLym1BW6NmoVLwKaWDahlAzUtG8ga73aA3rXQXHtRlIzm69HM3I1+KyBjJ5Q1dfBb" + "C2dyfE8gOeyPwODT3d3eyGDoImPWkVHBwZWjXCJ4VacxGwdnumNFz08Oe3A7XrUJWSnZdG" + "qmL5vNZmU7s18zn8U+jRIghtXbCWC3UMJ0NydhuitJmM46yS0nOTXzJDdFC8qsxDonln/+" + "D1XXXUg=" +) diff --git a/tests/unit/test_artifacts.py b/tests/unit/test_artifacts.py new file mode 100644 index 0000000..d553a21 --- /dev/null +++ b/tests/unit/test_artifacts.py @@ -0,0 +1,389 @@ +"""Tests for artifact creation (blueprints/artifacts/logic.py) and the owner-scoped API.""" + +from datetime import timedelta +from types import SimpleNamespace + +import pytest +from quart import Quart +from quart_jwt_extended import JWTManager, create_access_token +from tortoise import Tortoise + +from blueprints.artifacts import artifacts_blueprint +from blueprints.artifacts.logic import ARTIFACT_MAX_BYTES, create_artifact_from_output +from blueprints.artifacts.models import Artifact +from blueprints.conversation.models import Conversation, ConversationMessage +from blueprints.users.models import User + + +@pytest.fixture +async def db(): + await Tortoise.init( + db_url="sqlite://:memory:", + modules={ + "models": [ + "blueprints.users.models", + "blueprints.conversation.models", + "blueprints.artifacts.models", + ] + }, + ) + await Tortoise.generate_schemas() + yield + await Tortoise.close_connections() + + +@pytest.fixture +async def owner(db): + user = await User.create(username="owner", email="owner@example.com") + conversation = await Conversation.create(name="c", user=user) + message = await ConversationMessage.create( + text="[mealie_get_recipe result]\nSoup", + conversation=conversation, + speaker="tool", + ) + return user, conversation, message + + +def _output(artifact): + return SimpleNamespace(content="text", artifact=artifact) + + +RECIPE = {"kind": "mealie.recipe", "title": "Soup", "data": {"slug": "soup"}} + + +class TestCreateArtifact: + async def test_creates_and_links(self, owner): + user, conversation, message = owner + artifact = await create_artifact_from_output( + output=_output({**RECIPE, "source_ref": {"system": "mealie"}}), + tool="mealie_get_recipe", + user=user, + conversation=conversation, + message=message, + ) + assert artifact is not None + saved = await Artifact.get(id=artifact.id) + assert saved.message_id == message.id + assert saved.kind == "mealie.recipe" + assert saved.data == {"slug": "soup"} + assert saved.source_ref == {"system": "mealie"} + assert saved.version == 1 + + @pytest.mark.parametrize( + "artifact", [None, [{"page_content": "doc"}], {"title": "no kind"}, "str"] + ) + async def test_ignores_non_artifacts(self, owner, artifact): + user, conversation, message = owner + result = await create_artifact_from_output( + output=_output(artifact), + tool="simba_search", + user=user, + conversation=conversation, + message=message, + ) + assert result is None + assert await Artifact.all().count() == 0 + + async def test_plain_string_output(self, owner): + user, conversation, message = owner + result = await create_artifact_from_output( + output="just text", tool="t", user=user, conversation=conversation + ) + assert result is None + + async def test_drops_oversized(self, owner): + user, conversation, message = owner + big = {**RECIPE, "data": {"blob": "x" * (ARTIFACT_MAX_BYTES + 1)}} + result = await create_artifact_from_output( + output=_output(big), tool="t", user=user, conversation=conversation + ) + assert result is None + assert await Artifact.all().count() == 0 + + async def test_drops_unserialisable(self, owner): + user, conversation, message = owner + bad = {**RECIPE, "data": {"obj": object()}} + result = await create_artifact_from_output( + output=_output(bad), tool="t", user=user, conversation=conversation + ) + assert result is None + + async def test_message_delete_sets_null(self, owner): + user, conversation, message = owner + artifact = await create_artifact_from_output( + output=_output(RECIPE), + tool="t", + user=user, + conversation=conversation, + message=message, + ) + await message.delete() + assert (await Artifact.get(id=artifact.id)).message_id is None + + async def test_conversation_delete_cascades(self, owner): + user, conversation, message = owner + await create_artifact_from_output( + output=_output(RECIPE), tool="t", user=user, conversation=conversation + ) + await conversation.delete() + assert await Artifact.all().count() == 0 + + +@pytest.fixture +def app(): + app = Quart(__name__) + app.config["JWT_SECRET_KEY"] = "test-secret" + app.config["JWT_ACCESS_TOKEN_EXPIRES"] = timedelta(hours=1) + JWTManager(app) + app.register_blueprint(artifacts_blueprint) + return app + + +async def _headers(app, user): + async with app.app_context(): + token = create_access_token(identity=str(user.id)) + return {"Authorization": f"Bearer {token}"} + + +class TestArtifactRoutes: + async def _setup(self, owner): + user, conversation, message = owner + artifact = await create_artifact_from_output( + output=_output(RECIPE), + tool="mealie_get_recipe", + user=user, + conversation=conversation, + message=message, + ) + other = await User.create(username="other", email="other@example.com") + return user, other, conversation, message, artifact + + async def test_list_for_owner(self, app, owner): + user, _, conversation, message, artifact = await self._setup(owner) + resp = await app.test_client().get( + f"/api/artifacts?conversation_id={conversation.id}", + headers=await _headers(app, user), + ) + assert resp.status_code == 200 + body = await resp.get_json() + assert len(body) == 1 + assert body[0]["id"] == str(artifact.id) + assert body[0]["kind"] == "mealie.recipe" + assert body[0]["title"] == "Soup" + assert body[0]["tool"] == "mealie_get_recipe" + assert body[0]["message_id"] == str(message.id) + assert "created_at" in body[0] + assert "data" not in body[0] + + async def test_get_for_owner(self, app, owner): + user, _, conversation, message, artifact = await self._setup(owner) + resp = await app.test_client().get( + f"/api/artifacts/{artifact.id}", headers=await _headers(app, user) + ) + assert resp.status_code == 200 + body = await resp.get_json() + assert body["data"] == {"slug": "soup"} + assert body["source_ref"] is None + assert body["conversation_id"] == str(conversation.id) + + async def test_other_user_gets_404(self, app, owner): + _, other, conversation, _, artifact = await self._setup(owner) + headers = await _headers(app, other) + client = app.test_client() + resp = await client.get(f"/api/artifacts/{artifact.id}", headers=headers) + assert resp.status_code == 404 + resp = await client.get( + f"/api/artifacts?conversation_id={conversation.id}", headers=headers + ) + assert resp.status_code == 404 + + async def test_invalid_id_is_404(self, app, owner): + user, *_ = await self._setup(owner) + resp = await app.test_client().get( + "/api/artifacts/not-a-uuid", headers=await _headers(app, user) + ) + assert resp.status_code == 404 + + async def test_invalid_id_is_404_without_querying(self, app, owner, monkeypatch): + user, *_ = await self._setup(owner) + + async def _boom(*args, **kwargs): + raise AssertionError("queried with an invalid id") + + monkeypatch.setattr(Artifact, "get_or_none", _boom) + resp = await app.test_client().get( + "/api/artifacts/not-a-uuid", headers=await _headers(app, user) + ) + assert resp.status_code == 404 + + async def test_invalid_conversation_id_is_400_without_querying( + self, app, owner, monkeypatch + ): + user, *_ = await self._setup(owner) + + async def _boom(*args, **kwargs): + raise AssertionError("queried with an invalid id") + + monkeypatch.setattr(Conversation, "get_or_none", _boom) + resp = await app.test_client().get( + "/api/artifacts?conversation_id=garbage", headers=await _headers(app, user) + ) + assert resp.status_code == 400 + + async def test_requires_auth(self, app, owner): + await self._setup(owner) + resp = await app.test_client().get("/api/artifacts/whatever") + assert resp.status_code == 401 + + async def test_missing_conversation_id_is_400(self, app, owner): + user, *_ = await self._setup(owner) + resp = await app.test_client().get( + "/api/artifacts", headers=await _headers(app, user) + ) + assert resp.status_code == 400 + + +class TestConversationPayload: + async def test_tool_messages_gain_tool_label_artifact(self, owner): + from blueprints.conversation import conversation_blueprint + + user, conversation, message = owner + await ConversationMessage.create( + text="[web_search result]\nresults", + conversation=conversation, + speaker="tool", + ) + await ConversationMessage.create( + text="hello", conversation=conversation, speaker="user" + ) + artifact = await create_artifact_from_output( + output=_output(RECIPE), + tool="mealie_get_recipe", + user=user, + conversation=conversation, + message=message, + ) + + app = Quart(__name__) + app.config["JWT_SECRET_KEY"] = "test-secret" + JWTManager(app) + app.register_blueprint(conversation_blueprint) + resp = await app.test_client().get( + f"/api/conversation/{conversation.id}", headers=await _headers(app, user) + ) + assert resp.status_code == 200 + by_text = {m["text"]: m for m in (await resp.get_json())["messages"]} + + recipe = by_text["[mealie_get_recipe result]"] + assert recipe["tool"] == "mealie_get_recipe" + assert recipe["label"] == "Fetching recipe…" + assert recipe["artifact"] == { + "id": str(artifact.id), + "kind": "mealie.recipe", + "title": "Soup", + } + search = by_text["[web_search result]"] + assert search["tool"] == "web_search" + assert search["artifact"] is None + assert "tool" not in by_text["hello"] + + +class _FakeAgent: + def __init__(self, events): + self._events = events + + async def astream_events(self, payload, version, config): + for event in self._events: + yield event + + +class TestStreamQuery: + async def test_emits_v2_tool_events_and_saves_artifact(self, owner, monkeypatch): + import json + + import blueprints.conversation as conv + from langchain_core.messages import AIMessage, ToolMessage + + user, conversation, _ = owner + events = [ + { + "event": "on_tool_start", + "name": "mealie_get_recipe", + "run_id": "r1", + "data": {"input": {"recipe_slug": "soup", "config": {}}}, + }, + { + "event": "on_tool_start", + "name": "web_search", + "run_id": "r2", + "data": {"input": {"query": "cats"}}, + }, + { + "event": "on_tool_end", + "name": "mealie_get_recipe", + "run_id": "r1", + "data": { + "output": ToolMessage( + content="Soup recipe", + artifact=RECIPE, + tool_call_id="c1", + name="mealie_get_recipe", + ) + }, + }, + { + "event": "on_tool_error", + "name": "web_search", + "run_id": "r2", + "data": {"error": RuntimeError("boom")}, + }, + { + "event": "on_chain_end", + "name": "agent", + "data": {"output": {"messages": [AIMessage(content="Done")]}}, + }, + ] + monkeypatch.setattr(conv, "get_agent_for", lambda u: _FakeAgent(events)) + + async def _prompt(user_id): + return "system" + + monkeypatch.setattr(conv, "build_system_prompt_with_memories", _prompt) + + app = Quart(__name__) + app.config["JWT_SECRET_KEY"] = "test-secret" + JWTManager(app) + app.register_blueprint(conv.conversation_blueprint) + resp = await app.test_client().post( + "/api/conversation/stream-query", + json={"query": "soup?", "conversation_id": str(conversation.id)}, + headers=await _headers(app, user), + ) + body = (await resp.get_data()).decode() + frames = [ + json.loads(line[len("data: ") :]) + for line in body.splitlines() + if line.startswith("data: {") + ] + by_type = {} + for f in frames: + by_type.setdefault(f["type"], []).append(f) + + starts = by_type["tool_start"] + assert starts[0]["input"] == {"recipe_slug": "soup"} + assert starts[0]["label"] == "Fetching recipe…" + + ends = {f["run_id"]: f for f in by_type["tool_end"]} + artifact = await Artifact.get(tool="mealie_get_recipe") + assert ends["r1"]["status"] == "ok" + assert ends["r1"]["artifact"] == { + "id": str(artifact.id), + "kind": "mealie.recipe", + "title": "Soup", + } + assert ends["r2"]["status"] == "error" + assert ends["r2"]["artifact"] is None + assert by_type["response"][0]["message"] == "Done" + + tool_message = await ConversationMessage.get(id=artifact.message_id) + assert tool_message.text.startswith("[mealie_get_recipe result]") diff --git a/tests/unit/test_tool_events.py b/tests/unit/test_tool_events.py new file mode 100644 index 0000000..aa86e6a --- /dev/null +++ b/tests/unit/test_tool_events.py @@ -0,0 +1,110 @@ +"""Tests for SSE tool-event shaping and the tool label registry.""" + +from langchain_core.tools import BaseTool + +from blueprints.conversation import agents +from blueprints.conversation.tool_events import ( + INPUT_VALUE_MAX_CHARS, + summarize_tool_input, + tool_end_payload, + tool_name_from_message, + tool_start_payload, +) +from blueprints.conversation.tool_labels import TOOL_LABELS, tool_label + + +def _event(name="web_search", run_id="run-1", input=None): + return {"name": name, "run_id": run_id, "data": {"input": input or {}}} + + +class TestToolLabel: + def test_known_tool(self): + assert tool_label("web_search") == "Searching the web…" + + def test_fallback(self): + assert tool_label("mystery_tool") == "Using mystery_tool…" + + def test_mcp_prefixes(self): + assert tool_label("homebox_search_items") == "Checking home inventory…" + assert tool_label("bookstack_get_page") == "Checking the wiki…" + + def test_old_and_new_ynab_names(self): + for name in ( + "ynab_category_spending", + "ynab_insights", + "ynab_account_balances", + "ynab_spending_report", + "ynab_category_budget", + "ynab_recent_changes", + "ynab_sql", + ): + assert name in TOOL_LABELS + + def test_every_agent_tool_has_a_label(self): + names = [v.name for v in vars(agents).values() if isinstance(v, BaseTool)] + assert len(names) > 20 + assert [n for n in names if n not in TOOL_LABELS] == [] + + +class TestSummarizeToolInput: + def test_drops_config_and_none(self): + raw = {"content": "hi", "config": {"configurable": {}}, "date": None} + assert summarize_tool_input(raw) == {"content": "hi"} + + def test_truncates_long_strings(self): + out = summarize_tool_input({"query": "x" * 500}) + assert len(out["query"]) == INPUT_VALUE_MAX_CHARS + assert out["query"].endswith("…") + + def test_keeps_scalars(self): + assert summarize_tool_input({"days": 7, "flag": True, "f": 1.5}) == { + "days": 7, + "flag": True, + "f": 1.5, + } + + def test_serialises_nested_values(self): + assert summarize_tool_input({"tags": ["a", "b"]}) == {"tags": '["a", "b"]'} + + def test_non_dict_input(self): + assert summarize_tool_input("raw") == {"input": "raw"} + assert summarize_tool_input(None) == {} + + +class TestPayloads: + def test_tool_start(self): + payload = tool_start_payload(_event(input={"query": "cats"})) + assert payload == { + "type": "tool_start", + "run_id": "run-1", + "tool": "web_search", + "label": "Searching the web…", + "input": {"query": "cats"}, + } + + def test_tool_end_ok(self): + ref = {"id": "a1", "kind": "mealie.recipe", "title": "Soup"} + payload = tool_end_payload(_event(), 10.0, ref, now=10.25) + assert payload == { + "type": "tool_end", + "run_id": "run-1", + "tool": "web_search", + "label": "Searching the web…", + "status": "ok", + "ms": 250, + "artifact": ref, + } + + def test_tool_end_error_without_start(self): + payload = tool_end_payload(_event(), None, None, now=5.0, status="error") + assert payload["status"] == "error" + assert payload["ms"] == 0 + assert payload["artifact"] is None + + +class TestToolNameFromMessage: + def test_parses_header(self): + assert tool_name_from_message("[web_search result]\nbody") == "web_search" + + def test_unknown_header(self): + assert tool_name_from_message("something else") == "tool"