diff --git a/drizzle/0019_large_invaders.sql b/drizzle/0019_large_invaders.sql new file mode 100644 index 0000000..0580b65 --- /dev/null +++ b/drizzle/0019_large_invaders.sql @@ -0,0 +1,18 @@ +CREATE TABLE "folders" ( + "id" uuid PRIMARY KEY DEFAULT gen_random_uuid() NOT NULL, + "workspace_id" uuid NOT NULL, + "parent_id" uuid, + "name" text NOT NULL, + "created_at" timestamp with time zone DEFAULT now() NOT NULL, + "updated_at" timestamp with time zone DEFAULT now() NOT NULL, + "deleted_at" timestamp with time zone +); +--> statement-breakpoint +ALTER TABLE "sources" ADD COLUMN "folder_id" uuid;--> statement-breakpoint +ALTER TABLE "folders" ADD CONSTRAINT "folders_workspace_id_workspaces_id_fk" FOREIGN KEY ("workspace_id") REFERENCES "public"."workspaces"("id") ON DELETE cascade ON UPDATE no action;--> statement-breakpoint +ALTER TABLE "folders" ADD CONSTRAINT "folders_parent_id_folders_id_fk" FOREIGN KEY ("parent_id") REFERENCES "public"."folders"("id") ON DELETE restrict ON UPDATE no action;--> statement-breakpoint +CREATE INDEX "folders_workspace_parent_idx" ON "folders" USING btree ("workspace_id","parent_id");--> statement-breakpoint +CREATE UNIQUE INDEX "folders_workspace_parent_name_idx" ON "folders" USING btree ("workspace_id","parent_id","name") WHERE deleted_at IS NULL AND parent_id IS NOT NULL;--> statement-breakpoint +CREATE UNIQUE INDEX "folders_workspace_root_name_idx" ON "folders" USING btree ("workspace_id","name") WHERE deleted_at IS NULL AND parent_id IS NULL;--> statement-breakpoint +ALTER TABLE "sources" ADD CONSTRAINT "sources_folder_id_folders_id_fk" FOREIGN KEY ("folder_id") REFERENCES "public"."folders"("id") ON DELETE set null ON UPDATE no action;--> statement-breakpoint +CREATE INDEX "sources_workspace_folder_idx" ON "sources" USING btree ("workspace_id","folder_id") WHERE deleted_at IS NULL; \ No newline at end of file diff --git a/drizzle/meta/0019_snapshot.json b/drizzle/meta/0019_snapshot.json new file mode 100644 index 0000000..3c97a48 --- /dev/null +++ b/drizzle/meta/0019_snapshot.json @@ -0,0 +1,1136 @@ +{ + "id": "325d2bfd-d4cf-4c41-baf2-b93fc91afa5e", + "prevId": "01b8a746-ca75-4e79-adab-098a03a39e0e", + "version": "7", + "dialect": "postgresql", + "tables": { + "public.chat_messages": { + "name": "chat_messages", + "schema": "", + "columns": { + "id": { + "name": "id", + "type": "uuid", + "primaryKey": true, + "notNull": true, + "default": "gen_random_uuid()" + }, + "thread_id": { + "name": "thread_id", + "type": "uuid", + "primaryKey": false, + "notNull": true + }, + "role": { + "name": "role", + "type": "text", + "primaryKey": false, + "notNull": true + }, + "content": { + "name": "content", + "type": "text", + "primaryKey": false, + "notNull": true + }, + "citations": { + "name": "citations", + "type": "jsonb", + "primaryKey": false, + "notNull": false + }, + "artifacts": { + "name": "artifacts", + "type": "jsonb", + "primaryKey": false, + "notNull": false + }, + "created_at": { + "name": "created_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": true, + "default": "now()" + } + }, + "indexes": { + "chat_messages_thread_created_idx": { + "name": "chat_messages_thread_created_idx", + "columns": [ + { + "expression": "thread_id", + "isExpression": false, + "asc": true, + "nulls": "last" + }, + { + "expression": "created_at", + "isExpression": false, + "asc": true, + "nulls": "last" + } + ], + "isUnique": false, + "concurrently": false, + "method": "btree", + "with": {} + } + }, + "foreignKeys": { + "chat_messages_thread_id_chat_threads_id_fk": { + "name": "chat_messages_thread_id_chat_threads_id_fk", + "tableFrom": "chat_messages", + "tableTo": "chat_threads", + "columnsFrom": [ + "thread_id" + ], + "columnsTo": [ + "id" + ], + "onDelete": "cascade", + "onUpdate": "no action" + } + }, + "compositePrimaryKeys": {}, + "uniqueConstraints": {}, + "policies": {}, + "checkConstraints": {}, + "isRLSEnabled": false + }, + "public.chat_threads": { + "name": "chat_threads", + "schema": "", + "columns": { + "id": { + "name": "id", + "type": "uuid", + "primaryKey": true, + "notNull": true, + "default": "gen_random_uuid()" + }, + "workspace_id": { + "name": "workspace_id", + "type": "uuid", + "primaryKey": false, + "notNull": true + }, + "title": { + "name": "title", + "type": "text", + "primaryKey": false, + "notNull": false + }, + "demo_key": { + "name": "demo_key", + "type": "text", + "primaryKey": false, + "notNull": false + }, + "created_at": { + "name": "created_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": true, + "default": "now()" + }, + "updated_at": { + "name": "updated_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": true, + "default": "now()" + }, + "deleted_at": { + "name": "deleted_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": false + } + }, + "indexes": { + "chat_threads_workspace_updated_idx": { + "name": "chat_threads_workspace_updated_idx", + "columns": [ + { + "expression": "workspace_id", + "isExpression": false, + "asc": true, + "nulls": "last" + }, + { + "expression": "updated_at", + "isExpression": false, + "asc": false, + "nulls": "last" + } + ], + "isUnique": false, + "where": "deleted_at IS NULL", + "concurrently": false, + "method": "btree", + "with": {} + }, + "chat_threads_workspace_demo_key_idx": { + "name": "chat_threads_workspace_demo_key_idx", + "columns": [ + { + "expression": "workspace_id", + "isExpression": false, + "asc": true, + "nulls": "last" + }, + { + "expression": "demo_key", + "isExpression": false, + "asc": true, + "nulls": "last" + } + ], + "isUnique": true, + "concurrently": false, + "method": "btree", + "with": {} + } + }, + "foreignKeys": { + "chat_threads_workspace_id_workspaces_id_fk": { + "name": "chat_threads_workspace_id_workspaces_id_fk", + "tableFrom": "chat_threads", + "tableTo": "workspaces", + "columnsFrom": [ + "workspace_id" + ], + "columnsTo": [ + "id" + ], + "onDelete": "cascade", + "onUpdate": "no action" + } + }, + "compositePrimaryKeys": {}, + "uniqueConstraints": {}, + "policies": {}, + "checkConstraints": {}, + "isRLSEnabled": false + }, + "public.demo_source_visibilities": { + "name": "demo_source_visibilities", + "schema": "", + "columns": { + "id": { + "name": "id", + "type": "uuid", + "primaryKey": true, + "notNull": true, + "default": "gen_random_uuid()" + }, + "workspace_id": { + "name": "workspace_id", + "type": "uuid", + "primaryKey": false, + "notNull": true + }, + "demo_source_id": { + "name": "demo_source_id", + "type": "text", + "primaryKey": false, + "notNull": true + }, + "hidden_at": { + "name": "hidden_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": false + }, + "deleted_at": { + "name": "deleted_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": false + }, + "created_at": { + "name": "created_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": true, + "default": "now()" + }, + "updated_at": { + "name": "updated_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": true, + "default": "now()" + } + }, + "indexes": { + "demo_source_visibilities_workspace_source_idx": { + "name": "demo_source_visibilities_workspace_source_idx", + "columns": [ + { + "expression": "workspace_id", + "isExpression": false, + "asc": true, + "nulls": "last" + }, + { + "expression": "demo_source_id", + "isExpression": false, + "asc": true, + "nulls": "last" + } + ], + "isUnique": true, + "concurrently": false, + "method": "btree", + "with": {} + }, + "demo_source_visibilities_workspace_idx": { + "name": "demo_source_visibilities_workspace_idx", + "columns": [ + { + "expression": "workspace_id", + "isExpression": false, + "asc": true, + "nulls": "last" + } + ], + "isUnique": false, + "concurrently": false, + "method": "btree", + "with": {} + } + }, + "foreignKeys": { + "demo_source_visibilities_workspace_id_workspaces_id_fk": { + "name": "demo_source_visibilities_workspace_id_workspaces_id_fk", + "tableFrom": "demo_source_visibilities", + "tableTo": "workspaces", + "columnsFrom": [ + "workspace_id" + ], + "columnsTo": [ + "id" + ], + "onDelete": "cascade", + "onUpdate": "no action" + } + }, + "compositePrimaryKeys": {}, + "uniqueConstraints": {}, + "policies": {}, + "checkConstraints": {}, + "isRLSEnabled": false + }, + "public.folders": { + "name": "folders", + "schema": "", + "columns": { + "id": { + "name": "id", + "type": "uuid", + "primaryKey": true, + "notNull": true, + "default": "gen_random_uuid()" + }, + "workspace_id": { + "name": "workspace_id", + "type": "uuid", + "primaryKey": false, + "notNull": true + }, + "parent_id": { + "name": "parent_id", + "type": "uuid", + "primaryKey": false, + "notNull": false + }, + "name": { + "name": "name", + "type": "text", + "primaryKey": false, + "notNull": true + }, + "created_at": { + "name": "created_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": true, + "default": "now()" + }, + "updated_at": { + "name": "updated_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": true, + "default": "now()" + }, + "deleted_at": { + "name": "deleted_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": false + } + }, + "indexes": { + "folders_workspace_parent_idx": { + "name": "folders_workspace_parent_idx", + "columns": [ + { + "expression": "workspace_id", + "isExpression": false, + "asc": true, + "nulls": "last" + }, + { + "expression": "parent_id", + "isExpression": false, + "asc": true, + "nulls": "last" + } + ], + "isUnique": false, + "concurrently": false, + "method": "btree", + "with": {} + }, + "folders_workspace_parent_name_idx": { + "name": "folders_workspace_parent_name_idx", + "columns": [ + { + "expression": "workspace_id", + "isExpression": false, + "asc": true, + "nulls": "last" + }, + { + "expression": "parent_id", + "isExpression": false, + "asc": true, + "nulls": "last" + }, + { + "expression": "name", + "isExpression": false, + "asc": true, + "nulls": "last" + } + ], + "isUnique": true, + "where": "deleted_at IS NULL AND parent_id IS NOT NULL", + "concurrently": false, + "method": "btree", + "with": {} + }, + "folders_workspace_root_name_idx": { + "name": "folders_workspace_root_name_idx", + "columns": [ + { + "expression": "workspace_id", + "isExpression": false, + "asc": true, + "nulls": "last" + }, + { + "expression": "name", + "isExpression": false, + "asc": true, + "nulls": "last" + } + ], + "isUnique": true, + "where": "deleted_at IS NULL AND parent_id IS NULL", + "concurrently": false, + "method": "btree", + "with": {} + } + }, + "foreignKeys": { + "folders_workspace_id_workspaces_id_fk": { + "name": "folders_workspace_id_workspaces_id_fk", + "tableFrom": "folders", + "tableTo": "workspaces", + "columnsFrom": [ + "workspace_id" + ], + "columnsTo": [ + "id" + ], + "onDelete": "cascade", + "onUpdate": "no action" + }, + "folders_parent_id_folders_id_fk": { + "name": "folders_parent_id_folders_id_fk", + "tableFrom": "folders", + "tableTo": "folders", + "columnsFrom": [ + "parent_id" + ], + "columnsTo": [ + "id" + ], + "onDelete": "restrict", + "onUpdate": "no action" + } + }, + "compositePrimaryKeys": {}, + "uniqueConstraints": {}, + "policies": {}, + "checkConstraints": {}, + "isRLSEnabled": false + }, + "public.parsed_document_sync_leases": { + "name": "parsed_document_sync_leases", + "schema": "", + "columns": { + "id": { + "name": "id", + "type": "uuid", + "primaryKey": true, + "notNull": true, + "default": "gen_random_uuid()" + }, + "workspace_id": { + "name": "workspace_id", + "type": "uuid", + "primaryKey": false, + "notNull": true + }, + "source_id": { + "name": "source_id", + "type": "uuid", + "primaryKey": false, + "notNull": true + }, + "document_id": { + "name": "document_id", + "type": "text", + "primaryKey": false, + "notNull": true + }, + "revision_key": { + "name": "revision_key", + "type": "text", + "primaryKey": false, + "notNull": false + }, + "lease_token": { + "name": "lease_token", + "type": "text", + "primaryKey": false, + "notNull": true + }, + "acquired_at": { + "name": "acquired_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": true, + "default": "now()" + }, + "expires_at": { + "name": "expires_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": true + }, + "released_at": { + "name": "released_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": false + }, + "release_reason": { + "name": "release_reason", + "type": "text", + "primaryKey": false, + "notNull": false + }, + "created_at": { + "name": "created_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": true, + "default": "now()" + }, + "updated_at": { + "name": "updated_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": true, + "default": "now()" + } + }, + "indexes": { + "parsed_document_sync_leases_token_idx": { + "name": "parsed_document_sync_leases_token_idx", + "columns": [ + { + "expression": "lease_token", + "isExpression": false, + "asc": true, + "nulls": "last" + } + ], + "isUnique": true, + "concurrently": false, + "method": "btree", + "with": {} + }, + "parsed_document_sync_leases_active_idx": { + "name": "parsed_document_sync_leases_active_idx", + "columns": [ + { + "expression": "expires_at", + "isExpression": false, + "asc": true, + "nulls": "last" + } + ], + "isUnique": false, + "where": "released_at IS NULL", + "concurrently": false, + "method": "btree", + "with": {} + }, + "parsed_document_sync_leases_workspace_active_idx": { + "name": "parsed_document_sync_leases_workspace_active_idx", + "columns": [ + { + "expression": "workspace_id", + "isExpression": false, + "asc": true, + "nulls": "last" + } + ], + "isUnique": false, + "where": "released_at IS NULL", + "concurrently": false, + "method": "btree", + "with": {} + }, + "parsed_document_sync_leases_document_active_idx": { + "name": "parsed_document_sync_leases_document_active_idx", + "columns": [ + { + "expression": "document_id", + "isExpression": false, + "asc": true, + "nulls": "last" + } + ], + "isUnique": false, + "where": "released_at IS NULL", + "concurrently": false, + "method": "btree", + "with": {} + } + }, + "foreignKeys": { + "parsed_document_sync_leases_workspace_id_workspaces_id_fk": { + "name": "parsed_document_sync_leases_workspace_id_workspaces_id_fk", + "tableFrom": "parsed_document_sync_leases", + "tableTo": "workspaces", + "columnsFrom": [ + "workspace_id" + ], + "columnsTo": [ + "id" + ], + "onDelete": "cascade", + "onUpdate": "no action" + }, + "parsed_document_sync_leases_source_id_sources_id_fk": { + "name": "parsed_document_sync_leases_source_id_sources_id_fk", + "tableFrom": "parsed_document_sync_leases", + "tableTo": "sources", + "columnsFrom": [ + "source_id" + ], + "columnsTo": [ + "id" + ], + "onDelete": "cascade", + "onUpdate": "no action" + } + }, + "compositePrimaryKeys": {}, + "uniqueConstraints": {}, + "policies": {}, + "checkConstraints": {}, + "isRLSEnabled": false + }, + "public.source_parse_results": { + "name": "source_parse_results", + "schema": "", + "columns": { + "id": { + "name": "id", + "type": "uuid", + "primaryKey": true, + "notNull": true, + "default": "gen_random_uuid()" + }, + "source_id": { + "name": "source_id", + "type": "uuid", + "primaryKey": false, + "notNull": true + }, + "result_blob_url": { + "name": "result_blob_url", + "type": "text", + "primaryKey": false, + "notNull": false + }, + "snapshot_manifest_url": { + "name": "snapshot_manifest_url", + "type": "text", + "primaryKey": false, + "notNull": false + }, + "snapshot_manifest_key": { + "name": "snapshot_manifest_key", + "type": "text", + "primaryKey": false, + "notNull": false + }, + "revision_key": { + "name": "revision_key", + "type": "text", + "primaryKey": false, + "notNull": false + }, + "sync_status": { + "name": "sync_status", + "type": "text", + "primaryKey": false, + "notNull": false + }, + "sync_error": { + "name": "sync_error", + "type": "text", + "primaryKey": false, + "notNull": false + }, + "asset_urls": { + "name": "asset_urls", + "type": "jsonb", + "primaryKey": false, + "notNull": true + }, + "created_at": { + "name": "created_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": true, + "default": "now()" + }, + "updated_at": { + "name": "updated_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": true, + "default": "now()" + } + }, + "indexes": { + "source_parse_results_source_id_idx": { + "name": "source_parse_results_source_id_idx", + "columns": [ + { + "expression": "source_id", + "isExpression": false, + "asc": true, + "nulls": "last" + } + ], + "isUnique": false, + "concurrently": false, + "method": "btree", + "with": {} + } + }, + "foreignKeys": { + "source_parse_results_source_id_sources_id_fk": { + "name": "source_parse_results_source_id_sources_id_fk", + "tableFrom": "source_parse_results", + "tableTo": "sources", + "columnsFrom": [ + "source_id" + ], + "columnsTo": [ + "id" + ], + "onDelete": "cascade", + "onUpdate": "no action" + } + }, + "compositePrimaryKeys": {}, + "uniqueConstraints": { + "source_parse_results_source_id_unique": { + "name": "source_parse_results_source_id_unique", + "nullsNotDistinct": false, + "columns": [ + "source_id" + ] + } + }, + "policies": {}, + "checkConstraints": {}, + "isRLSEnabled": false + }, + "public.sources": { + "name": "sources", + "schema": "", + "columns": { + "id": { + "name": "id", + "type": "uuid", + "primaryKey": true, + "notNull": true, + "default": "gen_random_uuid()" + }, + "workspace_id": { + "name": "workspace_id", + "type": "uuid", + "primaryKey": false, + "notNull": true + }, + "title": { + "name": "title", + "type": "text", + "primaryKey": false, + "notNull": true + }, + "mime_type": { + "name": "mime_type", + "type": "text", + "primaryKey": false, + "notNull": true + }, + "size_bytes": { + "name": "size_bytes", + "type": "bigint", + "primaryKey": false, + "notNull": true + }, + "status": { + "name": "status", + "type": "text", + "primaryKey": false, + "notNull": true + }, + "failure_reason": { + "name": "failure_reason", + "type": "text", + "primaryKey": false, + "notNull": false + }, + "failure_stage": { + "name": "failure_stage", + "type": "text", + "primaryKey": false, + "notNull": false + }, + "knowhere_job_id": { + "name": "knowhere_job_id", + "type": "text", + "primaryKey": false, + "notNull": false + }, + "knowhere_document_id": { + "name": "knowhere_document_id", + "type": "text", + "primaryKey": false, + "notNull": false + }, + "staged_blob_pathname": { + "name": "staged_blob_pathname", + "type": "text", + "primaryKey": false, + "notNull": false + }, + "staged_blob_url": { + "name": "staged_blob_url", + "type": "text", + "primaryKey": false, + "notNull": false + }, + "original_blob_pathname": { + "name": "original_blob_pathname", + "type": "text", + "primaryKey": false, + "notNull": false + }, + "original_blob_url": { + "name": "original_blob_url", + "type": "text", + "primaryKey": false, + "notNull": false + }, + "demo_key": { + "name": "demo_key", + "type": "text", + "primaryKey": false, + "notNull": false + }, + "chunk_count": { + "name": "chunk_count", + "type": "integer", + "primaryKey": false, + "notNull": false + }, + "folder_id": { + "name": "folder_id", + "type": "uuid", + "primaryKey": false, + "notNull": false + }, + "created_at": { + "name": "created_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": true, + "default": "now()" + }, + "updated_at": { + "name": "updated_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": true, + "default": "now()" + }, + "deleted_at": { + "name": "deleted_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": false + } + }, + "indexes": { + "sources_workspace_created_idx": { + "name": "sources_workspace_created_idx", + "columns": [ + { + "expression": "workspace_id", + "isExpression": false, + "asc": true, + "nulls": "last" + }, + { + "expression": "created_at", + "isExpression": false, + "asc": false, + "nulls": "last" + } + ], + "isUnique": false, + "where": "deleted_at IS NULL", + "concurrently": false, + "method": "btree", + "with": {} + }, + "sources_workspace_status_idx": { + "name": "sources_workspace_status_idx", + "columns": [ + { + "expression": "workspace_id", + "isExpression": false, + "asc": true, + "nulls": "last" + }, + { + "expression": "status", + "isExpression": false, + "asc": true, + "nulls": "last" + } + ], + "isUnique": false, + "concurrently": false, + "method": "btree", + "with": {} + }, + "sources_workspace_demo_key_idx": { + "name": "sources_workspace_demo_key_idx", + "columns": [ + { + "expression": "workspace_id", + "isExpression": false, + "asc": true, + "nulls": "last" + }, + { + "expression": "demo_key", + "isExpression": false, + "asc": true, + "nulls": "last" + } + ], + "isUnique": true, + "concurrently": false, + "method": "btree", + "with": {} + }, + "sources_workspace_document_idx": { + "name": "sources_workspace_document_idx", + "columns": [ + { + "expression": "workspace_id", + "isExpression": false, + "asc": true, + "nulls": "last" + }, + { + "expression": "knowhere_document_id", + "isExpression": false, + "asc": true, + "nulls": "last" + } + ], + "isUnique": true, + "where": "knowhere_document_id IS NOT NULL AND deleted_at IS NULL", + "concurrently": false, + "method": "btree", + "with": {} + }, + "sources_workspace_folder_idx": { + "name": "sources_workspace_folder_idx", + "columns": [ + { + "expression": "workspace_id", + "isExpression": false, + "asc": true, + "nulls": "last" + }, + { + "expression": "folder_id", + "isExpression": false, + "asc": true, + "nulls": "last" + } + ], + "isUnique": false, + "where": "deleted_at IS NULL", + "concurrently": false, + "method": "btree", + "with": {} + } + }, + "foreignKeys": { + "sources_workspace_id_workspaces_id_fk": { + "name": "sources_workspace_id_workspaces_id_fk", + "tableFrom": "sources", + "tableTo": "workspaces", + "columnsFrom": [ + "workspace_id" + ], + "columnsTo": [ + "id" + ], + "onDelete": "cascade", + "onUpdate": "no action" + }, + "sources_folder_id_folders_id_fk": { + "name": "sources_folder_id_folders_id_fk", + "tableFrom": "sources", + "tableTo": "folders", + "columnsFrom": [ + "folder_id" + ], + "columnsTo": [ + "id" + ], + "onDelete": "set null", + "onUpdate": "no action" + } + }, + "compositePrimaryKeys": {}, + "uniqueConstraints": {}, + "policies": {}, + "checkConstraints": {}, + "isRLSEnabled": false + }, + "public.workspaces": { + "name": "workspaces", + "schema": "", + "columns": { + "id": { + "name": "id", + "type": "uuid", + "primaryKey": true, + "notNull": true, + "default": "gen_random_uuid()" + }, + "user_id": { + "name": "user_id", + "type": "text", + "primaryKey": false, + "notNull": true + }, + "namespace": { + "name": "namespace", + "type": "text", + "primaryKey": false, + "notNull": true + }, + "created_at": { + "name": "created_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": true, + "default": "now()" + } + }, + "indexes": { + "workspaces_user_id_idx": { + "name": "workspaces_user_id_idx", + "columns": [ + { + "expression": "user_id", + "isExpression": false, + "asc": true, + "nulls": "last" + } + ], + "isUnique": false, + "concurrently": false, + "method": "btree", + "with": {} + } + }, + "foreignKeys": {}, + "compositePrimaryKeys": {}, + "uniqueConstraints": { + "workspaces_user_id_unique": { + "name": "workspaces_user_id_unique", + "nullsNotDistinct": false, + "columns": [ + "user_id" + ] + }, + "workspaces_namespace_unique": { + "name": "workspaces_namespace_unique", + "nullsNotDistinct": false, + "columns": [ + "namespace" + ] + } + }, + "policies": {}, + "checkConstraints": {}, + "isRLSEnabled": false + } + }, + "enums": {}, + "schemas": {}, + "sequences": {}, + "roles": {}, + "policies": {}, + "views": {}, + "_meta": { + "columns": {}, + "schemas": {}, + "tables": {} + } +} \ No newline at end of file diff --git a/drizzle/meta/_journal.json b/drizzle/meta/_journal.json index 8439023..a3a4315 100644 --- a/drizzle/meta/_journal.json +++ b/drizzle/meta/_journal.json @@ -134,6 +134,13 @@ "when": 1789118032723, "tag": "0018_source_chunk_count", "breakpoints": true + }, + { + "idx": 19, + "version": "7", + "when": 1789124024330, + "tag": "0019_large_invaders", + "breakpoints": true } ] } \ No newline at end of file diff --git a/src/agent-harness/answer-context.test.ts b/src/agent-harness/answer-context.test.ts new file mode 100644 index 0000000..acc5053 --- /dev/null +++ b/src/agent-harness/answer-context.test.ts @@ -0,0 +1,203 @@ +import { describe, expect, it, vi } from "vitest" + +import { composeAnswerContext } from "./answer-context" +import type { EvidenceLedgerSnapshot } from "./types" + +describe("composeAnswerContext", () => { + it("composes retained chunks, full table HTML, images, memory, and the original question", async () => { + const readTableHtml = vi.fn().mockResolvedValue( + "
Q424.9
", + ) + const message = await composeAnswerContext({ + ledger: makeLedger(), + memoryItems: [ + { + ref: "mem:1", + itemId: "memory_1", + kind: "stance", + abstractL0: "关注毛利率", + overviewL1: "用户把毛利率当作核心指标。", + }, + ], + userText: "Q4 的收入是多少?", + readTableHtml, + }) + + expect(readTableHtml).toHaveBeenCalledWith( + "https://assets.example/tables/revenue.html?signature=valid", + ) + expect(message).toEqual({ + role: "user", + content: [ + { + type: "text", + text: [ + "## Evidence from Knowledge Base", + "", + "[pick 1] report.pdf › Summary", + '[artifact type="table" ref="asset:r1:result:1:table_chunk"]', + "Revenue increased in Q4.", + "
Q424.9
", + "", + '[image ref="asset:r1:result:1:image_chunk_1"]', + ].join("\n"), + }, + { + type: "image", + image: new URL("https://assets.example/images/revenue.png?signature=valid"), + }, + { + type: "text", + text: '[image ref="asset:r1:result:1:image_chunk_2"]', + }, + { + type: "image", + image: new URL("https://assets.example/images/margin.png?signature=valid"), + }, + { + type: "text", + text: [ + "## Fluid Memory", + "", + '[memory ref="mem:1" itemId="memory_1" kind="stance"]', + "关注毛利率", + "用户把毛利率当作核心指标。", + "", + "## User's Question", + "Q4 的收入是多少?", + ].join("\n"), + }, + ], + }) + }) + + it("omits knowledge and memory sections when neither search ran", async () => { + const message = await composeAnswerContext({ + ledger: { ...makeLedger(), retrievalCount: 0, retainedPicks: [] }, + memoryItems: [], + userText: "直接回答。", + }) + + expect(message).toEqual({ + role: "user", + content: [{ type: "text", text: "## User's Question\n直接回答。" }], + }) + }) + + it("notes when knowledge base search ran but found nothing to retain", async () => { + const message = await composeAnswerContext({ + ledger: { ...makeLedger(), retainedPicks: [] }, + memoryItems: [], + userText: "直接回答。", + }) + + expect(message).toEqual({ + role: "user", + content: [ + { + type: "text", + text: [ + "## Evidence from Knowledge Base", + "No knowledge base results were found for this search. Answer using your own knowledge only, and say so if relevant.", + "", + "## User's Question", + "直接回答。", + ].join("\n"), + }, + ], + }) + }) + + it("notes when memory search ran but found nothing", async () => { + const message = await composeAnswerContext({ + ledger: { ...makeLedger(), retrievalCount: 0, retainedPicks: [] }, + memoryItems: [], + memorySearchAttempted: true, + userText: "直接回答。", + }) + + expect(message).toEqual({ + role: "user", + content: [ + { + type: "text", + text: [ + "## Fluid Memory", + "No fluid memory results were found for this search.", + "", + "## User's Question", + "直接回答。", + ].join("\n"), + }, + ], + }) + }) +}) + +function makeLedger(): EvidenceLedgerSnapshot { + return { + retrievalCount: 1, + chunks: [ + { + ref: "r1:result:1", + kind: "result", + content: "Revenue increased in Q4.\n[Table: tables/revenue.html]", + contentPreview: "Revenue increased in Q4. [Table: tables/revenue.html]", + chunkType: "text", + score: 0.9, + source: { + documentId: "doc_1", + sourceFileName: "report.pdf", + sectionPath: "Summary", + }, + }, + ], + assets: [ + { + ref: "asset:r1:result:1:table_chunk", + chunkRef: "r1:result:1", + type: "table", + assetUrl: "https://assets.example/tables/revenue.html?signature=valid", + sourcePath: "tables/revenue.html", + source: { + documentId: "doc_1", + sourceFileName: "report.pdf", + sectionPath: "Summary", + }, + label: "report.pdf / Summary / tables/revenue.html / table", + }, + { + ref: "asset:r1:result:1:image_chunk_1", + chunkRef: "r1:result:1", + type: "image", + assetUrl: "https://assets.example/images/revenue.png?signature=valid", + sourcePath: "images/revenue.png", + source: { + documentId: "doc_1", + sourceFileName: "report.pdf", + sectionPath: "Summary", + }, + label: "report.pdf / Summary / images/revenue.png / image", + }, + { + ref: "asset:r1:result:1:image_chunk_2", + chunkRef: "r1:result:1", + type: "image", + assetUrl: "https://assets.example/images/margin.png?signature=valid", + sourcePath: "images/margin.png", + source: { + documentId: "doc_1", + sourceFileName: "report.pdf", + sectionPath: "Summary", + }, + label: "report.pdf / Summary / images/margin.png / image", + }, + ], + evidenceText: [], + stopReasons: [], + failureReasons: [], + decisionTraces: [], + retainedPicks: [1], + pendingRetention: null, + } +} diff --git a/src/agent-harness/answer-context.ts b/src/agent-harness/answer-context.ts new file mode 100644 index 0000000..3a6ff8e --- /dev/null +++ b/src/agent-harness/answer-context.ts @@ -0,0 +1,205 @@ +import type { ModelMessage } from "ai" + +import type { + EvidenceAsset, + EvidenceChunk, + EvidenceLedgerSnapshot, + MemorySearchItem, +} from "./types" + +export type ReadTableHtml = (assetUrl: string) => Promise +type UserContentParts = Exclude< + Extract["content"], + string +> + +export async function composeAnswerContext(input: { + readonly ledger: EvidenceLedgerSnapshot + readonly memoryItems: readonly MemorySearchItem[] + readonly memorySearchAttempted?: boolean + readonly userText: string + readonly readTableHtml?: ReadTableHtml +}): Promise { + const content: UserContentParts = [] + const retainedPicks = new Set(input.ledger.retainedPicks) + const assetsByChunkRef = groupAssetsByChunkRef(input.ledger.assets) + + const retainedChunks = input.ledger.chunks.flatMap((chunk, index) => { + const pick = index + 1 + return retainedPicks.has(pick) ? [{ chunk, pick }] : [] + }) + + if (retainedChunks.length === 0 && input.ledger.retrievalCount > 0) { + appendText( + content, + "## Evidence from Knowledge Base\nNo knowledge base results were found for this search. Answer using your own knowledge only, and say so if relevant.", + ) + } + + if (retainedChunks.length > 0) { + appendText(content, "## Evidence from Knowledge Base") + for (const { chunk, pick } of retainedChunks) { + const assets = assetsByChunkRef.get(chunk.ref) ?? [] + let chunkText = chunk.content + for (const table of assets.filter((asset) => asset.type === "table")) { + chunkText = await replaceTableWithHtml({ + chunk: { ...chunk, content: chunkText }, + asset: table, + readTableHtml: input.readTableHtml, + }) + } + + appendText( + content, + [ + formatChunkLabel(pick, chunk), + formatTableArtifactRefs(assets), + chunkText.trim(), + ] + .filter((part) => part.length > 0) + .join("\n"), + ) + + for (const image of assets.filter((asset) => asset.type === "image")) { + if (!image.assetUrl) { + throw new Error(`Retained image asset ${image.ref} has no signed URL.`) + } + appendText(content, `[image ref="${image.ref}"]`) + content.push({ + type: "image", + image: new URL(image.assetUrl), + }) + } + } + } + + if (input.memoryItems.length > 0) { + appendText(content, [ + "## Fluid Memory", + ...input.memoryItems.map(formatMemoryItem), + ].join("\n\n")) + } else if (input.memorySearchAttempted) { + appendText( + content, + "## Fluid Memory\nNo fluid memory results were found for this search.", + ) + } + + const userText = input.userText.trim() + if (userText) { + appendText(content, `## User's Question\n${userText}`) + } + + return { role: "user", content } +} + +function groupAssetsByChunkRef( + assets: readonly EvidenceAsset[], +): Map { + const grouped = new Map() + for (const asset of assets) { + const chunkAssets = grouped.get(asset.chunkRef) ?? [] + chunkAssets.push(asset) + grouped.set(asset.chunkRef, chunkAssets) + } + return grouped +} + +function formatTableArtifactRefs(assets: readonly EvidenceAsset[]): string { + return assets + .filter((asset) => asset.type === "table") + .map((asset) => `[artifact type="${asset.type}" ref="${asset.ref}"]`) + .join("\n") +} + +function appendText( + content: UserContentParts, + text: string, +): void { + const trimmed = text.trim() + if (!trimmed) return + const previous = content.at(-1) + if (previous?.type === "text") { + previous.text = `${previous.text}\n\n${trimmed}` + return + } + content.push({ type: "text", text: trimmed }) +} + +function formatChunkLabel(pick: number, chunk: EvidenceChunk): string { + const location = [chunk.source.sourceFileName, chunk.source.sectionPath] + .map((part) => part?.trim()) + .filter((part): part is string => Boolean(part)) + .join(" › ") + return location ? `[pick ${pick}] ${location}` : `[pick ${pick}]` +} + +function formatMemoryItem(item: MemorySearchItem): string { + return [ + `[memory ref="${item.ref}" itemId="${item.itemId}" kind="${item.kind}"]`, + item.abstractL0.trim(), + item.overviewL1.trim(), + ] + .filter((part) => part.length > 0) + .join("\n") +} + +async function replaceTableWithHtml(input: { + readonly chunk: EvidenceChunk + readonly asset: EvidenceAsset + readonly readTableHtml?: ReadTableHtml +}): Promise { + if (!input.asset.assetUrl) { + throw new Error(`Retained table asset ${input.asset.ref} has no signed URL.`) + } + if (!input.readTableHtml) { + throw new Error("No table HTML reader was provided for retained table evidence.") + } + + const html = await input.readTableHtml(input.asset.assetUrl) + const content = input.chunk.content + if (!content.trim()) return html + + const placeholders = getTablePlaceholderCandidates(input.chunk, input.asset) + const placeholderPattern = placeholders + .map(escapeRegExp) + .sort((left, right) => right.length - left.length) + .map((placeholder) => + `\\[(?:Table\\s*:\\s*)?${placeholder}\\]|${placeholder}`, + ) + .join("|") + const replaced = content.replace( + new RegExp(placeholderPattern, "gi"), + () => html, + ) + + if (replaced === content) { + throw new Error( + `Table placeholder for ${input.asset.ref} was not found in its chunk content.`, + ) + } + return replaced +} + +function getTablePlaceholderCandidates( + chunk: EvidenceChunk, + asset: EvidenceAsset, +): string[] { + const candidates = [ + asset.sourcePath, + chunk.filePath, + chunk.sourceChunkPath, + asset.assetUrl, + ] + const unique: string[] = [] + for (const candidate of candidates) { + const trimmed = candidate?.trim() + if (!trimmed || unique.includes(trimmed)) continue + unique.push(trimmed) + } + return unique +} + +function escapeRegExp(value: string): string { + return value.replace(/[.*+?^${}()|[\]\\]/g, "\\$&") +} diff --git a/src/agent-harness/index.ts b/src/agent-harness/index.ts index 40e843b..c2b33ee 100644 --- a/src/agent-harness/index.ts +++ b/src/agent-harness/index.ts @@ -1,3 +1,4 @@ +export * from "./answer-context" export * from "./image-highlights" export * from "./ledger" export * from "./knowhere-text" diff --git a/src/agent-harness/knowhere-text.test.ts b/src/agent-harness/knowhere-text.test.ts index af1094e..c07564e 100644 --- a/src/agent-harness/knowhere-text.test.ts +++ b/src/agent-harness/knowhere-text.test.ts @@ -23,10 +23,12 @@ describe("knowhereToolText", () => { expect(text).toContain('ref="r1:result:1"') expect(text).toContain('ref="asset:r1:result:1"') expect(text).toContain("Page one summary.") - expect(text).toContain("Call inspectImage") - expect(text).toContain("before finalize") + expect(text).not.toContain("inspectImage") + expect(text).not.toContain("asset_instruction") expect(text).not.toContain("") expect(text).not.toContain("Page one evidence.") + expect(text).not.toContain("referencedChunkCount") + expect(text).not.toContain("provenance-only-id") expect(text).not.toContain("https://assets.example/page-1.png") }) @@ -81,6 +83,12 @@ function makeSearchResponse(): RetrievalQueryResponse { }, }, ], - referencedChunks: [], + referencedChunks: [ + { + documentId: "doc_1", + chunkId: "provenance-only-id", + pageNums: [], + }, + ] as unknown as RetrievalQueryResponse["referencedChunks"], } } diff --git a/src/agent-harness/knowhere-text.ts b/src/agent-harness/knowhere-text.ts index 5f55ff7..c3a4cc6 100644 --- a/src/agent-harness/knowhere-text.ts +++ b/src/agent-harness/knowhere-text.ts @@ -20,9 +20,6 @@ type ErrorTextInput = { type KnowhereOperation = "search" -const assetInstruction = - "Notebook returned image/page asset refs. Call inspectImage with the asset refs you will cite before finalize so OCR/visual context and provenance boxes exist. Do not expose raw asset URLs." - export const knowhereToolText = { formatSearch(input: SearchTextInput): string { return wrapKnowhereBlock("search", [ @@ -31,7 +28,6 @@ export const knowhereToolText = { namespace: input.response.namespace, query: input.response.query, resultCount: String(input.response.results.length), - referencedChunkCount: String(input.response.referencedChunks.length), stopReason: input.response.stopReason ?? undefined, failureReason: input.response.failureReason ?? undefined, }), @@ -39,7 +35,6 @@ export const knowhereToolText = { // evidenceText — same bodies, no citeable refs, doubles context. formatEvidenceChunks(input.chunks, input.chunkPickStart), formatEvidenceAssets(input.assets), - formatAssetInstruction(input.assets), ]) }, @@ -88,7 +83,6 @@ function formatEvidenceChunks( sectionPath: chunk.source.sectionPath ?? undefined, sourceChunkPath: chunk.sourceChunkPath ?? undefined, filePath: chunk.filePath ?? undefined, - assetRef: chunk.assetRef, }), formatTextTag("content", chunk.content), "", @@ -119,11 +113,6 @@ function formatEvidenceAssets(assets: readonly EvidenceAsset[]): string { ].join("\n") } -function formatAssetInstruction(assets: readonly EvidenceAsset[]): string { - if (!assets.some((asset) => asset.type === "image")) return "" - return formatTextTag("asset_instruction", assetInstruction) -} - function formatTextTag(tagName: string, value: string): string { return [`<${tagName}>`, value, ``].join("\n") } diff --git a/src/agent-harness/ledger.test.ts b/src/agent-harness/ledger.test.ts index 6d3bdb3..5afd958 100644 --- a/src/agent-harness/ledger.test.ts +++ b/src/agent-harness/ledger.test.ts @@ -1,7 +1,8 @@ -import { describe, expect, it } from "vitest" +import { describe, expect, it, vi } from "vitest" import type { RetrievalQueryResponse } from "@ontos-ai/knowhere-sdk" import { createEvidenceLedger } from "./ledger" +import type { ResolveConnectedAssets } from "./types" describe("createEvidenceLedger", () => { it("normalizes retrieval chunks and media assets without treating candidates as final output", () => { @@ -177,17 +178,178 @@ describe("createEvidenceLedger", () => { }, }, ], - // Real agent_explore responses can return referencedChunks entries - // that only carry a summary id, with no chunkType/chunkId/documentId - // even though the SDK type declares those as required strings. + // Real agent_explore responses carry provenance IDs without evidence + // fields even though the SDK type declares chunkType as required. referencedChunks: [ - { summary: "8da0776b-c52b-5602-8579-25c421706f5f" }, + { + documentId: "doc_1", + chunkId: "8da0776b-c52b-5602-8579-25c421706f5f", + pageNums: [], + }, ] as unknown as RetrievalQueryResponse["referencedChunks"], }) expect(snapshot.chunks.map((chunk) => chunk.ref)).toEqual(["r1:result:1"]) expect(snapshot.chunks[0]?.content).toBe("Target BP is <130/80 mmHg.") }) + + it("opens a pending retention range for new chunks and keeps only retained picks", () => { + const ledger = createEvidenceLedger() + const snapshot = ledger.addRetrievalResponse(makeRetrievalResponse()) + + expect(snapshot.pendingRetention).toEqual({ startPick: 1, endPick: 3 }) + expect(snapshot.retainedPicks).toEqual([]) + expect(ledger.hasPendingRetention()).toBe(true) + + const retained = ledger.retainPicks([1, 3]) + expect(retained).toEqual({ ok: true, retainedPicks: [1, 3] }) + expect(ledger.hasPendingRetention()).toBe(false) + expect(ledger.isRetained(1)).toBe(true) + expect(ledger.isRetained(2)).toBe(false) + expect(ledger.isRetained(3)).toBe(true) + expect(ledger.snapshot().pendingRetention).toBeNull() + expect(ledger.snapshot().retainedPicks).toEqual([1, 3]) + }) + + it("does not open pending retention when a search adds no chunks", () => { + const ledger = createEvidenceLedger() + const snapshot = ledger.addRetrievalResponse({ + namespace: "notebook", + query: "empty", + routerUsed: "workflow_single_step", + answerText: null, + evidenceText: "No hits", + stopReason: "completed", + failureReason: null, + results: [], + referencedChunks: [], + }) + + expect(snapshot.chunks).toEqual([]) + expect(snapshot.pendingRetention).toBeNull() + expect(ledger.hasPendingRetention()).toBe(false) + expect(ledger.retainPicks([])).toMatchObject({ ok: false }) + }) + + it("rejects retain picks outside the latest search without changing state", () => { + const ledger = createEvidenceLedger() + ledger.addRetrievalResponse(makeRetrievalResponse()) + + expect(ledger.retainPicks([4])).toMatchObject({ + ok: false, + invalidPicks: [4], + }) + expect(ledger.hasPendingRetention()).toBe(true) + expect(ledger.snapshot().retainedPicks).toEqual([]) + }) + + it("locks unretained first-search picks after a later search is retained", () => { + const ledger = createEvidenceLedger() + ledger.addRetrievalResponse(makeRetrievalResponse()) + expect(ledger.retainPicks([])).toMatchObject({ ok: true, retainedPicks: [] }) + + ledger.addRetrievalResponse({ + ...makeRetrievalResponse(), + results: [ + { + content: "Second retrieval evidence.", + chunkType: "text", + score: 0.8, + source: { + documentId: "doc_2", + sourceFileName: "second.pdf", + sectionPath: "Second", + }, + }, + ], + referencedChunks: [], + }) + expect(ledger.pendingRetentionRange()).toEqual({ startPick: 4, endPick: 4 }) + expect(ledger.retainPicks([4])).toMatchObject({ + ok: true, + retainedPicks: [4], + }) + expect(ledger.isRetained(1)).toBe(false) + expect(ledger.isRetained(4)).toBe(true) + }) + + it("resolves every embedded table and image connected to retained text", async () => { + const ledger = createEvidenceLedger() + ledger.addRetrievalResponse({ + ...makeRetrievalResponse(), + results: [ + { + chunkId: "text_chunk", + content: "Comparison [tables/comparison.html]", + chunkType: "text", + score: 0.9, + metadata: { + connectTo: [ + { + target: "table_chunk", + relation: "embeds", + ref: "[tables/comparison.html]", + }, + { + target: "image_chunk_1", + relation: "embeds", + ref: "[images/comparison-a.jpg]", + }, + { + target: "image_chunk_2", + relation: "embeds", + ref: "[images/comparison-b.jpg]", + }, + ], + }, + source: { + documentId: "doc_1", + sourceFileName: "cardiology.pdf", + sectionPath: "Differential diagnosis", + }, + }, + ], + referencedChunks: [], + }) + ledger.retainPicks([1]) + const resolveConnectedAssets = vi.fn( + async (lookups) => + lookups.map((lookup) => ({ + ...lookup, + assetUrl: `https://assets.example/${lookup.chunkId}`, + })), + ) + + const snapshot = await ledger.resolveRetainedConnectedAssets( + resolveConnectedAssets, + ) + + expect(resolveConnectedAssets).toHaveBeenCalledWith([ + { documentId: "doc_1", chunkId: "table_chunk", type: "table" }, + { documentId: "doc_1", chunkId: "image_chunk_1", type: "image" }, + { documentId: "doc_1", chunkId: "image_chunk_2", type: "image" }, + ]) + expect(snapshot.assets).toEqual([ + expect.objectContaining({ + ref: "asset:r1:result:1:table_chunk", + chunkRef: "r1:result:1", + type: "table", + sourcePath: "tables/comparison.html", + }), + expect.objectContaining({ + ref: "asset:r1:result:1:image_chunk_1", + chunkRef: "r1:result:1", + type: "image", + sourcePath: "images/comparison-a.jpg", + }), + expect.objectContaining({ + ref: "asset:r1:result:1:image_chunk_2", + chunkRef: "r1:result:1", + type: "image", + sourcePath: "images/comparison-b.jpg", + }), + ]) + }) }) function makeRetrievalResponse(): RetrievalQueryResponse { diff --git a/src/agent-harness/ledger.ts b/src/agent-harness/ledger.ts index 9c80448..4cd3123 100644 --- a/src/agent-harness/ledger.ts +++ b/src/agent-harness/ledger.ts @@ -7,7 +7,10 @@ import type { EvidenceAsset, EvidenceChunk, EvidenceLedgerSnapshot, + PendingRetentionRange, + ResolveConnectedAssets, } from "./types" +import { hasReferencedChunkEvidence } from "./referenced-chunks" const contentPreviewLimit = 1_200 const imageExtensions = [".jpg", ".jpeg", ".png", ".gif", ".webp", ".svg"] as const @@ -20,6 +23,8 @@ type MutableLedger = { stopReasons: string[] failureReasons: string[] decisionTraces: unknown[] + retainedPicks: Set + pendingRetention: PendingRetentionRange | null } type EvidenceAssetCandidate = { @@ -29,6 +34,15 @@ type EvidenceAssetCandidate = { readonly label: string } +type ConnectedAssetCandidate = { + readonly chunkRef: string + readonly targetChunkId: string + readonly documentId: string + readonly type: EvidenceAsset["type"] + readonly sourcePath: string + readonly source: EvidenceChunk["source"] +} + type PageCitationAssetCandidate = { readonly pageNum: number readonly artifactRef?: string @@ -47,10 +61,13 @@ export function createEvidenceLedger() { stopReasons: [], failureReasons: [], decisionTraces: [], + retainedPicks: new Set(), + pendingRetention: null, } return { addRetrievalResponse(response: RetrievalQueryResponse): EvidenceLedgerSnapshot { + const chunkCountBefore = ledger.chunks.length ledger.retrievalCount += 1 const retrievalIndex = ledger.retrievalCount @@ -76,15 +93,9 @@ export function createEvidenceLedger() { }) response.referencedChunks.forEach((chunk, index) => { - // Knowhere's agent_explore router returns referencedChunks entries - // that may carry only a summary id with no chunkType/content - // (despite the SDK type declaring chunkType as required). Skip - // entries missing a usable chunkType: they have no real content - // (content is always "" here) and chunkType is required downstream - // (asset-type detection calls chunkType.toLowerCase()). - if (typeof chunk.chunkType !== "string" || chunk.chunkType.trim().length === 0) { - return - } + // ID-only references are provenance, not evidence. Structured page + // and media references remain available to the citation pipeline. + if (!hasReferencedChunkEvidence(chunk)) return const content = "" addChunk({ ledger, @@ -110,6 +121,121 @@ export function createEvidenceLedger() { }) }) + const chunkCountAfter = ledger.chunks.length + if (chunkCountAfter > chunkCountBefore) { + ledger.pendingRetention = { + startPick: chunkCountBefore + 1, + endPick: chunkCountAfter, + } + } + + return snapshot(ledger) + }, + + retainPicks(picks: readonly number[]): + | { readonly ok: true; readonly retainedPicks: readonly number[] } + | { + readonly ok: false + readonly message: string + readonly invalidPicks: readonly number[] + } { + const pending = ledger.pendingRetention + if (!pending) { + return { + ok: false, + message: + "retainEvidence can only be called after a search that returned new evidence.", + invalidPicks: [], + } + } + + const kept: number[] = [] + const invalidPicks: number[] = [] + for (const pick of picks) { + if ( + !Number.isInteger(pick) || + pick < pending.startPick || + pick > pending.endPick + ) { + if (!invalidPicks.includes(pick)) invalidPicks.push(pick) + continue + } + if (!kept.includes(pick)) kept.push(pick) + } + if (invalidPicks.length > 0) { + return { + ok: false, + message: [ + "retainEvidence picks must come from the latest search.", + `Invalid picks: ${invalidPicks.join(" ")}.`, + `Latest search picks: ${pending.startPick}-${pending.endPick}.`, + ].join(" "), + invalidPicks, + } + } + + for (const pick of kept) { + ledger.retainedPicks.add(pick) + } + ledger.pendingRetention = null + return { + ok: true, + retainedPicks: [...ledger.retainedPicks].sort((left, right) => left - right), + } + }, + + isRetained(pick: number): boolean { + return ledger.retainedPicks.has(pick) + }, + + hasPendingRetention(): boolean { + return ledger.pendingRetention !== null + }, + + pendingRetentionRange(): PendingRetentionRange | null { + return ledger.pendingRetention + }, + + async resolveRetainedConnectedAssets( + resolveConnectedAssets?: ResolveConnectedAssets, + ): Promise { + const candidates = getRetainedConnectedAssetCandidates(ledger) + if (candidates.length === 0) return snapshot(ledger) + if (!resolveConnectedAssets) { + throw new Error( + "No connected asset resolver was provided for retained evidence.", + ) + } + + const lookups = uniqueConnectedAssetLookups(candidates) + const resolved = await resolveConnectedAssets(lookups) + const assetUrlByLookup = new Map( + resolved.map((asset) => [connectedAssetLookupKey(asset), asset.assetUrl]), + ) + + for (const candidate of candidates) { + const lookupKey = connectedAssetLookupKey({ + documentId: candidate.documentId, + chunkId: candidate.targetChunkId, + type: candidate.type, + }) + const assetUrl = assetUrlByLookup.get(lookupKey) + if (!assetUrl) { + throw new Error( + `Connected ${candidate.type} chunk ${candidate.targetChunkId} was not resolved.`, + ) + } + ledger.assets.push({ + ref: `asset:${candidate.chunkRef}:${candidate.targetChunkId}`, + chunkRef: candidate.chunkRef, + type: candidate.type, + assetUrl, + sourcePath: candidate.sourcePath, + source: candidate.source, + label: formatConnectedAssetLabel(candidate), + }) + } + return snapshot(ledger) }, @@ -183,7 +309,7 @@ function addChunkFromResult(input: { function addChunk(input: { readonly ledger: MutableLedger - readonly chunk: Omit + readonly chunk: EvidenceChunk }): void { const asset = getEvidenceAssetCandidate(input.chunk) if (!asset) { @@ -192,23 +318,96 @@ function addChunk(input: { } const assetRef = `asset:${input.chunk.ref}` - const chunk: EvidenceChunk = { - ...input.chunk, - assetRef, - } - input.ledger.chunks.push(chunk) + input.ledger.chunks.push(input.chunk) input.ledger.assets.push({ ref: assetRef, - chunkRef: chunk.ref, + chunkRef: input.chunk.ref, type: asset.type, ...(asset.assetUrl ? { assetUrl: asset.assetUrl } : {}), ...(asset.sourcePath ? { sourcePath: asset.sourcePath } : {}), - ...(chunk.revisionKey ? { revisionKey: chunk.revisionKey } : {}), - source: chunk.source, + ...(input.chunk.revisionKey ? { revisionKey: input.chunk.revisionKey } : {}), + source: input.chunk.source, label: asset.label, }) } +function getRetainedConnectedAssetCandidates( + ledger: MutableLedger, +): ConnectedAssetCandidate[] { + return ledger.chunks.flatMap((chunk, index) => { + if (!ledger.retainedPicks.has(index + 1)) return [] + + const documentId = getTrimmedString(chunk.source.documentId) + const connections = chunk.metadata?.connectTo ?? chunk.metadata?.connect_to + if (!Array.isArray(connections)) return [] + + return connections.flatMap((connection): ConnectedAssetCandidate[] => { + if (!isRecord(connection) || connection.relation !== "embeds") return [] + const targetChunkId = getTrimmedString(connection.target) + const sourcePath = getConnectedAssetPath(connection.ref) + if (!targetChunkId || !sourcePath) return [] + if (!documentId) { + throw new Error( + `Retained evidence ${chunk.ref} has connected assets but no document ID.`, + ) + } + + return [{ + chunkRef: chunk.ref, + targetChunkId, + documentId, + type: sourcePath.startsWith("images/") ? "image" : "table", + sourcePath, + source: chunk.source, + }] + }) + }) +} + +function getConnectedAssetPath(value: unknown): string | null { + const ref = getTrimmedString(value) + if (!ref || !ref.startsWith("[") || !ref.endsWith("]")) return null + const path = ref.slice(1, -1).trim() + return path.startsWith("images/") || path.startsWith("tables/") ? path : null +} + +function uniqueConnectedAssetLookups( + candidates: readonly ConnectedAssetCandidate[], +) { + const lookups = new Map< + string, + { documentId: string; chunkId: string; type: EvidenceAsset["type"] } + >() + for (const candidate of candidates) { + const lookup = { + documentId: candidate.documentId, + chunkId: candidate.targetChunkId, + type: candidate.type, + } + lookups.set(connectedAssetLookupKey(lookup), lookup) + } + return [...lookups.values()] +} + +function connectedAssetLookupKey(input: { + readonly documentId: string + readonly chunkId: string + readonly type: EvidenceAsset["type"] +}): string { + return `${input.documentId}\u0000${input.type}\u0000${input.chunkId}` +} + +function formatConnectedAssetLabel(candidate: ConnectedAssetCandidate): string { + return [ + candidate.source.sourceFileName, + candidate.source.sectionPath, + candidate.sourcePath, + candidate.type, + ] + .filter((part): part is string => Boolean(part)) + .join(" / ") +} + function buildContentPreview(content: string): string { const normalized = content.replace(/\s+/g, " ").trim() if (normalized.length <= contentPreviewLimit) return normalized @@ -233,7 +432,7 @@ function isRenderableAsset(chunkType: string, assetUrl: string): boolean { } function getEvidenceAssetCandidate( - chunk: Omit, + chunk: EvidenceChunk, ): EvidenceAssetCandidate | null { const pageAsset = getPageCitationAssetCandidate(chunk) if (pageAsset) { @@ -253,7 +452,7 @@ function getEvidenceAssetCandidate( } function getPageCitationAssetCandidate( - chunk: Omit, + chunk: EvidenceChunk, ): EvidenceAssetCandidate | null { if (normalizeChunkType(chunk.chunkType) !== "page") return null @@ -296,7 +495,7 @@ function isImageAssetUrl(assetUrl: string): boolean { } function getAssetSourcePath( - chunk: Omit, + chunk: EvidenceChunk, assetUrl: string, ): string | null { const candidates = [ @@ -485,6 +684,8 @@ function snapshot(ledger: MutableLedger): EvidenceLedgerSnapshot { stopReasons: [...ledger.stopReasons], failureReasons: [...ledger.failureReasons], decisionTraces: [...ledger.decisionTraces], + retainedPicks: [...ledger.retainedPicks].sort((left, right) => left - right), + pendingRetention: ledger.pendingRetention, } } diff --git a/src/agent-harness/referenced-chunks.ts b/src/agent-harness/referenced-chunks.ts new file mode 100644 index 0000000..ccde61b --- /dev/null +++ b/src/agent-harness/referenced-chunks.ts @@ -0,0 +1,7 @@ +import type { RetrievalQueryResponse } from "@ontos-ai/knowhere-sdk" + +type ReferencedChunk = RetrievalQueryResponse["referencedChunks"][number] + +export function hasReferencedChunkEvidence(chunk: ReferencedChunk): boolean { + return typeof chunk.chunkType === "string" && chunk.chunkType.trim().length > 0 +} diff --git a/src/agent-harness/runtime.test.ts b/src/agent-harness/runtime.test.ts index d36ca1b..bc06251 100644 --- a/src/agent-harness/runtime.test.ts +++ b/src/agent-harness/runtime.test.ts @@ -1,11 +1,13 @@ import { describe, expect, it, vi } from "vitest" import type { RetrievalQueryResponse } from "@ontos-ai/knowhere-sdk" +import { MockLanguageModelV3 } from "ai/test" import { buildHarnessMessages, buildHarnessSystemPrompt, createHarnessTools, prepareHarnessStep, + runAgentHarness, sanitizeHarnessModelMessagesForStep, } from "./runtime" import { createEvidenceLedger } from "./ledger" @@ -13,25 +15,161 @@ import type { AgentTurnInput, ContextPolicy, HarnessToolCallTrace, - ImageInspectionRequest, IntentFrame, KnowhereToolRuntime, + MemorySearchItem, MemoryToolRuntime, OutputManifest, + ResolveConnectedAssets, } from "./types" describe("agent harness runtime", () => { - it("tells the agent to search fluid memory first and not treat every question as document retrieval", () => { + it("runs retention, resolves connected assets, assembles once, and finalizes", async () => { + const modelResults = [ + toolCallResult([ + { + toolCallId: "intent", + toolName: "declareIntent", + input: { + task: "compare", + dependsOnPreviousTurn: false, + retrievalNeeded: "yes", + targetModalities: ["text", "image", "table"], + constraints: {}, + groundingPolicy: "must_use_sources", + }, + }, + { + toolCallId: "context", + toolName: "setContextPolicy", + input: { + carryHistory: "none", + reason: "Self-contained request.", + activePriorTurnIds: [], + }, + }, + ]), + toolCallResult([ + { + toolCallId: "search", + toolName: "knowhere_search", + input: { query: "comparison", targetContent: "all" }, + }, + ]), + toolCallResult([ + { + toolCallId: "retain", + toolName: "retainEvidence", + input: { picks: [1] }, + }, + ]), + toolCallResult([ + { + toolCallId: "prepare", + toolName: "prepareAnswer", + input: {}, + }, + ]), + toolCallResult([ + { + toolCallId: "finalize", + toolName: "finalize", + input: { + text: "Comparison [[cite:1]].", + citations: [{ pick: 1 }], + memoryCitations: [], + artifacts: [ + { + type: "table", + ref: "asset:r1:result:1:table_chunk", + display: true, + reason: "Comparison table", + }, + { + type: "image", + ref: "asset:r1:result:1:image_chunk", + display: true, + reason: "Comparison image", + }, + ], + unresolved: [], + }, + }, + ]), + ] + let modelResultIndex = 0 + const model = new MockLanguageModelV3({ + supportedUrls: { "image/*": [/^https:\/\//] }, + doGenerate: async () => { + const result = modelResults[modelResultIndex] + modelResultIndex += 1 + if (!result) throw new Error("No mock model result remains.") + return result + }, + }) + const resolveConnectedAssets = vi.fn( + async (lookups) => + lookups.map((lookup) => ({ + ...lookup, + assetUrl: `https://assets.example/${lookup.chunkId}`, + })), + ) + const readTableHtml = vi + .fn() + .mockResolvedValue("
comparison
") + + const result = await runAgentHarness({ + model, + turn: makeTurnInput(), + knowhereTools: makeKnowhereTools( + vi.fn().mockResolvedValue(makeConnectedRetrievalResponse()), + ), + memoryTools: makeMemoryTools(), + resolveConnectedAssets, + readTableHtml, + }) + + expect(model.doGenerateCalls).toHaveLength(5) + expect(resolveConnectedAssets).toHaveBeenCalledWith([ + { documentId: "doc_1", chunkId: "table_chunk", type: "table" }, + { documentId: "doc_1", chunkId: "image_chunk", type: "image" }, + ]) + expect(readTableHtml).toHaveBeenCalledWith( + "https://assets.example/table_chunk", + ) + expect(JSON.stringify(model.doGenerateCalls[4]?.prompt)).toContain( + "
comparison
", + ) + expect(result.manifest.citations).toEqual([{ ref: "r1:result:1" }]) + expect(result.trace.ledger.assets).toHaveLength(2) + }) + + it("tells the agent to search memory and documents in parallel when both apply", () => { const prompt = buildHarnessSystemPrompt(makeTurnInput()) - expect(prompt).toContain("Call memory_search first") + expect(prompt).toContain("call them together in the same step") expect(prompt).toContain( - "Call knowhere_search only when memory is insufficient and groundingPolicy requires citing source documents", + "Call knowhere_search when groundingPolicy requires citing source documents", ) expect(prompt).toContain( "call knowhere_search again with a refined query", ) expect(prompt).toContain("Refine knowhere_search at most twice") + expect(prompt).toContain( + "After the second refined search, call prepareAnswer regardless of its result", + ) + expect(prompt).toContain( + "After each search that returns new evidence, call retainEvidence", + ) + expect(prompt).toContain( + "When calling knowhere_search after a previous search in this turn, set gapReason", + ) + expect(prompt).toContain( + "asset paths in retrieved chunks represent connected assets", + ) + expect(prompt).toContain( + "Do not refine the search solely because those assets have not been expanded yet", + ) expect(prompt).not.toContain("knowhere_list_documents") expect(prompt).not.toContain("knowhere_get_document_outline") expect(prompt).not.toContain("knowhere_read_chunks") @@ -70,15 +208,16 @@ describe("agent harness runtime", () => { expect(prompt).toContain("Do not collapse same-page citations") }) - it("requires inspectImage on cited page/image assets before finalize", () => { + it("tells the agent to read retained images in the assembled answer message", () => { const prompt = buildHarnessSystemPrompt(makeTurnInput()) expect(prompt).toContain( - "call inspectImage on the page/image assets you will cite before finalize", + "memory_search and knowhere_search are parallel retrieval sources", ) expect(prompt).toContain( - "Do not finalize cited page/image assets from chunk text alone", + "Images in retained evidence are embedded directly", ) + expect(prompt).not.toContain("inspectImage") }) it("passes only outer retrieval parameters to KNOWHERE without planning-tool gating", async () => { @@ -187,7 +326,11 @@ describe("agent harness runtime", () => { }) const firstResult = await executeTool(tools.knowhere_search, { query: "first" }) - const secondResult = await executeTool(tools.knowhere_search, { query: "second" }) + await retainLatestSearch(tools, ledger) + const secondResult = await executeTool(tools.knowhere_search, { + query: "second", + gapReason: "The first search lacked the second-source wording this query targets.", + }) expect(firstResult).toContain('pick="1"') expect(firstResult).toContain('ref="r1:result:1"') @@ -203,734 +346,260 @@ describe("agent harness runtime", () => { ]) }) - it("rejects image inspection before retrieval has returned image assets", async () => { - const inspectImages = vi.fn() + it("writes citation refs from ledger picks and rejects picks outside the ledger", async () => { + const ledger = createEvidenceLedger() + ledger.addRetrievalResponse(makeRetrievalResponse()) + const state: { + finalized?: boolean + finalizedManifest?: OutputManifest + } = {} const tools = createHarnessTools({ - state: {}, - ledger: createEvidenceLedger(), + state, + ledger, memoryTools: makeMemoryTools(), knowhereTools: makeKnowhereTools(), - inspectImages, recentTurns: [], }) + await retainLatestSearch(tools, ledger) - const result = await executeTool(tools.inspectImage, { - refs: ["asset:r1:result:1"], - question: "What text is visible?", + const rejected = await executeTool(tools.finalize, { + text: "Target is <130/80 mmHg [[cite:1]].", + citations: [{ pick: 99 }], + memoryCitations: [], + artifacts: [], + unresolved: [], }) - expect(result).toEqual({ + expect(rejected).toMatchObject({ ok: false, - message: - "Knowhere evidence tools must return image assets before inspectImage.", - inspected: [], - skipped: [ - { - ref: "asset:r1:result:1", - reason: "No Knowhere evidence is available yet.", - }, - ], + unknownPicks: [99], + }) + expect(String((rejected as { message: string }).message)).toContain( + "Available picks: 1-1", + ) + expect(state.finalized).not.toBe(true) + + const accepted = await executeTool(tools.finalize, { + text: "Target is <130/80 mmHg [[cite:1]].", + citations: [{ pick: 1 }], + memoryCitations: [], + artifacts: [], + unresolved: [], + }) + expect(accepted).toMatchObject({ + ok: true, + citations: [{ ref: "r1:result:1" }], }) - expect(inspectImages).not.toHaveBeenCalled() + expect(state.finalized).toBe(true) + expect(state.finalizedManifest?.citations).toEqual([{ ref: "r1:result:1" }]) }) - it("rejects image inspection refs that are unknown or not image assets", async () => { + it("maps finalize picks across successive searches", async () => { const ledger = createEvidenceLedger() - ledger.addRetrievalResponse(makeTableRetrievalResponse()) - const inspectImages = vi.fn() + ledger.addRetrievalResponse(makeRetrievalResponse()) + const state: { + finalizedManifest?: OutputManifest + } = {} const tools = createHarnessTools({ - state: {}, + state, ledger, memoryTools: makeMemoryTools(), knowhereTools: makeKnowhereTools(), - inspectImages, recentTurns: [], }) - - const result = await executeTool(tools.inspectImage, { - refs: ["asset:r1:result:1", "missing"], - question: "What is in these images?", - }) - - expect(result).toEqual({ - ok: false, - message: "No inspectable image asset refs were provided.", - inspected: [], - skipped: [ - { - ref: "asset:r1:result:1", - reason: "Ref is not an image asset.", - }, + await retainLatestSearch(tools, ledger) + ledger.addRetrievalResponse({ + ...makeRetrievalResponse(), + query: "second query", + results: [ { - ref: "missing", - reason: "Ref was not returned by Knowhere as an asset.", + content: "Second retrieval evidence.", + chunkType: "text", + score: 0.8, + source: { + documentId: "doc_2", + sourceFileName: "second.pdf", + sectionPath: "Second", + }, }, ], }) - expect(inspectImages).not.toHaveBeenCalled() + await retainLatestSearch(tools, ledger) + + const accepted = await executeTool(tools.finalize, { + text: "Second source [[cite:1]].", + citations: [{ pick: 2 }], + memoryCitations: [], + artifacts: [], + unresolved: [], + }) + + expect(accepted).toMatchObject({ + ok: true, + citations: [{ ref: "r2:result:1" }], + }) + expect(state.finalizedManifest?.citations).toEqual([{ ref: "r2:result:1" }]) }) - it("lets retrieval bound the number of images inspected in one call", async () => { - const ledger = createEvidenceLedger() - ledger.addRetrievalResponse(makeImageRetrievalResponse(7)) - const inspectImages = vi.fn(async (request: { - readonly assets: readonly { readonly ref: string; readonly label: string }[] - }) => ({ - analysis: "Compared all retrieved images.", - inspected: request.assets.map(({ ref, label }) => ({ ref, label })), - skipped: [], - })) + it("accepts free-form text with no citations, displayed artifacts, or unresolved gaps", async () => { + const state: { + finalizedManifest?: OutputManifest + finalized?: boolean + } = {} const tools = createHarnessTools({ - state: {}, - ledger, + state, + ledger: createEvidenceLedger(), memoryTools: makeMemoryTools(), knowhereTools: makeKnowhereTools(), - inspectImages, recentTurns: [], }) - const result = await executeTool(tools.inspectImage, { - refs: Array.from({ length: 7 }, (_, index) => `asset:r1:result:${index + 1}`), - question: "Compare these images.", + const result = await executeTool(tools.finalize, { + text: "Hi there! How can I help?", + citations: [], + memoryCitations: [], + artifacts: [], + unresolved: [], }) expect(result).toMatchObject({ ok: true }) - expect(inspectImages.mock.calls[0]?.[0].assets).toHaveLength(7) + expect(state.finalizedManifest?.text).toBe("Hi there! How can I help?") + expect(state.finalized).toBe(true) }) - it("calls the visual inspection capability with retrieved image ledger assets", async () => { - const ledger = createEvidenceLedger() - ledger.addRetrievalResponse(makeRetrievalResponse()) - const inspectImages = vi.fn().mockResolvedValue({ - analysis: "The image shows a Q4 revenue chart with a rising line.", - inspected: [ - { - ref: "asset:r1:result:1", - label: "report.pdf / images/chart.png / image", - }, - ], - skipped: [], - }) + it("accepts explicitly unresolved output without planning-tool gating", async () => { + const state: { + finalizedManifest?: OutputManifest + finalized?: boolean + memoryItems?: MemorySearchItem[] + } = {} const tools = createHarnessTools({ - state: {}, - ledger, + state, + ledger: createEvidenceLedger(), memoryTools: makeMemoryTools(), knowhereTools: makeKnowhereTools(), - inspectImages, recentTurns: [], }) - const result = await executeTool(tools.inspectImage, { - refs: ["asset:r1:result:1"], - question: "What does the chart show?", - }) + const manifest = { + text: "Answer.", + citations: [], + memoryCitations: [], + artifacts: [], + unresolved: ["No source evidence is available."], + } - expect(inspectImages).toHaveBeenCalledWith({ - question: "What does the chart show?", - assets: [ - { - ref: "asset:r1:result:1", - label: "report.pdf / images/chart.png / image", - assetUrl: "https://assets.example/chart.png", - sourcePath: "images/chart.png", - source: { - documentId: "doc_1", - sourceFileName: "report.pdf", - sectionPath: "images/chart.png", - }, - }, - ], - }) - expect(result).toEqual({ + expect(await executeTool(tools.finalize, manifest)).toMatchObject({ ok: true, - analysis: "The image shows a Q4 revenue chart with a rising line.", - inspected: [ - { - ref: "asset:r1:result:1", - label: "report.pdf / images/chart.png / image", - }, - ], - skipped: [], + text: "Answer.", }) + expect(state.finalizedManifest).toEqual(manifest) + expect(state.finalized).toBe(true) }) - it("calls the visual inspection capability with retrieved page citation assets", async () => { - const ledger = createEvidenceLedger() - ledger.addRetrievalResponse(makePageCitationRetrievalResponse()) - const inspectImages = vi.fn().mockResolvedValue({ - analysis: "The clause says the contractor pays 5000 yuan per occurrence.", - inspected: [ + it("returns memory refs from memory_search and stores memoryCitations on finalize", async () => { + const search = vi.fn().mockResolvedValue({ + query: "毛利率", + items: [ { - ref: "asset:r1:referenced:1", - label: - "Root / (6)现场工期进度管理方面的违约责任 / page_citation_assets/page-8.png / page", + ref: "mem:1", + itemId: "item_1", + kind: "stance", + abstractL0: "关注毛利率下滑", + overviewL1: "用户把毛利率当作核心观察指标。", }, ], - skipped: [], }) + const state: { + finalizedManifest?: OutputManifest + finalized?: boolean + memoryItems?: MemorySearchItem[] + } = {} const tools = createHarnessTools({ - state: {}, - ledger, - memoryTools: makeMemoryTools(), + state, + ledger: createEvidenceLedger(), + memoryTools: makeMemoryTools(search), knowhereTools: makeKnowhereTools(), - inspectImages, recentTurns: [], }) - const result = await executeTool(tools.inspectImage, { - refs: ["asset:r1:referenced:1"], - question: "What liquidated damages amount is visible in this clause?", + const searchText = await executeTool(tools.memory_search, { + query: "毛利率", + }) + const repeatedSearchText = await executeTool(tools.memory_search, { + query: "盈利能力", + }) + expect(searchText).toContain('') + expect(searchText).toContain('ref="mem:1"') + expect(searchText).toContain('itemId="item_1"') + expect(repeatedSearchText).toContain('status="error"') + expect(search).toHaveBeenCalledTimes(1) + expect(search).toHaveBeenNthCalledWith(1, { + query: "毛利率", + kinds: undefined, }) + expect(state.memoryItems).toEqual([ + { + ref: "mem:1", + itemId: "item_1", + kind: "stance", + abstractL0: "关注毛利率下滑", + overviewL1: "用户把毛利率当作核心观察指标。", + }, + ]) - expect(inspectImages).toHaveBeenCalledWith({ - question: "What liquidated damages amount is visible in this clause?", - assets: [ - { - ref: "asset:r1:referenced:1", - label: - "Root / (6)现场工期进度管理方面的违约责任 / page_citation_assets/page-8.png / page", - assetUrl: "https://assets.example/page-8.png", - sourcePath: "page_citation_assets/page-8.png", - revisionKey: "job_contract", - source: { - documentId: "doc_contract", - sourceFileName: null, - sectionPath: "Root / (6)现场工期进度管理方面的违约责任", - }, - }, + const manifest = { + text: "按已有记忆,毛利率是核心观察指标。", + citations: [], + memoryCitations: [ + { ref: "mem:1", itemId: "item_1", kind: "stance" as const }, ], + artifacts: [], + unresolved: [], + } + const invalidManifest = { + ...manifest, + memoryCitations: [ + { ref: "mem:1", itemId: "invented_item", kind: "stance" as const }, + ], + } + expect(await executeTool(tools.finalize, invalidManifest)).toMatchObject({ + ok: false, + invalidMemoryCitations: invalidManifest.memoryCitations, }) - expect(result).toMatchObject({ + expect(state.finalizedManifest).toBeUndefined() + + expect(await executeTool(tools.finalize, manifest)).toMatchObject({ ok: true, - analysis: "The clause says the contractor pays 5000 yuan per occurrence.", + memoryCitations: manifest.memoryCitations, }) + expect(state.finalizedManifest).toEqual(manifest) }) - it("inspects duplicate retrieved pages only once across namespaces", async () => { - const ledger = createEvidenceLedger() - const pageMetadata = { - pageNums: [4], - pageAssets: [ + it("exposes full prior-turn content through policy-approved readPriorTurn", async () => { + const state: { + contextPolicy?: ContextPolicy + priorTurnReads?: string[] + } = { + contextPolicy: { + carryHistory: "repair_previous", + reason: "The current request corrects the previous answer.", + activePriorTurnIds: ["turn_1"], + }, + priorTurnReads: [], + } + const tools = createHarnessTools({ + state, + ledger: createEvidenceLedger(), + memoryTools: makeMemoryTools(), + knowhereTools: makeKnowhereTools(), + recentTurns: [ { - pageNum: 4, - artifactRef: "page_citation_assets/page-4.png", - assetUrl: "https://assets.example/page-4.png", - contentType: "image/png", - }, - ], - } - ledger.addRetrievalResponse({ - namespace: "default,notebook", - query: "revenue", - routerUsed: "mapnav", - answerText: null, - evidenceText: "Revenue evidence", - stopReason: "completed", - failureReason: null, - results: [ - { - chunkId: "chunk_page_4", - content: "Revenue was $24.9B.", - chunkType: "page", - score: 0.9, - metadata: pageMetadata, - source: { - documentId: "doc_catalog", - sourceFileName: "original.pdf", - sectionPath: "FINANCIAL SUMMARY", - }, - }, - { - chunkId: "chunk_page_4", - content: "Revenue was $24.9B.", - chunkType: "page", - score: 0.9, - metadata: pageMetadata, - source: { - documentId: "doc_workspace", - sourceFileName: "TSLA-Q4-2025-Update.pdf", - sectionPath: "FINANCIAL SUMMARY", - }, - }, - ], - referencedChunks: [], - }) - const inspectImages = vi.fn().mockResolvedValue({ - analysis: "Revenue is visible in the financial summary.", - inspected: [ - { - ref: "asset:r1:result:1", - label: "original.pdf / page_citation_assets/page-4.png / page", - }, - ], - skipped: [], - highlights: [ - { - ref: "asset:r1:result:1", - regions: [{ x: 0.1, y: 0.2, w: 0.8, h: 0.1 }], - }, - ], - }) - const state: { - inspectedImageRefs?: string[] - } = {} - const tools = createHarnessTools({ - state, - ledger, - memoryTools: makeMemoryTools(), - knowhereTools: makeKnowhereTools(), - inspectImages, - recentTurns: [], - }) - - const inspection = await executeTool(tools.inspectImage, { - refs: ["asset:r1:result:1", "asset:r1:result:2"], - question: "Locate the cited revenue.", - }) - - expect(inspection).toMatchObject({ ok: true }) - expect(inspectImages.mock.calls[0]?.[0].assets).toHaveLength(1) - expect( - await executeTool(tools.finalize, { - text: "Revenue was $24.9B [[cite:1]] [[cite:2]].", - citations: [{ pick: 1 }, { pick: 2 }], - memoryCitations: [], - artifacts: [], - unresolved: [], - }), - ).toMatchObject({ ok: true }) - }) - - it("does not deduplicate unrelated documents with generic page paths", async () => { - const ledger = createEvidenceLedger() - ledger.addRetrievalResponse({ - namespace: "notebook", - query: "revenue", - routerUsed: "mapnav", - answerText: null, - evidenceText: "Revenue evidence", - stopReason: "completed", - failureReason: null, - results: [ - { - chunkId: "chunk_page_4", - content: "Revenue evidence from document A.", - chunkType: "page", - score: 0.9, - metadata: { - pageNums: [4], - pageAssets: [ - { - pageNum: 4, - artifactRef: "page_citation_assets/page-4.png", - assetUrl: "https://assets.example/doc-a/page-4.png", - contentType: "image/png", - }, - ], - }, - source: { - documentId: "doc_a", - sourceFileName: "a.pdf", - sectionPath: "Page 4", - }, - }, - { - chunkId: "chunk_page_4", - content: "Revenue evidence from document B.", - chunkType: "page", - score: 0.8, - metadata: { - pageNums: [4], - pageAssets: [ - { - pageNum: 4, - artifactRef: "page_citation_assets/page-4.png", - assetUrl: "https://assets.example/doc-b/page-4.png", - contentType: "image/png", - }, - ], - }, - source: { - documentId: "doc_b", - sourceFileName: "b.pdf", - sectionPath: "Page 4", - }, - }, - ], - referencedChunks: [], - }) - const inspectImages = vi.fn(async (request: ImageInspectionRequest) => ({ - analysis: "Inspected both pages.", - inspected: request.assets.map(({ ref, label }) => ({ ref, label })), - skipped: [], - })) - const tools = createHarnessTools({ - state: {}, - ledger, - memoryTools: makeMemoryTools(), - knowhereTools: makeKnowhereTools(), - inspectImages, - recentTurns: [], - }) - - expect( - await executeTool(tools.inspectImage, { - refs: ["asset:r1:result:1", "asset:r1:result:2"], - question: "Compare the cited revenue.", - }), - ).toMatchObject({ ok: true }) - expect(inspectImages.mock.calls[0]?.[0].assets).toHaveLength(2) - }) - - it("does not treat skipped image assets as successfully inspected", async () => { - const ledger = createEvidenceLedger() - ledger.addRetrievalResponse(makePageCitationRetrievalResponse()) - const state: { - inspectedImageRefs?: string[] - finalized?: boolean - } = {} - const tools = createHarnessTools({ - state, - ledger, - memoryTools: makeMemoryTools(), - knowhereTools: makeKnowhereTools(), - inspectImages: vi.fn().mockResolvedValue({ - analysis: "", - inspected: [], - skipped: [ - { - ref: "asset:r1:referenced:1", - reason: "The image asset was unavailable in Notebook storage.", - }, - ], - }), - recentTurns: [], - }) - - const inspection = await executeTool(tools.inspectImage, { - refs: ["asset:r1:referenced:1"], - question: "Locate the cited amount.", - }) - - expect(inspection).toMatchObject({ - ok: false, - inspected: [], - }) - expect(state.inspectedImageRefs).toEqual([]) - - const finalize = await executeTool(tools.finalize, { - text: "The amount is 5000 yuan [[cite:1]].", - citations: [{ pick: 1 }], - memoryCitations: [], - artifacts: [], - unresolved: [], - }) - expect(finalize).toMatchObject({ - ok: false, - inspectRefs: ["asset:r1:referenced:1"], - }) - expect(state.finalized).not.toBe(true) - }) - - it("writes citation refs from ledger picks and rejects picks outside the ledger", async () => { - const ledger = createEvidenceLedger() - ledger.addRetrievalResponse(makeRetrievalResponse()) - const state: { - finalized?: boolean - finalizedManifest?: OutputManifest - } = {} - const tools = createHarnessTools({ - state, - ledger, - memoryTools: makeMemoryTools(), - knowhereTools: makeKnowhereTools(), - recentTurns: [], - }) - - const rejected = await executeTool(tools.finalize, { - text: "Target is <130/80 mmHg [[cite:1]].", - citations: [{ pick: 99 }], - memoryCitations: [], - artifacts: [], - unresolved: [], - }) - - expect(rejected).toMatchObject({ - ok: false, - unknownPicks: [99], - }) - expect(String((rejected as { message: string }).message)).toContain( - "Available picks: 1-1", - ) - expect(state.finalized).not.toBe(true) - - const accepted = await executeTool(tools.finalize, { - text: "Target is <130/80 mmHg [[cite:1]].", - citations: [{ pick: 1 }], - memoryCitations: [], - artifacts: [], - unresolved: [], - }) - expect(accepted).toMatchObject({ - ok: true, - citations: [{ ref: "r1:result:1" }], - }) - expect(state.finalized).toBe(true) - expect(state.finalizedManifest?.citations).toEqual([{ ref: "r1:result:1" }]) - }) - - it("maps finalize picks across successive searches", async () => { - const ledger = createEvidenceLedger() - ledger.addRetrievalResponse(makeRetrievalResponse()) - ledger.addRetrievalResponse({ - ...makeRetrievalResponse(), - query: "second query", - results: [ - { - content: "Second retrieval evidence.", - chunkType: "text", - score: 0.8, - source: { - documentId: "doc_2", - sourceFileName: "second.pdf", - sectionPath: "Second", - }, - }, - ], - }) - const state: { - finalizedManifest?: OutputManifest - } = {} - const tools = createHarnessTools({ - state, - ledger, - memoryTools: makeMemoryTools(), - knowhereTools: makeKnowhereTools(), - recentTurns: [], - }) - - const accepted = await executeTool(tools.finalize, { - text: "Second source [[cite:1]].", - citations: [{ pick: 2 }], - memoryCitations: [], - artifacts: [], - unresolved: [], - }) - - expect(accepted).toMatchObject({ - ok: true, - citations: [{ ref: "r2:result:1" }], - }) - expect(state.finalizedManifest?.citations).toEqual([{ ref: "r2:result:1" }]) - }) - - it("accepts finalize output without planning-tool gating", async () => { - const state: { - finalizedManifest?: OutputManifest - finalized?: boolean - } = {} - const tools = createHarnessTools({ - state, - ledger: createEvidenceLedger(), - memoryTools: makeMemoryTools(), - knowhereTools: makeKnowhereTools(), - recentTurns: [], - }) - - const manifest = { - text: "Answer.", - citations: [], - memoryCitations: [], - artifacts: [], - unresolved: [], - } - - expect(await executeTool(tools.finalize, manifest)).toMatchObject({ - ok: true, - text: "Answer.", - }) - expect(state.finalizedManifest).toEqual(manifest) - expect(state.finalized).toBe(true) - }) - - it("returns memory refs from memory_search and stores memoryCitations on finalize", async () => { - const search = vi.fn().mockResolvedValue({ - query: "毛利率", - items: [ - { - ref: "mem:1", - itemId: "item_1", - kind: "stance", - abstractL0: "关注毛利率下滑", - overviewL1: "用户把毛利率当作核心观察指标。", - }, - ], - }) - const state: { - finalizedManifest?: OutputManifest - finalized?: boolean - } = {} - const tools = createHarnessTools({ - state, - ledger: createEvidenceLedger(), - memoryTools: makeMemoryTools(search), - knowhereTools: makeKnowhereTools(), - recentTurns: [], - }) - - const searchText = await executeTool(tools.memory_search, { - query: "毛利率", - }) - expect(searchText).toContain('') - expect(searchText).toContain('ref="mem:1"') - expect(searchText).toContain('itemId="item_1"') - expect(search).toHaveBeenCalledWith({ - query: "毛利率", - kinds: undefined, - }) - - const manifest = { - text: "按已有记忆,毛利率是核心观察指标。", - citations: [], - memoryCitations: [ - { ref: "mem:1", itemId: "item_1", kind: "stance" as const }, - ], - artifacts: [], - unresolved: [], - } - expect(await executeTool(tools.finalize, manifest)).toMatchObject({ - ok: true, - memoryCitations: manifest.memoryCitations, - }) - expect(state.finalizedManifest).toEqual(manifest) - }) - - it("rejects finalize of cited page images until inspectImage has run", async () => { - const ledger = createEvidenceLedger() - ledger.addRetrievalResponse(makePageCitationRetrievalResponse()) - const state: { - finalizedManifest?: OutputManifest - finalized?: boolean - inspectedImageRefs?: string[] - } = {} - const tools = createHarnessTools({ - state, - ledger, - memoryTools: makeMemoryTools(), - knowhereTools: makeKnowhereTools(), - inspectImages: vi.fn().mockResolvedValue({ - analysis: "The clause shows 5000 yuan per occurrence.", - inspected: [{ ref: "asset:r1:referenced:1", label: "page 8" }], - skipped: [], - highlights: [ - { - ref: "asset:r1:referenced:1", - regions: [{ x: 0.1, y: 0.2, w: 0.4, h: 0.15 }], - }, - ], - }), - recentTurns: [], - }) - const manifest = { - text: "The contractor pays 5000 yuan per occurrence [[cite:1]].", - citations: [{ pick: 1 }], - memoryCitations: [], - artifacts: [], - unresolved: [], - } - - const blocked = await executeTool(tools.finalize, manifest) - - expect(blocked).toMatchObject({ - ok: false, - inspectRefs: ["asset:r1:referenced:1"], - }) - expect(blocked).toEqual( - expect.objectContaining({ - message: expect.stringContaining("inspectImage"), - }), - ) - expect(state.finalized).not.toBe(true) - - await executeTool(tools.inspectImage, { - refs: ["asset:r1:referenced:1"], - question: "Locate the cited liquidated-damages amount on this page.", - }) - - expect(await executeTool(tools.finalize, manifest)).toMatchObject({ - ok: true, - text: manifest.text, - }) - expect(state.finalized).toBe(true) - }) - - it("requires inspection when the cited result shares a page asset with a referenced chunk", async () => { - const response = makePageCitationRetrievalResponse() - const ledger = createEvidenceLedger() - ledger.addRetrievalResponse({ - ...response, - results: [ - { - chunkId: "chunk_page_8", - content: "The contractor pays 5000 yuan per occurrence.", - chunkType: "page", - score: 0.9, - metadata: { pageNums: [8] }, - source: { - documentId: "doc_contract", - sourceFileName: "contract.pdf", - sectionPath: "Root / Liquidated damages", - }, - }, - ], - }) - const tools = createHarnessTools({ - state: {}, - ledger, - memoryTools: makeMemoryTools(), - knowhereTools: makeKnowhereTools(), - inspectImages: vi.fn(), - recentTurns: [], - }) - - const result = await executeTool(tools.finalize, { - text: "The contractor pays 5000 yuan [[cite:1]].", - citations: [{ pick: 1 }], - memoryCitations: [], - artifacts: [], - unresolved: [], - }) - - expect(result).toMatchObject({ - ok: false, - inspectRefs: ["asset:r1:referenced:1"], - }) - }) - - it("exposes full prior-turn content through policy-approved readPriorTurn", async () => { - const state: { - contextPolicy?: ContextPolicy - priorTurnReads?: string[] - } = { - contextPolicy: { - carryHistory: "repair_previous", - reason: "The current request corrects the previous answer.", - activePriorTurnIds: ["turn_1"], - }, - priorTurnReads: [], - } - const tools = createHarnessTools({ - state, - ledger: createEvidenceLedger(), - memoryTools: makeMemoryTools(), - knowhereTools: makeKnowhereTools(), - recentTurns: [ - { - id: "turn_1", - role: "assistant", - contentPreview: "Truncated preview...", - content: "The full earlier answer about the tax filing deadline.", - citationLabels: ["tax.pdf / deadline"], + id: "turn_1", + role: "assistant", + contentPreview: "Truncated preview...", + content: "The full earlier answer about the tax filing deadline.", + citationLabels: ["tax.pdf / deadline"], }, ], }) @@ -1094,17 +763,37 @@ describe("agent harness runtime", () => { expect(result.activeTools).toEqual([ "declareIntent", "setContextPolicy", - "inspectImage", "readPriorTurn", ]) expect(result.activeTools).not.toContain("finalize") - expect(result.activeTools).not.toContain("memory_search") + expect(result.activeTools).not.toContain("memory_search") + expect(result.activeTools).not.toContain("knowhere_search") + }) + + it("opens only memory_search after intent says retrieval may be needed", () => { + const result = prepareHarnessStep({ + stepNumber: 3, + intent: { + task: "answer", + dependsOnPreviousTurn: false, + retrievalNeeded: "maybe", + targetModalities: ["text"], + constraints: {}, + groundingPolicy: "can_use_context", + }, + messages: [], + }) + + expect(result.activeTools).toContain("memory_search") + expect(result.activeTools).toContain("prepareAnswer") + expect(result.activeTools).not.toContain("finalize") expect(result.activeTools).not.toContain("knowhere_search") }) - it("opens only memory_search after intent says retrieval may be needed", () => { + it("does not expose memory_search again after its one call", () => { const result = prepareHarnessStep({ - stepNumber: 3, + stepNumber: 4, + hasMemorySearch: true, intent: { task: "answer", dependsOnPreviousTurn: false, @@ -1116,9 +805,7 @@ describe("agent harness runtime", () => { messages: [], }) - expect(result.activeTools).toContain("memory_search") - expect(result.activeTools).toContain("finalize") - expect(result.activeTools).not.toContain("knowhere_search") + expect(result.activeTools).not.toContain("memory_search") }) it("keeps Knowhere tools closed for no_retrieval", () => { @@ -1135,7 +822,8 @@ describe("agent harness runtime", () => { messages: [], }) - expect(result.activeTools).toContain("finalize") + expect(result.activeTools).toContain("prepareAnswer") + expect(result.activeTools).not.toContain("finalize") expect(result.activeTools).not.toContain("memory_search") expect(result.activeTools).not.toContain("knowhere_search") }) @@ -1189,7 +877,7 @@ describe("agent harness runtime", () => { expect(result.activeTools).not.toContain("knowhere_grep_chunks") }) - it("reopens finalize after must_use_sources has called knowhere_search, including empty results", () => { + it("opens answer preparation after must_use_sources has called knowhere_search", () => { const result = prepareHarnessStep({ stepNumber: 4, hasKnowhereSearch: true, @@ -1204,138 +892,284 @@ describe("agent harness runtime", () => { messages: [], }) - expect(result.activeTools).toContain("finalize") + expect(result.activeTools).toContain("prepareAnswer") + expect(result.activeTools).not.toContain("finalize") expect(result.activeTools).toContain("knowhere_search") }) - it("forces image inspection before forced finalization when image assets are available", () => { + it("forces answer preparation after the second Knowhere refinement", () => { const result = prepareHarnessStep({ - stepNumber: 12, - hasUninspectedImageAssets: true, - messages: [ - { - role: "tool", - content: [ - { - type: "tool-result", - toolCallId: "call_1", - toolName: "knowhere_search", - output: { - type: "text", - value: - '', - }, - }, - ], - }, - ], + stepNumber: 8, + hasKnowhereSearch: true, + hasReachedKnowhereSearchLimit: true, + intent: { + task: "answer", + dependsOnPreviousTurn: false, + retrievalNeeded: "yes", + targetModalities: ["text"], + constraints: {}, + groundingPolicy: "must_use_sources", + }, + messages: [], }) - expect(result.activeTools).toEqual(["inspectImage"]) + expect(result.activeTools).toEqual(["prepareAnswer"]) expect(result.toolChoice).toEqual({ type: "tool", - toolName: "inspectImage", + toolName: "prepareAnswer", }) - expect(result.messages).toEqual([ - { - role: "tool", - content: [ - { - type: "tool-result", - toolCallId: "call_1", - toolName: "knowhere_search", - output: { - type: "text", - value: - '', - }, - }, - ], + }) + + it("replaces retrieval history with the assembled context for finalization", () => { + const answerContextMessage = { + role: "user" as const, + content: [{ type: "text" as const, text: "## User's Question\nQuestion" }], + } + const result = prepareHarnessStep({ + stepNumber: 5, + messages: [{ role: "user", content: "retrieval history" }], + answerContextMessage, + }) + + expect(result.messages).toEqual([answerContextMessage]) + expect(result.activeTools).toEqual(["finalize"]) + expect(result.toolChoice).toEqual({ type: "tool", toolName: "finalize" }) + }) + + it("forces retainEvidence after a search returns new evidence", () => { + const result = prepareHarnessStep({ + stepNumber: 4, + hasPendingRetention: true, + pendingRetentionRange: { startPick: 1, endPick: 3 }, + hasKnowhereSearch: true, + intent: { + task: "answer", + dependsOnPreviousTurn: false, + retrievalNeeded: "yes", + targetModalities: ["text"], + constraints: {}, + groundingPolicy: "must_use_sources", }, + messages: [], + }) + + expect(result.activeTools).toEqual(["retainEvidence"]) + expect(result.toolChoice).toEqual({ + type: "tool", + toolName: "retainEvidence", + }) + expect(result.messages).toEqual([ { role: "user", - content: expect.stringContaining("Call inspectImage now"), + content: expect.stringContaining("1-3"), }, ]) + expect(result.activeTools).not.toContain("finalize") + expect(result.activeTools).not.toContain("knowhere_search") }) - it("reserves the finalization step even when unused page assets remain", () => { - const result = prepareHarnessStep({ - stepNumber: 13, - hasUninspectedImageAssets: true, - messages: [ + it("prepares the answer only after the latest search has been retained", async () => { + const ledger = createEvidenceLedger() + ledger.addRetrievalResponse(makeRetrievalResponse()) + const state: { answerContextRequested?: boolean } = {} + const tools = createHarnessTools({ + state, + ledger, + memoryTools: makeMemoryTools(), + knowhereTools: makeKnowhereTools(), + recentTurns: [], + }) + + expect(await executeTool(tools.prepareAnswer, {})).toMatchObject({ ok: false }) + expect(state.answerContextRequested).not.toBe(true) + + await executeTool(tools.retainEvidence, { picks: [1] }) + expect(await executeTool(tools.prepareAnswer, {})).toEqual({ ok: true }) + expect(state.answerContextRequested).toBe(true) + }) + + it("rejects displayed artifacts whose evidence was not retained", async () => { + const ledger = createEvidenceLedger() + ledger.addRetrievalResponse(makeRetrievalResponse()) + const tools = createHarnessTools({ + state: {}, + ledger, + memoryTools: makeMemoryTools(), + knowhereTools: makeKnowhereTools(), + recentTurns: [], + }) + await executeTool(tools.retainEvidence, { picks: [] }) + + const result = await executeTool(tools.finalize, { + text: "See the chart.", + citations: [], + memoryCitations: [], + artifacts: [ { - role: "user", - content: "What is the penalty amount?", + type: "image", + ref: "asset:r1:result:1", + display: true, + reason: "Show the chart.", }, ], + unresolved: [], + }) + expect(result).toMatchObject({ + ok: false, + unretainedArtifactRefs: ["asset:r1:result:1"], }) + }) - expect(result.activeTools).toEqual(["finalize"]) - expect(result.toolChoice).toEqual({ - type: "tool", - toolName: "finalize", + it("rejects retainEvidence picks outside the latest search", async () => { + const ledger = createEvidenceLedger() + ledger.addRetrievalResponse(makeRetrievalResponse()) + const tools = createHarnessTools({ + state: {}, + ledger, + memoryTools: makeMemoryTools(), + knowhereTools: makeKnowhereTools(), + recentTurns: [], + }) + + const rejected = await executeTool(tools.retainEvidence, { picks: [2] }) + expect(rejected).toMatchObject({ + ok: false, + invalidPicks: [2], }) + expect(ledger.hasPendingRetention()).toBe(true) + + const accepted = await executeTool(tools.retainEvidence, { picks: [] }) + expect(accepted).toMatchObject({ ok: true, retainedPicks: [] }) + expect(ledger.hasPendingRetention()).toBe(false) + expect(ledger.isRetained(1)).toBe(false) }) - it("forces finalize at step 13 using existing tool results", () => { - const result = prepareHarnessStep({ - stepNumber: 13, - messages: [ + it("rejects finalize of unretained first-search picks after a later search", async () => { + const ledger = createEvidenceLedger() + const tools = createHarnessTools({ + state: {}, + ledger, + memoryTools: makeMemoryTools(), + knowhereTools: makeKnowhereTools(), + recentTurns: [], + }) + ledger.addRetrievalResponse(makeRetrievalResponse()) + await executeTool(tools.retainEvidence, { picks: [] }) + ledger.addRetrievalResponse({ + ...makeRetrievalResponse(), + query: "second query", + results: [ { - role: "tool", - content: [ - { - type: "tool-result", - toolCallId: "call_1", - toolName: "knowhere_search", - output: { - type: "text", - value: - '', - }, - providerOptions: { - google: { - thoughtSignature: "signature-1", - }, - }, - }, - ], + content: "Second retrieval evidence.", + chunkType: "text", + score: 0.8, + source: { + documentId: "doc_2", + sourceFileName: "second.pdf", + sectionPath: "Second", + }, }, ], }) + await executeTool(tools.retainEvidence, { picks: [2] }) - expect(result.activeTools).toEqual(["finalize"]) - expect(result.toolChoice).toEqual({ - type: "tool", - toolName: "finalize", + const rejected = await executeTool(tools.finalize, { + text: "First source [[cite:1]].", + citations: [{ pick: 1 }], + memoryCitations: [], + artifacts: [], + unresolved: [], }) - expect(result.messages).toEqual([ - { - role: "tool", - content: [ - { - type: "tool-result", - toolCallId: "call_1", - toolName: "knowhere_search", - output: { - type: "text", - value: - '', - }, - }, - ], - }, - { - role: "user", - content: expect.stringContaining( - "Use only the evidence and tool results already available", - ), - }, - ]) + expect(rejected).toMatchObject({ + ok: false, + unretainedPicks: [1], + }) + expect(String((rejected as { message: string }).message)).toContain( + "Unretained citation picks", + ) + expect(String((rejected as { message: string }).message)).not.toContain( + "Unknown citation picks", + ) + + const accepted = await executeTool(tools.finalize, { + text: "Second source [[cite:1]].", + citations: [{ pick: 2 }], + memoryCitations: [], + artifacts: [], + unresolved: [], + }) + expect(accepted).toMatchObject({ + ok: true, + citations: [{ ref: "r2:result:1" }], + }) + }) + + it("allows the first knowhere_search without gapReason and rejects a follow-up without it", async () => { + const search = vi + .fn() + .mockResolvedValue(makeRetrievalResponse()) + const tools = createHarnessTools({ + state: {}, + ledger: createEvidenceLedger(), + memoryTools: makeMemoryTools(), + knowhereTools: makeKnowhereTools(search), + recentTurns: [], + }) + + const first = await executeTool(tools.knowhere_search, { query: "first" }) + expect(first).toContain('status="ok"') + expect(search).toHaveBeenCalledWith( + expect.objectContaining({ query: "first" }), + ) + expect(search.mock.calls[0]?.[0]).not.toHaveProperty("gapReason") + + const rejected = await executeTool(tools.knowhere_search, { query: "second" }) + expect(rejected).toContain('status="error"') + expect(String(rejected)).toContain("gapReason") + expect(search).toHaveBeenCalledTimes(1) + + const accepted = await executeTool(tools.knowhere_search, { + query: "second", + gapReason: "The first search lacked the wording this query will add.", + }) + expect(accepted).toContain('status="ok"') + expect(search).toHaveBeenCalledWith( + expect.objectContaining({ + query: "second", + gapReason: "The first search lacked the wording this query will add.", + }), + ) + + const third = await executeTool(tools.knowhere_search, { + query: "third", + gapReason: "The second search still lacked the requested comparison.", + }) + expect(third).toContain('status="ok"') + + const fourth = await executeTool(tools.knowhere_search, { + query: "fourth", + gapReason: "Try one more query.", + }) + expect(fourth).toContain('status="error"') + expect(fourth).toContain("initial search and two refinements") + expect(search).toHaveBeenCalledTimes(3) }) + }) +function retainLatestSearch( + tools: { retainEvidence: unknown }, + ledger: ReturnType, +): Promise { + const pending = ledger.pendingRetentionRange() + if (!pending) return Promise.resolve(undefined) + const picks: number[] = [] + for (let pick = pending.startPick; pick <= pending.endPick; pick += 1) { + picks.push(pick) + } + return executeTool(tools.retainEvidence, { picks }) +} + function executeTool(tool: unknown, input: unknown): Promise { return (tool as { execute: (input: unknown) => Promise }).execute(input) } @@ -1396,25 +1230,39 @@ function makeRetrievalResponse(): RetrievalQueryResponse { } } -function makeTableRetrievalResponse(): RetrievalQueryResponse { +function makeConnectedRetrievalResponse(): RetrievalQueryResponse { return { namespace: "notebook", - query: "q4 table", - routerUsed: "workflow_single_step", + query: "comparison", + routerUsed: "agent_explore", answerText: null, - evidenceText: "Table evidence", + evidenceText: "Comparison evidence", stopReason: "answer_done", failureReason: null, results: [ { - content: "", - chunkType: "table", - score: 0.8, - assetUrl: "https://assets.example/tables/revenue.html", + chunkId: "text_chunk", + content: "Comparison [tables/comparison.html]", + chunkType: "text", + score: 0.9, + metadata: { + connectTo: [ + { + target: "table_chunk", + relation: "embeds", + ref: "[tables/comparison.html]", + }, + { + target: "image_chunk", + relation: "embeds", + ref: "[images/comparison.jpg]", + }, + ], + }, source: { documentId: "doc_1", - sourceFileName: "report.pdf", - sectionPath: "tables/revenue.html", + sourceFileName: "cardiology.pdf", + sectionPath: "Differential diagnosis", }, }, ], @@ -1422,61 +1270,34 @@ function makeTableRetrievalResponse(): RetrievalQueryResponse { } } -function makeImageRetrievalResponse(count: number): RetrievalQueryResponse { +function toolCallResult( + calls: readonly { + toolCallId: string + toolName: string + input: unknown + }[], +) { return { - namespace: "notebook", - query: "q4 images", - routerUsed: "workflow_single_step", - answerText: null, - evidenceText: "Image evidence", - stopReason: "answer_done", - failureReason: null, - results: Array.from({ length: count }, (_, index) => ({ - content: "", - chunkType: "image", - score: 0.8, - assetUrl: `https://assets.example/images/chart-${index + 1}.png`, - source: { - documentId: "doc_1", - sourceFileName: "report.pdf", - sectionPath: `images/chart-${index + 1}.png`, - }, + content: calls.map((call) => ({ + type: "tool-call" as const, + toolCallId: call.toolCallId, + toolName: call.toolName, + input: JSON.stringify(call.input), })), - referencedChunks: [], - } -} - -function makePageCitationRetrievalResponse(): RetrievalQueryResponse { - return { - namespace: "notebook", - query: "进度计划", - routerUsed: "workflow_single_step", - answerText: null, - evidenceText: "Root / (6)现场工期进度管理方面的违约责任", - stopReason: "answer_done", - failureReason: null, - results: [], - referencedChunks: [ - { - chunkId: "chunk_page_8", - documentId: "doc_contract", - chunkType: "page", - sectionPath: "Root / (6)现场工期进度管理方面的违约责任", - filePath: null, - jobId: "job_contract", - assetUrl: "https://assets.example/page-8.png", - metadata: { - pageNums: [8], - pageAssets: [ - { - pageNum: 8, - artifactRef: "page_citation_assets/page-8.png", - assetUrl: "https://assets.example/page-8.png", - contentType: "image/png", - }, - ], - }, + finishReason: { unified: "tool-calls" as const, raw: undefined }, + usage: { + inputTokens: { + total: undefined, + noCache: undefined, + cacheRead: undefined, + cacheWrite: undefined, }, - ], + outputTokens: { + total: undefined, + text: undefined, + reasoning: undefined, + }, + }, + warnings: [], } } diff --git a/src/agent-harness/runtime.ts b/src/agent-harness/runtime.ts index dbe91a5..4c0784b 100644 --- a/src/agent-harness/runtime.ts +++ b/src/agent-harness/runtime.ts @@ -7,39 +7,35 @@ import { } from "ai" import { z } from "zod" +import { composeAnswerContext, type ReadTableHtml } from "./answer-context" import { createEvidenceLedger } from "./ledger" -import { getCanonicalImageAssetKey } from "./image-asset-identity" import { knowhereToolText } from "./knowhere-text" import { memoryToolText } from "./memory-text" -import { mergeImageInspectionHighlights } from "./image-highlights" import type { AgentTurn, AgentTurnInput, ContextPolicy, - EvidenceAsset, - EvidenceChunk, EvidenceLedgerSnapshot, HarnessRunResult, HarnessToolCallTrace, HarnessTrace, - ImageInspectionAsset, - ImageInspectionHighlights, - ImageInspectionResponse, - InspectImages, IntentFrame, KnowhereSearchTargetContent, KnowhereToolRuntime, + MemoryCitation, MemorySearchKind, + MemorySearchItem, MemoryToolRuntime, OutputArtifactView, OutputCitation, OutputManifest, + ResolveConnectedAssets, } from "./types" import { memorySearchKinds } from "./types" -const defaultMaxSteps = 14 -const imageInspectionReminderStepNumber = 12 -const forcedFinalizationStepNumber = 13 +const defaultMaxSteps = 15 +const maxKnowhereRefinements = 2 +const maxKnowhereSearchAttempts = 1 + maxKnowhereRefinements type ToolLoopAgentSettings = ConstructorParameters[0] @@ -50,7 +46,8 @@ export type RunAgentHarnessInput = { readonly turn: AgentTurnInput readonly knowhereTools: KnowhereToolRuntime readonly memoryTools: MemoryToolRuntime - readonly inspectImages?: InspectImages + readonly resolveConnectedAssets?: ResolveConnectedAssets + readonly readTableHtml?: ReadTableHtml readonly maxSteps?: number } @@ -59,9 +56,12 @@ type HarnessToolState = { contextPolicy?: ContextPolicy finalizedManifest?: OutputManifest finalized?: boolean + answerContextRequested?: boolean + answerContextMessage?: ModelMessage + memoryItems?: MemorySearchItem[] + memorySearchAttempted?: boolean + knowhereSearchAttemptCount?: number priorTurnReads?: string[] - inspectedImageRefs?: string[] - imageHighlights?: ImageInspectionHighlights[] toolCalls?: HarnessToolCallTrace[] } @@ -91,13 +91,14 @@ const knowhereSearchTargetContentSchema = z.enum([ const knowhereSearchSchema = z.object({ query: z.string().min(1), includeDocumentIds: z.array(z.string().trim().min(1)).optional().describe( - "Only search these verified document IDs. Omit for unrestricted search; [] searches no documents. Use IDs from source context or previous search results, never filenames or guessed IDs.", + "Only search these verified document IDs. Omit for unrestricted search; [] searches no documents. Use IDs from previous search results, never filenames or guessed IDs.", ), excludeDocumentIds: z.array(z.string().trim().min(1)).optional().describe( "Exclude these verified document IDs. Exclusions take precedence over includeDocumentIds. If the ID is unknown, describe the document constraint in query instead.", ), targetContent: knowhereSearchTargetContentSchema.default("all"), purpose: z.string().optional(), + gapReason: z.string().optional(), topK: z.number().int().min(1).max(12).optional(), signalPaths: z.array(z.string().min(1)).max(8).optional(), filterMode: z.enum(["keep", "delete"]).optional(), @@ -206,36 +207,53 @@ export async function runAgentHarness( ledger, knowhereTools: input.knowhereTools, memoryTools: input.memoryTools, - inspectImages: input.inspectImages, recentTurns: input.turn.recentTurns, }) const agent = new ToolLoopAgent({ model: input.model, instructions: buildHarnessSystemPrompt(input.turn), tools, - prepareStep: ({ messages: stepMessages, stepNumber }) => - prepareHarnessStep({ + prepareStep: async ({ messages: stepMessages, stepNumber }) => { + if ( + state.answerContextRequested && + !state.answerContextMessage && + !ledger.hasPendingRetention() + ) { + await ledger.resolveRetainedConnectedAssets(input.resolveConnectedAssets) + state.answerContextMessage = await composeAnswerContext({ + ledger: ledger.snapshot(), + memoryItems: state.memoryItems ?? [], + memorySearchAttempted: state.memorySearchAttempted === true, + userText: input.turn.userText, + readTableHtml: input.readTableHtml, + }) + } + return prepareHarnessStep({ messages: stepMessages, stepNumber, intent: state.intent, - hasUninspectedImageAssets: - input.inspectImages !== undefined && - hasUninspectedImageAssets({ state, ledger }), - hasKnowhereSearch: (state.toolCalls ?? []).some( - (call) => call.tool === "knowhere_search", - ), - }), + hasKnowhereSearch: (state.knowhereSearchAttemptCount ?? 0) > 0, + hasReachedKnowhereSearchLimit: + (state.knowhereSearchAttemptCount ?? 0) >= maxKnowhereSearchAttempts, + hasMemorySearch: state.memorySearchAttempted === true, + hasPendingRetention: ledger.hasPendingRetention(), + pendingRetentionRange: ledger.pendingRetentionRange(), + answerContextMessage: state.answerContextMessage, + }) + }, stopWhen: [ () => state.finalized === true, stepCountIs(input.maxSteps ?? defaultMaxSteps), ], }) - const response = await agent.generate({ + await agent.generate({ messages: buildHarnessMessages(input.turn), }) - const manifest = - state.finalizedManifest ?? buildFallbackManifest(response.text.trim()) + if (!state.finalizedManifest) { + throw new Error("The agent did not finalize an output manifest.") + } + const manifest = state.finalizedManifest const ledgerSnapshot = ledger.snapshot() return { manifest, @@ -246,7 +264,7 @@ export async function runAgentHarness( finalized: state.finalized === true, priorTurnReads: [...(state.priorTurnReads ?? [])], toolCalls: [...(state.toolCalls ?? [])], - imageHighlights: [...(state.imageHighlights ?? [])], + imageHighlights: [], validationErrors: [], revisionsUsed: 0, }, @@ -256,7 +274,6 @@ export async function runAgentHarness( const alwaysAvailableTools = [ "declareIntent", "setContextPolicy", - "inspectImage", "readPriorTurn", ] as const @@ -280,46 +297,51 @@ const cognitionRetrievalTools = [] as const export function prepareHarnessStep(input: { readonly stepNumber: number readonly messages: readonly ModelMessage[] - readonly hasUninspectedImageAssets?: boolean readonly hasKnowhereSearch?: boolean + readonly hasReachedKnowhereSearchLimit?: boolean + readonly hasMemorySearch?: boolean + readonly hasPendingRetention?: boolean + readonly pendingRetentionRange?: { startPick: number; endPick: number } | null + readonly answerContextMessage?: ModelMessage readonly intent?: IntentFrame }): HarnessStepPreparation { const messages = sanitizeHarnessModelMessagesForStep(input.messages) - const shouldForceImageInspection = - input.hasUninspectedImageAssets === true && - input.stepNumber === imageInspectionReminderStepNumber + if (input.answerContextMessage) { + return { + messages: [input.answerContextMessage], + activeTools: ["finalize"], + toolChoice: { + type: "tool", + toolName: "finalize", + }, + } + } - if (shouldForceImageInspection) { + if (input.hasPendingRetention === true && input.pendingRetentionRange) { return { messages: [ ...messages, { role: "user", - content: buildImageInspectionReminderFeedback(), + content: buildRetainEvidenceFeedback(input.pendingRetentionRange), }, ], - activeTools: ["inspectImage"], + activeTools: ["retainEvidence"], toolChoice: { type: "tool", - toolName: "inspectImage", + toolName: "retainEvidence", }, } } - if (input.stepNumber >= forcedFinalizationStepNumber) { + if (input.hasReachedKnowhereSearchLimit === true) { return { - messages: [ - ...messages, - { - role: "user", - content: buildForcedFinalizationFeedback(), - }, - ], - activeTools: ["finalize"], + messages, + activeTools: ["prepareAnswer"], toolChoice: { type: "tool", - toolName: "finalize", + toolName: "prepareAnswer", }, } } @@ -329,10 +351,10 @@ export function prepareHarnessStep(input: { activeTools: selectHarnessActiveTools({ intent: input.intent, hasKnowhereSearch: input.hasKnowhereSearch === true, + hasMemorySearch: input.hasMemorySearch === true, }), - // finalize is the only output contract (see its tool description). Force - // a tool call every step so the model cannot end the turn with a bare - // text response that skips finalize's citation/artifact validation. + // Tool calls are required so the retrieval phase cannot end with bare + // text and skip the composed answer context or finalize validation. toolChoice: "required", } } @@ -340,12 +362,13 @@ export function prepareHarnessStep(input: { function selectHarnessActiveTools(input: { readonly intent?: IntentFrame readonly hasKnowhereSearch: boolean + readonly hasMemorySearch: boolean }): Array> { const tools: Array> = [ ...alwaysAvailableTools, ] - if (allowsFinalize(input)) { - tools.push("finalize") + if (allowsAnswerPreparation(input)) { + tools.push("prepareAnswer") } if (!allowsRetrieval(input.intent)) { return tools @@ -354,7 +377,7 @@ function selectHarnessActiveTools(input: { // memory_search and knowhere_search are peers: both open together once // retrieval is allowed. The agent decides which to call and in what // order — neither tool gates the other. - tools.push(...fluidRetrievalTools) + if (!input.hasMemorySearch) tools.push(...fluidRetrievalTools) if (input.intent?.groundingPolicy === "must_use_sources") { tools.push(...crystalRetrievalTools) } @@ -369,7 +392,7 @@ function allowsRetrieval(intent?: IntentFrame): boolean { ) } -function allowsFinalize(input: { +function allowsAnswerPreparation(input: { readonly intent?: IntentFrame readonly hasKnowhereSearch: boolean }): boolean { @@ -440,57 +463,24 @@ type ModelMessageForRole = Extract< { readonly role: TRole } > -function buildForcedFinalizationFeedback(): string { - return [ - "The retrieval step budget has been reached.", - "Do not search again.", - "Use only the evidence and tool results already available in this turn.", - "Call finalize now with the best supported answer.", - "If the existing evidence is insufficient, explain the gap in unresolved", - "instead of making unsupported claims.", - ].join("\n") -} - -function buildImageInspectionReminderFeedback(): string { +function buildRetainEvidenceFeedback(range: { + readonly startPick: number + readonly endPick: number +}): string { return [ - "Retrieved image/page assets are available and have not been inspected.", - "Call inspectImage now with the page/image asset refs you will cite.", - "Use a question that locates the cited evidence on those pages for OCR/visual context and provenance boxes.", - "Do not finalize until those cited image assets have been inspected.", - "Do not search again.", + "The latest search returned new evidence that has not been retained.", + `Call retainEvidence with the pick numbers you will keep from ${range.startPick}-${range.endPick}.`, + "An empty picks list means this search found nothing useful.", + "Picks not retained cannot be cited later.", + "Do not search or finalize until retainEvidence has been called.", ].join("\n") } -function hasUninspectedImageAssets(input: { - readonly state: HarnessToolState - readonly ledger: ReturnType -}): boolean { - const snapshot = input.ledger.snapshot() - const chunksByRef = new Map( - snapshot.chunks.map((chunk) => [chunk.ref, chunk] as const), - ) - const assetsByRef = new Map( - snapshot.assets.map((asset) => [asset.ref, asset] as const), - ) - const inspectedKeys = new Set( - (input.state.inspectedImageRefs ?? []).map((ref) => { - const asset = assetsByRef.get(ref) - return asset ? getCanonicalImageAssetKey(asset, chunksByRef) : ref - }), - ) - return snapshot.assets.some( - (asset) => - asset.type === "image" && - !inspectedKeys.has(getCanonicalImageAssetKey(asset, chunksByRef)), - ) -} - export function createHarnessTools(input: { readonly state: HarnessToolState readonly ledger: ReturnType readonly knowhereTools: KnowhereToolRuntime readonly memoryTools: MemoryToolRuntime - readonly inspectImages?: InspectImages readonly recentTurns: readonly AgentTurn[] }) { return { @@ -528,7 +518,7 @@ export function createHarnessTools(input: { memory_search: tool({ description: - "Search distilled fluid memory for this workspace. Returns tagged text with memory refs such as mem:1. Use this before Knowhere document search.", + "Search distilled fluid memory for this workspace. Returns tagged text with memory refs such as mem:1. It can run alongside Knowhere document search when both are useful.", inputSchema: memorySearchSchema, execute: async (request) => traceToolCall(input.state, { @@ -536,6 +526,7 @@ export function createHarnessTools(input: { inputSummary: summarizeMemorySearchRequest(request), execute: async () => { return await executeMemorySearch({ + state: input.state, memoryTools: input.memoryTools, request, }) @@ -544,9 +535,27 @@ export function createHarnessTools(input: { }), }), + retainEvidence: tool({ + description: + "Keep pick numbers from the latest Knowhere search that are useful. " + + "Call this after every search that returns new evidence. " + + "An empty picks list means that search found nothing useful. " + + "Picks not retained cannot be cited later.", + inputSchema: z.object({ + picks: z.array(z.number().int().positive()), + }), + execute: async ({ picks }) => + traceToolCall(input.state, { + toolName: "retainEvidence", + inputSummary: { picks }, + execute: async () => input.ledger.retainPicks(picks), + summarizeOutput: (output) => output, + }), + }), + knowhere_search: tool({ description: - "Search Knowhere for relevant Notebook evidence. Returns tagged text with pick numbers for finalize citations, evidence refs such as r1:result:1, and asset refs such as asset:r1:result:1.", + "Search Knowhere for relevant Notebook evidence. Returns tagged text with pick numbers for finalize citations, evidence refs such as r1:result:1, and connected image/table asset paths when present. After a previous search in this turn, set gapReason to what the previous search lacked and how this query will fill that gap.", inputSchema: knowhereSearchSchema, execute: async (request) => traceToolCall(input.state, { @@ -554,6 +563,7 @@ export function createHarnessTools(input: { inputSummary: summarizeKnowhereSearchRequest(request), execute: async () => executeKnowhereSearch({ + state: input.state, ledger: input.ledger, knowhereTools: input.knowhereTools, request, @@ -562,29 +572,6 @@ export function createHarnessTools(input: { }), }), - inspectImage: tool({ - description: - "Inspect cited Knowhere page/image asset refs for OCR, visual details, and provenance boxes. Call this after retrieval and before finalize whenever the answer cites page or image assets.", - inputSchema: z.object({ - refs: z.array(z.string().min(1)).min(1), - question: z.string().min(1), - }), - execute: async (request) => - traceToolCall(input.state, { - toolName: "inspectImage", - inputSummary: summarizeInspectImageRequest(request), - execute: async () => - inspectRetrievedImages({ - state: input.state, - ledger: input.ledger, - inspectImages: input.inspectImages, - refs: request.refs, - question: request.question, - }), - summarizeOutput: summarizeInspectImageOutput, - }), - }), - readPriorTurn: tool({ description: "Read the full text and citation labels of a specific prior turn by id " + @@ -644,6 +631,28 @@ export function createHarnessTools(input: { }), }), + prepareAnswer: tool({ + description: + "Finish retrieval and assemble retained Knowledge Base evidence, all fluid memory results, and the user's original question into one multimodal message for the answer step.", + inputSchema: z.object({}), + execute: async () => + traceToolCall(input.state, { + toolName: "prepareAnswer", + inputSummary: {}, + execute: async () => { + if (input.ledger.hasPendingRetention()) { + return { + ok: false as const, + message: "Call retainEvidence for the latest search before prepareAnswer.", + } + } + input.state.answerContextRequested = true + return { ok: true as const } + }, + summarizeOutput: (output) => output, + }), + }), + finalize: tool({ description: "Finalize the user-facing output manifest. This is the only final answer " + @@ -651,8 +660,8 @@ export function createHarnessTools(input: { "images/tables shown to the user. citations is the list of evidence picks " + "you used; each pick is the pick number on a Knowhere search chunk. " + "Notebook writes citation refs from the evidence ledger. " + - "Use memoryCitations for fluid memory refs. " + - "Cited page/image assets must be inspected with inspectImage first.", + "Use memoryCitations for fluid memory and copy ref, itemId, and kind " + + "exactly from the assembled Fluid Memory entries.", inputSchema: finalizeManifestSchema, execute: async (manifest) => traceToolCall(input.state, { @@ -664,38 +673,61 @@ export function createHarnessTools(input: { ledger: input.ledger, }) if (!resolvedCitations.ok) { + if ("unknownPicks" in resolvedCitations) { + return { + ok: false as const, + message: buildFinalizeRequiresPicksMessage({ + unknownPicks: resolvedCitations.unknownPicks, + ledger: input.ledger.snapshot(), + }), + unknownPicks: resolvedCitations.unknownPicks, + } + } return { ok: false as const, - message: buildFinalizeRequiresPicksMessage({ - unknownPicks: resolvedCitations.unknownPicks, + message: buildFinalizeRequiresRetainedPicksMessage({ + unretainedPicks: resolvedCitations.unretainedPicks, ledger: input.ledger.snapshot(), }), - unknownPicks: resolvedCitations.unknownPicks, + unretainedPicks: resolvedCitations.unretainedPicks, } } - const outputManifest: OutputManifest = { - text: manifest.text, - citations: resolvedCitations.citations, - memoryCitations: manifest.memoryCitations, - artifacts: manifest.artifacts, - unresolved: manifest.unresolved, + const invalidMemoryCitations = getInvalidMemoryCitations({ + citations: manifest.memoryCitations, + memoryItems: input.state.memoryItems ?? [], + }) + if (invalidMemoryCitations.length > 0) { + return { + ok: false as const, + message: + "memoryCitations must exactly match ref, itemId, and kind from this turn's Fluid Memory results.", + invalidMemoryCitations, + } } - const inspectRefs = getUninspectedCitedImageRefs({ - manifest: outputManifest, + const unretainedArtifactRefs = getUnretainedDisplayedArtifactRefs({ + artifacts: manifest.artifacts, ledger: input.ledger, - inspectedImageRefs: input.state.inspectedImageRefs ?? [], - inspectImagesAvailable: input.inspectImages !== undefined, }) - if (inspectRefs.length > 0) { + if (unretainedArtifactRefs.length > 0) { return { ok: false as const, - message: buildFinalizeRequiresInspectionMessage(inspectRefs), - inspectRefs, + message: buildFinalizeRequiresRetainedArtifactsMessage( + unretainedArtifactRefs, + ), + unretainedArtifactRefs, } } + const outputManifest: OutputManifest = { + text: manifest.text, + citations: resolvedCitations.citations, + memoryCitations: manifest.memoryCitations, + artifacts: manifest.artifacts, + unresolved: manifest.unresolved, + } + input.state.finalizedManifest = outputManifest input.state.finalized = true return { ok: true as const, ...outputManifest } @@ -706,175 +738,16 @@ export function createHarnessTools(input: { } as const } -async function inspectRetrievedImages(input: { - readonly state: HarnessToolState - readonly ledger: ReturnType - readonly inspectImages?: InspectImages - readonly refs: readonly string[] - readonly question: string -}): Promise< - | ({ readonly ok: true } & ImageInspectionResponse) - | { - readonly ok: false - readonly message: string - readonly inspected: readonly [] - readonly skipped: readonly { - readonly ref: string - readonly reason: string - }[] - } -> { - const refs = getUniqueTrimmedRefs(input.refs) - const question = input.question.trim() - - if (refs.length === 0 || question.length === 0) { - return { - ok: false, - message: "At least one image asset ref and a question are required.", - inspected: [], - skipped: [], - } - } - const snapshot = input.ledger.snapshot() - if (snapshot.chunks.length === 0 && snapshot.assets.length === 0) { - return { - ok: false, - message: - "Knowhere evidence tools must return image assets before inspectImage.", - inspected: [], - skipped: refs.map((ref) => ({ - ref, - reason: "No Knowhere evidence is available yet.", - })), - } - } - if (!input.inspectImages) { - return { - ok: false, - message: "Image inspection is not available for this turn.", - inspected: [], - skipped: refs.map((ref) => ({ - ref, - reason: "No image inspection capability is configured.", - })), - } - } - - const assetsByRef = new Map( - snapshot.assets.map((asset) => [asset.ref, asset] as const), - ) - const chunksByRef = new Map( - snapshot.chunks.map((chunk) => [chunk.ref, chunk] as const), - ) - const skipped: { - readonly ref: string - readonly reason: string - }[] = [] - const selectedAssets: ImageInspectionAsset[] = [] - const selectedKeys = new Set() - - for (const ref of refs) { - const asset = assetsByRef.get(ref) - if (!asset) { - skipped.push({ - ref, - reason: "Ref was not returned by Knowhere as an asset.", - }) - continue - } - if (asset.type !== "image") { - skipped.push({ - ref, - reason: "Ref is not an image asset.", - }) - continue - } - - const assetKey = getCanonicalImageAssetKey(asset, chunksByRef) - if (selectedKeys.has(assetKey)) { - skipped.push({ - ref, - reason: "Duplicate of another retrieved page selected for inspection.", - }) - continue - } - selectedKeys.add(assetKey) - selectedAssets.push({ - ref: asset.ref, - label: asset.label, - ...(asset.assetUrl ? { assetUrl: asset.assetUrl } : {}), - ...(asset.sourcePath ? { sourcePath: asset.sourcePath } : {}), - ...(asset.revisionKey ? { revisionKey: asset.revisionKey } : {}), - source: asset.source, - }) - } - - if (selectedAssets.length === 0) { - return { - ok: false, - message: "No inspectable image asset refs were provided.", - inspected: [], - skipped, - } - } - - const inspectedImageRefs = input.state.inspectedImageRefs ?? [] - - try { - const response = await input.inspectImages({ - question, - assets: selectedAssets, - }) - const selectedRefs = new Set(selectedAssets.map((asset) => asset.ref)) - const successfulRefs = response.inspected - .map((asset) => asset.ref) - .filter((ref) => selectedRefs.has(ref)) - input.state.inspectedImageRefs = [ - ...inspectedImageRefs, - ...successfulRefs.filter((ref) => !inspectedImageRefs.includes(ref)), - ] - input.state.imageHighlights = mergeImageInspectionHighlights( - input.state.imageHighlights, - response.highlights, - ) - if (successfulRefs.length === 0) { - return { - ok: false as const, - message: "Image inspection skipped every requested asset.", - inspected: [], - skipped: [...skipped, ...response.skipped], - } - } - return { - ok: true as const, - analysis: response.analysis, - inspected: response.inspected, - skipped: [...skipped, ...response.skipped], - } - } catch (error) { - return { - ok: false, - message: - error instanceof Error - ? `Image inspection failed: ${error.message}` - : "Image inspection failed.", - inspected: [], - skipped: selectedAssets.map((asset) => ({ - ref: asset.ref, - reason: "The image inspection request failed.", - })), - } - } -} - function resolveCitationPicks(input: { readonly citations: readonly { pick: number }[] readonly ledger: ReturnType }): | { ok: true; citations: OutputCitation[] } - | { ok: false; unknownPicks: number[] } { + | { ok: false; unknownPicks: number[] } + | { ok: false; unretainedPicks: number[] } { const chunks = input.ledger.snapshot().chunks const unknownPicks: number[] = [] + const unretainedPicks: number[] = [] const citations: OutputCitation[] = [] for (const citation of input.citations) { @@ -887,15 +760,39 @@ function resolveCitationPicks(input: { } continue } + if (!input.ledger.isRetained(citation.pick)) { + if (!unretainedPicks.includes(citation.pick)) { + unretainedPicks.push(citation.pick) + } + continue + } citations.push({ ref: chunk.ref }) } if (unknownPicks.length > 0) { return { ok: false, unknownPicks } } + if (unretainedPicks.length > 0) { + return { ok: false, unretainedPicks } + } return { ok: true, citations } } +function getInvalidMemoryCitations(input: { + readonly citations: readonly MemoryCitation[] + readonly memoryItems: readonly MemorySearchItem[] +}): MemoryCitation[] { + return input.citations.filter( + (citation) => + !input.memoryItems.some( + (item) => + item.ref === citation.ref && + item.itemId === citation.itemId && + item.kind === citation.kind, + ), + ) +} + function buildFinalizeRequiresPicksMessage(input: { readonly unknownPicks: readonly number[] readonly ledger: EvidenceLedgerSnapshot @@ -911,89 +808,66 @@ function buildFinalizeRequiresPicksMessage(input: { ].join(" ") } -function getUninspectedCitedImageRefs(input: { - readonly manifest: OutputManifest +function buildFinalizeRequiresRetainedPicksMessage(input: { + readonly unretainedPicks: readonly number[] + readonly ledger: EvidenceLedgerSnapshot +}): string { + const retained = + input.ledger.retainedPicks.length > 0 + ? `Retained picks: ${input.ledger.retainedPicks.join(" ")}.` + : "No evidence picks have been retained." + return [ + "Citations can only use retained evidence picks.", + `Unretained citation picks: ${input.unretainedPicks.join(" ")}.`, + retained, + "Call retainEvidence after each search, then finalize using only retained picks.", + ].join(" ") +} + +function getUnretainedDisplayedArtifactRefs(input: { + readonly artifacts: readonly OutputArtifactView[] readonly ledger: ReturnType - readonly inspectedImageRefs: readonly string[] - readonly inspectImagesAvailable: boolean }): string[] { - if (!input.inspectImagesAvailable) return [] - - const snapshot = input.ledger.snapshot() - const chunksByRef = new Map( - snapshot.chunks.map((chunk) => [chunk.ref, chunk] as const), - ) - const assetsByRef = new Map( - snapshot.assets.map((asset) => [asset.ref, asset] as const), - ) - const inspected = new Set(input.inspectedImageRefs) - const inspectedKeys = new Set( - input.inspectedImageRefs.map((ref) => { - const asset = assetsByRef.get(ref) - return asset ? getCanonicalImageAssetKey(asset, chunksByRef) : ref - }), - ) - const refs: string[] = [] - const seenKeys = new Set() - const addRef = (ref: string | null): void => { - if (!ref || inspected.has(ref)) return - const asset = assetsByRef.get(ref) - const key = asset ? getCanonicalImageAssetKey(asset, chunksByRef) : ref - if (seenKeys.has(key) || inspectedKeys.has(key)) return - seenKeys.add(key) - refs.push(ref) - } - - for (const citation of input.manifest.citations) { - addRef(resolveImageAssetRef(citation.ref, chunksByRef, assetsByRef)) - } - - for (const artifact of input.manifest.artifacts) { - if (artifact.type !== "image" || artifact.display === false) continue - addRef(resolveImageAssetRef(artifact.ref, chunksByRef, assetsByRef)) + for (const artifact of input.artifacts) { + if (artifact.display === false || artifact.type === "derived_table") continue + if (!isRefRetained(input.ledger, artifact.ref) && !refs.includes(artifact.ref)) { + refs.push(artifact.ref) + } } - return refs } -function resolveImageAssetRef( - ref: string, - chunksByRef: ReadonlyMap, - assetsByRef: ReadonlyMap, -): string | null { - const directAsset = assetsByRef.get(ref) - if (directAsset?.type === "image") return directAsset.ref - - const chunk = chunksByRef.get(ref) - if (!chunk) return null - - if (chunk.assetRef) { - const chunkAsset = assetsByRef.get(chunk.assetRef) - if (chunkAsset?.type === "image") return chunkAsset.ref - } +function buildFinalizeRequiresRetainedArtifactsMessage( + refs: readonly string[], +): string { + return [ + "Displayed artifacts can only use retained evidence.", + `Unretained artifact refs: ${refs.join(" ")}.`, + "Call retainEvidence after each search, then finalize using only retained evidence.", + ].join(" ") +} - if (!chunk.chunkId) return null - const sibling = Array.from(chunksByRef.values()).find( - (candidate) => - candidate.ref !== chunk.ref && - candidate.chunkId === chunk.chunkId && - candidate.source.documentId === chunk.source.documentId && - candidate.assetRef !== undefined && - assetsByRef.get(candidate.assetRef)?.type === "image", +function pickForRef( + snapshot: EvidenceLedgerSnapshot, + ref: string, +): number | null { + const chunkIndex = snapshot.chunks.findIndex((chunk) => chunk.ref === ref) + if (chunkIndex >= 0) return chunkIndex + 1 + const asset = snapshot.assets.find((candidate) => candidate.ref === ref) + if (!asset) return null + const parentIndex = snapshot.chunks.findIndex( + (chunk) => chunk.ref === asset.chunkRef, ) - return sibling?.assetRef ?? null + return parentIndex >= 0 ? parentIndex + 1 : null } -function buildFinalizeRequiresInspectionMessage( - inspectRefs: readonly string[], -): string { - return [ - "Cited page/image assets must be inspected before finalize.", - `Call inspectImage with refs: ${inspectRefs.join(" ")}.`, - "Use a question that locates the cited evidence on those pages.", - "Then call finalize again, using the inspection notes in the answer.", - ].join(" ") +function isRefRetained( + ledger: ReturnType, + ref: string, +): boolean { + const pick = pickForRef(ledger.snapshot(), ref) + return pick !== null && ledger.isRetained(pick) } type KnowhereToolOperation = "search" @@ -1002,14 +876,24 @@ type MemorySearchToolRequest = z.infer type KnowhereSearchToolRequest = z.infer async function executeMemorySearch(input: { + readonly state: HarnessToolState readonly memoryTools: MemoryToolRuntime readonly request: MemorySearchToolRequest }): Promise { + if (input.state.memorySearchAttempted) { + return memoryToolText.formatError({ + operation: "search", + message: "Fluid memory search already ran for this turn.", + }) + } + input.state.memorySearchAttempted = true + try { const response = await input.memoryTools.search({ query: input.request.query, kinds: input.request.kinds, }) + accumulateMemoryItems(input.state, response.items) return memoryToolText.formatSearch(response) } catch (error) { return memoryToolText.formatError({ @@ -1019,7 +903,22 @@ async function executeMemorySearch(input: { } } +function accumulateMemoryItems( + state: HarnessToolState, + items: readonly MemorySearchItem[], +): void { + const accumulated = state.memoryItems ?? [] + const itemIds = new Set(accumulated.map((item) => item.itemId)) + for (const item of items) { + if (itemIds.has(item.itemId)) continue + accumulated.push(item) + itemIds.add(item.itemId) + } + state.memoryItems = accumulated +} + async function executeKnowhereSearch(input: { + readonly state: HarnessToolState readonly ledger: ReturnType readonly knowhereTools: KnowhereToolRuntime readonly request: KnowhereSearchToolRequest @@ -1027,6 +926,23 @@ async function executeKnowhereSearch(input: { return executeKnowhereTextTool({ operation: "search", execute: async () => { + const attemptCount = input.state.knowhereSearchAttemptCount ?? 0 + const gapReason = input.request.gapReason?.trim() + if (attemptCount > 0 && !gapReason) { + return knowhereToolText.formatError({ + operation: "search", + message: + "This is a follow-up search. Set gapReason to what the previous search lacked and how this query will fill that gap.", + }) + } + if (attemptCount >= maxKnowhereSearchAttempts) { + return knowhereToolText.formatError({ + operation: "search", + message: + "Knowhere search already completed its initial search and two refinements. Call prepareAnswer now.", + }) + } + input.state.knowhereSearchAttemptCount = attemptCount + 1 const beforeSnapshot = input.ledger.snapshot() const response = await input.knowhereTools.search({ query: input.request.query, @@ -1038,6 +954,7 @@ async function executeKnowhereSearch(input: { : {}), targetContent: input.request.targetContent, purpose: input.request.purpose, + ...(gapReason ? { gapReason } : {}), topK: input.request.topK, signalPaths: input.request.signalPaths, filterMode: input.request.filterMode, @@ -1069,16 +986,6 @@ async function executeKnowhereTextTool(input: { } } -function getUniqueTrimmedRefs(refs: readonly string[]): string[] { - const normalizedRefs: string[] = [] - for (const ref of refs) { - const normalizedRef = ref.trim() - if (!normalizedRef || normalizedRefs.includes(normalizedRef)) continue - normalizedRefs.push(normalizedRef) - } - return normalizedRefs -} - async function traceToolCall(input: { readonly toolCalls?: HarnessToolCallTrace[] }, call: { @@ -1173,6 +1080,7 @@ function summarizeKnowhereSearchRequest(request: { readonly query: string readonly targetContent?: KnowhereSearchTargetContent readonly purpose?: string + readonly gapReason?: string readonly topK?: number readonly signalPaths?: readonly string[] readonly filterMode?: string @@ -1182,6 +1090,7 @@ function summarizeKnowhereSearchRequest(request: { query: request.query, targetContent: request.targetContent ?? "all", purpose: request.purpose, + gapReason: request.gapReason, topK: request.topK, signalPathCount: request.signalPaths?.length ?? 0, filterMode: request.filterMode, @@ -1211,30 +1120,6 @@ function countOccurrences(value: string, pattern: string): number { } } -function summarizeInspectImageRequest(request: { - readonly refs: readonly string[] - readonly question: string -}): unknown { - return { - refs: getUniqueTrimmedRefs(request.refs), - questionLength: request.question.trim().length, - } -} - -function summarizeInspectImageOutput(output: unknown): unknown { - if (!isRecord(output)) return output - return { - ok: output.ok, - analysisLength: - typeof output.analysis === "string" ? output.analysis.length : 0, - inspectedCount: Array.isArray(output.inspected) - ? output.inspected.length - : 0, - skippedCount: Array.isArray(output.skipped) ? output.skipped.length : 0, - message: output.message, - } -} - function summarizeReadPriorTurnOutput(output: unknown): unknown { if (!isRecord(output)) return output return { @@ -1305,17 +1190,21 @@ export function buildHarnessSystemPrompt(turn: AgentTurnInput): string { "1. Call declareIntent when it helps you plan the response. Capture constraints like a requested image/table count in constraints.desiredCount.", "2. Call setContextPolicy when prior turns may influence this turn.", "3. When the policy needs prior-turn detail (references or corrections), call readPriorTurn for the relevant ids.", - "4. Call memory_search first when known fluid memory may answer the request. Call knowhere_search when memory is insufficient and groundingPolicy requires citing source documents. If the returned evidence is not enough to answer, or the query needs to focus differently, call knowhere_search again with a refined query (different keywords / topK / targetContent) instead of trying to browse documents directly. Knowhere's own retrieval agent already navigates the corpus internally.", - "5. After Knowhere returns image/page asset refs, call inspectImage on the page/image assets you will cite before finalize. This supplies OCR/visual context and provenance boxes.", - "6. Inspect each unique cited page once; retrieval already bounds the available evidence set.", - "7. Call finalize with text, citations, artifacts, and unresolved issues when you are ready to answer.", + "4. Call memory_search for relevant fluid memory and knowhere_search when groundingPolicy requires source documents. When both are useful, call them together in the same step. If the returned evidence is not enough to answer, or the query needs to focus differently, call knowhere_search again with a refined query (different keywords / topK / targetContent) instead of trying to browse documents directly. Knowhere's own retrieval agent already navigates the corpus internally.", + "5. After a search returns new evidence, call retainEvidence with the pick numbers from that search you will keep, then decide whether to search again or prepare the answer. An empty picks list means this search found nothing useful. Picks not retained cannot be cited later.", + "6. When retrieval is complete, call prepareAnswer. The retained evidence, all fluid memory results, and the user's original question will be assembled into one multimodal message.", + "7. Read that assembled message once and call finalize with the answer, citations, artifacts, and unresolved issues.", "", "Retrieval rules:", - "- First use memory_search to see whether known fluid memory can answer directly.", - "- Call knowhere_search only when memory is insufficient and groundingPolicy requires citing source documents.", - "- Refine knowhere_search at most twice. If two refined searches still do not add new relevant evidence, call finalize and list the gap in unresolved.", + "- memory_search and knowhere_search are parallel retrieval sources. When both apply, call them together rather than making one wait for the other.", + "- Call memory_search once per turn. An empty result remains empty; do not retry it.", + "- Call knowhere_search when groundingPolicy requires citing source documents.", + "- Image/table asset paths in retrieved chunks represent connected assets that prepareAnswer will resolve into table HTML and image inputs. Do not refine the search solely because those assets have not been expanded yet; refine only when the source content needed to answer is missing.", + "- After each search that returns new evidence, call retainEvidence before searching again or finalizing.", + "- When calling knowhere_search after a previous search in this turn, set gapReason to what the previous search lacked and how the new query will fill that gap.", + "- Refine knowhere_search at most twice. After the second refined search, call prepareAnswer regardless of its result.", "- Do not treat every question as a document-retrieval task.", - "- For document-scoped searches, use includeDocumentIds/excludeDocumentIds only with verified IDs from source context or prior search results. If IDs are unknown, preserve the document requirement in query so Knowhere can locate it. Never invent IDs or substitute filenames. Exclusions win; an empty includeDocumentIds means no documents.", + "- For document-scoped searches, use includeDocumentIds/excludeDocumentIds only with verified IDs from prior search results. If IDs are unknown, preserve the document requirement in query so Knowhere can locate it. Never invent IDs or substitute filenames. Exclusions win; an empty includeDocumentIds means no documents.", "", "Context rules:", "- If the current user request is unrelated to prior turns, set carryHistory to none and do not reuse prior topics.", @@ -1326,19 +1215,16 @@ export function buildHarnessSystemPrompt(turn: AgentTurnInput): string { "- Final output is the OutputManifest passed to finalize, not freeform tool JSON or trailing text.", "- artifacts with display=true are the exact images/tables shown. Never display every candidate; honor constraints.desiredCount / maxCount.", "- Use type=derived_table only for tables you create from evidence; every derived_table.sourceRefs entry must reference evidence in the ledger.", - "- citations is a list of { pick }. pick is the 1-based number on the search chunk you are using. Notebook writes the citation list from those picks. Do not pass documentId or evidence refs as citations.", + "- citations is a list of { pick }. pick is the 1-based number on the search chunk you are using, and it must have been retained. Notebook writes the citation list from those picks. Do not pass documentId or evidence refs as citations.", "- Place [[cite:n]] immediately after the supported claim. n is the 1-based index into the citations array passed to finalize.", "- Write one marker per index: [[cite:1]] [[cite:3]] [[cite:5]]. Never group indices as [[cite:1, 3, 5]].", "- Do not write title/pN, [1], Markdown footnotes, or [Source N: ...] in the answer text. Notebook renders chips from [[cite:n]] and citation metadata.", "- Repeat [[cite:n]] when another claim uses the same page. Do not collapse same-page citations to one row.", "- If you have no supporting evidence pick, omit citations and list the gap in unresolved.", - "- inspectImage observations are inspection notes, not new source refs. Final citations and displayed image artifacts must use the original retrieved image asset refs.", - "- Do not finalize cited page/image assets from chunk text alone when inspectImage is available. Inspect those asset refs first, then write the answer using the inspection notes.", - "- If text evidence identifies a relevant page/image but does not include the exact fact, inspect the returned image asset for OCR/detail before saying the answer is unavailable.", + "- Images in retained evidence are embedded directly in the assembled user message. Read them there; do not call a separate image-inspection step.", "- If evidence is insufficient, list it in unresolved instead of fabricating facts.", `Surface: ${turn.surface}`, `Output capabilities: ${JSON.stringify(turn.outputCapabilities)}`, - turn.sourceContext ? `Searchable source context:\n${turn.sourceContext}` : "", ] .filter((line): line is string => line.length > 0) .join("\n") @@ -1370,15 +1256,4 @@ function formatRecentTurnIndex(turn: AgentTurnInput): string { }) return ["Recent turn index:", ...lines].join("\n") } - -function buildFallbackManifest(text: string): OutputManifest { - return { - text, - citations: [], - memoryCitations: [], - artifacts: [], - unresolved: text ? [] : ["The agent did not finalize an output manifest."], - } -} - export type { HarnessTrace } diff --git a/src/agent-harness/types.ts b/src/agent-harness/types.ts index b5a9416..1d7f9ae 100644 --- a/src/agent-harness/types.ts +++ b/src/agent-harness/types.ts @@ -68,7 +68,6 @@ export type AgentTurnInput = { readonly userText: string readonly recentTurns: readonly AgentTurn[] readonly localContext?: string - readonly sourceContext?: string readonly outputCapabilities: { readonly text: boolean readonly image: boolean @@ -94,6 +93,7 @@ export type KnowhereSearchRequest = Pick< readonly excludeDocumentIds?: string[] readonly targetContent?: KnowhereSearchTargetContent readonly purpose?: string + readonly gapReason?: string } export type KnowhereToolRuntime = { @@ -102,6 +102,20 @@ export type KnowhereToolRuntime = { ) => Promise } +export type ConnectedAssetLookup = { + readonly documentId: string + readonly chunkId: string + readonly type: "image" | "table" +} + +export type ResolvedConnectedAsset = ConnectedAssetLookup & { + readonly assetUrl: string +} + +export type ResolveConnectedAssets = ( + lookups: readonly ConnectedAssetLookup[], +) => Promise + export const memorySearchKinds = [ "indicator_pref", "stance", @@ -156,7 +170,6 @@ export type EvidenceChunk = { readonly sectionPath?: string | null } readonly revisionKey?: string | null - readonly assetRef?: string readonly assetUrl?: string } @@ -223,6 +236,11 @@ export type InspectImages = ( input: ImageInspectionRequest, ) => Promise +export type PendingRetentionRange = { + readonly startPick: number + readonly endPick: number +} + export type EvidenceLedgerSnapshot = { readonly retrievalCount: number readonly chunks: readonly EvidenceChunk[] @@ -231,6 +249,8 @@ export type EvidenceLedgerSnapshot = { readonly stopReasons: readonly string[] readonly failureReasons: readonly string[] readonly decisionTraces: readonly unknown[] + readonly retainedPicks: readonly number[] + readonly pendingRetention: PendingRetentionRange | null } export type OutputCitation = { diff --git a/src/app/api/demo-sources/materialize/route.test.ts b/src/app/api/demo-sources/materialize/route.test.ts index 3a3be24..618dc2c 100644 --- a/src/app/api/demo-sources/materialize/route.test.ts +++ b/src/app/api/demo-sources/materialize/route.test.ts @@ -242,6 +242,7 @@ function makeSource( originalBlobUrl: "https://example.com/tsla-q4-2025.pdf", demoKey: "demo-tsla-q4-2025", chunkCount: null, + folderId: null, createdAt: new Date("2026-05-10T00:00:00.000Z"), updatedAt: new Date("2026-05-10T00:00:00.000Z"), deletedAt: null, diff --git a/src/app/api/folders/[folderId]/route.ts b/src/app/api/folders/[folderId]/route.ts new file mode 100644 index 0000000..dc680e8 --- /dev/null +++ b/src/app/api/folders/[folderId]/route.ts @@ -0,0 +1,45 @@ +import type { NextRequest, NextResponse } from "next/server" + +import { folderRouteRequest } from "@/domains/folders/route-request" +import { folderRouteService } from "@/domains/folders/route-service" +import { withApiErrorResponse } from "@/lib/api-error-response" +import { nextRouteResponse } from "@/lib/next-route-response" + +type RouteContext = { + params: Promise<{ + folderId: string + }> +} + +export async function PATCH( + request: NextRequest, + context: RouteContext, +): Promise { + return withApiErrorResponse("folders:update", async () => { + const { folderId } = await context.params + const updateRequest = await folderRouteRequest.readUpdateFolder({ request }) + if (!updateRequest.ok) { + return nextRouteResponse.toNextResponse(updateRequest.result) + } + + return nextRouteResponse.toNextResponse( + await folderRouteService.updateFolder({ + folderId, + name: updateRequest.name, + parentId: updateRequest.parentId, + }), + ) + }) +} + +export async function DELETE( + _request: NextRequest, + context: RouteContext, +): Promise { + return withApiErrorResponse("folders:delete", async () => { + const { folderId } = await context.params + return nextRouteResponse.toNextResponse( + await folderRouteService.deleteFolder({ folderId }), + ) + }) +} diff --git a/src/app/api/folders/route.ts b/src/app/api/folders/route.ts new file mode 100644 index 0000000..cd9e6ad --- /dev/null +++ b/src/app/api/folders/route.ts @@ -0,0 +1,28 @@ +import type { NextRequest, NextResponse } from "next/server" + +import { folderRouteRequest } from "@/domains/folders/route-request" +import { folderRouteService } from "@/domains/folders/route-service" +import { withApiErrorResponse } from "@/lib/api-error-response" +import { nextRouteResponse } from "@/lib/next-route-response" + +export async function GET(): Promise { + return withApiErrorResponse("folders:list", async () => + nextRouteResponse.toNextResponse(await folderRouteService.listFolders()), + ) +} + +export async function POST(request: NextRequest): Promise { + return withApiErrorResponse("folders:create", async () => { + const createRequest = await folderRouteRequest.readCreateFolder({ request }) + if (!createRequest.ok) { + return nextRouteResponse.toNextResponse(createRequest.result) + } + + return nextRouteResponse.toNextResponse( + await folderRouteService.createFolder({ + name: createRequest.name, + parentId: createRequest.parentId, + }), + ) + }) +} diff --git a/src/app/api/sources/[sourceId]/chunks/route.test.ts b/src/app/api/sources/[sourceId]/chunks/route.test.ts index e38d9a8..3ff4188 100644 --- a/src/app/api/sources/[sourceId]/chunks/route.test.ts +++ b/src/app/api/sources/[sourceId]/chunks/route.test.ts @@ -312,6 +312,7 @@ describe("GET /api/sources/[sourceId]/chunks", () => { originalBlobUrl: "/api/demo-sources/demo-tsla-q4-2025/original", demoKey: "demo-tsla-q4-2025", chunkCount: null, + folderId: null, createdAt: new Date("2026-05-10T00:00:00.000Z"), updatedAt: new Date("2026-05-10T00:00:00.000Z"), deletedAt: null, @@ -847,6 +848,7 @@ function makeReadySource(overrides: Record) { originalBlobUrl: null, demoKey: null, chunkCount: null, + folderId: null, createdAt: new Date("2026-05-10T00:00:00.000Z"), updatedAt: new Date("2026-05-10T00:00:00.000Z"), deletedAt: null, diff --git a/src/app/api/sources/[sourceId]/route.test.ts b/src/app/api/sources/[sourceId]/route.test.ts index b9cbcd6..034d677 100644 --- a/src/app/api/sources/[sourceId]/route.test.ts +++ b/src/app/api/sources/[sourceId]/route.test.ts @@ -344,6 +344,7 @@ describe("PATCH /api/sources/[sourceId]", () => { "https://store.public.blob.vercel-storage.com/source-uploads/upload_1/document.pdf", demoKey: null, chunkCount: null, + folderId: null, createdAt: new Date("2026-05-10T00:00:00Z"), updatedAt: new Date("2026-05-10T00:00:00Z"), deletedAt: null, @@ -377,6 +378,7 @@ describe("PATCH /api/sources/[sourceId]", () => { "https://store.public.blob.vercel-storage.com/source-uploads/upload_1/document.pdf", demoKey: null, chunkCount: null, + folderId: null, createdAt: new Date("2026-05-10T00:00:00Z"), updatedAt: new Date("2026-05-10T00:00:00Z"), deletedAt: null, diff --git a/src/app/api/sources/[sourceId]/route.ts b/src/app/api/sources/[sourceId]/route.ts index 9d05d81..0104b42 100644 --- a/src/app/api/sources/[sourceId]/route.ts +++ b/src/app/api/sources/[sourceId]/route.ts @@ -31,7 +31,11 @@ export async function PATCH( const result = mutationRequest.mutation.kind === "archive" ? await sourceRouteService.archiveSource(mutationRequest.mutation.input) - : await sourceRouteService.retrySource(mutationRequest.mutation.input) + : mutationRequest.mutation.kind === "retry" + ? await sourceRouteService.retrySource(mutationRequest.mutation.input) + : await sourceRouteService.assignSourceFolder( + mutationRequest.mutation.input, + ) return nextRouteResponse.toNextResponse(result) } diff --git a/src/app/api/sources/route.test.ts b/src/app/api/sources/route.test.ts index 154efa2..edc83a2 100644 --- a/src/app/api/sources/route.test.ts +++ b/src/app/api/sources/route.test.ts @@ -76,6 +76,7 @@ const source: Source = { originalBlobUrl: null, demoKey: null, chunkCount: null, + folderId: null, createdAt: new Date("2026-05-10T00:00:00Z"), updatedAt: new Date("2026-05-10T00:00:00Z"), deletedAt: null, diff --git a/src/components/folder-destination-menu.tsx b/src/components/folder-destination-menu.tsx new file mode 100644 index 0000000..6ad20a4 --- /dev/null +++ b/src/components/folder-destination-menu.tsx @@ -0,0 +1,56 @@ +"use client" + +import type { ReactElement } from "react" +import { Folder } from "lucide-react" + +import { + DropdownMenuItem, + DropdownMenuLabel, + DropdownMenuSeparator, +} from "@/components/ui/dropdown-menu" +import { getFolderPath } from "@/domains/folders/tree" +import type { FolderView } from "@/domains/folders/types" + +export type FolderDestinationMenuProps = { + readonly folders: readonly FolderView[] + readonly excludedFolderIds?: ReadonlySet + readonly onSelect: (folderId: string | null) => void +} + +export function FolderDestinationMenu({ + folders, + excludedFolderIds, + onSelect, +}: FolderDestinationMenuProps): ReactElement { + const destinations = folders.filter( + (folder) => !excludedFolderIds?.has(folder.id), + ) + + return ( + <> + Move to + onSelect(null)}> + All sources + + {destinations.length > 0 ? : null} + {destinations.map((folder) => ( + onSelect(folder.id)} + > + + {formatFolderDestination(folders, folder.id)} + + ))} + + ) +} + +function formatFolderDestination( + folders: readonly FolderView[], + folderId: string, +): string { + return getFolderPath(folders, folderId) + .map((folder) => folder.name) + .join(" / ") +} diff --git a/src/components/folder-panel-state.ts b/src/components/folder-panel-state.ts new file mode 100644 index 0000000..16bffca --- /dev/null +++ b/src/components/folder-panel-state.ts @@ -0,0 +1,49 @@ +import { + getFolderPath, + listChildFolders, + listSourcesInFolder, +} from "@/domains/folders/tree" +import type { FolderView } from "@/domains/folders/types" +import type { SourceView } from "@/domains/sources/types" + +export type FolderListItem = + | { + readonly kind: "folder" + readonly folder: FolderView + } + | { + readonly kind: "source" + readonly source: SourceView + } + +export function getFolderListItems( + folders: readonly FolderView[], + sources: readonly SourceView[], + currentFolderId: string | null, +): FolderListItem[] { + return [ + ...listChildFolders(folders, currentFolderId).map( + (folder): FolderListItem => ({ kind: "folder", folder }), + ), + ...listSourcesInFolder(sources, currentFolderId).map( + (source): FolderListItem => ({ kind: "source", source }), + ), + ] +} + +export function getResolvedCurrentFolderId( + folders: readonly FolderView[], + currentFolderId: string | null, +): string | null { + if (!currentFolderId) return null + return folders.some((folder) => folder.id === currentFolderId) + ? currentFolderId + : null +} + +export function getFolderBreadcrumb( + folders: readonly FolderView[], + currentFolderId: string | null, +): FolderView[] { + return getFolderPath(folders, currentFolderId) +} diff --git a/src/components/folder-row.tsx b/src/components/folder-row.tsx new file mode 100644 index 0000000..0f749df --- /dev/null +++ b/src/components/folder-row.tsx @@ -0,0 +1,115 @@ +"use client" + +import type { ReactElement } from "react" +import { Folder, MoreHorizontal, Pencil, Trash2 } from "lucide-react" + +import { FolderDestinationMenu } from "@/components/folder-destination-menu" +import { + DropdownMenu, + DropdownMenuContent, + DropdownMenuItem, + DropdownMenuSeparator, + DropdownMenuSub, + DropdownMenuSubContent, + DropdownMenuSubTrigger, + DropdownMenuTrigger, +} from "@/components/ui/dropdown-menu" +import { Spinner } from "@/components/ui/spinner" +import { getDescendantFolderIds } from "@/domains/folders/tree" +import type { FolderView } from "@/domains/folders/types" + +export type FolderRowProps = { + readonly folder: FolderView + readonly folders: readonly FolderView[] + readonly isDeleting?: boolean + readonly isNarrow?: boolean + readonly onDelete: (folderId: string) => void + readonly onMove: (folderId: string, parentId: string | null) => void + readonly onOpen: (folderId: string) => void + readonly onRename: (folderId: string) => void +} + +export function FolderRow({ + folder, + folders, + isDeleting = false, + isNarrow = false, + onDelete, + onMove, + onOpen, + onRename, +}: FolderRowProps): ReactElement { + const excludedFolderIds = new Set(getDescendantFolderIds(folders, folder.id)) + + return ( +
+ + + + + + + onRename(folder.id)}> + + Rename + + + Move to… + + onMove(folder.id, parentId)} + /> + + + + onDelete(folder.id)}> + + Delete + + + +
+ ) +} diff --git a/src/components/source-row.tsx b/src/components/source-row.tsx index 473453f..4f95196 100644 --- a/src/components/source-row.tsx +++ b/src/components/source-row.tsx @@ -2,10 +2,18 @@ import type { ReactElement } from "react"; import Link from "next/link"; -import { FileText, ListTree, Plus, RotateCcw, Trash2 } from "lucide-react"; +import { FileText, FolderInput, ListTree, Plus, RotateCcw, Trash2 } from "lucide-react"; +import { FolderDestinationMenu } from "@/components/folder-destination-menu"; import { Checkbox } from "@/components/ui/checkbox"; +import { + DropdownMenu, + DropdownMenuContent, + DropdownMenuTrigger, +} from "@/components/ui/dropdown-menu"; import { Spinner } from "@/components/ui/spinner"; +import { canAssignSourceToFolder } from "@/domains/folders/tree"; +import type { FolderView } from "@/domains/folders/types"; import type { SourceView } from "@/domains/sources/types"; export type SourceRowProps = { @@ -20,6 +28,9 @@ export type SourceRowProps = { readonly onRetryClick?: (sourceId: string) => void; readonly onSelect: () => void; readonly onToggleIncluded?: (sourceId: string, included: boolean) => void; + readonly folders?: readonly FolderView[]; + readonly isMoving?: boolean; + readonly onMoveToFolder?: (sourceId: string, folderId: string | null) => void; readonly source: SourceView; }; @@ -36,8 +47,12 @@ export function SourceRow({ chunkTreeHref, isArchiving, isRetrying = false, + folders = [], + isMoving = false, + onMoveToFolder, }: SourceRowProps): ReactElement { const isReady = source.status === "ready"; + const canMove = canAssignSourceToFolder(source) && onMoveToFolder !== undefined; const isBusy = source.status === "uploading" || source.status === "parsing"; const isFailed = source.status === "failed"; const canRetry = isFailed && source.originalFile !== undefined; @@ -146,6 +161,31 @@ export function SourceRow({ {isNarrow ? null : "Add"} )} + {canMove ? ( + + + + + + onMoveToFolder?.(source.id, folderId)} + /> + + + ) : null} {canRetry && onRetryClick ? ( + + +
{onLoginClick ? ( @@ -199,6 +346,25 @@ export function SourcesPanel({ sourceCountSnapshot={sourceCountSnapshot} /> )} + {onCreateFolder && !onLoginClick ? ( + + ) : null}
@@ -221,39 +387,70 @@ export function SourcesPanel({ {isNarrow ? null : "open library"}
+ {onOpenFolder ? ( + + ) : null} - {workspaceSources.length === 0 ? ( + {listItems.length === 0 ? ( ) : (
- {sourcePagination.sources.map((source) => ( - - onSelectSource?.( - sourcePanelState.getNextSelectedSourceId({ - sourceId: source.id, - }), - ) - } - onToggleIncluded={onToggleIncluded} - onArchiveClick={ - onArchiveSource ? setConfirmSourceId : undefined - } - onRetryClick={onRetrySource} - isArchiving={archivingSourceIdSet.has(source.id)} - isRetrying={retryingSourceIdSet.has(source.id)} - isNarrow={isNarrow} - /> - ))} + {sourcePagination.items.map((item) => + item.kind === "folder" ? ( + { + void onMoveFolder?.(folderId, parentId); + }} + onOpen={(folderId) => onOpenFolder?.(folderId)} + onRename={(folderId) => { + const folder = folderRows.find( + (candidate) => candidate.id === folderId, + ); + setFolderName(folder?.name ?? ""); + setFolderNameDialog({ kind: "rename", folderId }); + }} + /> + ) : ( + + onSelectSource?.( + sourcePanelState.getNextSelectedSourceId({ + sourceId: item.source.id, + }), + ) + } + onToggleIncluded={onToggleIncluded} + onArchiveClick={ + onArchiveSource ? setConfirmSourceId : undefined + } + onRetryClick={onRetrySource} + folders={folderRows} + isMoving={movingSourceIdSet.has(item.source.id)} + onMoveToFolder={onMoveSourceToFolder} + isArchiving={archivingSourceIdSet.has(item.source.id)} + isRetrying={retryingSourceIdSet.has(item.source.id)} + isNarrow={isNarrow} + /> + ), + )}
)}
- {workspaceSources.length > sourceListPageSize ? ( + {listItems.length > sourceListPageSize ? ( source.id === selectedSourceId, + const selectedIndex = items.findIndex( + (item) => item.kind === "source" && item.source.id === selectedSourceId, ); return selectedIndex >= 0 ? getSourcePageForIndex(selectedIndex) : null; } +function FolderBreadcrumb({ + folders, + onOpenFolder, +}: { + readonly folders: readonly FolderView[]; + readonly onOpenFolder: (folderId: string | null) => void; +}): ReactElement { + return ( + + + + {folders.length === 0 ? ( + All sources + ) : ( + { + event.preventDefault(); + onOpenFolder(null); + }} + > + All sources + + )} + + {folders.map((folder, index) => ( + + + {index === folders.length - 1 ? ( + {folder.name} + ) : ( + { + event.preventDefault(); + onOpenFolder(folder.id); + }} + > + {folder.name} + + )} + + ))} + + + ); +} + function getChunkTreeHref(source: SourceView): string | undefined { if (source.documentPresentation?.kind === "page-assets") return undefined; return source.documentId diff --git a/src/components/ui/breadcrumb.tsx b/src/components/ui/breadcrumb.tsx new file mode 100644 index 0000000..58858ea --- /dev/null +++ b/src/components/ui/breadcrumb.tsx @@ -0,0 +1,113 @@ +import { cn } from "@/lib/utils" +import { Slot } from "@radix-ui/react-slot" +import { ChevronRight, MoreHorizontal } from "lucide-react" +import * as React from "react" + +const Breadcrumb = React.forwardRef< + HTMLElement, + React.ComponentPropsWithoutRef<"nav"> & { + separator?: React.ReactNode + } +>(({ ...props }, ref) =>