Stream v2 tool events, save artifacts, enrich conversation payload
stream_query emits the new tool_start/tool_end payloads, maps on_tool_error to tool_end status=error, and saves an artifact linked to the tool message. GET /api/conversation/<id> adds tool, label and artifact to tool messages. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
This commit is contained in:
@@ -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,33 @@ async def get_conversation(conversation_id: str):
|
||||
if conversation is None:
|
||||
return jsonify({"error": "conversation not found"}), 404
|
||||
await conversation.fetch_related("messages")
|
||||
artifacts_by_message = {
|
||||
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(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,363 @@
|
||||
"""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_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]")
|
||||
Reference in New Issue
Block a user