Merge pull request 'Cowork A: SSE v2 tool events, artifacts table + API' (#60) from feat/cowork-stream into main
This commit was merged in pull request #60.
This commit is contained in:
@@ -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():
|
||||
|
||||
@@ -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("/<artifact_id>")
|
||||
@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))
|
||||
@@ -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,
|
||||
)
|
||||
@@ -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"
|
||||
@@ -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),
|
||||
|
||||
@@ -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"
|
||||
@@ -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 <tool>…"."""
|
||||
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}…"
|
||||
@@ -19,6 +19,7 @@ TORTOISE_CONFIG = {
|
||||
"blueprints.scheduled_messages.models",
|
||||
"blueprints.cron.models",
|
||||
"blueprints.beli.models",
|
||||
"blueprints.artifacts.models",
|
||||
"aerich.models",
|
||||
],
|
||||
"default_connection": "default",
|
||||
|
||||
@@ -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="
|
||||
)
|
||||
@@ -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]")
|
||||
@@ -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"
|
||||
Reference in New Issue
Block a user