diff --git a/.env.example b/.env.example index 8c5049f6..674fe660 100644 --- a/.env.example +++ b/.env.example @@ -209,4 +209,31 @@ LLM_KEY_ENCRYPTION_KEY= # (case-insensitive), their role_id is set to 'admin'. Idempotent — leave # set forever; only the FIRST signup with this email is promoted. To add # a second admin later, use the admin UI role management tab. -INITIAL_ADMIN_EMAIL= \ No newline at end of file +INITIAL_ADMIN_EMAIL= + +# KB retrieval knobs (issue #13 v3). All server-only — backend tool +# surface and mention resolver read these via lib/kb/env.ts. Defaults +# reflect the community survey in `.claude/13-kb-v3.md`: +# LangChain EnsembleRetriever / LlamaIndex QueryFusionRetriever / +# Haystack DocumentJoiner — per-leg 50, fused 5–10; Mastra = 3; +# Mastra + Vercel AI SDK RAG chunks = 512–1024 tokens. +# KB_MENTION_TOPK_DEFAULT — chunks per single @-mention. 5. +# KB_MENTION_TOPK_MAX — per-mention upper bound when user overrides. 20. +# KB_MENTION_TOKEN_BUDGET — total token cap across multi-mention turns. +# Per-mention topK is rebudgeted as +# ceil(BUDGET / (CHUNK_MAX_CHARS/4 * mentions)). +# 8192. +# KB_HYBRID_TOPK_DEFAULT — fused topK for the search_kb tool. 8. +# KB_HYBRID_TOPK_MAX — fused topK upper bound for the tool. 20. +# KB_CHUNK_MAX_CHARS — per-chunk truncation before stuffing into +# the LLM prompt; ~512 tokens at 4 chars/token. +# 2000. +# KB_RERANK_MIN_SCORE — minimum rerank relevance score threshold for filtering. +# Chunks below this score are filtered out. 0.4. +KB_MENTION_TOPK_DEFAULT=5 +KB_MENTION_TOPK_MAX=20 +KB_MENTION_TOKEN_BUDGET=8192 +KB_HYBRID_TOPK_DEFAULT=8 +KB_HYBRID_TOPK_MAX=20 +KB_CHUNK_MAX_CHARS=2000 +KB_RERANK_MIN_SCORE=0.4 \ No newline at end of file diff --git a/.github/workflows/CI.yml b/.github/workflows/CI.yml index 3bad4051..34fa2f74 100644 --- a/.github/workflows/CI.yml +++ b/.github/workflows/CI.yml @@ -74,7 +74,13 @@ jobs: # to bring the schema up — the build-time DB can be empty otherwise. services: postgres: - image: postgres:16-alpine + # pgvector/pgvector:pg16 (Debian-based, FROM postgres:16) ships + # the vector extension pre-installed — required by migration + # 0005_*.sql (CREATE EXTENSION vector). Stock postgres:16-alpine + # does not bundle it. Wire-format (POSTGRES_*, PGDATA, init + # scripts) is unchanged from upstream so all 6 build/test steps + # keep working without further edits. + image: pgvector/pgvector:pg16 env: POSTGRES_USER: postgres POSTGRES_PASSWORD: postgres @@ -109,7 +115,13 @@ jobs: timeout-minutes: 10 services: postgres: - image: postgres:16-alpine + # pgvector/pgvector:pg16 (Debian-based, FROM postgres:16) ships + # the vector extension pre-installed — required by migration + # 0005_*.sql (CREATE EXTENSION vector). Stock postgres:16-alpine + # does not bundle it. Wire-format (POSTGRES_*, PGDATA, init + # scripts) is unchanged from upstream so all 6 build/test steps + # keep working without further edits. + image: pgvector/pgvector:pg16 env: POSTGRES_USER: postgres POSTGRES_PASSWORD: postgres diff --git a/.gitignore b/.gitignore index fd25936b..314035a0 100644 --- a/.gitignore +++ b/.gitignore @@ -31,6 +31,9 @@ yarn-debug.log* yarn-error.log* .pnpm-debug.log* +# v1 KB JSON store — temporary local storage, replaced by DB tables in v2 +/.kb-store/ + # env files (can opt-in for committing if needed) .env* !.env.test diff --git a/CLAUDE.md b/CLAUDE.md index 32754387..90537009 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -4,23 +4,24 @@ Guidance for Claude Code in this repo. Features, layout, env vars, tech stack ## Docs index -| Topic | File | -| ----------------------------------------------- | ----------------------- | -| Features, quick-start, layout, env vars, deps | `README.md` | -| Marketing landing page sections + assets | `docs/LANDING.md` | -| Every HTTP endpoint under `app/api/` | `docs/APIS.md` | -| Memory + thread-summarize design | `docs/MEMORY.md` | -| Observability panel design + retention | `docs/OBSERVABILITY.md` | -| LangGraph tool inventory + frontend card wiring | `docs/TOOLS.md` | -| Interrupt-driven tool flow contract | `docs/INTERRUPT.md` | -| Attachments backing (R2 + presign) design | `docs/ATTACHMENTS.md` | -| Auth setup, OAuth, troubleshooting | `docs/AUTH.md` | -| DB schema, ownership, indexes | `docs/DB.md` | -| Provider registry (round-robin + withFallbacks) | `docs/PROVIDERS.md` | -| Credit / LLM-call quota system | `docs/CREDIT.md` | -| Admin UI + endpoints for providers/users/roles | `docs/ADMIN.md` | -| Production deploys (Docker, env, Caddy) | `docs/DEPLOY.md` | -| CI/CD, Docker, deploys | `docs/CI.md` | +| Topic | File | +| ------------------------------------------------ | ------------------------ | +| Features, quick-start, layout, env vars, deps | `README.md` | +| Marketing landing page sections + assets | `docs/LANDING.md` | +| Every HTTP endpoint under `app/api/` | `docs/APIS.md` | +| Memory + thread-summarize design | `docs/MEMORY.md` | +| Knowledge base, ingestion, hybrid search, Rerank | `docs/KNOWLEDGE_BASE.md` | +| Observability panel design + retention | `docs/OBSERVABILITY.md` | +| LangGraph tool inventory + frontend card wiring | `docs/TOOLS.md` | +| Interrupt-driven tool flow contract | `docs/INTERRUPT.md` | +| Attachments backing (R2 + presign) design | `docs/ATTACHMENTS.md` | +| Auth setup, OAuth, troubleshooting | `docs/AUTH.md` | +| DB schema, ownership, indexes | `docs/DB.md` | +| Provider registry (round-robin + withFallbacks) | `docs/PROVIDERS.md` | +| Credit / LLM-call quota system | `docs/CREDIT.md` | +| Admin UI + endpoints for providers/users/roles | `docs/ADMIN.md` | +| Production deploys (Docker, env, Caddy) | `docs/DEPLOY.md` | +| CI/CD, Docker, deploys | `docs/CI.md` | API changes update `docs/APIS.md` in the same commit (rule 1). Tool changes update `docs/TOOLS.md` (rule 10). @@ -75,7 +76,7 @@ Non-negotiable. Every change. ## Things to know before editing -- Graph id `agent` is in `langgraph.json`, `LANGGRAPH_ASSISTANT_ID` (`.env.example`, surfaced to client via `window.__CONFIG__` — see rule #12), and `unstable_createLangGraphStream({ assistantId })`. Keep aligned. +- Graph id `agent` is in `langgraph.json`, `LANGGRAPH_ASSISTANT_ID` (`.env.example`, surfaced to client via `window.__CONFIG__` — see rule #12), and `unstable_createLangGraphStream({ assistantId })`. Keep aligned. The `kbAgent` assistant id (KB ingestion, `langgraph.json` + `lib/kb/ingest.ts:fireIngestionRun` dispatch) follows the same sync rule — there is no env-driven alias for it, but renaming it requires updating `langgraph.json` AND `fireIngestionRun` together. - `app/api/[..._path]/route.ts` proxy uses `runtime = "nodejs"` (was edge) — `withAuth` needs Node `net` for Postgres session reads. - `components.json` declares a `@assistant-ui` registry at `https://r.assistant-ui.com/{name}.json` for `shadcn`-style component adds. - `feat/*` branches: `git fetch origin main` and merge if main moved before committing — see [[feature-branch-tracks-main]]. diff --git a/README.md b/README.md index 38f4453d..6cf6d676 100644 --- a/README.md +++ b/README.md @@ -16,6 +16,7 @@ A self-hostable chat app (this repo: `langgraph-app`) that streams tokens from a - **Crypto sub-agent**: price, NFT holdings (5-chain gallery via Alchemy Portfolio), and a simulated swap flow against an auto-funded Mock Coin balance. - **Observability panel**: every LLM / Tool / Chain / Node span is captured by a `BaseCallbackHandler` and persisted to a `observability_spans` Postgres table. Each assistant message shows an icon button that opens a per-turn waterfall — duration, token usage, nested parent/child spans. The list endpoint is server-transformed (panel never carries the raw collector payload); per-row click lazy-loads the full span via a dedicated detail endpoint. See [docs/OBSERVABILITY.md](docs/OBSERVABILITY.md). - **Chat attachments**: assistant-ui's `AttachmentAdapter` plus a presigned PUT to Cloudflare R2 — the browser uploads bytes directly to R2, nothing traverses Next.js. Lazy-register on missing env (mirrors DENO / ALCHEMY). See [docs/ATTACHMENTS.md](docs/ATTACHMENTS.md) for the key convention, messageId-deferred decision, and Content-Disposition XSS guard. +- **Knowledge Base & Hybrid Search**: PDF ingestion, page screenshot rendering, text chunking, and pgvector-backed three-leg RRF (Keyword, Vector, Tag) hybrid search combined with semantic Reranking (Cohere/Jina), `@` mention resolution (with automatic fallback to full markdown when chunks are not ready), dynamic budget scaling, and iterative search. See [docs/KNOWLEDGE_BASE.md](docs/KNOWLEDGE_BASE.md). - **Per-LLM-call credit quota**: every successful call is metered against a UTC-aligned rolling-window cap read from `role.creditLimit` / `role.windowHours`. Enforcement lives at the `/api/[..._path]` proxy — when the cap is hit, the proxy synthesizes a `show_credit_card` SSE stream and the chat UI renders the credit-limit-reached card inline. The call log backs a per-user history (Settings → Credits) and an admin-managed rate config. See [docs/CREDIT.md](docs/CREDIT.md). - **Admin console**: a single `/admin` page with three tabs — Providers (registry + encrypted API keys + per-model rates), Roles (credit caps + window length), Users (role assignment, ban with immediate session revoke, delete). The first admin is bootstrapped via `INITIAL_ADMIN_EMAIL`. See [docs/ADMIN.md](docs/ADMIN.md). @@ -310,6 +311,7 @@ Test database stays isolated from dev — never put production-like data in `lan - [`docs/APIS.md`](docs/APIS.md) — HTTP endpoint reference. Update whenever a route under `app/api/` changes. - [`docs/LANDING.md`](docs/LANDING.md) — marketing landing at `/`: per-section file map, asset inventory, motion keyframes, route-group naming rule, frontend test layout. - [`docs/MEMORY.md`](docs/MEMORY.md) — memory + thread-summarize design: dual-graph topology, `` + `` recall, `save_memory` RFC 6902 patches, store-anchored trigger window math, Memory tab UI, security stance. +- [`docs/KNOWLEDGE_BASE.md`](docs/KNOWLEDGE_BASE.md) — knowledge base design: ingestion pipeline, three-leg RRF hybrid search (Keyword, Vector, Tag), semantic Reranking and score filtering, `@` mention resolution (Meta vs Full Markdown mode), mention budgeting, and iterative search. - [`docs/OBSERVABILITY.md`](docs/OBSERVABILITY.md) — observability panel design: callback handler wiring, `observability_spans` schema, server-side transform + aggregate, lazy-loaded row detail, security/redaction, retention config, and curl examples. - [`docs/TOOLS.md`](docs/TOOLS.md) — LangGraph tool inventory and frontend card wiring. Update whenever a tool or card is added/removed/rerouted. - [`docs/INTERRUPT.md`](docs/INTERRUPT.md) — interrupt-driven tool flows (ask_location, connect_wallet, place_crypto_order, get_order_status) — the two runtime paths the cards can take. diff --git a/app/admin/admin-tabs.tsx b/app/admin/admin-tabs.tsx index 218de54d..b7bc11a1 100644 --- a/app/admin/admin-tabs.tsx +++ b/app/admin/admin-tabs.tsx @@ -10,6 +10,7 @@ import { Button } from "@/components/ui/button"; import { Input } from "@/components/ui/input"; import { Badge } from "@/components/ui/badge"; import { Card, CardContent, CardDescription, CardHeader, CardTitle } from "@/components/ui/card"; +import { Checkbox } from "@/components/ui/checkbox"; import { Separator } from "@/components/ui/separator"; import { Switch } from "@/components/ui/switch"; import { @@ -32,11 +33,16 @@ import { Tooltip, TooltipContent, TooltipProvider, TooltipTrigger } from "@/comp import { toast } from "sonner"; type PublicProviderApiKey = { name: string }; +type ModelKind = "chat" | "ocr" | "embed" | "extract" | "rerank"; type PublicModel = { name: string; enabled: boolean; inputPer1k: number; outputPer1k: number; + // ponytail: server-side default ["chat"] means an older row may not + // carry the field at all. Fall back to chat at render time so the + // table shows a value for every model, regardless of seed-vs-new. + kind?: ModelKind[]; }; type PublicProviderRow = { id: string; @@ -56,6 +62,12 @@ type RoleRow = { createdAt: string; updatedAt: string; }; + +// ponytail: module-scope constant so add-mode `initialKind ?? ["chat"]` +// keeps the same array reference across renders — `useEffect` dep +// comparison uses Object.is, and a fresh literal each render would loop +// forever (Next dev StrictMode catches it as "Maximum update depth"). +const DEFAULT_KIND: ModelKind[] = ["chat"]; // ponytail: server-side join shape — user + role name inlined so the // table can render "Admin" / "User" labels without a second round-trip. // `roleName` is null when the FK points at a missing role (defensive — @@ -360,8 +372,9 @@ function ModelsSection({ providerId, models }: { providerId: string; models: Pub - - + + + {models.length === 0 ? ( - ) : ( models.map((m) => ( - + +
NameEnabledNameStatusKind Input / 1k Output / 1k @@ -370,19 +383,28 @@ function ModelsSection({ providerId, models }: { providerId: string; models: Pub
+ No models configured.
{m.name}{m.name} {m.enabled ? "Enabled" : "Disabled"} +
+ {(m.kind ?? ["chat"]).map((k) => ( + + {k.charAt(0).toUpperCase() + k.slice(1)} + + ))} +
+ 		{m.inputPer1k} {m.outputPer1k} @@ -511,11 +533,16 @@ function ModelDialog( const initialEnabled = isEdit ? props.model.enabled : true; const initialIn = isEdit ? String(props.model.inputPer1k) : "0.001"; const initialOut = isEdit ? String(props.model.outputPer1k) : "0.002"; + // ponytail: backend defaults kind to ["chat"] on POST when omitted, + // so the add-mode local seed mirrors the server contract and we don't + // need to send `kind` in the body. Edit-mode seeds from existing row. + const initialKind: ModelKind[] = isEdit ? (props.model.kind ?? DEFAULT_KIND) : DEFAULT_KIND; const [name, setName] = useState(initialName); const [enabled, setEnabled] = useState(initialEnabled); const [inputPer1k, setInputPer1k] = useState(initialIn); const [outputPer1k, setOutputPer1k] = useState(initialOut); + const [kind, setKind] = useState(initialKind); const [saving, start] = useTransition(); useEffect(() => { @@ -523,7 +550,18 @@ function ModelDialog( setEnabled(initialEnabled); setInputPer1k(initialIn); setOutputPer1k(initialOut); - }, [initialName, initialEnabled, initialIn, initialOut]); + setKind(initialKind); + }, [initialName, initialEnabled, initialIn, initialOut, initialKind]); + + const toggleKind = (k: ModelKind, checked: boolean) => { + setKind((prev) => { + const next = checked ? Array.from(new Set([...prev, k])) : prev.filter((x) => x !== k); + // ponytail: at least one kind must stay selected — unchecking the + // last box is a no-op so the user can never save a model with + // kind = [] (which the API rejects with 400 anyway). + return next.length === 0 ? prev : next; + }); + }; const save = () => { if (!name.trim()) { @@ -541,8 +579,8 @@ function ModelDialog( : `/api/admin/providers/${encodeURIComponent(props.providerId)}/models`; const method = isEdit ? "PATCH" : "POST"; const body = isEdit - ? { name: name.trim(), enabled, inputPer1k: inp, outputPer1k: out } - : { name: name.trim(), enabled, inputPer1k: inp, outputPer1k: out }; + ? { name: name.trim(), enabled, inputPer1k: inp, outputPer1k: out, kind } + : { name: name.trim(), enabled, inputPer1k: inp, outputPer1k: out, kind }; start(async () => { const r = await jsonFetch(path, { method, body: JSON.stringify(body) }); if (!r.ok) { @@ -580,15 +618,6 @@ function ModelDialog( className="font-mono" /> - + +
+ Kind +
+ {(["chat", "ocr", "embed", "extract", "rerank"] as const).map((k) => ( + + ))} +
+

+ At least one kind is required — chat for general inference, ocr for PDF vision, embed + for KB chunk vectors, extract for structured-output triples from chunks, rerank for + secondary retrieval reranking. +

+
); diff --git a/app/api/admin/providers/[id]/models/[modelName]/route.ts b/app/api/admin/providers/[id]/models/[modelName]/route.ts index 567ae14b..d2651df3 100644 --- a/app/api/admin/providers/[id]/models/[modelName]/route.ts +++ b/app/api/admin/providers/[id]/models/[modelName]/route.ts @@ -1,32 +1,21 @@ import { NextResponse } from "next/server"; import { eq } from "drizzle-orm"; -import { z } from "zod"; import { db } from "@/db/client"; import { provider } from "@/lib/provider/schema"; import { stripProviderSecrets } from "@/lib/provider/admin"; +import { modelPatchSchema } from "@/lib/credit/zod"; import { invalidateModelCache } from "@/lib/provider/model-registry"; import { withAuth } from "@/lib/auth/with-auth"; type ModelParams = { id: string; modelName: string }; -// ponytail: the input side requires enabled / inputPer1k / outputPer1k, so -// a partial PATCH has to lift those out of `.partial()` and re-require them -// individually — otherwise a PATCH with `{}` would no-op and the caller has -// no signal that "no fields were sent". -const ModelPatchBody = z.object({ - name: z.string().min(1).max(128).optional(), - enabled: z.boolean().optional(), - inputPer1k: z.number().min(0).optional(), - outputPer1k: z.number().min(0).optional(), -}); - export const PATCH = withAuth({ role: "admin" }, async (req, { params }) => { const [existing] = await db.select().from(provider).where(eq(provider.id, params.id)); if (!existing) return NextResponse.json({ code: "NOT_FOUND" }, { status: 404 }); const json = await req.json().catch(() => ({})); - const parsed = ModelPatchBody.safeParse(json); + const parsed = modelPatchSchema.safeParse(json); if (!parsed.success) { return NextResponse.json({ code: "BAD_REQUEST", error: parsed.error.issues }, { status: 400 }); } diff --git a/app/api/kb/documents/[id]/observability/route.ts b/app/api/kb/documents/[id]/observability/route.ts new file mode 100644 index 00000000..3fc3f00a --- /dev/null +++ b/app/api/kb/documents/[id]/observability/route.ts @@ -0,0 +1,50 @@ +import { desc, eq } from "drizzle-orm"; +import { NextResponse } from "next/server"; + +import { withAuth } from "@/lib/auth/with-auth"; +import { db } from "@/db/client"; +import { findKbDocumentById } from "@/lib/kb/queries"; +import { kbObservability } from "@/lib/kb/schema"; + +// ponytail: Settings → KB → doc row → Activity icon → popover data +// source. Reads the kb_observability table directly (no SDK call) so +// the popover sees runs from BOTH ingest paths: +// - standalone (Settings upload / reprocess): threadId = docId-派生 +// - chat (mainAgent → kbAgent subgraph): threadId = chat thread +// Every kbAgent invocation inserts a row in prepareKBDataNode with +// (docId, threadId, parentMessageId, source, mode, created_at), so the +// popover's "View runs" list IS the union. Per-run LangGraph status +// (running/pending/success/error) used to come from runs.list; that +// info lives in observability_spans now, surfaced via the sheet. + +type Params = { id: string }; + +export const GET = withAuth(async (_req, { user, params }) => { + // ponytail: rule #9 — cross-user doc id is 404, not 401/403, so + // callers can't enumerate which doc ids exist. + const doc = await findKbDocumentById(user.id, params.id); + if (!doc) { + return NextResponse.json({ code: "NOT_FOUND" }, { status: 404 }); + } + + const rows = await db + .select() + .from(kbObservability) + .where(eq(kbObservability.docId, doc.id)) + .orderBy(desc(kbObservability.createdAt)) + .limit(50); + + const runs = rows.map((r) => ({ + runId: r.runId, + threadId: r.threadId, + parentMessageId: r.parentMessageId, + source: r.source, + mode: r.mode, + createdAt: r.createdAt.toISOString(), + })); + + return NextResponse.json({ + doc_id: doc.id, + runs, + }); +}); diff --git a/app/api/kb/documents/[id]/reprocess/route.ts b/app/api/kb/documents/[id]/reprocess/route.ts new file mode 100644 index 00000000..6c9d068f --- /dev/null +++ b/app/api/kb/documents/[id]/reprocess/route.ts @@ -0,0 +1,236 @@ +import { NextResponse } from "next/server"; + +import { withAuth } from "@/lib/auth/with-auth"; +import { fireIngestionRun } from "@/lib/kb/ingest"; +import { getAttachmentForUser } from "@/lib/attachments/queries"; +import { + deleteKbChunksByDocumentId, + findKbDocumentById, + markFailedKbChunksRetryingByDocumentId, + resetKbDocumentForReprocess, + updateKbDocumentStatus, + withKbTx, +} from "@/lib/kb/queries"; + +// ponytail: Settings → KB → per-row "Refresh" button. Re-runs OCR + +// chunk + embed against the existing attachment for a doc that's +// already in the DB. The Settings UI shows status via 2s polling, so +// we fire-and-forget the run and return 202. +// +// Two modes via `?chunksOnly=true|false` query string: +// - default (full): wipe doc row back to "pending" + clear +// chunks, dispatch kbAgent which re-runs +// PDF render + OCR + chunk + embed. +// - chunksOnly=true: only clear chunks, leave doc row at its +// terminal status (success/failed), dispatch +// kbAgent with `mode: "chunksOnly"` so OCR +// is skipped and pages[].markdown is reused. +// +// Status guards: +// - status='pending' / 'parsing' → 409 PROCESSING — a run is already +// in flight, double-clicking the refresh button shouldn't kick off +// a parallel pipeline against the same attachment. +// Cross-user docs → 404 (no existence leak, same convention as +// /api/threads). +// +// ponytail: chunksOnly doc-status guard is stricter — we need a +// doc whose OCR already landed (status='success' or 'failed' with +// pages[].markdown populated). 409 NOT_READY surfaces the gap so +// the client can fall back to the default "Refresh" (full +// reprocess) which seeds pages. + +type Params = { id: string }; + +export const POST = withAuth(async (req, { user, params }) => { + // ponytail: withAuth hands us a plain `Request`; parse the query + // string with stdlib URL instead of NextRequest.nextUrl. + const modeParam = new URL(req.url).searchParams.get("mode"); + const chunksOnly = new URL(req.url).searchParams.get("chunksOnly") === "true"; + + let mode: "full" | "chunksOnly" | "retryFailed" | "retryFailedChunks" = "full"; + if (modeParam === "chunksOnly" || chunksOnly) { + mode = "chunksOnly"; + } else if (modeParam === "retryFailed") { + mode = "retryFailed"; + } else if (modeParam === "retryFailedChunks") { + // ponytail: 4th reprocess mode — only re-run chunks whose status + // is 'failed'. Successful chunks keep their id/embedding/ + // entities verbatim; doc.status stays 'success' so the live UI + // doesn't flicker to "Indexing" while chunks heal in the + // background. Surfaces the partial-chunk-failure case without + // paying OCR cost again. + mode = "retryFailedChunks"; + } + + const doc = await findKbDocumentById(user.id, params.id); + if (!doc) { + return NextResponse.json({ code: "NOT_FOUND" }, { status: 404 }); + } + if (doc.status === "pending" || doc.status === "parsing") { + return NextResponse.json({ code: "PROCESSING" }, { status: 409 }); + } + + // ponytail: chunksOnly and retryFailed both bypass re-rendering pages, + // requiring doc.pages to be populated first. + if (mode === "chunksOnly" || mode === "retryFailed") { + const pages = (doc.pages ?? []) as Array<{ markdown?: string }>; + if (pages.length === 0) { + return NextResponse.json( + { code: "NOT_READY", reason: "doc has no pages; run full reprocess first" }, + { status: 409 }, + ); + } + + if (mode === "chunksOnly") { + const hasUsableMarkdown = pages.some((p) => (p.markdown ?? "").trim().length > 0); + if (!hasUsableMarkdown) { + return NextResponse.json( + { code: "NOT_READY", reason: "doc has no pages[].markdown; run full reprocess first" }, + { status: 409 }, + ); + } + } + + // ponytail: wipe chunks inside one tx + await withKbTx(async (tx) => { + await deleteKbChunksByDocumentId(tx, doc.id); + }); + + if (mode === "retryFailed") { + // update document status to parsing so live UI gets polling feedback! + await updateKbDocumentStatus(user.id, doc.id, { + status: "parsing", + errorMessage: null, + }); + } + + try { + await fireIngestionRun({ + userId: user.id, + attachment: { + r2Key: "chunks-only-no-op", + contentType: doc.contentType, + name: doc.title, + }, + docId: doc.id, + title: doc.title, + source: "kb-reprocess", + mode, + }); + } catch (err) { + console.error( + `POST /api/kb/documents/[id]/reprocess?mode=${mode}: fireIngestionRun failed`, + err, + ); + } + + return NextResponse.json( + { docId: doc.id, chunksOnly: mode === "chunksOnly", mode }, + { status: 202 }, + ); + } + + // ponytail: retryFailedChunks — UPDATE failed chunks in place + // (status='parsing', clear error_message + entities, keep id/ + // ordinal/embedding/content). doc.status STAYS 'success' the + // whole time — flipping it to 'parsing' would flicker the UI + // badge and confuse the user ("did my doc break again?"). The + // chunk-level "Indexed N/M" rollup is the progress signal; + // doc.status reflects the macro OCR + chunk pipeline state, + // which hasn't actually changed. + // + // Why UPDATE instead of DELETE+INSERT: the DELETE+INSERT design + // had a race where pageToMarkdownNode skips under chunksOnly / + // retryFailed, so doc.pages is empty, fullMarkdown is empty, the + // IIFE inside generateChunkEmbedNode throws — and the DELETE has + // already run, leaving the doc with N-K chunks and no recovery + // path. In-place UPDATE makes the row indestructible. + if (mode === "retryFailedChunks") { + if (doc.status !== "success") { + return NextResponse.json( + { + code: "NOT_READY", + reason: "doc has no successfully indexed chunks yet; run full reprocess first", + }, + { status: 409 }, + ); + } + + await withKbTx(async (tx) => { + // ponytail: reset failed chunks to status='parsing' so the + // IIFE can find them and UPDATE them back to success/failed + // per-row. The row's id, ordinal, embedding, and content are + // preserved verbatim — embedding API is deterministic so the + // old vector is still valid for KB search even if the + // entity-extract LLM fails on this run. + await markFailedKbChunksRetryingByDocumentId(tx, doc.id); + }); + + try { + await fireIngestionRun({ + userId: user.id, + attachment: { + r2Key: "chunks-only-no-op", + contentType: doc.contentType, + name: doc.title, + }, + docId: doc.id, + title: doc.title, + source: "kb-reprocess", + mode, + }); + } catch (err) { + console.error( + `POST /api/kb/documents/[id]/reprocess?mode=retryFailedChunks: fireIngestionRun failed`, + err, + ); + } + + return NextResponse.json({ docId: doc.id, mode }, { status: 202 }); + } + + // ponytail: validate attachment BEFORE the destructive tx — a + // 409 ATTACHMENT_MISSING must be side-effect-free. Previously the + // order was delete-then-check; two clicks of "reprocess" against + // a doc whose R2 source was missing would wipe its chunks first + // and then surface the missing-attachment error, leaving the user + // with permanent data loss and no recovery path. + if (!doc.attachmentId) { + // ponytail: no attachment means no R2 file to re-OCR. The row + // stays at its terminal status — the user can re-upload or delete. + // Surface the gap rather than swallowing it. + return NextResponse.json({ code: "ATTACHMENT_MISSING" }, { status: 409 }); + } + + const attachment = await getAttachmentForUser(doc.attachmentId, user.id); + if (!attachment) { + return NextResponse.json({ code: "ATTACHMENT_MISSING" }, { status: 409 }); + } + + // ponytail: full reprocess — clear stale chunks + flip the doc + // row's status back to "pending" inside one tx so the Settings UI + // sees a clean state if the deletion succeeds but the dispatch + // step below fails. + await withKbTx(async (tx) => { + await deleteKbChunksByDocumentId(tx, doc.id); + await resetKbDocumentForReprocess(user.id, doc.id); + }); + + try { + await fireIngestionRun({ + userId: user.id, + attachment, + docId: doc.id, + title: doc.title, + source: "kb-reprocess", + }); + } catch (err) { + // ponytail: row is back to pending + chunks are wiped. A failed + // dispatch leaves the row in pending and the next Settings poll + // (or another reprocess click) can retry — same UX as a fresh + // upload that hit a transient network blip. + console.error("POST /api/kb/documents/[id]/reprocess: fireIngestionRun failed", err); + } + + return NextResponse.json({ docId: doc.id }, { status: 202 }); +}); diff --git a/app/api/kb/documents/[id]/route.ts b/app/api/kb/documents/[id]/route.ts new file mode 100644 index 00000000..a9269942 --- /dev/null +++ b/app/api/kb/documents/[id]/route.ts @@ -0,0 +1,53 @@ +import { NextResponse } from "next/server"; + +import { withAuth } from "@/lib/auth/with-auth"; +import { + deleteKbDocumentForUser, + findKbChunksContentByDocumentId, + findKbDocumentById, + type KbChunkPreview, +} from "@/lib/kb/queries"; + +// ponytail: Settings → KB → doc detail (right pane). Returns the +// kb_document row + slim chunk preview (no 1536-dim embedding, no +// generated tsv column) so the UI can show parsed content without +// paying ~6 KB per chunk in the payload. + +export const GET = withAuth<{ id: string }>(async (_req, { user, params }) => { + const doc = await findKbDocumentById(user.id, params.id); + if (!doc) { + return NextResponse.json({ code: "NOT_FOUND" }, { status: 404 }); + } + + const chunks: KbChunkPreview[] = + doc.status === "success" ? await findKbChunksContentByDocumentId(user.id, doc.id) : []; + + return NextResponse.json({ + doc: { + id: doc.id, + title: doc.title, + status: doc.status, + errorMessage: doc.errorMessage, + contentType: doc.contentType, + attachmentId: doc.attachmentId, + folderId: doc.folderId, + contentHash: doc.contentHash, + pages: doc.pages, + createdAt: doc.createdAt.toISOString(), + updatedAt: doc.updatedAt.toISOString(), + }, + chunks, + }); +}); + +// ponytail: Settings → KB → doc delete. Cascades to kb_chunk via +// `document_id ... ON DELETE cascade`. R2 objects (kb-tmp/* page PNGs +// and the source PDF in attachments/) are NOT deleted — they live in +// R2 forever; v3 retention sweep can clean them up. +export const DELETE = withAuth<{ id: string }>(async (_req, { user, params }) => { + const deleted = await deleteKbDocumentForUser(user.id, params.id); + if (!deleted) { + return NextResponse.json({ code: "NOT_FOUND" }, { status: 404 }); + } + return new NextResponse(null, { status: 204 }); +}); diff --git a/app/api/kb/documents/route.ts b/app/api/kb/documents/route.ts new file mode 100644 index 00000000..e7ca4b1b --- /dev/null +++ b/app/api/kb/documents/route.ts @@ -0,0 +1,119 @@ +import { NextResponse } from "next/server"; + +import { withAuth } from "@/lib/auth/with-auth"; +import { listKbDocumentsByFolder, listKbDocumentsGroupedWithAttachment } from "@/lib/kb/queries"; + +// ponytail: Settings → KB tab list. Per-user scoped at the query layer +// (the helper filters by userId from withAuth). Returns folders + their +// docs in one round-trip, with attachmentUrl joined in for the "View +// source" link. No pagination yet (KB volume per user is O(tens of docs)). +// +// v3: also serves the @-mention composer. When `?mention=1` is set, +// returns a flat list of `status='success'` docs (the popover shows +// only ingest-ready docs). The flat mode skips the folder grouping and +// limits each doc to the fields the popover needs. +// +// v4: `?folderId=` scopes the doc payload to a single folder. All +// folders are still listed (the sidebar needs them), but only the +// targeted folder gets its `documents` array populated — other folders +// return `documents: []`. The frontend's `anyInflight` poll can then +// skip the JOIN cost for every other folder the user owns. + +type GroupedDoc = { + id: string; + title: string; + status: string; + errorMessage: string | null; + contentType: string; + attachmentId: string | null; + attachmentUrl: string | null; + createdAt: string; + updatedAt: string; + totalChunks?: number; + successChunks?: number; + failedChunks?: number; + pendingChunks?: number; + parsingChunks?: number; + totalPages?: number; + failedPages?: number; + pendingPages?: number; + parsingPages?: number; +}; + +export const GET = withAuth(async (req: Request, { user }) => { + const { searchParams } = new URL(req.url); + const mentionMode = searchParams.get("mention") === "1"; + const folderIdParam = searchParams.get("folderId"); + + try { + if (mentionMode) { + // ponytail: composer popover surface. Now grouped by folder so the + // popover can drill into folders as categories. Empty folders are + // dropped (a folder with zero ingest-ready docs has nothing to + // offer the user). Folders with zero docs at all are also dropped + // — keeps the popover focused on actionable choices. + // + // ponytail: mention mode is always cross-folder — the popover + // shows docs from every folder the user owns, so we never pass + // `folderId` here even if the URL has one. + const groups = await listKbDocumentsGroupedWithAttachment(user.id); + const folders = groups + .map(({ folder, documents }) => { + const successDocs = documents + .filter((d) => d.status === "success") + .map((d) => ({ + id: d.id, + title: d.title, + status: d.status, + })); + if (successDocs.length === 0) return null; + return { + id: folder.id, + name: folder.name, + docCount: successDocs.length, + docs: successDocs, + }; + }) + .filter((g): g is NonNullable => g !== null); + return NextResponse.json({ folders }); + } + + const groups = await listKbDocumentsGroupedWithAttachment( + user.id, + folderIdParam && folderIdParam.length > 0 ? folderIdParam : null, + ); + return NextResponse.json({ + groups: groups.map(({ folder, documents }) => ({ + folder: { id: folder.id, name: folder.name }, + documents: documents.map( + (d): GroupedDoc => ({ + id: d.id, + title: d.title, + status: d.status, + errorMessage: d.errorMessage, + contentType: d.contentType, + attachmentId: d.attachmentId, + attachmentUrl: d.attachmentUrl, + createdAt: d.createdAt.toISOString(), + updatedAt: d.updatedAt.toISOString(), + totalChunks: d.totalChunks, + successChunks: d.successChunks, + failedChunks: d.failedChunks, + pendingChunks: d.pendingChunks, + parsingChunks: d.parsingChunks, + totalPages: d.totalPages, + failedPages: d.failedPages, + pendingPages: d.pendingPages, + parsingPages: d.parsingPages, + }), + ), + })), + }); + } catch (err) { + console.error("GET /api/kb/documents failed", err); + return NextResponse.json({ code: "INTERNAL" }, { status: 500 }); + } +}); + +// keep unused-import referenced for tree-shake tests +void listKbDocumentsByFolder; diff --git a/app/api/kb/folders/[id]/route.ts b/app/api/kb/folders/[id]/route.ts new file mode 100644 index 00000000..528f88ad --- /dev/null +++ b/app/api/kb/folders/[id]/route.ts @@ -0,0 +1,101 @@ +import { NextResponse } from "next/server"; +import { z } from "zod"; + +import { withAuth } from "@/lib/auth/with-auth"; +import { + deleteKbFolderForUser, + findKbChunksByFolderId, + findKbFolderById, + findKbFolderByName, + listKbDocumentsByFolder, + updateKbFolderNameForUser, +} from "@/lib/kb/queries"; + +// ponytail: Settings → KB → folder detail (combined graph). Returns the +// kb_folder row + chunks content of all documents in the folder. +export const GET = withAuth<{ id: string }>(async (_req, { user, params }) => { + const folder = await findKbFolderById(user.id, params.id); + if (!folder) { + return NextResponse.json({ code: "NOT_FOUND" }, { status: 404 }); + } + + const chunks = await findKbChunksByFolderId(user.id, folder.id); + + return NextResponse.json({ + folder: { + id: folder.id, + name: folder.name, + createdAt: folder.createdAt.toISOString(), + }, + chunks, + }); +}); + +// ponytail: Settings → KB → folder rename. Same UNIQUE(user_id, name) +// guard as POST /api/kb/folders — duplicate names surface as 409 +// DUPLICATE, missing folder as 404. No body shape change; the +// `name` field is the only thing the UI edits. +const PatchBody = z.object({ + name: z.string().min(1).max(64).trim(), +}); + +export const PATCH = withAuth<{ id: string }>(async (req, { user, params }) => { + const body = PatchBody.safeParse(await req.json().catch(() => null)); + if (!body.success) { + return NextResponse.json({ code: "INVALID_NAME" }, { status: 400 }); + } + const { name } = body.data; + + const folder = await findKbFolderById(user.id, params.id); + if (!folder) { + return NextResponse.json({ code: "NOT_FOUND" }, { status: 404 }); + } + + // Same name → no-op success. + if (name === folder.name) { + return NextResponse.json({ folder }, { status: 200 }); + } + + // ponytail: race-safe duplicate check. If another tab creates a + // folder with the same name in between the SELECT and the UPDATE, + // the UPDATE will fail with 23505 — catch + re-read + 409. + const dup = await findKbFolderByName(user.id, name); + if (dup && dup.id !== folder.id) { + return NextResponse.json({ code: "DUPLICATE" }, { status: 409 }); + } + + try { + const updated = await updateKbFolderNameForUser(user.id, folder.id, name); + if (!updated) { + return NextResponse.json({ code: "NOT_FOUND" }, { status: 404 }); + } + return NextResponse.json({ folder: updated }, { status: 200 }); + } catch (err) { + if ((err as { code?: string }).code === "23505") { + return NextResponse.json({ code: "DUPLICATE" }, { status: 409 }); + } + console.error("PATCH /api/kb/folders/[id] failed", err); + return NextResponse.json({ code: "INTERNAL" }, { status: 500 }); + } +}); + +// ponytail: Settings → KB → folder delete. The DB has `folder_id ... +// ON DELETE RESTRICT` on kb_document, so Postgres refuses to drop a +// folder that still has docs. We surface that as 409 NON_EMPTY with +// a doc count, so the UI can render "delete its 3 docs first" instead +// of a generic FK-violation message. +export const DELETE = withAuth<{ id: string }>(async (_req, { user, params }) => { + const folder = await findKbFolderById(user.id, params.id); + if (!folder) { + return NextResponse.json({ code: "NOT_FOUND" }, { status: 404 }); + } + const docs = await listKbDocumentsByFolder(user.id, folder.id); + if (docs.length > 0) { + return NextResponse.json({ code: "NON_EMPTY", docCount: docs.length }, { status: 409 }); + } + const deleted = await deleteKbFolderForUser(user.id, folder.id); + if (!deleted) { + return NextResponse.json({ code: "NOT_FOUND" }, { status: 404 }); + } + return new NextResponse(null, { status: 204 }); +}); diff --git a/app/api/kb/folders/route.ts b/app/api/kb/folders/route.ts new file mode 100644 index 00000000..444ebd58 --- /dev/null +++ b/app/api/kb/folders/route.ts @@ -0,0 +1,49 @@ +import { NextResponse } from "next/server"; +import { randomUUID } from "node:crypto"; +import { z } from "zod"; + +import { withAuth } from "@/lib/auth/with-auth"; +import { findKbFolderByName, insertKbFolder } from "@/lib/kb/queries"; + +// ponytail: Settings → KB → "New Folder" modal endpoint. Same UNIQUE +// (user_id, name) constraint as the auto-created "Attachments" folder, +// so concurrent creates collapse via the 23505 retry path inside +// ensureDefaultKbFolder. We don't use that helper here because the UI +// needs to know whether the folder already exists (409) vs. was just +// created (201) — that distinction is lost behind ensureDefaultKbFolder. + +const Schema = z.object({ + name: z.string().min(1).max(64).trim(), +}); + +export const POST = withAuth(async (req, { user }) => { + const body = Schema.safeParse(await req.json().catch(() => null)); + if (!body.success) { + return NextResponse.json({ code: "INVALID_NAME" }, { status: 400 }); + } + const { name } = body.data; + + const existing = await findKbFolderByName(user.id, name); + if (existing) { + return NextResponse.json({ code: "DUPLICATE", folder: existing }, { status: 409 }); + } + + try { + const folder = await insertKbFolder({ + id: `f-${randomUUID()}`, + userId: user.id, + name, + }); + return NextResponse.json({ folder }, { status: 201 }); + } catch (err) { + // Race with another tab creating the same folder: 23505 → re-read. + if ((err as { code?: string }).code === "23505") { + const again = await findKbFolderByName(user.id, name); + if (again) { + return NextResponse.json({ code: "DUPLICATE", folder: again }, { status: 409 }); + } + } + console.error("POST /api/kb/folders failed", err); + return NextResponse.json({ code: "INTERNAL" }, { status: 500 }); + } +}); diff --git a/app/api/kb/upload/route.ts b/app/api/kb/upload/route.ts new file mode 100644 index 00000000..e9f9368b --- /dev/null +++ b/app/api/kb/upload/route.ts @@ -0,0 +1,105 @@ +import { NextResponse } from "next/server"; +import { randomUUID } from "node:crypto"; +import { z } from "zod"; + +import { withAuth } from "@/lib/auth/with-auth"; +import { fireIngestionRun } from "@/lib/kb/ingest"; +import { getAttachmentForUser } from "@/lib/attachments/queries"; +import { findKbDocumentByContentHash, findKbFolderById, insertKbDocument } from "@/lib/kb/queries"; + +// ponytail: Settings → KB → "Add Doc". Frontend uploads the file via +// the existing /api/attachments/presign → PUT → confirm flow first, +// then POSTs the resulting attachmentId + a target folderId here. +// Backend creates a kb_document row (status=pending) and kicks off +// the kbAgent graph — registered as a top-level assistant in +// langgraph.json so the synthetic "ingest this file" thread skips the +// mainAgent router + renameThreadAgent LLM calls. Shared with +// POST /api/kb/documents/[id]/reprocess via lib/kb/ingest. +// +// The run is fire-and-forget: we return 202 with the docId. The +// frontend polls GET /api/kb/documents and watches the row's status +// flip pending → parsing → success. + +const Schema = z.object({ + folderId: z.string().min(1), + attachmentId: z.string().min(1), + title: z.string().min(1).max(256).optional(), +}); + +export const POST = withAuth(async (req, { user }) => { + const body = Schema.safeParse(await req.json().catch(() => null)); + if (!body.success) { + return NextResponse.json({ code: "INVALID" }, { status: 400 }); + } + const { folderId, attachmentId, title } = body.data; + + // 1. Verify attachment belongs to the caller and is uploaded. + const attachment = await getAttachmentForUser(attachmentId, user.id); + if (!attachment) { + return NextResponse.json({ code: "ATTACHMENT_NOT_FOUND" }, { status: 404 }); + } + if (attachment.status !== "uploaded") { + return NextResponse.json({ code: "ATTACHMENT_NOT_UPLOADED" }, { status: 409 }); + } + + // 2. Verify target folder belongs to the caller. + const folder = await findKbFolderById(user.id, folderId); + if (!folder) { + return NextResponse.json({ code: "FOLDER_NOT_FOUND" }, { status: 404 }); + } + + // 3. PRIMARY dedup: if a doc with this contentHash already exists, + // re-fire ingestion if the previous attempt failed/stalled. + const contentHash = attachment.sha256 ?? `r2key:${attachment.r2Key}`; + const existing = await findKbDocumentByContentHash(user.id, contentHash); + if (existing) { + if ( + existing.status === "pending" || + existing.status === "failed" || + existing.status === "parsing" + ) { + try { + await fireIngestionRun({ + userId: user.id, + attachment, + docId: existing.id, + title: title ?? existing.title, + }); + } catch (err) { + console.error("POST /api/kb/upload: fireIngestionRun failed", err); + } + } + return NextResponse.json({ doc: existing, deduped: true }, { status: 200 }); + } + + // 4. Create the kb_document row (status=pending) so the UI has + // something to show immediately and a target to update when the run + // lands. + const docId = `d-${randomUUID()}`; + const doc = await insertKbDocument({ + id: docId, + userId: user.id, + folderId, + attachmentId, + title: title ?? attachment.name, + contentType: attachment.contentType, + contentHash, + status: "pending", + errorMessage: null, + }); + + // 5. Fire-and-forget kbAgent run. + try { + await fireIngestionRun({ + userId: user.id, + attachment, + docId: doc.id, + title: doc.title, + }); + } catch (err) { + // The row is already created; the user can retry from the UI. + console.error("POST /api/kb/upload: fireIngestionRun failed", err); + } + + return NextResponse.json({ doc }, { status: 202 }); +}); diff --git a/app/api/threads/[id]/observability/[parentMessageId]/route.ts b/app/api/threads/[id]/observability/[parentMessageId]/route.ts index f594ae6c..0608ddb6 100644 --- a/app/api/threads/[id]/observability/[parentMessageId]/route.ts +++ b/app/api/threads/[id]/observability/[parentMessageId]/route.ts @@ -39,11 +39,12 @@ export const GET = withAuth(async (_req, { user, params }) => { // pending set is empty most of the time. fetchInFlightRuns(params.id, params.parentMessageId), ]); + // ponytail: server-side transform. CapturedSpan → SpanData runs here // so the panel never has to import transformCapturedToSpanData and // the wire carries only what the waterfall needs. const spans = transformCapturedToSpanData(capturedSpans); - const aggregate = aggregateRoot(capturedSpans); + const aggregate = aggregateRoot(capturedSpans, spans); const stepIdToRawSpanId = buildStepIdToRawSpanId(capturedSpans); return NextResponse.json({ thread_id: params.id, diff --git a/app/api/threads/[id]/observability/route.ts b/app/api/threads/[id]/observability/route.ts index 76983eeb..1162111b 100644 --- a/app/api/threads/[id]/observability/route.ts +++ b/app/api/threads/[id]/observability/route.ts @@ -30,7 +30,7 @@ export const GET = withAuth(async (_req, { user, params }) => { await markRunningAsFailed(params.id); const capturedSpans = await getSpansByThreadId(params.id); const spans = transformCapturedToSpanData(capturedSpans); - const aggregate = aggregateRoot(capturedSpans); + const aggregate = aggregateRoot(capturedSpans, spans); const stepIdToRawSpanId = buildStepIdToRawSpanId(capturedSpans); return NextResponse.json({ thread_id: params.id, diff --git a/app/assistant.tsx b/app/assistant.tsx index 16329a1b..08c7c53b 100644 --- a/app/assistant.tsx +++ b/app/assistant.tsx @@ -11,7 +11,7 @@ import { import { useLangGraphRuntime } from "@assistant-ui/react-langgraph"; import { Client } from "@langchain/langgraph-sdk"; import { ThreadListPrimitive } from "@assistant-ui/react"; -import { Brain, MenuIcon, PanelLeftIcon, PlusIcon, ShareIcon } from "lucide-react"; +import { BookOpen, Brain, MenuIcon, PanelLeftIcon, PlusIcon, ShareIcon } from "lucide-react"; import { BrandMark } from "@/components/brand-mark"; @@ -48,6 +48,17 @@ const memoryLink = [ icon: , visibility: "authenticated" as const, }, + // ponytail: Knowledge Base lives behind the Better Auth settings + // shim (kbSettingsPlugin registers the path "knowledge-base"). Add + // it next to Memory so the sidebar's UserButton has a KB shortcut — + // the Settings tab also surfaces it but most users reach KB from + // here. + { + label: "Knowledge Base", + href: "/settings/knowledge-base", + icon: , + visibility: "authenticated" as const, + }, ]; const Sidebar: FC<{ collapsed?: boolean }> = ({ collapsed }) => { @@ -418,6 +429,7 @@ const ThreadUrlShadow: FC = () => { // e.g. immediately after +New before adapter.initialize resolves; // we treat that as "user wants /chat for now"). Real uuid → push // /chat/. No preconditions, no reads of aUI's state. +// const writeUrlForThread = (remoteId: string | undefined): void => { if (typeof window === "undefined") return; const target = remoteId ? `/chat/${remoteId}` : "/chat"; diff --git a/app/auth-shell.tsx b/app/auth-shell.tsx index 66e563bf..9aeb3341 100644 --- a/app/auth-shell.tsx +++ b/app/auth-shell.tsx @@ -7,6 +7,7 @@ import { AuthProvider as BetterAuthUIProvider } from "@/components/auth/auth-pro import { authClient } from "@/lib/auth/client"; import { memorySettingsPlugin } from "@/components/auth/settings/memory-tab"; import { creditSettingsPlugin } from "@/components/auth/settings/credit-tab"; +import { kbSettingsPlugin } from "@/components/auth/settings/kb-tab"; export function AuthShell({ children }: { children: ReactNode }) { const router = useRouter(); @@ -16,7 +17,7 @@ export function AuthShell({ children }: { children: ReactNode }) { basePaths={{ auth: "/login", settings: "/settings" }} socialProviders={["github", "google"]} multipleAccountsPerProvider={false} - plugins={[memorySettingsPlugin, creditSettingsPlugin]} + plugins={[memorySettingsPlugin, creditSettingsPlugin, kbSettingsPlugin]} emailAndPassword={{ minPasswordLength: 8, confirmPassword: true, diff --git a/app/globals.css b/app/globals.css index cc94b358..60cb5f35 100644 --- a/app/globals.css +++ b/app/globals.css @@ -342,3 +342,85 @@ margin-left: 0.25rem; margin-right: 0.25rem; } + +/* ponytail: hide tab label text on mobile so the settings tab strip + doesn't overflow. Uses Tailwind's `sr-only` technique (clip-path + + absolute position) so the text stays in the accessibility tree — + screen readers still announce the full tab label, sighted users see + only the icon. Above 768px the rule doesn't apply and the span + shows inline next to the icon. */ +@media (max-width: 767px) { + [role="tab"] > span:not(.sr-only) { + position: absolute; + width: 1px; + height: 1px; + padding: 0; + margin: -1px; + overflow: hidden; + clip: rect(0, 0, 0, 0); + white-space: nowrap; + } +} + +/* Lexical editor styling overrides — remove default borders, outlines, and box shadows + to keep the input seamless and transparent inside the parent composer shell. */ +.aui-lexical-editor, +.aui-lexical-editor:focus, +.aui-lexical-editor:focus-visible, +.aui-lexical-input, +.aui-lexical-input:focus, +.aui-lexical-input:focus-visible { + border: none !important; + outline: none !important; + box-shadow: none !important; +} + +.aui-lexical-editor { + position: relative; +} + +.aui-lexical-placeholder { + position: absolute !important; + top: 0; + left: 0; + padding: inherit; + pointer-events: none; + user-select: none; + color: hsl(var(--muted-foreground)); + opacity: 0.6; +} + +.aui-lexical-editor span[contenteditable="false"], +.aui-lexical-input span[contenteditable="false"] { + display: inline-flex !important; + vertical-align: middle; +} + +.aui-lexical-input p, +.aui-lexical-editor p, +.aui-user-message-content { + margin: 0 !important; +} + +.aui-lexical-input p > *, +.aui-lexical-editor p > *, +.aui-user-message-content > * { + vertical-align: middle; +} + +.aui-user-message-content .aui-directive-chip:not(:first-child) { + margin-left: 0.25rem; +} +.aui-user-message-content .aui-directive-chip:not(:last-child) { + margin-right: 0.25rem; +} + +/* ponytail: override react-force-graph default tooltip wrapper styling + so that our custom HTML card renders natively without double + backgrounds, paddings, or black boxes wrapping it. */ +.graph-tooltip { + background: transparent !important; + border: none !important; + box-shadow: none !important; + padding: 0 !important; +} diff --git a/backend/agent.ts b/backend/agent.ts index 2af0d241..7c219806 100644 --- a/backend/agent.ts +++ b/backend/agent.ts @@ -1,63 +1,97 @@ import { START, END, StateGraph } from "@langchain/langgraph"; +import { HumanMessage } from "@langchain/core/messages"; import { triggerBackgroundAgentNode } from "@/backend/node/trigger-background-agent-node"; import { capturingHandler, creditTrackingHandler } from "@/backend/callbacks"; import { renameThreadAgentNode } from "@/backend/node/rename-thread-agent-node"; +import { prepareDataNode } from "@/backend/node/prepare-data-node"; import { weatherAgent } from "@/backend/agent/weather-agent"; import { chatAgent } from "@/backend/agent/chat-agent"; import { cryptoAgent } from "@/backend/agent/crypto-agent"; import { codeAgent } from "@/backend/agent/code-agent"; +import { kbAgent } from "@/backend/agent/kb-agent"; import { routerAgentNode } from "@/backend/node/router-agent-node"; import { checkpointer } from "@/backend/checkpointer"; import { store } from "@/backend/store"; import { RouterAgentState } from "@/backend/state"; import { getThreadTitle } from "@/lib/threads/queries"; import { DEFAULT_THREAD_TITLE } from "@/lib/constants"; +import { isFilePart } from "@/lib/kb/extract"; -// After the router speaks, decide which sub-agent gets the turn. -// Falls back to chatAgent if the router hasn't run yet or its -// decision didn't make it into state. -function routeToSubAgent({ - routerDecision, -}: { - routerDecision?: { next: "weatherAgent" | "chatAgent" | "cryptoAgent" | "codeAgent" }; -}): "weatherAgent" | "chatAgent" | "cryptoAgent" | "codeAgent" { - return routerDecision?.next ?? "chatAgent"; +// After the router speaks, decide which sub-agent gets the turn AND +// whether to fan out to renameThreadAgent in parallel. Falls back to +// chatAgent if the router hasn't run yet or its decision didn't make +// it into state. kbAgent routes back to the router after stamping a +// `kb_ref` sibling onto the PDF file part, so a SECOND router pass +// picks the final sub-agent (chat / weather / etc.) — the PDF is +// still in the message (file part preserved) but the kb_ref sibling +// marks it as already-ingested, so the PDF-short-circuit no longer +// fires. On that second pass the rename fanout also fires. +// +// Returning an ARRAY of destinations makes LangGraph run the listed +// nodes in parallel; returning a single string is the normal one-shot +// routing case. So we collapse the sub-agent pick + rename fanout into +// a single conditional edge to satisfy langgraph's "one condition per +// source node" rule. +function hasPendingFilePart(messages: { content: unknown }[]): boolean { + return messages.some( + (m) => + m instanceof HumanMessage && + Array.isArray(m.content) && + m.content.some((item) => isFilePart(item) && !item?.kb_ref), + ); } -// ponytail: renameThreadAgent only needs to run once per thread — the -// first time the user sends a message. After that, threads.title is -// already set; re-invoking the LLM every turn (interrupt + resume, -// regenerate, follow-up) wastes tokens. Query the title from the DB -// before entering; if it has been replaced from the default placeholder, -// skip the node entirely. The conditional edge wires both branches from -// START, so renameThreadAgent is never even entered (no callback, no -// span) when the LLM-generated title already exists. -async function shouldRenameRouter( - _state: unknown, +type SubAgent = "weatherAgent" | "chatAgent" | "cryptoAgent" | "codeAgent" | "kbAgent"; + +async function routeAndMaybeRename( + state: { messages: { content: unknown }[]; routerDecision?: { next: SubAgent } }, config: { configurable?: { thread_id?: string } }, -): Promise<"renameThreadAgent" | typeof END> { +): Promise { + const subAgent: SubAgent = state.routerDecision?.next ?? "chatAgent"; + // ponytail: renameThreadAgent only needs to run once per thread — the + // first time the user sends a message. After that, threads.title is + // already set; re-invoking the LLM every turn (interrupt + resume, + // regenerate, follow-up) wastes tokens. + // + // Skip while a raw file part is present (kbAgent hasn't rewritten it + // yet) — for PDF uploads the SECOND routerAgent pass after kbAgent + // sees the kb_ref in place and fires rename. + if (hasPendingFilePart(state.messages)) return subAgent; const threadId = config.configurable?.thread_id; - if (typeof threadId !== "string" || !threadId) return END; + if (typeof threadId !== "string" || !threadId) return subAgent; const title = await getThreadTitle(threadId); // ponytail: the column has `notNull().default(DEFAULT_THREAD_TITLE)` // ("New Chat"), so title is always a non-null string in the DB. The // "auto-rename not yet run" signal is `title === DEFAULT_THREAD_TITLE`; // anything else is the LLM-generated title from a prior turn. - if (typeof title === "string" && title !== DEFAULT_THREAD_TITLE) return END; - return "renameThreadAgent"; + if (typeof title === "string" && title !== DEFAULT_THREAD_TITLE) return subAgent; + // Array return = parallel fanout (langgraph runs both nodes). + return ["renameThreadAgent", subAgent]; } export const builder = new StateGraph(RouterAgentState) + .addNode("prepareData", prepareDataNode) .addNode("routerAgent", routerAgentNode) .addNode("chatAgent", chatAgent) .addNode("weatherAgent", weatherAgent) .addNode("cryptoAgent", cryptoAgent) .addNode("codeAgent", codeAgent) + .addNode("kbAgent", kbAgent) .addNode("triggerBackgroundAgent", triggerBackgroundAgentNode) .addNode("renameThreadAgent", renameThreadAgentNode) - // Topology: - // START ──▶ routerAgent ──▶ (sub-agent) ──▶ triggerBackgroundAgent ──▶ END - // START ─────────────────────────────────▶ renameThreadAgent (parallel, leaf) + // Topology (issue #13 v3): + // START ──▶ prepareData ──▶ routerAgent ──┬──▶ (sub-agent | kbAgent) ──▶ triggerBackgroundAgent ──▶ END + // └──▶ renameThreadAgent (terminal, no outgoing edge needed) + // + // kbAgent loops back to routerAgent (NOT through prepareData) after + // stamping a `kb_ref` sibling onto the PDF file part, so a SECOND + // router pass picks the final sub-agent (chat / weather / etc.) — + // the file part is preserved (not replaced), but the kb_ref sibling + // marks it as already-ingested and the PDF-short-circuit no longer + // fires. We deliberately skip prepareData on the second pass so + // the ToolMessage we injected at the start of the turn doesn't + // re-fire (it'd be a no-op anyway since directives only appear in + // the original HumanMessage, but re-running is wasted DB work). // // ask_location's picker card is owned by the weather subgraph // (see backend/agent/weather-agent.ts + components/tool-ui/ask-location). @@ -66,34 +100,41 @@ export const builder = new StateGraph(RouterAgentState) // write_code's editor card is owned by the code subgraph // (see backend/agent/code-agent.ts + components/tool-ui/code). // - // renameThreadAgent runs as a parallel leaf off the main response - // path (END). The graph invocation only returns after ALL active - // branches complete, but the chat stream ends on the END branch, - // so the user sees no rename latency. - // - // triggerBackgroundAgent is the chat's last node before END. It fires + // triggerBackgroundAgent is the chat's last sub-agent step. It fires // the `background_agent` graph (registered separately in // langgraph.json) and returns `{}` immediately — that graph does // `last_message_at` touch + threadSummarizeNode work on its own // thread. See backend/node/trigger-background-agent-node.ts for the // fire-and-forget pattern; see backend/background-agent.ts for // what the background graph runs. - .addEdge(START, "routerAgent") - .addConditionalEdges("routerAgent", routeToSubAgent, [ - "weatherAgent", - "chatAgent", - "cryptoAgent", - "codeAgent", - ]) + // + // routeAndMaybeRename collapses the sub-agent pick + rename fanout + // into one conditional edge (langgraph forbids two conditions on the + // same source node). When the function returns an array of two + // names, langgraph runs them in parallel. The rename fires only + // when the messages are clean (no raw file part) AND the thread + // title is still the default placeholder — see routeAndMaybeRename. + .addEdge(START, "prepareData") + .addEdge("prepareData", "routerAgent") + .addConditionalEdges("routerAgent", routeAndMaybeRename, { + chatAgent: "chatAgent", + weatherAgent: "weatherAgent", + cryptoAgent: "cryptoAgent", + codeAgent: "codeAgent", + kbAgent: "kbAgent", + renameThreadAgent: "renameThreadAgent", + }) .addEdge("chatAgent", "triggerBackgroundAgent") .addEdge("weatherAgent", "triggerBackgroundAgent") .addEdge("cryptoAgent", "triggerBackgroundAgent") .addEdge("codeAgent", "triggerBackgroundAgent") - .addEdge("triggerBackgroundAgent", END) - .addConditionalEdges(START, shouldRenameRouter, { - renameThreadAgent: "renameThreadAgent", - __end__: END, - }); + // ponytail: kbAgent loops back to the router — after it stamps a + // `kb_ref` sibling onto the PDF file part, the router's + // PDF-short-circuit no longer fires (the PDF is still in the + // message but it's marked as already-ingested via the sibling) and + // the router routes to the final sub-agent (chatAgent, etc.). + .addEdge("kbAgent", "routerAgent") + .addEdge("triggerBackgroundAgent", END); // ponytail: one handler per process (per module), shared across all // concurrent runs AND across every Pregel that wires it via withConfig. diff --git a/backend/agent/chat-agent.ts b/backend/agent/chat-agent.ts index df3d5760..ceb05f13 100644 --- a/backend/agent/chat-agent.ts +++ b/backend/agent/chat-agent.ts @@ -3,14 +3,15 @@ import { ToolNode, toolsCondition } from "@langchain/langgraph/prebuilt"; import type { BaseMessage } from "@langchain/core/messages"; import type { RunnableConfig } from "@langchain/core/runnables"; import { getChatModel } from "@/backend/model"; -import { ALL_TOOLS } from "@/backend/tool"; +import { CHAT_TOOLS } from "@/backend/tool"; import { CHAT_AGENT_PROMPT } from "@/backend/prompt/system"; import { CommonAgentState } from "@/backend/state"; import { buildSystemMessageWithMemory, loadThreadSummariesForPrompt, - trimMessagesForInvoke, + prepareMessagesForInvoke, } from "@/backend/memory/template"; +import { extractUserId } from "@/backend/memory/recall"; import { subgraphCheckpointerConfig } from "@/backend/checkpointer"; // Chat agent gets every tool — the router already decided whether this @@ -37,13 +38,18 @@ async function chatModelNode({ messages }: { messages: BaseMessage[] }, config?: // context-loss one). state.messages is NEVER touched — UI + // checkpointer read from it directly. const threads = await loadThreadSummariesForPrompt(config); - const history = trimMessagesForInvoke(messages, threads?.summaries ?? []); + const userId = extractUserId(config); + const history = await prepareMessagesForInvoke( + messages, + threads?.summaries ?? [], + userId ?? undefined, + ); const sysMsg = await buildSystemMessageWithMemory(CHAT_AGENT_PROMPT, config, threads); const response = await ( await getChatModel() ) - .bindTools(ALL_TOOLS) + .bindTools(CHAT_TOOLS) .invoke([sysMsg, ...history], config); return { messages: [response] }; @@ -53,7 +59,7 @@ function chatModelRoute(state: { messages: BaseMessage[] }) { return toolsCondition(state) === END ? END : "chatTools"; } -const chatToolNode = new ToolNode(ALL_TOOLS); +const chatToolNode = new ToolNode(CHAT_TOOLS); const builder = new StateGraph(CommonAgentState) .addNode("chatModel", chatModelNode) diff --git a/backend/agent/code-agent.ts b/backend/agent/code-agent.ts index a054ab05..7940b3f1 100644 --- a/backend/agent/code-agent.ts +++ b/backend/agent/code-agent.ts @@ -9,8 +9,9 @@ import { CommonAgentState } from "@/backend/state"; import { buildSystemMessageWithMemory, loadThreadSummariesForPrompt, - trimMessagesForInvoke, + prepareMessagesForInvoke, } from "@/backend/memory/template"; +import { extractUserId } from "@/backend/memory/recall"; import { subgraphCheckpointerConfig } from "@/backend/checkpointer"; // Code sub-agent: model ↔ tools loop. write_code proposes code that @@ -30,7 +31,11 @@ async function codeModelNode({ messages }: { messages: BaseMessage[] }, config?: // and drop the original turns from the input array. state.messages // is NEVER touched. const threads = await loadThreadSummariesForPrompt(config); - const history = trimMessagesForInvoke(messages, threads?.summaries ?? []); + const history = await prepareMessagesForInvoke( + messages, + threads?.summaries ?? [], + extractUserId(config) ?? undefined, + ); const sysMsg = await buildSystemMessageWithMemory(CODE_AGENT_PROMPT, config, threads); const response = await ( await getChatModel() diff --git a/backend/agent/crypto-agent.ts b/backend/agent/crypto-agent.ts index d74c17e1..1def2e3e 100644 --- a/backend/agent/crypto-agent.ts +++ b/backend/agent/crypto-agent.ts @@ -9,8 +9,9 @@ import { CommonAgentState } from "@/backend/state"; import { buildSystemMessageWithMemory, loadThreadSummariesForPrompt, - trimMessagesForInvoke, + prepareMessagesForInvoke, } from "@/backend/memory/template"; +import { extractUserId } from "@/backend/memory/recall"; import { subgraphCheckpointerConfig } from "@/backend/checkpointer"; // Crypto sub-agent: mirrors the weather subgraph. The model ↔ tools @@ -26,7 +27,11 @@ async function cryptoModelNode({ messages }: { messages: BaseMessage[] }, config // and drop the original turns from the input array. state.messages // is NEVER touched. const threads = await loadThreadSummariesForPrompt(config); - const history = trimMessagesForInvoke(messages, threads?.summaries ?? []); + const history = await prepareMessagesForInvoke( + messages, + threads?.summaries ?? [], + extractUserId(config) ?? undefined, + ); const sysMsg = await buildSystemMessageWithMemory(CRYPTO_AGENT_PROMPT, config, threads); const response = await ( await getChatModel() diff --git a/backend/agent/kb-agent.ts b/backend/agent/kb-agent.ts new file mode 100644 index 00000000..32b6903b --- /dev/null +++ b/backend/agent/kb-agent.ts @@ -0,0 +1,1399 @@ +import { END, START, StateGraph } from "@langchain/langgraph"; +import { HumanMessage, SystemMessage, type BaseMessage } from "@langchain/core/messages"; +import { MarkdownTextSplitter } from "@langchain/textsplitters"; +import PQueue from "p-queue"; +import { randomUUID } from "node:crypto"; +import { z } from "zod"; +import type { RunnableConfig } from "@langchain/core/runnables"; +import { getEmbeddingModel, getExtractModel, getOcrModel } from "@/backend/model"; +import { + KB_OCR_PAGE_PROMPT, + KB_ENTITY_EXTRACTION_SYSTEM_PROMPT, + KB_ENTITY_ALIGNMENT_SYSTEM_PROMPT, +} from "@/backend/prompt/system"; +import { creditTrackingHandler, capturingHandler } from "@/backend/callbacks"; +import { lastHumanMessageId } from "@/lib/langgraph/last-human-message-id"; +import { checkpointer, subgraphCheckpointerConfig } from "@/backend/checkpointer"; +import { store } from "@/backend/store"; +import { + KbAgentState, + type KbAgentStateShape, + type PageResult, + type ProcessedFile, +} from "@/backend/state"; +import { screenshotPdf } from "@/lib/kb/screenshot"; +import { extractPdfText } from "@/lib/kb/text"; +import { + ensureDefaultKbFolder, + findKbDocumentByContentHash, + findKbDocumentByAttachmentId, + findKbDocumentById, + findKbChunksByDocumentId, + insertKbChunks, + insertKbDocument, + insertKbObservability, + markAllKbChunksParsingForDocInTx, + markKbChunkFailed, + markKbChunkSuccess, + updateKbChunkForFailure, + updateKbChunkForSuccess, + updateKbChunkGraphData, + updateKbDocumentStatus, + withKbTx, +} from "@/lib/kb/queries"; +import { findAttachmentByR2Key } from "@/lib/attachments/queries"; +import { extractAllPdfParts, isFilePart, stampKbRefOnFilename } from "@/lib/kb/extract"; +import { invalidateKbDoc } from "@/lib/kb/cache"; +import { EMBEDDING_DIM } from "@/lib/kb/schema"; +import { r2KeyFromPublicUrl, uploadKbImage, getR2PublicBaseUrl, getObject } from "@/lib/r2/client"; +import { KB_OCR_CONCURRENCY, KB_ENTITY_CONCURRENCY } from "@/lib/constants"; + +const KB_CHUNK_SIZE = 1024; +const KB_CHUNK_OVERLAP = 200; + +// ponytail: v3 KB ingest subgraph — per-doc state. Compiled once at +// module load, wired into agent.ts as `kbAgent`. Sits between +// RouterNode ("PDF → kbAgent") and the sub-agents. +// +// Flow: +// START → prepareKBData → splitFilePage → pageToMarkdown → rewriteMessages ─┬─▶ END +// └─▶ generateChunkEmbed → END +// (non-blocking, triggers generateChunkEmbedNode) +// +// Every PDF file part in EVERY HumanMessage gets one of three +// outcomes: +// 1. `kb_ref` sibling stamped onto the file part (success, dedup, +// failed with a docId, or parsing). The file part is PRESERVED — +// the sibling just marks it as ingested so the next router pass +// skips re-processing. The resolve layer (lib/kb/resolve.ts) +// replaces the file part with resolved text at LLM-invoke time. +// 2. file part stripped (unknown attachment, can't even dedup). +// 3. carried over as a non-PDF file (images etc. — preserved). +// +// After one kbAgent invocation there are zero UNSTAMPED PDF file +// parts left in state.messages — every PDF either carries a kb_ref +// sibling or has been stripped. extractAllPdfParts / hasUnprocessedPdf +// filter on `!p.kb_ref` so the second router pass won't re-dispatch +// kbAgent. + +// ponytail: debug toggle. When true, generateChunkEmbedNode bails out +// before firing its per-row entity-LLM + write-back arms. Use to +// isolate the INSERT path vs the LLM path when reproducing a +// partial-pipeline failure — flip in source, restart backend, reprocess +// the target doc, then revert. +const SKIP_CHUNK_TO_ENTRIES = false; + +const ocrPageSchema = z.object({ + markdown: z + .string() + .describe( + "Clean markdown extraction of this PDF page. " + + "Preserve headings, lists, code blocks, tables, and inline formatting. " + + "Return an empty string if the page is blank or contains only decorative images. " + + "Output ONLY the markdown — no preamble, no commentary, no code fences.", + ), +}); + +const lightRagSchema = z.object({ + entities: z + .array( + z.object({ + name: z.string().describe("Entity name (e.g., person, tech stack, system component)"), + type: z.string().describe("Category of the entity (e.g., Person, Tool, Concept)"), + description: z.string().describe("Brief description of this entity in the current context"), + }), + ) + .describe("All distinct entities mentioned in the text"), + relationships: z + .array( + z.object({ + source: z.string().describe("Source entity name"), + target: z.string().describe("Target entity name"), + relation: z.string().describe("The action or logical connection between them"), + type: z.string().describe("Alias for the relation (used by some models)"), + description: z.string().describe("Detailed explanation of this relationship"), + }), + ) + .describe("Directed relationships connecting the extracted entities"), + themes: z + .array(z.string()) + .describe( + "3 to 5 high-level macroscopic keywords or core concepts summarizing this chunk's main point", + ), +}); + +function makeError(message: string): Partial { + return { status: "failed", errorMessage: message, processedFiles: [] }; +} + +// --------------------------------------------------------------------------- +// Node 1: prepareKBDataNode — DB queries + dedup + insert kb_documents row +// --------------------------------------------------------------------------- + +async function prepareKBDataNode( + state: KbAgentStateShape, + config?: { + configurable?: { + userId?: string; + mode?: "full" | "chunksOnly" | "retryFailed" | "retryFailedChunks"; + docId?: string; + forceRerun?: boolean; + thread_id?: string; + source?: "kb-upload" | "kb-reprocess" | "chat"; + run_id?: string; + }; + }, +): Promise> { + // ponytail: chunksOnly / retryFailed / retryFailedChunks dispatch + // from POST /reprocess. We bypass the entire attachment/file-part + // lookup chain and reuse the existing kb_document row's pages[].markdown: + // - splitFileToPageNode: r2Key=null → filter drops it → no-op + // - pageToMarkdownNode: retryFailed runs OCR only on failed pages; + // chunksOnly / retryFailedChunks skip OCR entirely (chunksOnly + // early-returns in pageToMarkdownNode; retryFailedChunks never + // reaches it because generateChunkEmbedNode intercepts on + // state.mode='retryFailedChunks'). + // - rewriteMessagesNode: guarded → no messages rewrite + // - generateChunkEmbedNode: pipelineStatus="new" → runs. Each + // mode then branches internally on state.mode. + // dispatchable mode precedence: config.configurable (per-run + // override set by fireIngestionRun) wins over state.mode (default + // "full" in the schema). + const mode = config?.configurable?.mode ?? state.mode ?? "full"; + + const userId = config?.configurable?.userId ?? state.userId; + + if (!userId) return makeError("user not provided"); + + // ponytail: source discriminator for kb_observability. Standalone + // callers (ingest.ts from Settings upload / reprocess) pass it + // explicitly via config.configurable. The chat subgraph path + // (mainAgent → kbAgent) inherits configurable from the parent + // graph, which doesn't set source — default to 'chat' in that + // case. Standalone path is locked to ingest.ts so the default + // won't mislabel an upload as chat; if a third caller shows up + // later without setting source, audit logs will surface it. + const source = config?.configurable?.source ?? "chat"; + + // ponytail: parent_message_id for kb_observability + spans link. + // Both paths converge on state.messages[-1] — synthetic in standalone, + // user chat msg in chat path. lastHumanMessageId walks backward so + // it survives any extra tool/AI messages the router injected. + const parentMessageId = lastHumanMessageId(state.messages); + if (!parentMessageId) { + return makeError("kbAgent prepareKBDataNode: no HumanMessage to anchor parent_message_id on"); + } + + const runId = config?.configurable?.run_id ?? null; + + // ponytail: thread_id is required for kb_observability rows + span + // lookups, but the rest of prepareKBDataNode (kb_document insert, + // dedup, file processing) doesn't need it. Standalone path always + // sets it via ingest.ts; chat subgraph inherits it from the parent + // graph's configurable; direct unit-test invocations may omit it. + // When missing, skip the kb_observability inserts below rather than + // failing the whole graph run — observability is a side feature, + // not a precondition. + const threadId = config?.configurable?.thread_id ?? null; + + if (mode === "chunksOnly" || mode === "retryFailed" || mode === "retryFailedChunks") { + // ponytail: chunksOnly / retryFailed requires an explicit docId either from + // config.configurable.docId (per-run) or state.docId. fallback + // to the explicit dispatch path. fail closed if neither set. + const targetDocId = config?.configurable?.docId ?? state.docId; + if (!targetDocId) { + return makeError(`${mode} requires docId`); + } + const doc = await findKbDocumentById(userId, targetDocId); + if (!doc) return makeError(`doc ${targetDocId} not found`); + if (threadId) { + await insertKbObservability({ + docId: doc.id, + threadId, + parentMessageId, + runId, + source, + mode, + }); + } + if (doc.status !== "success" && doc.status !== "failed" && doc.status !== "parsing") { + return makeError( + `${mode} requires settled doc or parsing doc, got status='${doc.status}'. Run full reprocess first.`, + ); + } + const pages = (doc.pages ?? []) as PageResult[]; + + // ponytail: stub FilePart is required by ProcessedFile.shape but + // rewriteMessagesNode skips the stamp pass under mode= + // "chunksOnly" / "retryFailed" so the values are never read. url/data empty → + // resolveKbRefs won't try to look up a public R2 path that + // doesn't exist for this synthetic dispatch. + const stubFilePart = { type: "file" as const, url: "", data: "", metadata: {} as never }; + return { + userId, + mode, + docId: doc.id, + pagesByDocId: { [doc.id]: pages }, + processedFiles: [ + { + messageIndex: -1, + filePart: stubFilePart as never, + docId: doc.id, + attachmentId: doc.attachmentId, + // r2Key=null → splitFileToPageNode filter rejects this entry + // (its filter chains `&& p.r2Key !== null`), so PDF rendering + // + screenshot + image upload are skipped. + r2Key: null, + title: doc.title, + contentHash: doc.contentHash, + // pipelineStatus="new" drives both routeAfterRewrite (pushes + // generateChunkEmbed) AND generateChunkEmbedNode's own filter + // (only acts on "new" entries). + pipelineStatus: "new", + errorMessage: null, + existingStatus: doc.status, + }, + ], + // preserve doc's terminal status (or force parsing for live UI if retryFailed) + status: mode === "retryFailed" ? "parsing" : doc.status, + errorMessage: null, + }; + } + + const pdfs = extractAllPdfParts(state.messages); + if (pdfs.length === 0) return makeError("no PDF file parts found"); + + const base = getR2PublicBaseUrl(); + + // ponytail: per-PDF processing runs in parallel — each PDF is + // independent, failures isolated to one entry, and the LRU on + // findKbDocumentByContentHash makes repeat lookups free within a + // single invocation. + const processed = await Promise.all( + pdfs.map(async ({ messageIndex, filePart }): Promise => { + const url = filePart.url || filePart.data; + const r2Key = r2KeyFromPublicUrl(url, base); + try { + const attachment = await findAttachmentByR2Key(userId, r2Key); + if (!attachment) { + return { + messageIndex, + filePart, + docId: null, + attachmentId: null, + r2Key, + title: null, + contentHash: null, + pipelineStatus: "unknown", + errorMessage: "attachment not found", + }; + } + const contentHash = attachment.sha256 ?? `r2key:${attachment.r2Key}`; + + let existing = await findKbDocumentByContentHash(userId, contentHash); + if (!existing) existing = await findKbDocumentByAttachmentId(userId, attachment.id); + if (existing) { + // ponytail: dedup short-circuit only when the row already + // ran the pipeline to completion (`success`/`failed`) — a + // stale `pending` row means a prior kbAgent run never + // landed its status writes (the t- prefix bug, dropped + // dispatch, etc.), and a second dispatch should re-process + // the file pointing at the SAME row id so the row actually + // flips to `success`. Falling through to the fresh-create + // branch would insert a NEW docId and leave the stale row + // stuck forever — instead, reuse existing.id with a `new` + // pipelineStatus so splitFileToPageNode writes back to it. + const forceRerun = config?.configurable?.forceRerun ?? false; + if (!forceRerun && (existing.status === "success" || existing.status === "failed")) { + return { + messageIndex, + filePart, + docId: existing.id, + attachmentId: attachment.id, + r2Key: attachment.r2Key, + title: attachment.name, + contentHash, + pipelineStatus: "dedup", + errorMessage: existing.errorMessage, + existingStatus: existing.status, + }; + } + // pending/parsing rows: reuse the row, re-process the file + return { + messageIndex, + filePart, + docId: existing.id, + attachmentId: attachment.id, + r2Key: attachment.r2Key, + title: attachment.name, + contentHash, + pipelineStatus: "new", + errorMessage: null, + existingStatus: existing.status, + }; + } + + const docId = `d-${randomUUID()}`; + return { + messageIndex, + filePart, + docId, + attachmentId: attachment.id, + r2Key: attachment.r2Key, + title: attachment.name, + contentHash, + pipelineStatus: "new", + errorMessage: null, + }; + } catch (err) { + return { + messageIndex, + filePart, + docId: null, + attachmentId: null, + r2Key, + title: null, + contentHash: null, + pipelineStatus: "failed", + errorMessage: (err as Error).message, + }; + } + }), + ); + + // ponytail: persist a "parsing" row for every new doc NOW so the + // Settings UI sees the doc immediately (2s poll picks it up), and so + // a later OCR / chunk failure still leaves a row in kb_documents — + // resolveKbRefs then renders "[Failed: ...]" instead of silently + // dropping the document context. + const folder = await ensureDefaultKbFolder(userId, "Attachments"); + const newDocs = processed.filter( + (p) => p.pipelineStatus === "new" && p.docId !== null && p.attachmentId !== null, + ); + + await Promise.allSettled( + newDocs.map(async (pf) => { + try { + await insertKbDocument({ + id: pf.docId!, + userId, + folderId: folder.id, + attachmentId: pf.attachmentId!, + title: pf.title ?? "untitled", + contentType: "application/pdf", + contentHash: pf.contentHash!, + status: "parsing", + errorMessage: null, + }); + } catch (err) { + // 23505 unique_violation: a row with this id already exists + // (the dedup-pending branch reuses the existing docId so the + // OCR pipeline writes its status back to that row). Flip the + // row from `pending` to `parsing` to surface progress. + // ponytail: Drizzle wraps PostgresError in DrizzleQueryError — + // top-level `err.code` is undefined; the actual pg code lives + // on `err.cause.code`. Same lookup pattern as + // ensureDefaultKbFolder above (lib/kb/queries.ts:77). + const code = + (err as { code?: string }).code ?? (err as { cause?: { code?: string } }).cause?.code; + if (code === "23505") { + try { + await updateKbDocumentStatus(userId, pf.docId!, { + status: "parsing", + errorMessage: null, + }); + } catch (statusErr) { + console.error( + `kbAgent prepareKBDataNode: recovery UPDATE failed for ${pf.docId}`, + statusErr, + ); + } + return; + } + console.error(`kbAgent prepareKBDataNode: insertKbDocument failed for ${pf.docId}`, err); + } finally { + if (threadId) { + await insertKbObservability({ + docId: pf.docId!, + threadId, + parentMessageId, + runId, + source, + mode, + }); + } + } + }), + ); + + return { + userId, + processedFiles: processed, + status: "parsing", + }; +} + +// --------------------------------------------------------------------------- +// Node 2: splitFileToPageNode — PDF rendering + text extraction + R2 upload +// --------------------------------------------------------------------------- + +async function splitFileToPageNode(state: KbAgentStateShape): Promise> { + const newDocs = state.processedFiles.filter( + (p) => p.pipelineStatus === "new" && p.docId !== null && p.r2Key !== null, + ); + + const pagesByDocId: Record = {}; + const updatedProcessed = state.processedFiles.map((p) => ({ ...p })); + + for (const pf of newDocs) { + try { + const pdfBytes = await getObject(pf.r2Key!); + const [rendered, extracted] = await Promise.all([ + screenshotPdf({ pdfBytes, dpi: 250 }), + extractPdfText({ pdfBytes }), + ]); + const textByPage = Object.fromEntries(extracted.map((e) => [e.pageIndex, e.text])); + const pages: PageResult[] = await Promise.all( + rendered.map(async (p) => { + const key = `kb-tmp/${state.userId}/${pf.docId}/page-${p.pageIndex}.png`; + const imageUrl = await uploadKbImage({ key, body: p.png }); + return { + pageIndex: p.pageIndex, + imageUrl, + markdown: "", + referenceText: textByPage[p.pageIndex] ?? "", + status: "pending", + }; + }), + ); + pagesByDocId[pf.docId!] = pages; + if (state.userId && pf.docId) { + await updateKbDocumentStatus(state.userId, pf.docId, { + status: "parsing", + pages, + }); + } + } catch (err) { + // ponytail: render failure flips this PDF to "failed" — keep the + // docId so the rewritten HumanMessage still carries a kb_ref + // sibling, and persist the failure on the row so the [Failed: ...] + // placeholder resolves correctly in resolveKbRefs. + const idx = state.processedFiles.indexOf(pf); + if (idx >= 0) { + updatedProcessed[idx] = { + ...updatedProcessed[idx], + pipelineStatus: "failed", + errorMessage: (err as Error).message, + }; + } + + console.error("kbAgent splitFileToPageNode", err); + + if (state.userId && pf.docId) { + try { + await updateKbDocumentStatus(state.userId, pf.docId, { + status: "failed", + errorMessage: (err as Error).message, + }); + } catch (statusErr) { + console.error( + `kbAgent splitFileToPageNode: updateKbDocumentStatus failed for ${pf.docId}`, + statusErr, + ); + } + } + } + } + + return { pagesByDocId, processedFiles: updatedProcessed }; +} + +// --------------------------------------------------------------------------- +// Node 3: pageToMarkdownNode — OCR + fullMarkdown + fire-and-forget chunk +// --------------------------------------------------------------------------- + +async function pageToMarkdownNode(state: KbAgentStateShape) { + // ponytail: chunksOnly dispatch reuses doc.pages[].markdown as-is, + // so OCR is by definition out of scope. Returning the original + // pagesByDocId + processedFiles without any work keeps the graph + // edges valid while skipping a full re-render that the user + // explicitly asked NOT to do. + if (state.mode === "chunksOnly") { + return { + pagesByDocId: state.pagesByDocId, + processedFiles: state.processedFiles, + }; + } + + const ocrModel = await getOcrModel(); + + const system = new SystemMessage(KB_OCR_PAGE_PROMPT); + const structured = ocrModel.withStructuredOutput(ocrPageSchema, { + method: "jsonSchema", + strict: true, + }); + + // ponytail: one p-queue across ALL docs — caps total apimart + // concurrency at OCR_CONCURRENCY regardless of how many PDFs were + // in flight. Per-doc pages still complete in order (Promise.all per + // doc preserves it). + const queue = new PQueue({ concurrency: KB_OCR_CONCURRENCY }); + + const newDocs = state.processedFiles.filter( + (p) => + p.pipelineStatus === "new" && p.docId !== null && state.pagesByDocId[p.docId] !== undefined, + ); + + const updatedPagesByDocId: Record = { ...state.pagesByDocId }; + const updatedProcessed = state.processedFiles.map((p) => ({ ...p })); + + const results = await Promise.allSettled( + newDocs.map((pf) => + queue.add(async () => { + const pages = state.pagesByDocId[pf.docId!]; + const controller = new AbortController(); + let hasFailed = false; + + const ocrResults = await Promise.all( + pages.map(async (p) => { + if ( + state.mode === "retryFailed" && + (p.markdown ?? "").trim().length > 0 && + !p.errorMessage + ) { + return p; + } + + if (hasFailed || controller.signal.aborted) { + return { + ...p, + markdown: "", + status: "failed" as const, + errorMessage: + "Bypassed: OCR aborted due to another page failure in the same document", + }; + } + + const contentParts: Array<{ type: string; [key: string]: unknown }> = [ + { type: "image_url", image_url: { url: p.imageUrl } }, + ]; + if (p.referenceText?.trim()) { + contentParts.push({ + type: "text", + text: `Reference text extracted directly from the PDF (may contain layout noise — trust the image for structure):\n\n${p.referenceText}`, + }); + } + try { + if (hasFailed || controller.signal.aborted) { + return { + ...p, + markdown: "", + status: "failed" as const, + errorMessage: + "Bypassed: OCR aborted due to another page failure in the same document", + }; + } + const out = (await structured.invoke( + [system, new HumanMessage({ content: contentParts })], + { tags: ["nostream"], signal: controller.signal }, + )) as z.infer; + return { + ...p, + markdown: out.markdown.trim(), + status: "success" as const, + errorMessage: undefined, + }; + } catch (err) { + hasFailed = true; + controller.abort(); + console.error( + `kbAgent pageToMarkdownNode: OCR failed for doc ${pf.docId} page ${p.pageIndex}:`, + err, + ); + const isAborted = + err instanceof Error && + (err.name === "AbortError" || err.message?.toLowerCase().includes("abort")); + const msg = isAborted + ? "Bypassed: OCR aborted due to another page failure in the same document" + : err instanceof Error + ? err.message + : String(err); + return { ...p, markdown: "", status: "failed" as const, errorMessage: msg }; + } + }), + ); + return { docId: pf.docId!, pages: ocrResults, messageIndex: pf.messageIndex }; + }), + ), + ); + + results.forEach((r, i) => { + const pf = newDocs[i]; + if (r.status === "fulfilled") { + updatedPagesByDocId[r.value.docId] = r.value.pages; + } else { + const idx = state.processedFiles.indexOf(pf); + if (idx >= 0) { + updatedProcessed[idx] = { + ...updatedProcessed[idx], + pipelineStatus: "failed", + errorMessage: (r.reason as Error).message, + }; + } + } + }); + + const successfulDocIds: string[] = []; + const failedNewDocs: ProcessedFile[] = []; + + for (let i = 0; i < updatedProcessed.length; i++) { + const pf = updatedProcessed[i]; + if (pf.pipelineStatus !== "new" || !pf.docId) continue; + // ponytail: retryFailedChunks skips the page-level OCR check — + // pages[] is empty for chunksOnly-style dispatches (route + // already verified the doc reached a terminal status), and + // even when pages are populated this loop is checking the + // wrong axis (OCR), not the chunk axis the retry is fixing. + // Without this guard, the loop flips the stub entry to + // pipelineStatus='failed' on pages.length===0, which then + // makes routeAfterRewrite return END before generateChunkEmbed + // runs. Routing the retry through the same check would also + // downgrade chunks that legitimately have all-pages-markdown. + if (state.mode === "retryFailedChunks") { + successfulDocIds.push(pf.docId); + continue; + } + const pages = updatedPagesByDocId[pf.docId] ?? []; + const hasAnyFailedPage = pages.some((p) => !!p.errorMessage || !(p.markdown ?? "").trim()); + if (pages.length > 0 && !hasAnyFailedPage) { + successfulDocIds.push(pf.docId); + } else { + const pageErrors = pages + .map((p) => p.errorMessage) + .filter((e): e is string => !!e && e.length > 0); + const uniqueErrors = Array.from(new Set(pageErrors)); + const combinedError = + uniqueErrors.length === 1 + ? uniqueErrors[0] + : uniqueErrors.length > 1 + ? `OCR failed on some pages: ${uniqueErrors.join("; ")}` + : "some pages have empty markdown after OCR"; + updatedProcessed[i] = { + ...pf, + pipelineStatus: "failed", + errorMessage: combinedError, + }; + failedNewDocs.push(updatedProcessed[i]); + } + } + + const failedOcrDocs = updatedProcessed.filter( + (p) => + p.pipelineStatus === "failed" && + p.docId !== null && + state.pagesByDocId[p.docId] !== undefined && + !failedNewDocs.some((orig) => orig.docId === p.docId), + ); + const allFailedNew = [...failedNewDocs, ...failedOcrDocs]; + + if (state.userId) { + const userId = state.userId; + await Promise.allSettled([ + ...allFailedNew.map(async (p) => { + await updateKbDocumentStatus(userId, p.docId!, { + status: "failed", + errorMessage: p.errorMessage, + pages: updatedPagesByDocId[p.docId!] ?? null, + }); + }), + ...successfulDocIds.map(async (docId) => { + await updateKbDocumentStatus(userId, docId, { + status: "success", + pages: updatedPagesByDocId[docId], + }); + }), + ]); + } + + return { pagesByDocId: updatedPagesByDocId, processedFiles: updatedProcessed }; +} + +// --------------------------------------------------------------------------- + +// --------------------------------------------------------------------------- +// Node 5: rewriteMessagesNode — stamp kb_ref on file parts + compute status +// --------------------------------------------------------------------------- + +async function rewriteMessagesNode(state: KbAgentStateShape): Promise> { + // ponytail: chunksOnly dispatch has no messages to stamp — the + // synthetic fireIngestionRun HumanMessage (or empty payload) never + // carries a chat-context file part. Skip the stamp pass entirely + // and forward state.messages untouched. The empty filePart.url/data + // we constructed in prepareKBDataNode means fileToDoc below stays + // empty; this guard avoids the unnecessary iteration + isHumanLike + // rebuild cost. + if (state.mode === "chunksOnly" || state.mode === "retryFailed") { + return { + messages: state.messages, + status: state.status, + errorMessage: null, + }; + } + + const fileToDoc = new Map(); + for (const pf of state.processedFiles) { + if (pf.docId) { + const url = pf.filePart.url || pf.filePart.data; + fileToDoc.set(url, { docId: pf.docId, attachmentId: pf.attachmentId }); + } + } + + const messages = state.messages.map((m): BaseMessage => { + // ponytail: match BOTH HumanMessage instance AND plain + // `{type:"human", content:[...]}` rehydration form (see + // lib/kb/extract.ts isHumanLike comment — the standalone + // runs.create path produces the plain-object form after the + // MessagesValue reducer round-trips). Without the type fallback + // the rewrite skips the message and the kb_ref sibling never + // lands on the file part. + const mType = (m as { type?: unknown }).type; + const isHuman = m instanceof HumanMessage || mType === "human"; + if (!isHuman || !Array.isArray(m.content)) return m; + let changed = false; + const newContent: unknown[] = []; + for (const part of m.content) { + if (isFilePart(part)) { + // ponytail: already-stamped parts carry through untouched. + if (part.kb_ref) { + newContent.push(part); + continue; + } + const url = part.url || part.data; + const matched = fileToDoc.get(url); + if (!matched) { + // non-PDF or unknown/failed PDF with no docId → drop + changed = true; + continue; + } + changed = true; + // ponytail: stamp BOTH the kb_ref sibling AND the filename + // prefix on the same write. stampKbRefOnFilename is idempotent + // so a re-stamp is a no-op. + const baseFilename = + typeof part.filename === "string" + ? part.filename + : typeof part.metadata?.filename === "string" + ? part.metadata.filename + : undefined; + const stampedFilename = stampKbRefOnFilename(baseFilename, matched.docId); + newContent.push({ + ...part, + kb_ref: { + docId: matched.docId, + ...(matched.attachmentId ? { attachmentId: matched.attachmentId } : {}), + }, + filename: stampedFilename, + metadata: { ...part.metadata, filename: stampedFilename }, + }); + continue; + } + newContent.push(part); + } + if (!changed) return m; + return new HumanMessage({ content: newContent as never, id: m.id }); + }); + + // ponytail: status follows the loudest outcome. If anything failed + // (OCR / render) the run is "failed" overall. Otherwise "success". + const hasFailure = state.processedFiles.some((p) => p.pipelineStatus === "failed"); + const allUnknown = state.processedFiles.every((p) => p.pipelineStatus === "unknown"); + const newDocCount = state.processedFiles.filter((p) => p.pipelineStatus === "new").length; + const dedupCount = state.processedFiles.filter((p) => p.pipelineStatus === "dedup").length; + + let status: KbAgentStateShape["status"] = "success"; + let errorMessage: string | null = null; + if (allUnknown) { + status = "failed"; + errorMessage = "no PDF could be processed"; + } else if (hasFailure) { + status = "failed"; + const firstFailure = state.processedFiles.find((p) => p.pipelineStatus === "failed"); + errorMessage = firstFailure?.errorMessage ?? "kbAgent failed"; + } else if (newDocCount === 0 && dedupCount === 0) { + status = "failed"; + errorMessage = "no PDF could be processed"; + } + + // ponytail: dedup row status sync. When kbAgent dedupes onto an + // existing kb_document row whose prior run stalled (status=pending), + // the row never gets a write — splitFileToPage / pageToMarkdown + // filter updates to `pipelineStatus === "new"`. Read each dedup'd + // docId's CURRENT row state and write it back. Best-effort: a + // mid-write DB hiccup doesn't fail the run; the chat dedup + // contract just needs the row to eventually converge. + if (state.userId) { + const dedupRows = state.processedFiles.filter( + (p) => p.pipelineStatus === "dedup" && p.docId !== null, + ); + await Promise.allSettled( + dedupRows.map(async (pf) => { + const row = await findKbDocumentById(state.userId!, pf.docId!); + if (!row) return; + // Forward the current row state so the user sees the settled + // status (success / failed) — that's what they uploaded for. + // Skip write if it's already at terminal state (avoid + // rewriting a success to success). Only sync forward when + // the row is currently pending/parsing. + // ponytail: forward whatever the dedup target's actual current + // state is, including intermediate (`parsing`/`pending`). + // Re-running kbAgent on the dedup target SHOULD surface a + // real progress signal even if the prior run stalled. The + // row is read-fresh here so a row that was previously + // `pending` but is now `success` gets flipped to `success`. + await updateKbDocumentStatus(state.userId!, pf.docId!, { + status: row.status, + errorMessage: row.errorMessage, + pages: row.pages ?? null, + }); + }), + ); + } + + return { messages, status, errorMessage }; +} + +// --------------------------------------------------------------------------- +// ponytail: normalizeLightRagOut — pure mapping from the LLM's +// structured output to the shape persisted on kb_chunk. Shared by +// generateChunkEmbedNode's main path AND its retryFailedChunks +// branch so the two can never drift (e.g. one day the schema adds +// a new field and the other path forgets to map it). +// --------------------------------------------------------------------------- +function normalizeLightRagOut(out: z.infer) { + return { + entities: (out.entities ?? []).map((e) => ({ + name: e.name, + type: e.type, + description: e.description ?? "", + })), + relationships: (out.relationships ?? []).map((r) => ({ + source: r.source, + target: r.target, + relation: r.relation || r.type || "", + description: r.description ?? "", + })), + themes: out.themes ?? [], + }; +} + +// --------------------------------------------------------------------------- +// ponytail: resolveEntityAliasesForDoc — extracted from +// generateChunkEmbedNode's IIFE so the alignment pass is unit-testable +// in isolation. NOT a LangGraph node (would need polling for chunks +// to land — out of scope). Errors are swallowed (best-effort, same +// as the original inline try/catch). +// --------------------------------------------------------------------------- +export async function resolveEntityAliasesForDoc(args: { + userId: string; + docId: string; + documentTitle: string; + config?: RunnableConfig; +}): Promise { + const { userId, docId, documentTitle, config } = args; + try { + // 1. Fetch all successfully saved chunks for this document + const dbChunks = await findKbChunksByDocumentId(userId, docId); + const successChunks = dbChunks.filter((c) => c.status === "success"); + + // 2. Gather all unique entities + const allEntityNames = new Set(); + for (const c of successChunks) { + for (const e of c.entities ?? []) { + if (e.name) { + allEntityNames.add(e.name.trim()); + } + } + } + + const entityList = Array.from(allEntityNames); + // ponytail: skip when no entities OR only a single unique name — + // LLM alignment can't fold a singleton. Saves the roundtrip. + if (entityList.length <= 1) return; + + // 3. Define structured output schema for the alignment map + const alignmentSchema = z.object({ + mappings: z + .array( + z.object({ + original: z.string().describe("The original entity name variation found in the list"), + canonical: z.string().describe("The resolved canonical standard name to merge into"), + }), + ) + .describe("A list of name mappings to resolve aliases and variants"), + }); + + const extractModel = await getExtractModel(); + const systemMsg = new SystemMessage(KB_ENTITY_ALIGNMENT_SYSTEM_PROMPT); + const humanMsg = new HumanMessage( + `Document Title: ${documentTitle || "Unknown Document"}\n` + + `Extracted Entities List:\n${JSON.stringify(entityList, null, 2)}`, + ); + + // 4. Call LLM to find alignments + const alignmentResult = (await extractModel + .withStructuredOutput(alignmentSchema, { method: "jsonSchema", strict: true }) + .invoke([systemMsg, humanMsg], { ...config, tags: ["nostream"] })) as z.infer< + typeof alignmentSchema + >; + + // 5. Create mapping dictionary + const nameMap = new Map(); + if (alignmentResult?.mappings) { + for (const m of alignmentResult.mappings) { + const orig = m.original.trim(); + const canon = m.canonical.trim(); + if (orig && canon && orig !== canon) { + nameMap.set(orig.toLowerCase(), canon); + } + } + } + + // 6. Update database if any mappings found + if (nameMap.size === 0) return; + for (const c of successChunks) { + let chunkUpdated = false; + + // Standardize entities list + const updatedEntities = (c.entities ?? []).map((e) => { + const match = nameMap.get(e.name.trim().toLowerCase()); + if (match) { + chunkUpdated = true; + return { ...e, name: match }; + } + return e; + }); + + // Standardize relationships list + const updatedRelationships = (c.relationships ?? []).map((r) => { + let relUpdated = false; + let source = r.source; + let target = r.target; + + const matchSrc = nameMap.get(r.source.trim().toLowerCase()); + if (matchSrc) { + source = matchSrc; + relUpdated = true; + } + const matchTgt = nameMap.get(r.target.trim().toLowerCase()); + if (matchTgt) { + target = matchTgt; + relUpdated = true; + } + + if (relUpdated) { + chunkUpdated = true; + return { ...r, source, target }; + } + return r; + }); + + // Write-back to DB if this chunk had aligned elements + if (chunkUpdated) { + await updateKbChunkGraphData(c.id, updatedEntities, updatedRelationships); + } + } + } catch (alignErr) { + console.error( + `kbAgent resolveEntityAliasesForDoc: alignment failed for doc ${docId}:`, + alignErr, + ); + } +} + +// --------------------------------------------------------------------------- +// Node 6: generateChunkEmbedNode — chunk + embed + entity + insert +// ponytail: registered as a LangGraph node but the heavy work runs inside +// a fire-and-forget IIFE so the node returns in milliseconds and never +// blocks the main kbAgent or the RAG chat loop. RunnableConfig is +// captured for callback propagation into the entity-extract LLM call. +// --------------------------------------------------------------------------- + +async function generateChunkEmbedNode( + state: KbAgentStateShape, + config?: RunnableConfig, +): Promise> { + console.log( + `[kbAgent] Entering generateChunkEmbedNode, files=`, + state.processedFiles.map((p) => ({ docId: p.docId, status: p.pipelineStatus })), + ); + + // ponytail: the chunk + embed + entity-extract pass is the slow leg + // of the pipeline (per-row LLM calls). The chat path (mainAgent → + // kbAgent subgraph) wants it fire-and-forget so the graph node + // returns in milliseconds and the chat reply flow is unblocked. + // The Settings standalone path (source='kb-upload' / 'kb-reprocess') + // wants it awaited — the caller (`POST /api/kb/upload`, + // `POST /api/kb/documents/[id]/reprocess`) already polls the row's + // status, but awaiting here makes the route's 202 contract more + // honest: the entity-extract pass has either landed or thrown by + // the time the route returns, so a poll-then-fetch chunks sees the + // final state on the first try. + // + // Default: wait. Only the chat subgraph path fires kbAgent + // fire-and-forget so the reply flow isn't blocked. `source` defaults + // to "chat" when unset (matches line 176), which means direct + // callers (tests, future invokers) that don't stamp a source also + // get fire-and-forget — opt in by setting `source` to anything else, + // or override explicitly via `waitForChunks`. + const source = (config?.configurable as { source?: string } | undefined)?.source ?? "chat"; + const waitFromConfig = (config?.configurable as { waitForChunks?: boolean } | undefined) + ?.waitForChunks; + const waitForChunks = waitFromConfig ?? source !== "chat"; + + if (state.userId) { + const pendingChunks: Array> = []; + for (const pf of state.processedFiles) { + if (pf.pipelineStatus === "new" && pf.docId) { + const docId = pf.docId; + const userId = state.userId; + + console.log( + `[kbAgent] Starting ${waitForChunks ? "blocking" : "background"} chunking task for docId=${docId}`, + ); + const work = (async () => { + try { + // ponytail: retryFailedChunks is an input-source branch, + // not a parallel pipeline. The route handler has marked + // failed chunks as status='parsing' (preserving id/ + // ordinal/embedding/content). We skip OCR/markdown/ + // splitter/embedder, read those rows directly, and feed + // them into the same entity-extract loop + write-back + // the main path uses. This way the two paths share + // normalization / error handling / queue / concurrency + // — there is no "second IIFE" to drift out of sync. + // + // Type-wise: chunkInputs[] is the shared shape produced + // by either path. From the entity-extract loop's POV, + // it doesn't care whether the rows were just split or + // read from the DB. + type ChunkInput = { id: string; ordinal: number; content: string }; + + const isRetryFailedChunks = state.mode === "retryFailedChunks"; + + const doc = await findKbDocumentById(userId, docId); + if (!doc) { + throw new Error(`Document ${docId} not found`); + } + const docTitle = doc.title ?? "Unknown Document"; + + // ponytail: entity-extract LLM routes through the extract + // pool so admin can flag a cheaper model (e.g. gpt-4o-mini) + // for this work without forcing the same model on the + // extractModel default. Falls back to the extractModel + // pool when no extract-tagged model is registered (see + // getExtractModel). + const extractModel = await getExtractModel(); + const entityQueue = new PQueue({ concurrency: KB_ENTITY_CONCURRENCY }); + + // ponytail: input source. + // - retryFailedChunks: SELECT WHERE status='parsing' + // (route already marked them). id/ordinal/content + // come from the existing row; embedding is preserved + // verbatim on the row. + // - other modes: MarkdownTextSplitter over + // pages[].markdown, then fresh chunkIds + embeddings. + // The two paths converge here into a single + // `chunkInputs` array — everything downstream + // (entity-extract, normalize, write-back) is identical. + let chunkInputs: ChunkInput[]; + let totalChunksForPrompt: number; + + if (isRetryFailedChunks) { + const existing = await findKbChunksByDocumentId(userId, docId); + const retryTargets = existing.filter((c) => c.status === "parsing"); + totalChunksForPrompt = existing.length; + console.log( + `[kbAgent] retryFailedChunks: docId=${docId}, ${retryTargets.length} to re-extract (out of ${totalChunksForPrompt})`, + ); + chunkInputs = retryTargets.map((c) => ({ + id: c.id, + ordinal: c.ordinal, + content: c.content, + })); + if (chunkInputs.length === 0) { + invalidateKbDoc(userId, docId); + return; + } + } else { + const pages = (doc.pages ?? []) as Array<{ + pageIndex: number; + imageUrl: string; + markdown: string; + }>; + const fullMarkdown = pages + .map((p) => p.markdown) + .filter((m) => m && m.length > 0) + .join("\n\n"); + console.log( + `[kbAgent] Background task: loaded docId=${docId}, pages count=${pages.length}, fullMarkdown length=${fullMarkdown.length}`, + ); + if (!fullMarkdown) { + throw new Error(`Document ${docId} has no markdown content extracted yet`); + } + + const embedder = await getEmbeddingModel(); + const lengthSplitter = new MarkdownTextSplitter({ + chunkSize: KB_CHUNK_SIZE, + chunkOverlap: KB_CHUNK_OVERLAP, + }); + + const splitDocs = await lengthSplitter.createDocuments([fullMarkdown]); + const texts = splitDocs.map((d) => d.pageContent); + const embeddings = await embedder.embedDocuments(texts); + + // ponytail: schema expects vector(1024) (kb_chunk.embedding + + // HNSW index). If the embedder returns anything else, pgvector + // rejects every insert with 22P02 — caught too late to be useful. + // Fail fast with a single clear sentence instead. + const actualDim = embeddings[0]?.length ?? 0; + if (actualDim !== EMBEDDING_DIM) { + throw new Error( + `embedding dimension mismatch: schema expects ${EMBEDDING_DIM}, embedder returned ${actualDim}. Update lib/kb/schema.ts EMBEDDING_DIM + run the matching ALTER COLUMN migration.`, + ); + } + + const chunkIds = texts.map(() => `c-${randomUUID()}`); + + totalChunksForPrompt = texts.length; + chunkInputs = texts.map((text, i) => ({ + id: chunkIds[i]!, + ordinal: i, + content: text, + })); + + // ponytail: 3-stage chunk lifecycle for fresh chunks — + // INSERT all rows at status='pending', then + // markAllParsingForDocInTx in the same tx so polling + // either sees 'parsing' or nothing, never a stuck + // 'pending' frame. + await withKbTx(async (tx) => { + await insertKbChunks( + tx, + texts.map((text, i) => ({ + id: chunkIds[i]!, + documentId: docId, + ordinal: i, + content: text, + embedding: embeddings[i] ?? [], + entities: [], + // status defaults to 'pending' on insert. + })) as never, + ); + await markAllKbChunksParsingForDocInTx(tx, docId); + }); + console.log( + `[kbAgent] Background task: successfully inserted ${texts.length} chunks for docId=${docId}`, + ); + } + + // ponytail: per-row, streaming write-back. Each task writes + // ITS OWN row the moment its entity LLM resolves — no + // `Promise.all` to await siblings, no `results[]` buffer. + // The 2s UI poll sees chunks flip from `parsing` → + // `success`/`failed` one by one as each LLM lands. + // Otherwise the row status lives in two states + // (`parsing` for ~30s then `success` everywhere) and the + // preview's "Indexed N/N, K failed" never moves while + // the pipeline is still grinding. Each task wraps its + // DB writes in a per-row try/catch so a single row's + // UPDATE rejection can't crash the queued task and + // silence its siblings. + + // ponytail: debug bail-out kept behind a const toggle so + // oxlint's no-unreachable stays quiet. Flip to `true` to + // isolate one of the Promise.allSettled arms (entity-LLM + // vs insert vs per-row write-back) when reproducing a + // partial-pipeline failure. + if (SKIP_CHUNK_TO_ENTRIES) { + return; + } + await Promise.allSettled( + chunkInputs.map((chunk) => + entityQueue.add(async (): Promise => { + const chunkId = chunk.id; + const ordinal = chunk.ordinal; + const text = chunk.content; + + const systemMessage = new SystemMessage(KB_ENTITY_EXTRACTION_SYSTEM_PROMPT); + const humanMessage = new HumanMessage( + `Context Document Title: [${docTitle}]\n` + + `Chunk: [${ordinal + 1} / ${totalChunksForPrompt}]\n\n` + + `Text to extract:\n${text}`, + ); + + try { + const out = (await extractModel + .withStructuredOutput(lightRagSchema, { method: "jsonSchema", strict: true }) + .invoke([systemMessage, humanMessage], { + ...config, + tags: ["nostream"], + })) as z.infer; + // write-back: entities + status='success' in one go + // so the row never sits at success with a blank + // entities field. + await Promise.allSettled([ + updateKbChunkForSuccess(chunkId, normalizeLightRagOut(out)), + markKbChunkSuccess(chunkId), + ]); + } catch (err) { + // entity-extract LLM failure (or per-row DB + // write) — surface as kb_chunk.status='failed' + // + errorMessage. kb_document stays success + // (Step 3 contract). + const msg = err instanceof Error ? (err as any).message : String(err); + console.error( + `kbAgent generateChunkEmbedNode: chunk ${chunkId} failed (doc ${docId} ordinal ${ordinal}): ${msg}`, + err as any, + ); + try { + await Promise.allSettled([ + updateKbChunkForFailure(chunkId, msg), + markKbChunkFailed(chunkId, msg), + ]); + } catch (writeErr) { + console.error( + `kbAgent generateChunkEmbedNode: failed-row write-back itself errored for chunk ${chunkId}:`, + writeErr, + ); + } + } + }), + ), + ); + + // ponytail: invalidate the doc cache so the next + // poll-from-DB read picks up the fresh chunk counts. + // For retryFailedChunks this is also where the run + // ends — no entity alignment pass (no new entities + // were introduced; existing aligned entities are + // untouched because the chunk content didn't change). + invalidateKbDoc(userId, docId); + + // ponytail: alignment pass only runs for modes that + // introduce new entities from scratch. retryFailedChunks + // re-extracts entities for already-existing chunks — + // alignment over those would be a no-op (entities + // already share names) AND would pointlessly spend LLM + // tokens. full / chunksOnly / retryFailed all produce + // fresh chunks where alignment is genuinely useful. + if (!isRetryFailedChunks) { + await resolveEntityAliasesForDoc({ + userId, + docId, + documentTitle: docTitle, + config, + }); + } + } catch (err) { + // ponytail: WHOLE-batch failure (embedding dim mismatch, + // DB write rejection). kb_document.status stays at + // 'success' (it was flipped by imageToMarkdownNode) — + // chunks are a downstream derived store, and a chunk + // pipeline crash shouldn't downgrade the doc itself. + // The Settings UI surfaces "0/47 chunks indexed" via + // the chunk count roll-up, and the user can rebuild + // chunks via Reprocess > "Only rebuild chunks". + const pgErr = err as Error & { + code?: string; + detail?: string; + hint?: string; + }; + const reason = pgErr.code + ? `${pgErr.code}: ${pgErr.detail ?? pgErr.message}${pgErr.hint ? ` (${pgErr.hint})` : ""}` + : pgErr.message; + console.error( + `kbAgent generateChunkEmbedNode: batch failure for doc ${docId}: ${reason}`, + pgErr, + ); + } + })(); + pendingChunks.push(work); + } + } + if (waitForChunks && pendingChunks.length > 0) { + // ponytail: await the whole-doc IIFEs in parallel — each one is + // already independent (per-row entity-extract queue is internal + // to that doc). Waiting lets the route's 202 contract land with + // chunks indexed; one doc's failure doesn't fail its siblings + // because the IIFE catches its own errors. + await Promise.allSettled(pendingChunks); + } + } + return {}; +} + +// --------------------------------------------------------------------------- +// Graph builder + dual compilation +// --------------------------------------------------------------------------- + +function routeAfterRewrite(state: KbAgentStateShape): string | typeof END { + const hasNew = state.processedFiles.some((p) => p.pipelineStatus === "new"); + console.log( + `[kbAgent] routeAfterRewrite: hasNew=${hasNew}, files=`, + state.processedFiles.map((p) => ({ docId: p.docId, status: p.pipelineStatus })), + ); + if (hasNew) { + console.log(`[kbAgent] Routing to generateChunkEmbed`); + return "generateChunkEmbed"; + } + console.log(`[kbAgent] Routing to END`); + return END; +} + +const builder = new StateGraph(KbAgentState) + .addNode("prepareKBData", prepareKBDataNode) + .addNode("splitFilePage", splitFileToPageNode) + .addNode("pageToMarkdown", pageToMarkdownNode) + .addNode("rewriteMessages", rewriteMessagesNode) + .addNode("generateChunkEmbed", generateChunkEmbedNode) + .addEdge(START, "prepareKBData") + .addEdge("prepareKBData", "splitFilePage") + .addEdge("splitFilePage", "pageToMarkdown") + .addEdge("pageToMarkdown", "rewriteMessages") + .addConditionalEdges("rewriteMessages", routeAfterRewrite, { + generateChunkEmbed: "generateChunkEmbed", + __end__: END, + }) + .addEdge("generateChunkEmbed", END); + +// ponytail: TWO compiled graphs from the same builder. +// - `kbAgent` (in-process subgraph): empty subgraphCheckpointerConfig. +// mainAgent (`backend/agent.ts`) calls `.addNode("kbAgent", kbAgent)` +// and the parent graph's checkpointer governs persistence. The +// existing direct-node-call tests (tests/backend/kb-agent.test.ts) +// bypass the parent and never set a thread_id — they rely on this +// config being empty so no checkpoint write is attempted. +// - `graph` (standalone top-level assistant registered in +// langgraph.json): gets the global checkpointer + store + callbacks +// so observability + credit tracking + per-thread persistence work +// for the synthetic "ingest this file" runs dispatched from +// `lib/kb/ingest.fireIngestionRun()`. +// +// Both keep `name: "kbAgent"` so the runtime identifies the standalone +// assistant under the expected key. +export const kbAgent = builder.compile({ + name: "kbAgent", + ...subgraphCheckpointerConfig, +}); + +const standaloneCompiled = builder.compile({ + name: "kbAgent", + checkpointer, + store, +}); + +type WithConfigPregel = (config: Record) => typeof standaloneCompiled; +export const graph = (standaloneCompiled.withConfig as unknown as WithConfigPregel)({ + callbacks: [capturingHandler, creditTrackingHandler], +}); +void END; diff --git a/backend/agent/weather-agent.ts b/backend/agent/weather-agent.ts index 39fdeafa..d6c96b51 100644 --- a/backend/agent/weather-agent.ts +++ b/backend/agent/weather-agent.ts @@ -9,8 +9,9 @@ import { CommonAgentState } from "@/backend/state"; import { buildSystemMessageWithMemory, loadThreadSummariesForPrompt, - trimMessagesForInvoke, + prepareMessagesForInvoke, } from "@/backend/memory/template"; +import { extractUserId } from "@/backend/memory/recall"; import { subgraphCheckpointerConfig } from "@/backend/checkpointer"; // Weather agent: a focused sub-agent that owns the RAG-style weather @@ -32,7 +33,11 @@ async function weatherModelNode( // and drop the original turns from the input array. state.messages // is NEVER touched. const threads = await loadThreadSummariesForPrompt(config); - const history = trimMessagesForInvoke(messages, threads?.summaries ?? []); + const history = await prepareMessagesForInvoke( + messages, + threads?.summaries ?? [], + extractUserId(config) ?? undefined, + ); const sysMsg = await buildSystemMessageWithMemory(WEATHER_AGENT_PROMPT, config, threads); const response = await ( diff --git a/backend/background-agent.ts b/backend/background-agent.ts index edd679a1..7d44f88a 100644 --- a/backend/background-agent.ts +++ b/backend/background-agent.ts @@ -44,7 +44,7 @@ export async function touchLastMessageNode( if (typeof threadId === "string" && threadId.length > 0) { await touchLastMessageAt(threadId); } - // Empty state update — the messages reducer on CommonAgentState would + // Empty state update — the messages reducer on BackgroundAgentState would // turn `[]` into a no-op (no replace, no add). Same side-effect-only // contract as the original afterAgentNode. return { messages: [] }; diff --git a/backend/memory/template.ts b/backend/memory/template.ts index f23ab2bf..6e3a0c72 100644 --- a/backend/memory/template.ts +++ b/backend/memory/template.ts @@ -11,6 +11,7 @@ import { } from "@/backend/memory/recall"; import { MEMORY_AUGMENTED_PROMPT_TEMPLATE } from "@/backend/prompt/system"; import { formatSummaryText } from "@/lib/langgraph/format-summary"; +import { resolveKbRefs } from "@/lib/kb/resolve"; // ponytail: the system prompt carries TWO dynamic blocks: // = the user's saved profile (keys + values), plus OAuth @@ -106,11 +107,24 @@ export function formatThreadsForPrompt(threads: ThreadSummariesPayload): string // Pure function — `template.test.ts` pins every branch (no summary, // no humans, single summary, multiple summaries, last human covered, // out-of-order store rows, tool interleaving preserved). -export function trimMessagesForInvoke( +// +// v2 (issue #13): kb_ref resolution happens here too — the LLM only +// ever sees resolved text. Async because resolveKbRefs awaits the +// LRU-cached DB lookup. +// +// v3 (issue #13): resolveKbMentions leaves the `:kb-doc[…]{id=…}` / +// `:kb-folder[…]{id=…}` directive tokens in the HumanMessage text +// (the LLM reads them and calls `search_kb` itself with the right +// filter). The resolver only injects a synthetic search_kb +// ToolMessage for the rare 0-chunk-fallback case — see +// `lib/kb/resolve-mentions.ts`. +export async function prepareMessagesForInvoke( messages: BaseMessage[], summaries: ThreadSummariesPayload["summaries"], -): BaseMessage[] { - const noSystem = messages.filter((m) => !(m instanceof SystemMessage)); + userId?: string, +): Promise { + const resolved = userId ? await resolveKbRefs(messages, userId) : messages; + const noSystem = resolved.filter((m) => !(m instanceof SystemMessage)); const humanIndices: number[] = []; for (let i = 0; i < noSystem.length; i++) { if (noSystem[i] instanceof HumanMessage) humanIndices.push(i); @@ -119,11 +133,13 @@ export function trimMessagesForInvoke( for (const s of summaries) { if (s.endMessageIndex > maxEnd) maxEnd = s.endMessageIndex; } + // ponytail: no summary OR no human turns in the array → nothing to // trim. Returning noSystem (not messages) still drops a stray // SystemMessage if one slipped in — the strip pass is unconditional. if (maxEnd < 0 || humanIndices.length === 0) return noSystem; const trimTo = maxEnd + 1 < humanIndices.length ? humanIndices[maxEnd + 1] : noSystem.length; + return noSystem.slice(trimTo); } diff --git a/backend/model.ts b/backend/model.ts index 8aea0c40..bfbf2ac7 100644 --- a/backend/model.ts +++ b/backend/model.ts @@ -1,12 +1,18 @@ -import { ChatOpenAI } from "@langchain/openai"; +import { ChatOpenAI, OpenAIEmbeddings } from "@langchain/openai"; +import type { Embeddings } from "@langchain/core/embeddings"; import { getChatModelFromDB, + getEmbeddingModelFromDB, + getExtractModelFromDB, + getOcrModelFromDB, + getRerankModelFromDB, invalidateModelCache, type GetChatModelOpts, + type RerankModel, } from "@/lib/provider/model-registry"; -function buildEnvModel(): ChatOpenAI { +function buildEnvChatModel(): ChatOpenAI { return new ChatOpenAI({ model: process.env.OPENAI_MODEL ?? "gpt-4o-mini", apiKey: process.env.OPENAI_API_KEY, @@ -25,6 +31,28 @@ function buildEnvModel(): ChatOpenAI { }); } +// ponytail: ocr reuses the chat-model env vars today (vision-capable +// chat models handle image_url content for OCR). When a non-chat +// vision upstream lands (e.g. a vision-only OCR service) this +// builder splits off. +function buildEnvOcrModel(): ChatOpenAI { + return buildEnvChatModel(); +} + +function buildEnvEmbeddingModel(): Embeddings { + return new OpenAIEmbeddings({ + model: process.env.OPENAI_EMBEDDING_MODEL ?? "text-embedding-3-small", + apiKey: process.env.OPENAI_API_KEY, + ...(process.env.OPENAI_BASE_URL + ? { + configuration: { + baseURL: process.env.OPENAI_BASE_URL, + }, + } + : {}), + }); +} + /** * Canonical entry point for runtime chat-model lookup. Tries the DB-backed * registry first (with LRU caching + admin CUD invalidation); on miss / @@ -40,7 +68,58 @@ export async function getChatModel(opts: GetChatModelOpts = {}): Promise { + try { + return (await getOcrModelFromDB(opts)) as ChatOpenAI; + } catch { + return buildEnvOcrModel(); + } +} + +/** + * Embedding model entry point. Same fallback chain as chat: registry + * first, env on miss. Returns Embeddings interface (not chat) — caller + * chains `.embedDocuments` / `.embedQuery`, not `.invoke`. + */ +export async function getEmbeddingModel(opts: GetChatModelOpts = {}): Promise { + try { + return await getEmbeddingModelFromDB(opts); + } catch { + return buildEnvEmbeddingModel(); + } +} + +/** + * ponytail: structured-output extraction (KB chunk → entity / relationship / + * theme triples) routes to `kind="extract"` models. Today's chat-LLM + * call sites for this work resolve here. `getExtractModelFromDB` falls + * back to the chat pool when no extract-tagged model is registered, + * so the wrapper doesn't need its own env fallback — by the time the + * inner fallback fires, we've already degraded to chat. + */ +export async function getExtractModel(opts: GetChatModelOpts = {}): Promise { + return (await getExtractModelFromDB(opts)) as ChatOpenAI; +} + +/** + * Retrieve the rerank model. If none is configured in DB, returns null, + * allowing caller to gracefully skip reranking. + */ +export async function getRerankModel(opts: GetChatModelOpts = {}): Promise { + try { + return await getRerankModelFromDB(opts); + } catch { + return null; } } diff --git a/backend/node/prepare-data-node.ts b/backend/node/prepare-data-node.ts new file mode 100644 index 00000000..3449c5f6 --- /dev/null +++ b/backend/node/prepare-data-node.ts @@ -0,0 +1,26 @@ +import type { RunnableConfig } from "@langchain/core/runnables"; +import type { BaseMessage } from "@langchain/core/messages"; + +// ponytail: per-turn data-prep node. Currently a pass-through — no +// pre-LLM transformation happens here. The KB @-mention flow is +// driven entirely by the LLM reading the directive token +// (':kb-document[label]{documentId=…}' / ':kb-folder[label]{folderId=…}') +// from the HumanMessage text and calling search_kb / list_documents +// with the right filter. +// +// The node is reserved for future per-turn Message transforms: +// stripping tool_call_id duplicates, hydrating attachments, or +// compressing long history. Right now there's nothing to do, so +// messages pass through unchanged. The node still exists because (a) +// the router reads message SHAPE (PDFs → kbAgent via hasUnprocessedPdf, +// tool calls → resume) and downstream prepareMessagesForInvoke expects +// a `messages` field on the returned partial state, and (b) the +// router-loop edge (kbAgent → routerAgent) bypasses this node by +// design — we don't want to re-inject after kbAgent has stamped +// `kb_ref` onto the PDF. +export async function prepareDataNode( + state: { messages: BaseMessage[] }, + _config?: RunnableConfig, +): Promise<{ messages: BaseMessage[] }> { + return { messages: state.messages }; +} diff --git a/backend/node/rename-thread-agent-node.ts b/backend/node/rename-thread-agent-node.ts index 00ce4585..f1414815 100644 --- a/backend/node/rename-thread-agent-node.ts +++ b/backend/node/rename-thread-agent-node.ts @@ -2,6 +2,7 @@ import { BaseMessage, HumanMessage, SystemMessage } from "@langchain/core/messag import { renameThread } from "@/lib/threads/queries"; import { getChatModel } from "@/backend/model"; import { RENAME_THREAD_PROMPT } from "@/backend/prompt/system"; +import { stripFileParts } from "@/lib/kb/extract"; export async function renameThreadAgentNode( state: { messages: BaseMessage[] }, @@ -12,7 +13,9 @@ export async function renameThreadAgentNode( const response = await ( await getChatModel() - ).invoke([new SystemMessage(RENAME_THREAD_PROMPT), firstUserMessage], { tags: ["nostream"] }); + ).invoke([new SystemMessage(RENAME_THREAD_PROMPT), stripFileParts(firstUserMessage)], { + tags: ["nostream"], + }); const trimmed = (typeof response.content === "string" ? response.content : "").trim(); const threadId = config.configurable?.thread_id; diff --git a/backend/node/router-agent-node.ts b/backend/node/router-agent-node.ts index 3bf6eff1..3cd5a47d 100644 --- a/backend/node/router-agent-node.ts +++ b/backend/node/router-agent-node.ts @@ -4,18 +4,23 @@ import { z } from "zod"; import { getChatModel } from "@/backend/model"; import { ROUTER_AGENT_PROMPT } from "@/backend/prompt/system"; +import { hasUnprocessedPdf, stripFileParts } from "@/lib/kb/extract"; +import { prepareMessagesForInvoke } from "@/backend/memory/template"; +import { extractUserId } from "@/backend/memory/recall"; -// Router agent: inspects the latest user message and decides which -// sub-agent should handle the turn. Output is a zod-validated object -// so the parser never has to handle malformed JSON — the schema is -// the contract. +// ponytail: v3 router. Two short-circuits and a fallback: +// 1. ANY HumanMessage has an unprocessed PDF → route to kbAgent. +// 2. Otherwise → resolve kb_refs + trim, ask the LLM. // -// Method: `functionCalling` with the `nostream` invocation tag. The -// schema is registered as a tool so the model emits a `tool_call` on -// the AIMessage; we discard the AIMessage and return only the parsed -// `routerDecision` to the state. `tags: ["nostream"]` keeps the run's -// token stream free of the router's internal reasoning. +// RouterNode is intentionally DB-free: kbAgent owns the contentHash +// dedup probe. The router only inspects message shape. + const RouteDecisionSchema = z.object({ + next: z.enum(["weatherAgent", "chatAgent", "cryptoAgent", "codeAgent", "kbAgent"]), +}); + +// delete RouteDecisionSchema kbAgent +const InvokeRouteDecisionSchema = z.object({ next: z.enum(["weatherAgent", "chatAgent", "cryptoAgent", "codeAgent"]), }); @@ -25,21 +30,37 @@ export async function routerAgentNode( state: { messages: BaseMessage[] }, config?: RunnableConfig, ): Promise<{ routerDecision: RouterDecision }> { - // ponytail: router is a yes/no classifier on the CURRENT turn. Full - // history is a token-cost move AND can distract the classifier into - // routing off a stale topic. The trailing HumanMessage is always the - // current turn — the router runs before any AI reply for this turn - // exists. const lastUserMessage = state.messages.findLast((m) => m instanceof HumanMessage); + + // Short-circuit: any HumanMessage has an unprocessed PDF → kbAgent. + // kbAgent now processes every PDF across every HumanMessage in one + // invocation, so the router only needs to know "is there still + // work to do?" — not which turn owns it. + if (hasUnprocessedPdf(state.messages)) { + return { routerDecision: { next: "kbAgent" } }; + } + const system = new SystemMessage(ROUTER_AGENT_PROMPT); - const invokeMessages = lastUserMessage ? [system, lastUserMessage] : [system]; + const userId = extractUserId(config); + const trimmed = await prepareMessagesForInvoke(state.messages, [], userId ?? undefined); + + const trimmedClean = trimmed.map(stripFileParts); + const lastClean = lastUserMessage ? stripFileParts(lastUserMessage) : null; + + const invokeMessages = lastClean + ? [system, lastClean, ...trimmedClean.filter((m) => m.id !== lastClean.id)] + : [system, ...trimmedClean]; + // LLM route — schema now includes kbAgent for completeness, but + // the explicit short-circuit above means we never reach this with a + // new PDF. const decision = (await ( await getChatModel() ) - .withStructuredOutput(RouteDecisionSchema, { + .withStructuredOutput(InvokeRouteDecisionSchema, { name: "route_decision", method: "jsonSchema", + strict: true, }) .invoke(invokeMessages, { ...config, diff --git a/backend/node/thread-summarize-node.ts b/backend/node/thread-summarize-node.ts index ee0a688b..819e9ff1 100644 --- a/backend/node/thread-summarize-node.ts +++ b/backend/node/thread-summarize-node.ts @@ -7,6 +7,7 @@ import { summaryOutputSchema } from "@/lib/langgraph/summary-schema"; import { getChatModel } from "@/backend/model"; import { THREAD_SUMMARIZE_PROMPT } from "@/backend/prompt/system"; import { HumanMessage, BaseMessage } from "@langchain/core/messages"; +import { prepareMessagesForInvoke } from "@/backend/memory/template"; function isHumanMessage(m: BaseMessage | ExcerptMessage): boolean { return m instanceof HumanMessage || m.type === "human"; @@ -21,6 +22,17 @@ type ExcerptMessage = { type Config = { configurable?: { userId?: unknown; thread_id?: unknown } }; +type TranscriptMsg = { + role: "user" | "assistant" | "tool"; + content: string; + tool_calls?: unknown; +}; + +type ThreadLine = { + ref: string; + messages: TranscriptMsg[]; +}; + // ponytail: the LLM produces ordered Q&A entries with refs to the // #N labels we generated in the prompt. ref strings are the labels // (e.g. "#1", "#2-#4") — the original BaseMessage.id values are @@ -107,51 +119,55 @@ function normalizeRole(t: string | undefined): "user" | "assistant" | "tool" { } } -// ponytail: JSONL output — one line per human turn in the THREAD, 1-indexed -// globally. The LLM reads each line as a self-contained record and emits -// OUTPUT entries whose `refs` map 1:1 back to these id strings (matches -// SummaryEntry.startMessageIndex..endMessageIndex byte-for-byte in the new -// 1-indexed world — Memory tab display "messages [3..5]" → humanIndex 2..4 -// under the cumulative formula, and the LLM's `refs: ["#3"..."#5"]` reuses -// that exact id, so the read path stays structural). Replacing the prior -// markdown "#N\nUser: ...\nAssistant: ..." format — role labels in plain -// text collided with content containing ":" and tool_calls had to be -// appended as ad-hoc "[tool_call X]" trailers that the model often -// ignored, producing the meta-question paraphrase ("User said … what was -// the assistant's response?"). Structured input ↔ structured output -// eliminates the prose↔JSON translation step on the model side. -function renderTranscript(excerpt: Array, startHumanIdx: number): string { - const lines: string[] = []; +// ponytail: per-turn transcript — one ThreadLine per human turn in the +// THREAD, 1-indexed globally. Each line is a self-contained record +// (`{ref, messages[]}`) that the LLM reads as a separate user message +// in the chat-conversation; the leading text part ("This turn ref is +// #N") carries the byte-for-byte label the model must copy verbatim +// into OUTPUT `refs`. Matches SummaryEntry.startMessageIndex..endMessageIndex +// — Memory tab display "messages [3..5]" → humanIndex 2..4 under the +// cumulative formula, and the LLM's `refs: ["#3"..."#5"]` reuses the +// exact label, so the read path stays structural. Replacing the prior +// markdown "#N\nUser: ...\nAssistant: ..." format — role labels in +// plain text collided with content containing ":" and tool_calls had +// to be appended as ad-hoc "[tool_call X]" trailers that the model +// often ignored, producing the meta-question paraphrase ("User said +// … what was the assistant's response?"). Splitting each line into +// its own user message gives the model a clean per-turn boundary to +// index against. +function renderTranscript(excerpt: Array, startHumanIdx: number): ThreadLine[] { + const lines: ThreadLine[] = []; let humanCount = 0; - let current: { id: string; messages: unknown[] } | null = null; + let current: ThreadLine | null = null; const flush = () => { - if (current !== null) lines.push(JSON.stringify(current)); + if (current !== null) lines.push(current); current = null; }; for (const m of excerpt) { if (isHumanMessage(m)) { flush(); humanCount++; - current = { id: `#${startHumanIdx + humanCount}`, messages: [] }; + current = { ref: `#${startHumanIdx + humanCount}`, messages: [] }; } if (!current) continue; - const msg: { role: string; content: unknown; tool_calls?: unknown } = { + const msg: TranscriptMsg = { role: normalizeRole(m.type), content: stringifyContent(m.content) || "", }; - // ponytail: carry tool_calls as a first-class field. Models trained on - // log-style JSONL iterate it naturally and don't drop it like they - // drop trailing "[tool_call X]" prose — the prior format's missing - // tool_call lines were the root cause of the meta-question paraphrase - // failures in #3..#5 chunks (see issue notes 2026-07-04). + // ponytail: carry tool_calls as a first-class field on the line + // object so the LLM can attribute the matching tool message's data + // to the same turn (rather than narrating "called X" with no + // source). Model trained on log-style JSON iterates it naturally. if (Array.isArray(m.tool_calls) && m.tool_calls.length > 0) { msg.tool_calls = m.tool_calls; } + current.messages.push(msg); } flush(); - return lines.join("\n"); + + return lines; } // ponytail: STORE-ANCHORED trigger (replaces the prior stateless @@ -232,7 +248,7 @@ export function computeCumulativeWindow( // These would have been END'd by the conditional edge, but a tick // can race — re-deriving here is the safety belt. export async function threadSummarizeNode( - state: { messages?: Array }, + state: { messages?: BaseMessage[] }, config: Config, ): Promise<{ messages: never[] }> { const userId = config.configurable?.userId; @@ -240,7 +256,12 @@ export async function threadSummarizeNode( if (typeof userId !== "string" || userId.length === 0) return { messages: [] }; if (typeof threadId !== "string" || threadId.length === 0) return { messages: [] }; - const messages = (state.messages ?? []) as Array; + const messages = (await prepareMessagesForInvoke( + state.messages ?? [], + [], + userId ?? undefined, + )) as Array; + const humanIndices: number[] = []; for (let i = 0; i < messages.length; i++) { const m = messages[i]; @@ -270,7 +291,7 @@ export async function threadSummarizeNode( // user question captures its assistant reply. Slicing on // humanIndices[endIdx] alone stops at the user message itself — the // AI/tool messages immediately following would be dropped, leaving the - // last JSONL entry as a Q with no A. + // last line as a Q with no A. // // Unknown / orphan roles are dropped from the LLM-facing transcript // (KEEPABLE_TYPES gate) — the original messages still live in @@ -294,18 +315,14 @@ export async function threadSummarizeNode( // state.messages and the checkpointer serves them when needed. const humanMessageIds = excerpt.filter((m) => isHumanMessage(m)).map((m) => m.id ?? ""); - // ponytail: token counts before/after compression. Recorded in the - // SummaryEntry for future analytics + UI stats. - // countTokensApproximately is the same heuristic LangChain's - // summarizationMiddleware uses for its token-budget gate (4 chars - // ≈ 1 token). We don't gate on a hard budget — turn-based trigger - // is primary — but the numbers let a future token-based second - // pass skip work and let the UI render "compressed 420 → 80 tokens" - // without a separate re-tokenize call. // ponytail: token counts before/after compression. Recorded in the // SummaryEntry for future analytics + UI stats — local char-based // estimate (~4 chars/token) is good enough for analytics; the - // trigger itself is turn-count-based, not token-budget-gated. + // trigger itself is turn-count-based, not token-budget-gated. A + // future token-based secondary pass can swap in @langchain/core's + // real counting without changing the call sites, and the UI can + // render "compressed 420 → 80 tokens" without a separate + // re-tokenize call. const tokenCountBefore = estimateTokensFromExcerpt(excerpt); const transcript = renderTranscript(excerpt, startIdx); @@ -319,11 +336,23 @@ export async function threadSummarizeNode( out = await ( await getChatModel() ) - .withStructuredOutput(summaryOutputSchema, { method: "jsonSchema" }) + .withStructuredOutput(summaryOutputSchema, { method: "jsonSchema", strict: true }) .invoke( [ { role: "system", content: THREAD_SUMMARIZE_PROMPT }, - { role: "user", content: transcript }, + ...transcript.map((item) => ({ + role: "user", + content: [ + { + type: "text", + text: `This turn ref is ${item.ref}`, + }, + { + type: "text", + text: JSON.stringify(item), + }, + ], + })), ] as never, { tags: ["nostream"], diff --git a/backend/node/trigger-background-agent-node.ts b/backend/node/trigger-background-agent-node.ts index d2b3effc..deb65162 100644 --- a/backend/node/trigger-background-agent-node.ts +++ b/backend/node/trigger-background-agent-node.ts @@ -19,6 +19,7 @@ // spans land in the same observability row set as the chat invoke. import { langGraphClient } from "@/lib/langgraph/client"; import { lastHumanMessageId } from "@/lib/langgraph/last-human-message-id"; +import type { BaseMessage } from "@langchain/core/messages"; type ScheduleConfig = { configurable?: { @@ -28,13 +29,13 @@ type ScheduleConfig = { }; type ScheduleState = { - messages?: unknown[]; + messages?: BaseMessage[]; }; type PreparedCall = { userId: string; threadId: string; - messages: unknown[]; + messages: BaseMessage[]; // ponytail: last HumanMessage id from the chat invoke — the same // id assistant-ui stamps on the user message. The observability // API uses it to scope in-flight runs to the current turn (the @@ -42,7 +43,10 @@ type PreparedCall = { parentMessageId: string | null; }; -function readBackgroundCall(state: ScheduleState, config: ScheduleConfig): PreparedCall | null { +async function readBackgroundCall( + state: ScheduleState, + config: ScheduleConfig, +): Promise { const userId = config.configurable?.userId; const threadId = config.configurable?.thread_id; if (typeof userId !== "string" || userId.length === 0) return null; @@ -51,7 +55,7 @@ function readBackgroundCall(state: ScheduleState, config: ScheduleConfig): Prepa return { userId, threadId, - messages: state.messages ?? [], + messages: state.messages as BaseMessage[], parentMessageId: lastHumanMessageId(state.messages), }; } @@ -83,7 +87,7 @@ export async function triggerBackgroundAgentNode( state: ScheduleState, config: ScheduleConfig, ): Promise> { - const prepared = readBackgroundCall(state, config); + const prepared = await readBackgroundCall(state, config); if (!prepared) return {}; // ponytail: must await so SDK rejections propagate to the catch diff --git a/backend/prompt/system.ts b/backend/prompt/system.ts index a96e5821..4f5c5d1a 100644 --- a/backend/prompt/system.ts +++ b/backend/prompt/system.ts @@ -13,6 +13,8 @@ export const CHAT_AGENT_PROMPT = `You are ${APP_NAME}, a careful and direct AI a GOALS: - Give the user a correct, complete answer. If you are unsure, state clearly what information you lack, and ask the user a specific question to clarify — never invent facts, numbers, citations, or tool outputs. - Use the available tools whenever the answer depends on current information, a specific URL, or anything you cannot reliably recall. +- [KNOWLEDGE BASE] When the answer depends on information you can't reliably recall (the user's own docs, a specific PDF, prior research, a fact you might be wrong about), reach for a tool — don't guess. Tool priority when the user might have their own content: (1) 'search_kb' — the KB outranks the web for anything the user has uploaded. The 'query' arg is optional: pass the user's natural-language question to rank by relevance, OR omit it (or pass an empty string) to dump the full filtered scope. Narrow the scope with a ':kb-document[label]{documentId=…}' / ':kb-folder[label]{folderId=…}' directive from the message; (2) iterate on the query if results are thin; (3) only THEN 'search_web' / 'fetch_url' for the gap. Don't search when you can answer from built-in knowledge — small talk, definitions, coding patterns, well-known facts don't need a tool. + - [Iterative Search]: If a search returns empty results or the retrieved chunks are insufficient to fully answer the query, DO NOT give up. You MUST refine your query terms (using different keywords, synonyms, translating between English and Chinese, or relaxing the filters if too restrictive) and retry the search. You can perform up to 3 consecutive search attempts in a single turn. Carefully evaluate if you have enough information to answer; if not, retry or state what is missing. - Match the user's language. If they write in Chinese, reply in Chinese; English, reply in English; otherwise match the dominant language in the conversation. - [MEMORY] When the conversation yields a durable fact worth recalling in a future session — from the user's own statements or from a tool result that captures user input — save it to memory using the 'save_memory' tool. @@ -174,6 +176,43 @@ ON FAILURE: - If execute_code returns \`{ ok: false, error }\`: read the error, fix the code, retry. If the fix is non-trivial, call write_code first. - After 3 failed attempts on the same problem: stop. Tell the user what went wrong in one sentence and ask if they want to try a different approach.`; +// Dropped into the ocrNode inside the kb sub-agent. Asks the OCR +// model to read ONE rendered PDF page (passed in as an image_url) +// and emit the page's content as clean markdown. Per-page, so the +// prompt has to stay generic — no document-level structure, just +// "what's on this page?". Runs at OCR_CONCURRENCY=5 (see kb-agent.ts). +// +export const KB_OCR_PAGE_PROMPT = `You are a precise document digitizer. Your task is to convert a single PDF page image into clean, accurate Markdown. + +## Inputs +- **Image** (always present): The rendered PDF page. This is your primary source for both text content and visual layout. +- **Reference Text** (optional, appears after the image): Raw text programmatically extracted from the PDF's text layer. + *WARNING: The reference text is often severely fragmented, out of order, and displaced due to PDF multi-column layout extraction. For example, dates, titles, or subtitles visible in a specific card on the image might be extracted at a completely different place in the reference text. It is NOT a reliable guide for reading order or layout structure.* + +## Rules + +### Segmentation & Layout — follow the IMAGE +- Analyze the **Image** to determine how the content is grouped, partitioned, and structured. +- Do NOT use the reference text to segment or organize the content. The layout, section divisions, reading order, and block structure must come entirely from the visual flow of the image. +- Use heading levels (#, ##, ###) matching the visual hierarchy in the image. +- Preserve lists (bullet / numbered), tables (using GFM table syntax), and code blocks (\`\`\`) matching their visual representations in the image. + +### Completeness — transcribing EVERYTHING without omission +- Convert **ALL** readable text and visible content from the Image. +- Do NOT summarize, skip, truncate, or paraphrase any sections of the page. +- Ensure every sentence, paragraph, table row, and cell visible in the image is completely translated into the markdown output. +- If text is clearly visible in the image (e.g., job titles, dates, or company names) but is missing from its expected place in the reference text, you **must** transcribe it fully based on what you see in the image (and look for it elsewhere in the reference text if needed). + +### Character Disambiguation — use the REFERENCE TEXT +- The reference text is provided **solely as a lookup reference** to help you resolve or verify individual characters that are hard to recognize visually (especially rare CJK characters like 焯 vs 炜, proper nouns, technical terms, and numbers). +- Because the reference text is often out of order, do not expect it to align spatially with the image. Search the **entire** reference text to find the correct spelling/character for a given visual section. +- Do NOT copy the whitespace, line breaks, or block grouping of the reference text. + +### Edge cases +- If the page is blank or contains only decorative images with no readable text, return an empty string. +- Do not add headings, summaries, or commentary that are not present in the image. +- Output ONLY the Markdown content — no preamble, no explanation, no code fences wrapping the output.`; + // ponytail: shared system-prompt skeleton — wraps the per-agent base // prompt (CHAT_AGENT_PROMPT, WEATHER_AGENT_PROMPT, etc.) with the // user-memory + past-thread context blocks. Renders as {{base}} + @@ -244,43 +283,107 @@ OBJECTIVE Produce the smallest set of self-contained Q&A entries that capture: the topic being asked, the substance of the answer, and any concrete data the tools returned. Skip filler. The entries MUST cover every #N exactly once (or mark it as skipped). INPUT -JSONL — one line per human turn in the THREAD, 1-indexed globally ("#1" is the very first User message in this thread, "#3" is the fourth, etc.). Each line is a JSON object: {"id": "#N", "messages": [...]}. Lines are separated by a single newline; do NOT wrap the whole payload in an array. +Each user message in this conversation represents one human turn and contains two text parts: + 1. A short label: "This turn ref is #N" — use this #N verbatim in OUTPUT refs. + 2. A JSON object: {"ref": "#N", "messages": [...]} -Inside each line: - - "id": the #N label, byte-for-byte the value the model must put in OUTPUT refs. This is the SAME numbering used by SummaryEntry.startMessageIndex..endMessageIndex and by the Memory tab's "messages [start..end]" header. - - "messages": the ordered list of this turn's messages. Each message has: - - "role": "user" | "assistant" | "tool" (assistant covers both "ai" and "assistant"; tool covers ToolMessage results). - - "content": the message text. tool results are stringified JSON objects — read them as data, not as chat prose. - - "tool_calls" (assistant only, optional): array of {name, args} describing which tools the assistant invoked that turn. The matching tool message's "content" IS the answer's data — surface it verbatim. Don't narrate the call ("called get_weather", "queried the API"); treat the tool as an implementation detail of how the data was sourced, not part of the answer itself. +#N is 1-indexed globally across the entire thread (not slice-local). It maps byte-for-byte to the Memory tab's "messages [start..end]" header. -Example shape (covering #1..#2): -{"id":"#1","messages":[{"role":"user","content":"hello"},{"role":"assistant","content":"hi! how can I help?"}]} -{"id":"#2","messages":[{"role":"user","content":"weather in BJ"},{"role":"assistant","content":"","tool_calls":[{"name":"get_weather","args":{"loc":"BJ"}}]},{"role":"tool","content":"{...}"}]} +"messages" is the ordered list of messages in this turn. Each item has: + - "role": "user" | "assistant" | "tool" + - "content": the message text. tool results are often stringified JSON — read them as data, not as chat prose. + - "tool_calls" (assistant only, optional): array of {name, args} describing which tools were invoked. Read the matching tool message's "content" as source data — summarize its key outcome, do not reproduce it verbatim. + +Example (two turns sent as two separate user messages): +Part 1: "This turn ref is #1" Part 2: {"ref":"#1","messages":[{"role":"user","content":"hello"},{"role":"assistant","content":"hi! how can I help?"}]} +Part 1: "This turn ref is #2" Part 2: {"ref":"#2","messages":[{"role":"user","content":"weather in BJ"},{"role":"assistant","content":"","tool_calls":[{"name":"get_weather","args":{"loc":"BJ"}}]},{"role":"tool","content":"{"temp":32}"}]} OUTPUT (strict JSON, no prose before or after) { "entries": [ { - "question": "", - "answer": "", + "question": "", + "answer": "", "refs": ["#1"] } ] } INSTRUCTIONS -- One entry covers ONE topic or ONE resolved question. Group consecutive turns on the same topic into one entry; use refs to list every covered turn. -- For consecutive labels, abbreviate refs: ["#1", "#2", "#3"] → ["#1-#3"]. Do NOT abbreviate non-consecutive. -- Order entries chronologically (matching the #N labels). -- Preserve concrete facts the user or tools shared — numbers, names, places, IDs, URLs, command outputs — verbatim when they fit. -- Skip turns that carry no information (greetings, "ok", empty tool errors, system chatter). Do not emit entries with empty questions or answers. -- Take a third-party observer's voice. Q names the topic being asked; A states the substantive content of the answer. Skip the interaction scaffolding — no meta-verbs (提供了/请求了/询问了/回应了/请选择), no first/second-person pronouns (我/你/我们/您) in your own prose. When roles need to be named for clarity, use third-person tags (用户/助手). Verbatim quotes from the original transcript are the only place first/second-person text may appear. +- One entry covers ONE topic or ONE coherent task. Group related turns (consecutive or not) into a single entry; list every covered ref in \`refs\`. +- For consecutive refs, abbreviate: ["#1", "#2", "#3"] → ["#1-#3"]. Do NOT abbreviate non-consecutive refs. +- Order entries chronologically by the earliest ref they contain. +- If the assistant called tools to get the answer, briefly name the tool and summarize the key outcome (e.g., "通过 get_weather 查询到北京气温为 32°C" — one clause, not a data dump). Do not reproduce raw JSON, long lists, or intermediate steps verbatim. +- Skip turns that carry no information (greetings, "ok", empty errors, system chatter). Do not emit entries with empty questions or answers. +- Write from a third-party observer's perspective. Do not use first/second-person pronouns (我/你/我们/您) in your own prose; use (用户/助手) when roles need to be named. Verbatim quotes from the transcript are the only place first/second-person text may appear. CONSTRAINTS - Match the dominant language of the transcript. If the transcript mixes languages, match the language the user used most. -- Each entry is a self-contained Q&A — the future reader sees only the entry, not the surrounding context. Don't refer to "this thread", "the conversation", the assistant, the user, or anyone's first/second-person voice; describe only the concrete content being exchanged. +- Each entry is self-contained — the future reader sees only the entry, not the surrounding context. Do not refer to "this thread" or "the conversation"; describe only the concrete content exchanged. SELF-CHECK before emitting -- Every #N from 1 to last is referenced exactly once across all entries (or marked skipped). +- Every #N in the input is referenced exactly once across all entries (or intentionally skipped). - All refs use the #N form, never bare numbers. - JSON is valid (no trailing commas, no comments).`; + +// ponytail: knowledge base entity/relation extraction system prompt constant +export const KB_ENTITY_EXTRACTION_SYSTEM_PROMPT = `You are a top-tier GraphRAG extraction algorithm. Your goal is to extract low-level knowledge graph elements (entities and relationships) and high-level macro themes from the provided text block. + +## Contextual Inputs +You will be provided with: +1. **Context Document Title**: The main subject or name of the document (e.g., candidate's name for a resume). +2. **Section Path**: The structural path within the document (e.g., "Work Experience > Company A"). +3. **Text to extract**: The raw text content of the page chunk. + +## Extraction Rules + +### 1. Core Reference Bridging (Anti-Isolation) +- If the text describes experiences, projects, or attributes of an implicit subject (e.g., "he", "she", "the author", "the candidate"), you MUST resolve this reference and explicitly link it to the **Context Document Title** as the source or target entity. +- Avoid leaving relationship lines floating without a connection to the core subject. + +### 2. Entity Name Normalization +- Avoid generic pronoun entities (e.g., "he", "she", "the company", "the project") as standalone nodes. +- Always map aliases, abbreviations, and informal references to their full, standardized name. + +### 3. Technology & Concept Alignment +- Standardize technology names and tools to their official casings (e.g., use "jQuery" instead of "jquery", "React" instead of "react framework"). +- Group identical terms to avoid generating duplicate nodes with minor casing or spelling variances. + +### 4. Themes +- Themes should be macroscopic abstractions or key topics (e.g., "Web3", "Frontend Development") summarizing the chunk's intent. + +## Output Format +You MUST output a valid JSON object matching the following structure. Do NOT rename any fields, and ensure description and relation properties are included for all elements: + +\`\`\`json +{ + "entities": [ + { + "name": "Name of the entity (e.g., React)", + "type": "Category of the entity (e.g., Technology)", + "description": "Short description of the entity in the current context" + } + ], + "relationships": [ + { + "source": "Source entity name", + "target": "Target entity name", + "relation": "The action or logical connection (e.g., USED_IN)", + "description": "Short description of how they are related" + } + ], + "themes": [ + "High-level theme topic" + ] +} +\`\`\` +`; + +export const KB_ENTITY_ALIGNMENT_SYSTEM_PROMPT = `You are a specialized entity resolution and alignment algorithm. Given a list of entity names extracted from a document, your goal is to identify synonyms, aliases, acronyms, minor typos, spelling variations, and generic references that refer to the same logical entity, and resolve them to a single canonical standard name. + +## Core Rules: +1. **Implicit Subject Resolution**: Identify names that refer to the main subject of the document (e.g., "Owner", "User", "Author", variations of the candidate's name) and resolve them to the primary canonical name of that person (usually the Document Title). +2. **Technological/Conceptual Alignment**: Group together variations of technology names, frameworks, or tools (e.g., "react", "React.js", "ReactJS" -> "React"; "jquery", "jQuery" -> "jQuery"). +3. **Company/Organization Alignment**: Resolve variations of company names (e.g., "ArcBlock Inc", "Arcblock", "ArcBlock" -> "ArcBlock"). +4. **Output Mappings**: Produce a mapping dictionary that maps each original entity name variation to its resolved canonical name. Only include mappings where the original name is different from the canonical name. Do not map unrelated entities. +`; diff --git a/backend/state.ts b/backend/state.ts index b25bcf4b..45c88e09 100644 --- a/backend/state.ts +++ b/backend/state.ts @@ -1,5 +1,8 @@ import { z } from "zod"; import { StateSchema, MessagesValue } from "@langchain/langgraph"; +import type { BaseMessage } from "@langchain/core/messages"; + +import type { FilePart } from "@/lib/kb/extract"; // ponytail: userMessageCount is *not* on RouterAgentState because // deriving it from state.messages in the summarize node (one filter @@ -10,10 +13,98 @@ import { StateSchema, MessagesValue } from "@langchain/langgraph"; export const RouterAgentState = new StateSchema({ messages: MessagesValue, routerDecision: z.object({ - next: z.enum(["weatherAgent", "chatAgent", "cryptoAgent", "codeAgent"]), + next: z.enum(["weatherAgent", "chatAgent", "cryptoAgent", "codeAgent", "kbAgent"]), }), }); export const CommonAgentState = new StateSchema({ messages: MessagesValue, }); + +// --------------------------------------------------------------------------- +// KB ingest subgraph state +// --------------------------------------------------------------------------- + +export type PageResult = { + pageIndex: number; + imageUrl: string; + markdown: string; + /** Native text extracted from the PDF text layer by mupdf. Empty for scanned/image-only pages. */ + referenceText?: string; + errorMessage?: string; + // ponytail: per-page 4-stage status mirroring kbChunkStatusEnum. + // Written by pageToMarkdownNode when OCR succeeds/fails; pending + // is the default for a freshly-screenshot page whose markdown + // hasn't been produced yet. Legacy rows (no status field) read as + // "success" when markdown is non-empty, "failed" when errorMessage + // is set, "pending" when both are empty — preserves existing UI + // behaviour for docs ingested before this field existed. + status?: "pending" | "parsing" | "success" | "failed"; +}; + +// Per-file record. One entry per PDF file part found across every +// HumanMessage. Drives every node — prepareKBDataNode fills it, +// splitFileToPageNode uploads images + extracts reference text, +// pageToMarkdownNode updates page markdown, rewriteMessagesNode +// uses it to rewrite HumanMessages. filePart.data is the join key +// when matching back to the original HumanMessage content. +export type ProcessedFile = { + messageIndex: number; + filePart: FilePart; + docId: string | null; + attachmentId: string | null; + r2Key: string | null; + title: string | null; + contentHash: string | null; + // "new" = docId freshly generated, needs OCR + chunk + insert. + // "dedup" = existing docId, skip the heavy pipeline. + // "failed" = OCR failed (or empty markdown); docId may or may not + // exist in DB — resolve layer shows [Failed: ...] for + // the file part's kb_ref prefix, or strips the file + // part entirely if no docId was ever written. + // "unknown" = attachment row missing, no docId at all. + pipelineStatus: "new" | "dedup" | "failed" | "unknown"; + errorMessage: string | null; + // ponytail: when pipelineStatus === "dedup", this is the row's + // CURRENT status read from the kb_documents table at dispatch time. + // kbAgent's terminal node (`rewriteMessagesNode`) writes this back + // to the row so a previous kbAgent run that landed `success` is + // visible to the user even when the dispatch path went through a + // re-upload dedup short-circuit. Optional for the `new` / `failed` + // / `unknown` branches that produce their own row status. + existingStatus?: "pending" | "parsing" | "success" | "failed"; +}; + +export const KbAgentState = new StateSchema({ + // From parent — populated by RouterNode at invoke time. + messages: z.array(z.custom()), + userId: z.string().nullable().default(null), + // ponytail: "full" = original OCR + chunk + embed pipeline. + // "chunksOnly" = skip the OCR chain (prepareKBData reads an + // existing doc row whose pages[].markdown is reused) — only the + // chunk + embed + entity stage lands. Populated by + // `fireIngestionRun` from `config.configurable` when invoked by + // `POST /api/kb/documents/[id]/reprocess?chunksOnly=true`. The + // kb_documents row stays at its terminal status (no reset). + mode: z.enum(["full", "chunksOnly", "retryFailed", "retryFailedChunks"]).default("full"), + // ponytail: for chunksOnly dispatch, this is the target docId + // (the row whose pages[].markdown will be re-chunked). Ignored in + // full mode (prepareKBData figures the docId out per file part). + docId: z.string().nullable().default(null), + // Internal. + pagesByDocId: z.record(z.string(), z.array(z.custom())).default({}), + processedFiles: z.array(z.custom()).default([]), + status: z.enum(["pending", "parsing", "success", "failed"]).default("pending"), + errorMessage: z.string().nullable().default(null), +}); + +export type KbAgentStateShape = { + messages: BaseMessage[]; + userId: string | null; + mode: "full" | "chunksOnly" | "retryFailed" | "retryFailedChunks"; + docId: string | null; + pagesByDocId: Record; + processedFiles: ProcessedFile[]; + status: "pending" | "parsing" | "success" | "failed"; + errorMessage: string | null; +}; diff --git a/backend/tool/index.ts b/backend/tool/index.ts index d705b2c4..0fea4978 100644 --- a/backend/tool/index.ts +++ b/backend/tool/index.ts @@ -11,6 +11,7 @@ import { getOrderStatusTool } from "@/backend/tool/crypto/get-order-status"; import { getNftHoldingsTool } from "@/backend/tool/crypto/get-nft-holdings"; import { saveMemoryTool } from "@/backend/tool/memory/save-memory-tool"; import { executeCodeTool, writeCodeTool } from "@/backend/tool/code"; +import { listDocumentsTool, searchKbTool } from "@/backend/tool/kb"; // ponytail: keep the tool list in one place so the graph binds it from a // single source. Adding a tool = drop a file + add one line here. @@ -25,8 +26,12 @@ import { executeCodeTool, writeCodeTool } from "@/backend/tool/code"; // Tools that need a third-party key (search_web → JINA_API_KEYS, // get_NFT_holdings → ALCHEMY_API_KEY) are gated: they return `null` // when the key is missing, and the spreads below skip them. `fetch_url` -// is unconditional because r.jina.ai accepts unauthenticated requests -// on the free tier (lower rate limit, no key needed). +// is unconditional because r.jina.ai accepts unauthenticated requests on +// the free tier (lower rate limit, no key needed). +// +// KB tools (issue #13 v3): +// - search_kb — gated on pgvector extension (rule #10). +// - list_documents — pure SQL, always available. export const WEATHER_TOOLS = [askLocationTool, geocodeLocationTool, getWeatherTool, saveMemoryTool]; @@ -47,19 +52,18 @@ export const CRYPTO_TOOLS = [ // fallback runs at click-time. export const CODE_TOOLS = [writeCodeTool, ...(executeCodeTool ? [executeCodeTool] : [])]; -export const ALL_TOOLS = [ +// ponytail: KB tools — search_kb throws at runtime when pgvector is +// missing (the tool is still registered so the LLM sees a consistent +// tool surface; missing-extension produces a clean error message instead +// of a 500). list_documents is unconditional. +export const KB_TOOLS = [searchKbTool, listDocumentsTool]; + +export const CHAT_TOOLS = [ fetchUrl, ...(searchWeb ? [searchWeb] : []), - askLocationTool, - geocodeLocationTool, - getWeatherTool, - getCryptoPriceTool, - getFxRateTool, - connectWalletTool, - placeCryptoOrderTool, - getOrderStatusTool, - ...(getNftHoldingsTool ? [getNftHoldingsTool] : []), - saveMemoryTool, + ...WEATHER_TOOLS, + ...CRYPTO_TOOLS, + ...KB_TOOLS, ]; export { @@ -74,4 +78,6 @@ export { placeCryptoOrderTool, getOrderStatusTool, getNftHoldingsTool, + searchKbTool, + listDocumentsTool, }; diff --git a/backend/tool/kb/format.ts b/backend/tool/kb/format.ts new file mode 100644 index 00000000..3829a2a8 --- /dev/null +++ b/backend/tool/kb/format.ts @@ -0,0 +1,35 @@ +import type { HybridSearchResult } from "@/lib/kb/search"; + +import type { KbSearchDocument, KbToolResult } from "./types"; + +function truncate(s: string, max: number): string { + if (s.length <= max) return s; + return s.slice(0, max) + "…"; +} + +// ponytail: build the dual-purpose ToolMessage payload. `content` is the +// LLM string with `[1] [2] …` markers; `documents` is the full structured +// rows for the UI. Same data, two views. +export function formatSearchResult( + results: HybridSearchResult[], + chunkMaxChars: number, +): KbToolResult { + if (results.length === 0) { + return { content: "", documents: [], empty: true }; + } + const documents: KbSearchDocument[] = results.map((r) => ({ + chunkId: r.chunkId, + documentId: r.documentId, + docTitle: r.docTitle, + pageNumbers: r.pageNumbers, + content: r.content, + rrfScore: r.rrfScore, + legsHit: r.legsHit, + })); + // Truncate each chunk to keep the LLM prompt within budget; the full + // chunk is in `documents` for the UI to fetch if needed. + const content = documents + .map((d, i) => `[${i + 1}] ${truncate(d.content, chunkMaxChars)}`) + .join("\n\n"); + return { content, documents, empty: false }; +} diff --git a/backend/tool/kb/index.ts b/backend/tool/kb/index.ts new file mode 100644 index 00000000..7a067cb4 --- /dev/null +++ b/backend/tool/kb/index.ts @@ -0,0 +1,21 @@ +// ponytail: barrel re-export. The tools, helpers, types, and pgvector +// gate are all imported by the rest of the app via this single entry +// point (`@/backend/tool/kb` resolves to this file). Individual +// sub-modules stay private to the package. + +export { isPgVectorAvailable, _resetPgVectorCache } from "./pgvector"; +export { setKbToolUserId, thisUserId } from "./user-id"; +export { formatSearchResult } from "./format"; +export { searchKbTool } from "./search-kb"; +export { + LIST_DOCUMENTS_STATUSES, + listDocumentsTool, + listKbDocumentsForUser, + listKbFoldersForUser, +} from "./list-documents"; +export type { + KbSearchDocument, + KbToolResult, + ListDocumentsArgs, + ListDocumentsResult, +} from "./types"; diff --git a/backend/tool/kb/list-documents.ts b/backend/tool/kb/list-documents.ts new file mode 100644 index 00000000..ec6aabf3 --- /dev/null +++ b/backend/tool/kb/list-documents.ts @@ -0,0 +1,177 @@ +import { tool, type StructuredTool } from "@langchain/core/tools"; +import { asc, eq } from "drizzle-orm"; +import { z } from "zod"; + +import { db } from "@/db/client"; +import { kbFolder } from "@/lib/kb/schema"; +import { listKbDocumentsGroupedWithAttachment } from "@/lib/kb/queries"; +import { extractUserId } from "@/backend/memory/recall"; + +import { thisUserId } from "./user-id"; +import type { + ListDocumentsArgs, + ListDocumentsDoc, + ListDocumentsFolder, + ListDocumentsResult, +} from "./types"; + +// ponytail: the four statuses are mirrored in the DB column enum + the +// API mention endpoint. Single source of truth here. +export const LIST_DOCUMENTS_STATUSES = ["success", "failed", "parsing", "pending"] as const; + +const listDocumentsSchema = z.object({ + folderId: z + .string() + .optional() + .describe( + "Restrict to a specific folder. Copy the value verbatim from the " + + "':kb-folder[label]{folderId=...}' directive in the user message " + + "(optional — omit for all folders).", + ), + status: z + .enum(["success", "failed", "parsing", "pending"]) + .optional() + .describe("Filter by ingest status. Defaults to 'success' — only fully-ingested docs."), + titleQuery: z.string().optional().describe("Case-insensitive substring match on document title."), + page: z.number().int().min(1).optional().describe("1-indexed page number. Default 1."), + pageSize: z + .number() + .int() + .min(1) + .max(100) + .optional() + .describe("Per-folder cap on documents returned. Default 20, max 100."), +}) satisfies z.ZodType; + +// ponytail: chat-tool view of a KB document. Strip the heavy fields +// (attachment URL, raw pages array, folderId) the LLM / chat card +// doesn't need — just the badge fields + a stable identifier. +function toListDoc(row: { + id: string; + title: string; + status: "success" | "failed" | "parsing" | "pending"; + errorMessage: string | null; + createdAt: Date; + totalPages?: number; + failedPages?: number; + parsingPages?: number; + pendingPages?: number; + totalChunks?: number; + successChunks?: number; + failedChunks?: number; + pendingChunks?: number; + parsingChunks?: number; +}): ListDocumentsDoc { + const totalPages = row.totalPages ?? 0; + const pending = row.pendingPages ?? 0; + const parsing = row.parsingPages ?? 0; + const failed = row.failedPages ?? 0; + const success = Math.max(totalPages - pending - parsing - failed, 0); + return { + id: row.id, + title: row.title, + status: row.status, + errorMessage: row.errorMessage, + createdAt: row.createdAt.toISOString(), + totalPages, + successPages: success, + failedPages: failed, + parsingPages: parsing, + pendingPages: pending, + totalChunks: row.totalChunks ?? 0, + successChunks: row.successChunks ?? 0, + failedChunks: row.failedChunks ?? 0, + pendingChunks: row.pendingChunks ?? 0, + parsingChunks: row.parsingChunks ?? 0, + }; +} + +// ponytail: Title Case the folder name for the LLM-facing content +// so the model doesn't have to deal with all-caps user-typed names +// like "ARCBLOCK". The chat card displays the original name as-is +// — only the LLM string is normalised. +function titleCase(s: string): string { + return s + .split(/(\s+|-|_)/) + .map((part) => + part.match(/\s+|-|_/) ? part : part.charAt(0).toUpperCase() + part.slice(1).toLowerCase(), + ) + .join(""); +} + +export async function listKbDocumentsForUser( + args: ListDocumentsArgs & { userId: string }, +): Promise { + const page = args.page ?? 1; + const pageSize = Math.min(args.pageSize ?? 20, 100); + const status = args.status ?? "success"; + const titleQuery = args.titleQuery?.trim(); + + const groups = await listKbDocumentsGroupedWithAttachment(args.userId, args.folderId ?? null); + + const folders: ListDocumentsFolder[] = []; + let total = 0; + for (const g of groups) { + let docs = g.documents; + if (status) docs = docs.filter((d) => d.status === status); + if (titleQuery) { + const q = titleQuery.toLowerCase(); + docs = docs.filter((d) => d.title.toLowerCase().includes(q)); + } + docs = docs.slice(0, pageSize); + folders.push({ id: g.folder.id, name: g.folder.name, documents: docs.map(toListDoc) }); + total += docs.length; + } + + return { folders, total, page, pageSize }; +} + +// ponytail: LLM-facing summary string. Kept terse — folder names + doc +// titles + status + a short error tail. The structured `folders` +// field carries the badge data for the chat card; the LLM doesn't +// need chunk/page counts to decide what to do next. +function buildListContent(folders: ListDocumentsFolder[]): string { + const parts: string[] = []; + for (const f of folders) { + if (f.documents.length === 0) continue; + parts.push( + `[Folder "${titleCase(f.name)}" (${f.documents.length} ${f.documents.length === 1 ? "document" : "documents"})]`, + ); + for (const d of f.documents) { + const err = d.status === "failed" && d.errorMessage ? ` — ${d.errorMessage}` : ""; + parts.push(` - ${d.title} (id=${d.id}, status=${d.status}${err})`); + } + } + return parts.length > 0 ? parts.join("\n") : "No documents matched."; +} + +export const listDocumentsTool: StructuredTool = tool( + async (args, config) => { + const userId = extractUserId(config) ?? thisUserId(); + const result = await listKbDocumentsForUser({ ...args, userId }); + return JSON.stringify({ + ...result, + content: buildListContent(result.folders), + empty: result.folders.every((f) => f.documents.length === 0), + }); + }, + { + name: "list_documents", + description: + "List the user's knowledge-base documents grouped by folder. Supports filtering " + + "by folder, ingest status, and title substring. Returns each document's id, " + + "title, status, and per-doc page/chunk progress counts. Use when the user " + + "asks what's in their KB or wants to navigate to a specific doc. " + + "If the user @-mentioned a folder, copy the id from the " + + "':kb-folder[label]{folderId=...}' directive into the folderId arg.", + schema: listDocumentsSchema, + }, +); + +// ponytail: re-exported so the @mention composer API route can pull the +// user's folder list without going through queries.ts. Not a tool. +export async function listKbFoldersForUser( + userId: string, +): Promise> { + return db.select().from(kbFolder).where(eq(kbFolder.userId, userId)).orderBy(asc(kbFolder.name)); +} diff --git a/backend/tool/kb/pgvector.ts b/backend/tool/kb/pgvector.ts new file mode 100644 index 00000000..d89bf124 --- /dev/null +++ b/backend/tool/kb/pgvector.ts @@ -0,0 +1,26 @@ +import { sql } from "drizzle-orm"; + +import { db } from "@/db/client"; + +// ponytail: cache the extension check at module load. The check is a +// single SQL round-trip against pg_extension; once known, every tool +// import uses the cached value. The test suite can stub the cache via +// `_resetPgVectorCache()`. + +let cachedExtensionAvailable: boolean | null = null; + +export async function isPgVectorAvailable(): Promise { + if (cachedExtensionAvailable !== null) return cachedExtensionAvailable; + const rows = await db.execute<{ ok: number }>(sql` + SELECT 1 AS ok FROM pg_extension WHERE extname = 'vector' + `); + const result = Array.isArray(rows) + ? rows + : ((rows as { rows?: Array<{ ok: number }> }).rows ?? []); + cachedExtensionAvailable = result.length > 0; + return cachedExtensionAvailable; +} + +export function _resetPgVectorCache(value: boolean | null = null): void { + cachedExtensionAvailable = value; +} diff --git a/backend/tool/kb/search-kb.ts b/backend/tool/kb/search-kb.ts new file mode 100644 index 00000000..48cc4e62 --- /dev/null +++ b/backend/tool/kb/search-kb.ts @@ -0,0 +1,92 @@ +import { tool, type StructuredTool } from "@langchain/core/tools"; +import { z } from "zod"; + +import { getKbEnv } from "@/lib/kb/env"; +import { hybridSearch } from "@/lib/kb/search"; +import { extractUserId } from "@/backend/memory/recall"; + +import { formatSearchResult } from "./format"; +import { isPgVectorAvailable } from "./pgvector"; +import { thisUserId } from "./user-id"; + +// ponytail: hybrid RRF over BM25 + pgvector + entity-tag. Returns the +// top-K chunks with `[1] [2] …` markers for the LLM to cite inline. +// Gated on pgvector — a missing extension throws a clear error rather +// than crashing the tool (so the LLM tool surface stays stable). + +const searchKbSchema = z.object({ + query: z + .string() + .optional() + .default("") + .describe( + "Two modes. (1) Ranked search — pass 5-10 space-separated entries " + + "(NOT a verbatim copy of the user's question): build from (a) the " + + "user's question, broken apart (never pass the meta-question " + + "itself, e.g. 'what is this', 'summarize please', '这是什么', as " + + "a single entry), and (b) the @-directive label if present " + + "(split the file/folder name on spaces, dashes, underscores, or " + + "dots and treat each piece as an entry). Aim for 5-10 entries " + + "total; fewer is fine if the question is narrow; don't pad with " + + "repeats or filler. (2) Full scope dump — OMIT query (or pass an " + + "empty string) when the user wants everything in the filtered " + + "scope, e.g. 'summarize @doc', 'extract all clauses from @folder', " + + "'list contents of @doc'. Returns every chunk in documentId, " + + "every chunk in folderId, or — with no other filter — the user's " + + "most recent chunks.", + ), + folderId: z + .string() + .optional() + .describe( + "Filter results to documents within this specific folder ID. " + + "Copy the value verbatim from the ':kb-folder[label]{folderId=...}' " + + "directive in the user message (optional).", + ), + documentId: z + .string() + .optional() + .describe( + "Filter results to this specific document ID only. " + + "Copy the value verbatim from the ':kb-document[label]{documentId=...}' " + + "directive in the user message (optional).", + ), +}); + +export const searchKbTool: StructuredTool = tool( + async ({ query, folderId, documentId }, config) => { + if (!(await isPgVectorAvailable())) { + throw new Error("search_kb unavailable: pgvector extension is not installed on the database"); + } + const userId = extractUserId(config) ?? thisUserId(); + const env = getKbEnv(); + + // ponytail: hybridSearch owns the query -> embed -> search + // pipeline (qvec auto-embedded if not pre-computed; embed + // failures fall back to the BM25 + tag legs only). Empty query + // returns the full filtered scope (capped at 1000). When ranked + // retrieval returns 0 with a scope filter, hybridSearch itself + // transparently retries with an empty query for the same scope + // (see "path A fallback" in lib/kb/search.ts). + const results = await hybridSearch({ userId, query, folderId, documentId }); + return JSON.stringify(formatSearchResult(results, env.chunkMaxChars)); + }, + { + name: "search_kb", + description: + "Search the user's knowledge base (uploaded PDFs / docs) using hybrid " + + "BM25 + vector + entity-tag retrieval. Returns the most relevant " + + "chunks with `[1]`, `[2]`, ... markers the LLM can cite inline. Use " + + "when the user references their KB or asks about content they've uploaded. " + + "Pass a natural-language query to rank by relevance, OR omit query " + + "(pass an empty string) to dump the full filtered scope — useful " + + "for 'summarize @doc', 'extract all clauses from @folder'. " + + "If the user @-mentioned a doc or folder, narrow the search by copying " + + "the id from the ':kb-document[label]{documentId=...}' or " + + "':kb-folder[label]{folderId=...}' directive in the message. " + + "If results are empty or insufficient, retry with a fresh query " + + "(rephrased keywords, synonyms, English↔Chinese, or relaxed filters) " + + "— up to 3 attempts per turn before falling back to search_web.", + schema: searchKbSchema, + }, +); diff --git a/backend/tool/kb/types.ts b/backend/tool/kb/types.ts new file mode 100644 index 00000000..58a9524a --- /dev/null +++ b/backend/tool/kb/types.ts @@ -0,0 +1,75 @@ +import type { HybridSearchResult } from "@/lib/kb/search"; + +// ponytail: ToolMessage shape returned by every KB search tool. Both +// fields describe the SAME payload from two perspectives: +// - `content` — the LLM-facing string with `[1] [2] …` markers baked +// in. The model emits inline citations by copying these. +// Community consensus (LangChain / LlamaIndex / Haystack): +// hide scores from the LLM, use them only for ranking. +// - `documents` — the UI-facing structured array for Sources cards. +// Carries `rrfScore` + `legsHit` so the frontend can +// show debug badges. +export type KbSearchDocument = { + chunkId: string; + documentId: string; + docTitle: string; + pageNumbers: number[]; + content: string; + rrfScore: number; + // ponytail: mirrors HybridSearchResult.legsHit in lib/kb/search.ts. + legsHit: Array<"kw" | "vec" | "tag" | "full">; +}; + +export type KbToolResult = { + content: string; + documents: KbSearchDocument[]; + empty: boolean; +}; + +// ponytail: `list_documents` returns a folder-grouped shape so the +// chat-side card can render each folder as a collapsible section and +// the LLM can see the folder structure in `content`. `documents[]` +// carries the same per-doc status counts (totalPages / successPages / +// totalChunks / …) that the Settings → KB DocStatusBadge + +// ChunksStatusBadge read — same source of truth for both surfaces. + +export type ListDocumentsArgs = { + folderId?: string; + status?: "success" | "failed" | "parsing" | "pending"; + titleQuery?: string; + page?: number; + pageSize?: number; +}; + +export type ListDocumentsDoc = { + id: string; + title: string; + status: "success" | "failed" | "parsing" | "pending"; + errorMessage: string | null; + createdAt: string; + totalPages: number; + successPages: number; + failedPages: number; + parsingPages: number; + pendingPages: number; + totalChunks: number; + successChunks: number; + failedChunks: number; + pendingChunks: number; + parsingChunks: number; +}; + +export type ListDocumentsFolder = { + id: string; + name: string; + documents: ListDocumentsDoc[]; +}; + +export type ListDocumentsResult = { + folders: ListDocumentsFolder[]; + total: number; + page: number; + pageSize: number; +}; + +export type { HybridSearchResult }; diff --git a/backend/tool/kb/user-id.ts b/backend/tool/kb/user-id.ts new file mode 100644 index 00000000..e06cd767 --- /dev/null +++ b/backend/tool/kb/user-id.ts @@ -0,0 +1,23 @@ +// ponytail: userId is set per-turn by chat-agent.chatModelNode +// (chat-agent.ts: `setKbToolUserId(userId)` before the model is +// invoked). The LangGraph `ToolNode` runs in a separate node; threading +// the userId across that hop cleanly requires reading RunnableConfig +// (which the `tool()` factory's typed signature drops). Module-local +// capture is the simplest path that works with the existing bindTools + +// ToolNode pattern. Cleared on every chat-agent invocation, set fresh, +// then the tool fires within that same execution. + +let currentUserId = ""; + +export function setKbToolUserId(userId: string): void { + currentUserId = userId; +} + +export function thisUserId(): string { + if (!currentUserId) { + throw new Error( + "KB tool: userId not set — caller must invoke setKbToolUserId() before the tool runs", + ); + } + return currentUserId; +} diff --git a/components/assistant-ui/attachment.tsx b/components/assistant-ui/attachment.tsx index b32429fc..6adf9270 100644 --- a/components/assistant-ui/attachment.tsx +++ b/components/assistant-ui/attachment.tsx @@ -1,9 +1,9 @@ "use client"; import { type PropsWithChildren, useEffect, useMemo, useState, type FC } from "react"; -import { XIcon, PlusIcon, FileText, Loader2Icon } from "lucide-react"; +import { XIcon, PlusIcon, FileText, BookOpen, Loader2Icon } from "lucide-react"; import { AttachmentPrimitive, ComposerPrimitive, useAuiState, useAui } from "@assistant-ui/react"; -import type { CompleteAttachment, ThreadUserMessagePart } from "@assistant-ui/react"; +import type { CompleteAttachment } from "@assistant-ui/react"; import { useShallow } from "zustand/shallow"; import { Tooltip, TooltipContent, TooltipTrigger } from "@/components/ui/tooltip"; import { Dialog, DialogTitle, DialogContent, DialogTrigger } from "@/components/ui/dialog"; @@ -11,6 +11,9 @@ import { Avatar, AvatarImage, AvatarFallback } from "@/components/ui/avatar"; import { TooltipIconButton } from "@/components/assistant-ui/tooltip-icon-button"; import { cn } from "@/lib/utils"; import { useUploadStore } from "@/lib/attachments/upload-store"; +import { extractKbRefFromFilename, stripKbRefFromFilename } from "@/lib/kb/extract"; + +const KB_SETTINGS_PATH = "/settings/knowledge-base"; // ponytail: client-side object URL for a pending File. Drops on unmount. const useFileSrc = (file: File | undefined) => { @@ -147,7 +150,12 @@ const AttachmentPreviewDialogWithSdk: FC = ({ children }) => return {children}; }; -const AttachmentThumb: FC<{ src?: string }> = ({ src }) => { +type IconComponent = FC<{ className?: string }>; + +const AttachmentThumb: FC<{ src?: string; fallbackIcon?: IconComponent }> = ({ + src, + fallbackIcon: FallbackIcon = FileText, +}) => { return ( = ({ src }) => { className="aui-attachment-tile-image object-cover" /> - + ); @@ -262,22 +270,6 @@ const AttachmentUI: FC = () => { ); }; -// ponytail: image parts in message.content carry only the URL — the -// SDK's `contentToParts` drops `filename` on the round trip through -// `image_url`. R2 keys look like `u//-`, so -// the last URL segment is the original filename with the uuid prefix -// stripped. -function filenameFromImageUrl(url: string): string { - try { - const last = new URL(url).pathname.split("/").pop() ?? ""; - const decoded = decodeURIComponent(last); - const stripped = decoded.replace(/^[0-9a-f-]{36}-/, ""); - return stripped || "image"; - } catch { - return "image"; - } -} - // ponytail: prop-based message-path card. Renders one attachment for // the message-list. No SDK attachment runtime here — the attachment // data is plain JS, sourced from message.content by the parent. @@ -291,7 +283,23 @@ const MessageAttachmentCard: FC = ({ attachment }) = [attachment], ); const isImage = attachment.type === "image"; + // ponytail: kbAgent stamps a `[kb:]` prefix onto the file + // part's filename (both `filename` and `metadata.filename`). We + // parse it once and use it to (a) override the click target + // (deep-link into /settings/... rather than open the file preview + // dialog) and (b) rename the tile label from "File" to "KB + // document". The user-facing filename has the prefix stripped so + // the bracket is invisible — `attachment.name` is the SDK's view of + // the post-round-trip filename, so this is the only place we need + // to scrub. + const kbRefDocId = useMemo(() => { + const ref = extractKbRefFromFilename(attachment.name); + return ref?.docId ?? null; + }, [attachment.name]); + const isKbDoc = kbRefDocId !== null; + const displayName = useMemo(() => stripKbRefFromFilename(attachment.name), [attachment.name]); const typeLabel = useMemo(() => { + if (isKbDoc) return "KB document"; switch (attachment.type) { case "image": return "Image"; @@ -302,7 +310,21 @@ const MessageAttachmentCard: FC = ({ attachment }) = default: return attachment.type; } - }, [attachment.type]); + }, [attachment.type, isKbDoc]); + + const activate = () => { + if (kbRefDocId) { + // ponytail: open in a new tab so the chat stays put — users + // typically want to peek at the KB doc + come back without + // losing scroll/state. router.push would unmount the chat + // thread and force a re-fetch of all messages on return. + window.open( + `${KB_SETTINGS_PATH}?doc=${encodeURIComponent(kbRefDocId)}`, + "_blank", + "noopener,noreferrer", + ); + } + }; return ( @@ -313,20 +335,40 @@ const MessageAttachmentCard: FC = ({ attachment }) = isImage && "aui-attachment-root-message only:*:first:size-24", )} > - + {isKbDoc ? (
{ + if (e.key === "Enter" || e.key === " ") { + e.preventDefault(); + activate(); + } + }} > - +
- + ) : ( + + +
+ +
+ + + )} - {attachment.name} + {displayName} ); }; @@ -351,64 +393,102 @@ const MessageAttachmentCard: FC = ({ attachment }) = // and delete the message.content parser below. The visual is identical // (same `MessageAttachmentCard` shape), so no UX change is needed at // switchover. +// ponytail: image parts in message.content carry only the URL — the +// SDK's `contentToParts` drops `filename` on the round trip through +// `image_url`. R2 keys look like `u//-`, so +// the last URL segment is the original filename with the uuid prefix +// stripped. +function filenameFromImageUrl(url: string): string { + try { + const last = new URL(url).pathname.split("/").pop() ?? ""; + const decoded = decodeURIComponent(last); + const stripped = decoded.replace(/^[0-9a-f-]{36}-/, ""); + return stripped || "image"; + } catch { + return "image"; + } +} + +function asRecord(part: unknown): Record { + return part as Record; +} + +// ponytail: pure projection of message.content → CompleteAttachment[]. +// File tiles always come out as `type: "file"`; the KB marker rides +// on the file part's filename as a `[kb:]` prefix (both +// `filename` and `metadata.filename`) — the SDK's `contentToParts` +// round-trip preserves that filename, so the projection just copies +// it through. MessageAttachmentCard parses the prefix off +// `attachment.name` to decide whether to deep-link into /settings/ +// knowledge-base and to strip the prefix before showing the filename +// in the tooltip. Image parts never carry the prefix so they always +// render plain. +function buildUserMessageAttachments(parts: readonly unknown[]): CompleteAttachment[] { + const seen = new Set(); + const out: CompleteAttachment[] = []; + for (const part of parts) { + const r = asRecord(part); + const type = r.type; + + if (type === "image" && typeof r.image === "string") { + const url = r.image; + const name = filenameFromImageUrl(url); + if (seen.has(url)) continue; + seen.add(url); + out.push({ + id: url, + type: "image", + name, + contentType: "image", + status: { type: "complete" }, + content: [{ type: "image", image: url, filename: name }], + }); + } else if (type === "file" && typeof r.data === "string") { + // ponytail: SDK's contentToParts defaults filename to "file" + // when the source part lacks `metadata.filename`, but the type + // still allows undefined. Fall back to "file" so we always + // have a non-empty id / name. + const fileName = (typeof r.filename === "string" && r.filename) || "file"; + const mimeType = typeof r.mimeType === "string" ? r.mimeType : ""; + const kbRef = extractKbRefFromFilename(fileName); + // ponytail: dedupe on docId when the file is KB-tagged, otherwise + // on filename. A retried / duplicate upload with the same name + // but a different docId would otherwise show two tiles; the same + // file uploaded twice with no docId change stays as one tile. + const key = kbRef ? `kb_ref:${kbRef.docId}` : fileName; + if (seen.has(key)) continue; + seen.add(key); + out.push({ + id: key, + type: "file", + name: fileName, + contentType: mimeType, + status: { type: "complete" }, + content: [ + { + type: "file" as const, + data: r.data, + mimeType, + filename: fileName, + }, + ], + }); + } + // legacy standalone `{ type: "kb_ref", docId }` parts (older + // threads) are dropped here — the parent message still has the + // file part that produced the original upload, and that file + // part will carry the kb_ref prefix on any subsequent re-ingest. + } + return out; +} + export const UserMessageAttachments: FC = () => { const content = useAuiState((s) => s.message.content); - const attachments = useMemo(() => { - const seen = new Set(); - const result: CompleteAttachment[] = []; - for (const part of content) { - let key: string | undefined; - let complete: CompleteAttachment | undefined; - if (part.type === "image") { - const imagePart = part; - const name = filenameFromImageUrl(imagePart.image); - key = imagePart.image; - complete = { - id: imagePart.image, - type: "image", - name, - contentType: "image", - status: { type: "complete" }, - content: [ - { - type: "image", - image: imagePart.image, - filename: name, - }, - ], - }; - } else if (part.type === "file") { - const filePart = part as Extract; - // ponytail: SDK's contentToParts defaults filename to "file" - // when the source part lacks `metadata.filename`, but the type - // still allows undefined. Fall back to "file" so we always - // have a non-empty id / name. - const fileName = filePart.filename || "file"; - key = fileName; - complete = { - id: fileName, - type: "file", - name: fileName, - contentType: filePart.mimeType, - status: { type: "complete" }, - content: [ - { - type: "file", - data: filePart.data, - mimeType: filePart.mimeType, - filename: fileName, - }, - ], - }; - } - if (complete && key && !seen.has(key)) { - seen.add(key); - result.push(complete); - } - } - return result; - }, [content]); + const attachments = useMemo( + () => buildUserMessageAttachments(content), + [content], + ); if (attachments.length === 0) return null; diff --git a/components/assistant-ui/directive-chip.tsx b/components/assistant-ui/directive-chip.tsx new file mode 100644 index 00000000..1fd97992 --- /dev/null +++ b/components/assistant-ui/directive-chip.tsx @@ -0,0 +1,105 @@ +"use client"; + +import { FileTextIcon, FolderIcon } from "lucide-react"; +import type { ReactElement } from "react"; +import { kbMentionFormatter } from "./kb-mention-formatter"; + +type DirectiveType = string; + +// ponytail: shared chip visual used by both the read-only +// DirectiveText (user message bubble) and the editable +// DirectiveComposerInput (composer typing area). aUI's +// kbMentionFormatter.parse is the single source of +// truth — both call sites pass segments into `renderDirectiveSegments` +// so what the user types and what they see after send stay in sync. +// +// `renderDirectiveSegments` is a pure render-prop: given segments, it +// returns a flat array of React elements (text spans + chip spans). +// Unit-tested directly; no DOM mounting required. + +export const DIRECTIVE_CHIP_CLASS = [ + "aui-directive-chip", + "inline-flex", + "items-center", + "gap-1", + "rounded", + "px-1", + "py-0.5", + "align-middle", + "text-xs", + "font-medium", + "border", +].join(" "); + +/** Returns the icon component for a given directive type. */ +export function getChipIcon(directiveType: DirectiveType) { + if (directiveType === "kb-folder") return FolderIcon; + if (directiveType === "kb-document" || directiveType === "kb-doc") return FileTextIcon; + return null; +} + +/** Returns the Tailwind color classes for a given directive type. */ +export function getChipColorClass(directiveType: DirectiveType): string { + if (directiveType === "kb-folder") { + return "bg-indigo-500/10 dark:bg-indigo-500/20 text-indigo-600 dark:text-indigo-400 border-indigo-500/20 dark:border-indigo-400/30"; + } + return "bg-emerald-500/10 dark:bg-emerald-500/20 text-emerald-600 dark:text-emerald-400 border-emerald-500/20 dark:border-emerald-400/30"; +} + +/** Renders a single directive chip span. Used by both the static renderer and the Lexical chip. */ +export function DirectiveChipSpan({ + directiveType, + label, + directiveId, +}: { + directiveType: DirectiveType; + label: string; + directiveId?: string; +}): ReactElement { + const Icon = getChipIcon(directiveType); + const colorClass = getChipColorClass(directiveType); + return ( + + {Icon ? : null} + {label} + + ); +} + +export function renderDirectiveSegments( + text: string, + // ponytail: only used by the composer overlay; in message bubbles + // (DirectiveText) we always want chips parsed. During IME + // composition we pass composing=true to skip directive parsing and + // render the raw buffer verbatim — otherwise the chips flicker as + // the user types pinyin. + options?: { composing?: boolean }, +): ReactElement[] { + if (options?.composing) { + // ponytail: during IME composition, the overlay shows the raw + // text exactly as the textarea holds it. Any directive parse would + // flash chips as the user types pinyin. Visually inert during + // composition; chips re-appear on compositionend. + return [{text}]; + } + const segments = kbMentionFormatter.parse(text); + return segments.map((segment, i) => { + if (segment.kind === "text") { + // ponytail: preserve whitespace + line breaks verbatim — + // the composer emits them and the user expects them. + return {segment.text}; + } + return ( + + ); + }); +} diff --git a/components/assistant-ui/directive-text.tsx b/components/assistant-ui/directive-text.tsx new file mode 100644 index 00000000..29b65089 --- /dev/null +++ b/components/assistant-ui/directive-text.tsx @@ -0,0 +1,21 @@ +"use client"; + +import { type TextMessagePartComponent } from "@assistant-ui/react"; + +import { renderDirectiveSegments } from "@/components/assistant-ui/directive-chip"; + +// ponytail: render the user-message text content with `:kb-document[…]` +// and `:kb-folder[…]` directives as inline chips instead of raw text. +// aUI ships `DirectiveText` (install via shadcn), but its default +// formatter + plain `` styling doesn't match our visual +// language. The chip rendering itself lives in `directive-chip.tsx` +// so the editable composer (DirectiveComposerInput) and this +// read-only bubble renderer share one source of truth. +// +// `unstable_defaultDirectiveFormatter.parse` returns alternating +// `text` and `mention` segments. We render text as plain spans and +// mentions as chips with the right icon. + +export const DirectiveText: TextMessagePartComponent = ({ text }) => { + return <>{renderDirectiveSegments(text)}; +}; diff --git a/components/assistant-ui/kb-mention-formatter.ts b/components/assistant-ui/kb-mention-formatter.ts new file mode 100644 index 00000000..77777372 --- /dev/null +++ b/components/assistant-ui/kb-mention-formatter.ts @@ -0,0 +1,69 @@ +import { + type Unstable_DirectiveFormatter, + type Unstable_DirectiveSegment, +} from "@assistant-ui/react"; + +// ponytail: the key in the brace group matches the search_kb / +// list_documents parameter name so the LLM can copy the value +// directly into the tool call. Doc directive: {documentId=…}; +// folder directive: {folderId=…}. No more generic {id=…} — the +// LLM no longer has to guess which arg to pass. + +export const kbMentionFormatter: Unstable_DirectiveFormatter = { + serialize(item) { + const key = item.type === "kb-folder" ? "folderId" : "documentId"; + return `:${item.type}[${item.label}]{${key}=${item.id}}`; + }, + parse(text) { + // ponytail: accept both {documentId=…} / {folderId=…} and the + // older {id=…} so existing transcript lines (and tests pinned to + // the old format) keep rendering. The LLM only sees the wire + // form, so newly serialised chips will all use the typed keys. + const regex = + /:([\w-]{1,64})\[([^\]\n]{1,1024})\](?:\{(?:documentId|folderId|id)=([^}\n]{1,1024})\})?/g; + const segments: Unstable_DirectiveSegment[] = []; + let lastIndex = 0; + let match: RegExpExecArray | null; + + while ((match = regex.exec(text)) !== null) { + const matchIndex = match.index; + if (matchIndex > lastIndex) { + segments.push({ + kind: "text", + text: text.slice(lastIndex, matchIndex), + }); + } + + const rawType = match[1]!; + const type = rawType === "kb-doc" ? "kb-document" : rawType; + const label = match[2]!; + const explicitId = match[3]; + const id = explicitId || label; // Fallback to label if no explicit id is provided + + segments.push({ + kind: "mention", + type, + id, + label, + }); + + lastIndex = regex.lastIndex; + } + + if (lastIndex < text.length) { + segments.push({ + kind: "text", + text: text.slice(lastIndex), + }); + } + + if (segments.length === 0) { + segments.push({ + kind: "text", + text, + }); + } + + return segments; + }, +}; diff --git a/components/assistant-ui/kb-mention.tsx b/components/assistant-ui/kb-mention.tsx new file mode 100644 index 00000000..ee6dcaff --- /dev/null +++ b/components/assistant-ui/kb-mention.tsx @@ -0,0 +1,376 @@ +"use client"; + +import { + ComposerPrimitive, + unstable_useMentionAdapter, + unstable_useTriggerPopoverScopeContextOptional, + type Unstable_Mention, + type Unstable_MentionCategory, +} from "@assistant-ui/react"; +import { + ChevronLeftIcon, + ChevronRightIcon, + FileTextIcon, + FolderIcon, + SparklesIcon, +} from "lucide-react"; +import { memo, useCallback, useEffect, useRef, useState, useMemo } from "react"; +import { createPortal } from "react-dom"; + +import { cn } from "@/lib/utils"; +import { kbMentionFormatter } from "./kb-mention-formatter"; + +// ponytail: @-mention adapter for KB folders + docs (issue #13 v3). +// Reads `/api/kb/documents?mention=1` (returns folders grouped with +// their success-only docs) and exposes them as categories for +// assistant-ui's `TriggerPopover`. Each folder is a category; each +// category's items are docs in that folder. The first item in each +// category is a synthetic `type: "kb-folder"` row labeled "All in +// " — picking it inserts a `:kb-folder[label]{name=id}` +// directive that the backend resolver expands to every success doc. +// +// Why the synthetic row instead of a header button: aUI's +// `TriggerPopoverItem` already covers selection + keyboard nav + +// directive insertion. A custom header button would mean re-implementing +// the directive formatter, ARIA wiring, and keyboard nav — pure +// overhead. Putting it as the first item gives us the same UX with +// the existing primitives. + +type Doc = { id: string; title: string; status: string }; +type FolderGroup = { + id: string; + name: string; + docCount: number; + docs: Doc[]; +}; +type MentionPayload = { folders?: FolderGroup[] }; + +async function fetchMentionPayload(): Promise { + const res = await fetch("/api/kb/documents?mention=1", { credentials: "include" }); + if (!res.ok) return []; + const body = (await res.json()) as MentionPayload; + return body.folders ?? []; +} + +function KbDocIcon({ className }: { className?: string }) { + return ; +} + +function KbFolderIcon({ className }: { className?: string }) { + return ; +} + +function KbAllInFolderIcon({ className }: { className?: string }) { + return ; +} + +const KB_ICON_MAP = { + "kb-doc": KbDocIcon, + "kb-folder": KbFolderIcon, + "kb-folder-all": KbAllInFolderIcon, +}; + +// formatMentionCategories: each folder becomes its own category. +// The first item in each category is a synthetic "kb-folder" row that +// lets the user mention the entire folder. The remaining items are +// individual docs ("kb-document"). +export function formatMentionCategories( + folders: readonly FolderGroup[], +): Unstable_MentionCategory[] { + return folders.map((f) => ({ + id: f.id, + label: f.name, + items: [ + // First item: select the whole folder + { + id: f.id, + type: "kb-folder", + label: f.name, + description: `All ${f.docCount} doc${f.docCount === 1 ? "" : "s"}`, + icon: "kb-folder", + metadata: { folderId: f.id, folderName: f.name }, + } satisfies Unstable_Mention, + // Subsequent items: individual documents + ...f.docs.map( + (d): Unstable_Mention => ({ + id: d.id, + type: "kb-document", + label: d.title, + description: d.status, + icon: "kb-doc", + metadata: { docId: d.id, parentFolderId: f.id }, + }), + ), + ], + })); +} + +// ponytail: KB popover hook — bundles all state + side effects. +// Returns the aUI adapter bundle plus folders/isLoading/refetch so +// callers can either use the pre-built popover or compose their own. +export function useKbMention() { + const [folders, setFolders] = useState([]); + const [error, setError] = useState(null); + const [isLoading, setIsLoading] = useState(true); + // ponytail: optimistic update — only the FIRST fetch shows the + // skeleton. Subsequent refetches (popover reopens) keep the existing + // data on screen and update silently in the background. Avoids the + // skeleton flash every time the user types `@` after the first time. + const hasInitialDataRef = useRef(false); + + const refetch = useCallback(() => { + let cancelled = false; + if (!hasInitialDataRef.current) { + setIsLoading(true); + } + fetchMentionPayload() + .then((g) => { + if (!cancelled) { + setFolders(g); + setError(null); + hasInitialDataRef.current = true; + } + }) + .catch((e: Error) => { + if (!cancelled) setError(e.message); + }) + .finally(() => { + if (!cancelled) setIsLoading(false); + }); + return () => { + cancelled = true; + }; + }, []); + + // ponytail: no mount refetch — the first popover open drives the + // fetch via PopoverOpenWatcher. The skeleton guard (scope.open check + // in KbMentionSkeleton) keeps the initial isLoading=true from + // rendering an inline skeleton in the composer root. + + // ponytail: convert API payload to aUI's categories shape. + // Each folder is its own top-level category; its docs are category items. + // The library auto-manages navigation: show categories → click folder → + // load category items (docs). TriggerPopoverBack returns to folder list. + const categories = useMemo(() => formatMentionCategories(folders), [folders]); + + const bundle = unstable_useMentionAdapter({ + categories, + iconMap: KB_ICON_MAP, + fallbackIcon: KbFolderIcon, + formatter: kbMentionFormatter, + }); + + // ponytail: aUI strips items before passing categories to the + // popover UI (Unstable_TriggerCategory only has id+label). Pass the + // per-folder docCount separately so the dropdown header can read + // "Research · 5 docs". memoized — bare object literal would defeat + // React.memo on KbMentionPopover. + const docCountByFolderId = useMemo(() => { + const out: Record = {}; + for (const f of folders) out[f.id] = f.docCount; + return out; + }, [folders]); + + return { bundle, folders, error, isLoading, docCountByFolderId, refetch }; +} + +// ponytail: debounce onOpen. aUI's resource can flicker scope.open +// during a close (true → false → true → false) — the bare guard +// `if (isOpen && !wasOpenRef.current)` fires onOpen on every false→true +// transition, so a single close can trigger two refetches. We hold +// onOpen for 100ms of stable-open before firing; if scope.open flips +// back to false in that window, the timeout is cleared and onOpen +// never runs. +const PopoverOpenWatcher = ({ onOpen }: { onOpen: () => void }) => { + const scope = unstable_useTriggerPopoverScopeContextOptional(); + const wasOpenRef = useRef(false); + const timeoutRef = useRef | null>(null); + + useEffect(() => { + const isOpen = scope?.open ?? false; + + if (timeoutRef.current !== null) { + clearTimeout(timeoutRef.current); + timeoutRef.current = null; + } + + if (isOpen) { + if (!wasOpenRef.current) { + timeoutRef.current = setTimeout(() => { + wasOpenRef.current = true; + timeoutRef.current = null; + onOpen(); + }, 100); + } + } else { + wasOpenRef.current = false; + } + }, [scope?.open, onOpen]); + + useEffect( + () => () => { + if (timeoutRef.current !== null) clearTimeout(timeoutRef.current); + }, + [], + ); + + return null; +}; + +// Renders the popover header: +// - When a folder is active (drill-down): shows back button + folder name +// - When at root level: shows "Knowledge Base" title +const KbMentionPopoverHeader = () => { + const scope = unstable_useTriggerPopoverScopeContextOptional(); + if (!scope || !scope.open) return null; + + if (scope.activeCategoryId) { + return ( + + + {"Back"} + + ); + } + + return ( +
+ Knowledge Base +
+ ); +}; + +// ponytail: must guard on scope.open — aUI's Unstable_TriggerPopover +// renders children WITHOUT the positioning wrapper when closed +// (`t15 = open ?
: children` in TriggerPopover.js). The popover +// is portaled to .aui-composer-root (flex-col), so an unguarded +// skeleton would render as an in-flow child there and push the input +// up by ~30px on every popover close. Same pattern as +// KbMentionPopoverHeader / TriggerPopoverCategories / TriggerPopoverItems. +const KbMentionSkeleton = () => { + const scope = unstable_useTriggerPopoverScopeContextOptional(); + if (!scope?.open) return null; + return ( +
+
+
+
+ ); +}; + +// ponytail: wrapped in React.memo + self-contained. The parent Composer +// re-renders on every aUI store update (typing / running / disabled), so +// without memo this component re-runs the JSX construction, the portal +// call, and the PopoverOpenWatcher's effect chain on every store tick. +// State lives inside (useKbMention) — no props to thread through. Memo +// still helps: external re-renders without internal state changes skip. +export const KbMentionPopover = memo(function KbMentionPopover() { + const { bundle, docCountByFolderId, folders, isLoading, refetch } = useKbMention(); + const [portalContainer, setPortalContainer] = useState(null); + + useEffect(() => { + const container = document.querySelector(".aui-composer-root"); + setPortalContainer(container); + }, []); + + if (!isLoading && (!folders || folders.length === 0)) return null; + + // ponytail: wrap the popover JSX in useMemo. aUI's Unstable_TriggerPopover + // toggles a wrapper div on open/close (`t15 = open ?
: children`), + // and React treats a parent-type change as unmount + remount of the + // subtree. That remount resets the PopoverOpenWatcher's wasOpenRef and + // fires onOpen on the next render. Keeping the children reference stable + // (memoized on the real content deps) means the wrapper doesn't toggle + // when the popover re-renders for unrelated reasons — only when the + // content actually changes. + const popoverNode = useMemo( + () => ( + + + {/* Directive must always be mounted so the popover can open when @ is typed */} + + + {isLoading ? ( + + ) : ( + <> + + {(categories) => + categories.map((category) => ( + + + + + {category.label} + + + {docCountByFolderId[category.id] ?? 0} docs + + + + + )) + } + + + {(items) => + items.map((item, index) => ( + + + {item.type === "kb-folder" ? ( + + ) : ( + + )} + + {item.label} + + {item.description && item.type === "kb-folder" ? ( + + {item.description} + + ) : null} + + + )) + } + + + )} + + ), + [bundle.adapter, bundle.directive.formatter, isLoading, docCountByFolderId, refetch], + ); + + if (!portalContainer) return null; + + return createPortal(popoverNode, portalContainer); +}); + +// ponytail: helper hook for components that need the folder payload +// directly (e.g. the DirectiveText chip renderer can use folder icons +// when the directive type is "kb-folder"). +export type { FolderGroup, Doc }; diff --git a/components/assistant-ui/lexical-directive-chip.tsx b/components/assistant-ui/lexical-directive-chip.tsx new file mode 100644 index 00000000..f5e6e7aa --- /dev/null +++ b/components/assistant-ui/lexical-directive-chip.tsx @@ -0,0 +1,18 @@ +import { FC } from "react"; +import { DirectiveChipSpan } from "./directive-chip"; + +type DirectiveChipProps = { + directiveId: string; + directiveType: string; + label: string; +}; + +export const LexicalDirectiveChip: FC = ({ + directiveId, + directiveType, + label, +}) => { + return ( + + ); +}; diff --git a/components/assistant-ui/thread.tsx b/components/assistant-ui/thread.tsx index 9944017d..1d2dc5a1 100644 --- a/components/assistant-ui/thread.tsx +++ b/components/assistant-ui/thread.tsx @@ -3,7 +3,11 @@ import { ComposerAttachments, UserMessageAttachments, } from "@/components/assistant-ui/attachment"; +import { DirectiveText } from "@/components/assistant-ui/directive-text"; +import { LexicalComposerInput, DirectiveNode } from "@assistant-ui/react-lexical"; +import { LexicalDirectiveChip } from "@/components/assistant-ui/lexical-directive-chip"; import { MarkdownText } from "@/components/assistant-ui/markdown-text"; +import { KbMentionPopover } from "@/components/assistant-ui/kb-mention"; import { Reasoning, ReasoningContent, @@ -19,7 +23,7 @@ import { } from "@/components/assistant-ui/tool-group"; import { TooltipIconButton } from "@/components/assistant-ui/tooltip-icon-button"; import { ObservabilityButton } from "@/components/observability/button"; -import { ObservabilitySheet } from "@/components/observability/sheet"; +import { ChatObservabilitySheet } from "@/components/observability/sheet"; import { ObservabilitySheetProvider } from "@/components/observability/sheet-context"; import { WorkingIndicator } from "@/components/assistant-ui/working-indicator"; import { Button } from "@/components/ui/button"; @@ -38,7 +42,10 @@ import { ThreadPrimitive, type ToolCallMessagePartComponent, useAuiState, + useAui, } from "@assistant-ui/react"; +import { useLexicalComposerContext } from "@lexical/react/LexicalComposerContext"; +import { COMMAND_PRIORITY_LOW, PASTE_COMMAND, $getNodeByKey, $createTextNode } from "lexical"; import { ArrowDownIcon, ArrowUpIcon, @@ -57,6 +64,7 @@ import { import { createContext, useContext, + useEffect, type ComponentType, type FC, type PropsWithChildren, @@ -100,7 +108,7 @@ export const Thread: FC = ({ components = EMPTY_COMPONENTS }) => { - + ); @@ -216,6 +224,62 @@ const ThreadSuggestionItem: FC = () => { ); }; +const LexicalPasteFilesPlugin: FC = () => { + const [editor] = useLexicalComposerContext(); + const aui = useAui(); + + useEffect(() => { + return editor.registerCommand( + PASTE_COMMAND, + (event: ClipboardEvent) => { + const files = event.clipboardData?.files; + if (files && files.length > 0) { + event.preventDefault(); + Promise.all(Array.from(files).map((file) => aui.composer().addAttachment(file))).catch( + (err) => { + console.error("Failed to paste files:", err); + }, + ); + return true; + } + return false; + }, + COMMAND_PRIORITY_LOW, + ); + }, [editor, aui]); + + return null; +}; + +const LexicalAutoSpacePlugin: FC = () => { + const [editor] = useLexicalComposerContext(); + + useEffect(() => { + return editor.registerMutationListener(DirectiveNode as any, (mutations) => { + mutations.forEach((mutation, key) => { + if (mutation === "created") { + editor.update(() => { + const node = $getNodeByKey(key); + if (node instanceof DirectiveNode) { + const nextSibling = node.getNextSibling(); + if ( + !nextSibling || + (nextSibling.getType() === "text" && !nextSibling.getTextContent().startsWith(" ")) + ) { + const spaceNode = $createTextNode(" "); + node.insertAfter(spaceNode); + (spaceNode as any).select(); + } + } + }); + } + }); + }); + }, [editor]); + + return null; +}; + const Composer: FC = () => { // ponytail: lock the textbox while r2-adapter.send() is in flight so // the user can't edit a message that's mid-upload. SDK keeps the @@ -232,6 +296,8 @@ const Composer: FC = () => { typeof window !== "undefined" && window.matchMedia("(pointer: fine)").matches && navigator.maxTouchPoints === 0; + // ponytail: KbMentionPopover manages its own state via useKbMention + // internally — no need to call the hook here or thread props down. return ( @@ -240,15 +306,29 @@ const Composer: FC = () => { className="border-border/60 data-[dragging=true]:border-ring focus-within:border-border dark:border-muted-foreground/15 dark:focus-within:border-muted-foreground/30 flex w-full flex-col gap-2 rounded-(--composer-radius) border bg-(--composer-bg) p-(--composer-padding) shadow-[0_4px_16px_-8px_rgba(0,0,0,0.08),0_1px_2px_rgba(0,0,0,0.04)] transition-[border-color,box-shadow] focus-within:shadow-[0_6px_24px_-8px_rgba(0,0,0,0.12),0_1px_2px_rgba(0,0,0,0.05)] data-[dragging=true]:border-dashed data-[dragging=true]:bg-[color-mix(in_oklab,var(--color-accent)_50%,var(--color-background))] dark:shadow-none" > - - + {/* ponytail: aUI's trigger popover contract (issue #13 v3): + TriggerPopoverRoot owns the registered trigger map + ARIA + context. The Input lives inside the root so the root's + composer state scope is shared. Unstable_TriggerPopover is + a SIBLING of the input — it renders the popover container + when `@` is typed. Its children are the directive + + categories/items render-prop components. Without those + children (and without `Directive`), the popover stays + closed (per the docstring on TriggerPopover). */} + + + + + + + +
@@ -533,6 +613,12 @@ const UserMessage: FC = () => { > null, File: () => null, }} @@ -576,9 +662,10 @@ const EditComposer: FC = () => { return ( -
diff --git a/components/auth/settings/credit-tab.tsx b/components/auth/settings/credit-tab.tsx index 95aeff0b..f19abcdd 100644 --- a/components/auth/settings/credit-tab.tsx +++ b/components/auth/settings/credit-tab.tsx @@ -15,7 +15,11 @@ export const creditSettingsPlugin = { label: ( <> - Credits + {/* ponytail: text is auto-hidden on mobile by the global + rule in app/globals.css ([role="tab"] > span:not(.sr-only)) + using the sr-only technique. Above md the rule doesn't + apply and the text shows inline. */} + Credits ), component: CreditHistory, diff --git a/components/auth/settings/kb-tab.tsx b/components/auth/settings/kb-tab.tsx new file mode 100644 index 00000000..8f9f89de --- /dev/null +++ b/components/auth/settings/kb-tab.tsx @@ -0,0 +1,27 @@ +import { BookOpen } from "lucide-react"; + +import { KbView } from "@/components/settings/kb-view"; + +// ponytail: KB settings plugin. Mirrors memory-tab.tsx — same shape +// (id + viewPaths + settingsTabs) so BetterAuthUIProvider renders the +// tab alongside Memory + Credit. +export const kbSettingsPlugin = { + id: "knowledge-base", + viewPaths: { settings: { "knowledge-base": "knowledge-base" } }, + settingsTabs: [ + { + view: "knowledge-base", + label: ( + <> + + {/* ponytail: text is auto-hidden on mobile by the global + rule in app/globals.css ([role="tab"] > span:not(.sr-only)) + using the sr-only technique. Above md the rule doesn't + apply and the text shows inline. */} + Knowledge Base + + ), + component: KbView, + }, + ], +}; diff --git a/components/auth/settings/memory-tab.tsx b/components/auth/settings/memory-tab.tsx index da0e17d1..f59714d9 100644 --- a/components/auth/settings/memory-tab.tsx +++ b/components/auth/settings/memory-tab.tsx @@ -16,7 +16,11 @@ export const memorySettingsPlugin = { label: ( <> - Memory + {/* ponytail: text is auto-hidden on mobile by the global + rule in app/globals.css ([role="tab"] > span:not(.sr-only)) + using the sr-only technique. Above md the rule doesn't + apply and the text shows inline. */} + Memory ), component: MemoryView, diff --git a/components/auth/settings/settings.tsx b/components/auth/settings/settings.tsx index ce763545..c9b38472 100644 --- a/components/auth/settings/settings.tsx +++ b/components/auth/settings/settings.tsx @@ -69,8 +69,10 @@ export function Settings({ className, view, path, hideNav }: SettingsProps) { } > - - {localization.settings.account} + {/* ponytail: wrap the label in a so the mobile CSS + rule ([role="tab"] > span:not(.sr-only)) can hide it. + Bare text nodes can't be selected by element selectors. */} + {localization.settings.account} - - {localization.settings.security} + {localization.settings.security} {plugins.flatMap( diff --git a/components/brand-mark.tsx b/components/brand-mark.tsx index 1f833c3c..baef995f 100644 --- a/components/brand-mark.tsx +++ b/components/brand-mark.tsx @@ -1,11 +1,10 @@ import type { FC } from "react"; import Link from "next/link"; -import { MessageSquareTextIcon } from "lucide-react"; import { cn } from "@/lib/utils"; import { APP_NAME } from "@/lib/constants"; -const TEXT_CLASS = "text-foreground/90 ml-2 text-sm font-medium whitespace-nowrap"; +const TEXT_CLASS = "text-foreground/90 text-sm font-medium whitespace-nowrap"; // ponytail: app logo block — icon + APP_NAME in a single rounded row. // The collapse animation lives on the text node (opacity-0 when @@ -19,7 +18,6 @@ export const BrandMark: FC<{ collapsed?: boolean; className?: string }> = ({ className, }) => (
- {APP_NAME} diff --git a/components/landing/cta.tsx b/components/landing/cta.tsx index 067fa75d..f1615d93 100644 --- a/components/landing/cta.tsx +++ b/components/landing/cta.tsx @@ -27,7 +27,7 @@ export const Cta: FC = ({ signedIn }) => { // `from` angle rotates via `cta-marquee` (see globals.css) so // the warm overflow breathes instead of standing still. The // radial stays put to anchor the layout. -
+
= ({ signedIn }) => { }} />
-
+

Read the code. Run it. Skip the demo.

diff --git a/components/landing/features.tsx b/components/landing/features.tsx index 2301378a..3a833222 100644 --- a/components/landing/features.tsx +++ b/components/landing/features.tsx @@ -110,9 +110,9 @@ const BENTO: BentoCard[] = [ codePreview: ``, }, { - title: "Cross-conversation memory", + title: "Memory + Knowledge Base", description: - "User facts and recent threads surface automatically. \nThe model sees them in a prepended system block; the Memory tab lets you review and delete.", + "User facts surface in the system block; PDFs go through a per-doc pipeline (OCR → chunk → embed → entity) and become a hybrid-searchable index. Both are reviewable and deletable from settings.", icon: , hue: "violet", span: "wide", diff --git a/components/landing/how-it-works.tsx b/components/landing/how-it-works.tsx index cf06fd75..b9856488 100644 --- a/components/landing/how-it-works.tsx +++ b/components/landing/how-it-works.tsx @@ -6,6 +6,7 @@ import type { FC } from "react"; import { BackgroundSplitDemo } from "@/components/landing/motion/background-split-demo"; import { HumanInTheLoopDemo } from "@/components/landing/motion/human-in-the-loop-demo"; +import { KbExplainerDemo } from "@/components/landing/motion/kb-explainer-demo"; import { MemoryRecallDemo } from "@/components/landing/motion/memory-recall-demo"; import { ObservabilityWaterfallDemo } from "@/components/landing/motion/observability-waterfall-demo"; import { StreamingTokensDemo } from "@/components/landing/motion/streaming-tokens-demo"; @@ -63,6 +64,31 @@ export const HowItWorks: FC = () => ( > + + +

+ The KB agent renders each page, OCRs the markdown, splits and embeds the chunks, + then runs an LLM to extract entities and relationships. +

+

+ Hybrid Search: three legs run + in one round-trip — BM25 for + the exact term, pgvector cosine + for semantic closeness, and an{" "} + entity-overlap leg that walks + the graph from your query node. The three are fused with{" "} + RRF (Reciprocal Rank Fusion). +

+ + } + reverse + > + +
@@ -71,7 +97,10 @@ export const HowItWorks: FC = () => ( type ExplainRowProps = { eyebrow: string; title: string; - body: string; + // ponytail: ReactNode so a row can pass multi-paragraph copy + // (e.g. the KB row splits ingest / search / rerank / scopes + // into separate

blocks). Plain-string bodies still work. + body: React.ReactNode; reverse?: boolean; children: React.ReactNode; }; @@ -86,7 +115,10 @@ const ExplainRow: FC = ({ eyebrow, title, body, reverse, childr

{eyebrow}

{title}

-

{body}

+ {/* ponytail: a div (not

) so callers can pass a Fragment of + multiple

blocks for longer copy. Space-y-3 keeps the + rhythm consistent whether the body is a string or 4 paragraphs. */} +

{body}
{children} diff --git a/components/landing/motion/kb-explainer-demo.tsx b/components/landing/motion/kb-explainer-demo.tsx new file mode 100644 index 00000000..a2e7d913 --- /dev/null +++ b/components/landing/motion/kb-explainer-demo.tsx @@ -0,0 +1,499 @@ +"use client"; + +// ponytail: KB explainer — two stacked panels in one card. +// Top: PDF → OCR → chunk → embed pipeline. Plays once on scroll +// into view (stages light up + arrows draw in sequence). +// Bottom: entity graph traversal. Same scroll trigger; plays +// after the pipeline settles. +// +// Why both: the pipeline shows the WRITE path (how a PDF becomes a +// searchable index); the graph shows the READ path (how a query +// traverses entities + relationships). Together they cover the +// round trip — that's the story the How-it-works row tells. + +import { m, useInView, useReducedMotion } from "motion/react"; +import { Fragment, useEffect, useRef, useState } from "react"; +import { + ArrowRightIcon, + BlocksIcon, + FileTextIcon, + LayersIcon, + ScanTextIcon, + SearchIcon, + SparklesIcon, + WorkflowIcon, +} from "lucide-react"; + +import { cn } from "@/lib/utils"; + +type Stage = { + id: string; + label: string; + icon: typeof FileTextIcon; + // ponytail: literal Tailwind colors so JIT picks them up — dynamic + // text-${color}-500 / bg-${color}-500 would silently no-op. + lit: string; +}; + +const STAGES: Stage[] = [ + { + id: "pdf", + label: "PDF", + icon: FileTextIcon, + lit: "text-rose-500 border-rose-500/40 bg-rose-500/10", + }, + { + id: "ocr", + label: "OCR", + icon: ScanTextIcon, + lit: "text-amber-500 border-amber-500/40 bg-amber-500/10", + }, + { + id: "chunk", + label: "Chunk", + icon: LayersIcon, + lit: "text-emerald-500 border-emerald-500/40 bg-emerald-500/10", + }, + { + id: "embed", + label: "Embed", + icon: BlocksIcon, + lit: "text-sky-500 border-sky-500/40 bg-sky-500/10", + }, + { + id: "entity", + label: "Entity", + icon: WorkflowIcon, + lit: "text-violet-500 border-violet-500/40 bg-violet-500/10", + }, +]; + +type GraphNode = { + id: string; + label: string; + x: number; + y: number; + lit: string; + type: "query" | "person" | "tech" | "concept" | "project"; + z?: number; // 0..1 for perspective scaling (front-side larger) +}; + +type GraphEdge = [string, string]; + +// ponytail: spherical layout — Fibonacci-sphere distribution around +// the center reads as a 3D ball at a glance. ~20 nodes is dense +// enough to feel like a "real" graph without crowding a 320px-wide +// demo card. Z coordinate drives the SVG radius so back-side nodes +// sit smaller (perspective). +// +// Node colors cycle through a fixed palette — palette is the +// visual story, individual labels matter less. Query node is the +// primary-colored "?" in the middle; everything else is an entity. +// ponytail: solid fill only (no stroke) — the previous +// `fill-X stroke-X` pair rendered the line through the circle's +// outline. Without a stroke the node reads as a solid dot and the +// edge ends cleanly behind it. +const PALETTE = [ + "fill-rose-500", + "fill-amber-500", + "fill-emerald-500", + "fill-sky-500", + "fill-violet-500", + "fill-orange-500", + "fill-fuchsia-500", + "fill-teal-500", +]; + +// ponytail: deterministic pseudo-random edges so the graph looks +// organic but reproducible. Each new node gets 2-3 random edges to +// earlier nodes — keeps mean degree low so the layout stays legible. +function buildGraph(nodeCount: number) { + const nodes: GraphNode[] = []; + // Query node centered + nodes.push({ + id: "q", + label: "?", + x: 100, + y: 100, + lit: "fill-primary", + type: "query", + }); + const cx = 100; + const cy = 100; + const sphereR = 75; + // Fibonacci sphere — gives even angular distribution + const goldenAngle = Math.PI * (3 - Math.sqrt(5)); + for (let i = 0; i < nodeCount - 1; i++) { + const idx = i + 1; + const yNorm = 1 - (idx / (nodeCount - 1)) * 2; // -1..1 + const radiusAtY = Math.sqrt(1 - yNorm * yNorm); + const theta = goldenAngle * idx; + const x3d = Math.cos(theta) * radiusAtY; + const z3d = Math.sin(theta) * radiusAtY; // -1..1 + // Map to 2D + scale Z to perspective radius + nodes.push({ + id: `n${idx}`, + label: "", + x: cx + x3d * sphereR, + y: cy + yNorm * sphereR, + lit: PALETTE[(idx - 1) % PALETTE.length] ?? PALETTE[0]!, + type: "concept", + // z normalized 0..1 (front-side larger) + z: (z3d + 1) / 2, + }); + } + // Edges: each non-query node connects to its 2 nearest neighbours + // (by Euclidean distance) — produces an organic-looking web. The + // edge set is canonical (i,j) regardless of which node added it + // first, so dedup is a Set keyed on sorted-pair ids. + const edgeKeys = new Set(); + const edges: GraphEdge[] = []; + for (let i = 1; i < nodes.length; i++) { + const a = nodes[i]!; + const dists: Array<{ id: string; d: number }> = []; + for (let j = 1; j < nodes.length; j++) { + if (i === j) continue; + const b = nodes[j]!; + dists.push({ id: b.id, d: Math.hypot(a.x - b.x, a.y - b.y) }); + } + dists.sort((x, y) => x.d - y.d); + for (const n of dists.slice(0, 2)) { + const key = [a.id, n.id].sort().join("|"); + if (edgeKeys.has(key)) continue; + edgeKeys.add(key); + edges.push([a.id, n.id]); + } + } + // Hook the query node to the 3 closest entities so the traversal + // visual reads as "query hits the graph" + const queryDists = nodes + .slice(1) + .map((n) => ({ id: n.id, d: Math.hypot(n.x - cx, n.y - cy) })) + .sort((a, b) => a.d - b.d) + .slice(0, 3); + for (const q of queryDists) { + edges.push(["q", q.id]); + } + return { nodes, edges }; +} + +const { nodes: NODES, edges: EDGES } = buildGraph(20); + +const NODE_R = 5; + +// ponytail: 2-hop BFS from the query node to find the pulse set. +// Pre-computed (NODES is static) so the demo runs without a real +// graph search — keeps the file readable and unit-testable. +const PULSE_IDS = new Set(["q"]); +const hop1 = new Set(); +const hop2 = new Set(); +for (const [from, to] of EDGES) { + if (from === "q") hop1.add(to); +} +for (const [from, to] of EDGES) { + if (hop1.has(from)) hop2.add(to); +} +for (const id of hop1) PULSE_IDS.add(id); +for (const id of hop2) PULSE_IDS.add(id); + +// ponytail: timing — pipeline first, graph second. The pipeline is +// the "how a doc gets in" story; the graph is the "how a query +// finds it" story. Sequencing keeps the eye on the ingest path +// before the read path lights up. +const STAGE_DURATION = 200; +const PIPELINE_QUERY_DELAY = STAGES.length * STAGE_DURATION + 200; +const PIPELINE_RETRIEVE_DELAY = PIPELINE_QUERY_DELAY + 300; +const GRAPH_NODE_DELAY = PIPELINE_RETRIEVE_DELAY + 400; +const GRAPH_PULSE_DELAY = GRAPH_NODE_DELAY + NODES.length * 100 + 200; + +export const KbExplainerDemo = () => { + const reduced = useReducedMotion(); + const ref = useRef(null); + const inView = useInView(ref, { once: true, amount: 0.3 }); + const [step, setStep] = useState(reduced ? STAGES.length + NODES.length + 4 : 0); + + useEffect(() => { + if (reduced) return; + if (!inView) { + setStep(0); + return; + } + let cancelled = false; + const timers: number[] = []; + // Pipeline stages + for (let i = 1; i <= STAGES.length; i++) { + timers.push( + window.setTimeout(() => !cancelled && setStep((s) => Math.max(s, i)), STAGE_DURATION * i), + ); + } + // Query card + timers.push( + window.setTimeout( + () => !cancelled && setStep((s) => Math.max(s, STAGES.length + 1)), + PIPELINE_QUERY_DELAY, + ), + ); + // Retrieved chunk card + timers.push( + window.setTimeout( + () => !cancelled && setStep((s) => Math.max(s, STAGES.length + 2)), + PIPELINE_RETRIEVE_DELAY, + ), + ); + // Graph nodes + for (let i = 1; i <= NODES.length; i++) { + timers.push( + window.setTimeout( + () => !cancelled && setStep((s) => Math.max(s, STAGES.length + 2 + i)), + GRAPH_NODE_DELAY + 100 * i, + ), + ); + } + // Graph pulse + timers.push( + window.setTimeout( + () => !cancelled && setStep((s) => Math.max(s, STAGES.length + 2 + NODES.length + 1)), + GRAPH_PULSE_DELAY, + ), + ); + return () => { + cancelled = true; + timers.forEach((id) => window.clearTimeout(id)); + }; + }, [inView, reduced]); + + const pipelineLit = Math.min(step, STAGES.length); + const queryVisible = step > STAGES.length; + const retrieveVisible = step > STAGES.length + 1; + const graphLit = Math.max(0, Math.min(step - STAGES.length - 2, NODES.length)); + const graphPulse = step > STAGES.length + 2 + NODES.length; + + return ( +
+ {/* === Pipeline row === */} +
+
+ KB ingest pipeline + per document +
+ +
+
+ {STAGES.map((stage, i) => ( + + + {i < STAGES.length - 1 && ( + i + 1 ? { opacity: 1, scale: 1 } : { opacity: 0, scale: 0.6 } + } + transition={{ duration: 0.25, ease: "easeOut" }} + aria-hidden + > + + + )} + + ))} +
+
+ + {/* Retrieval layer — only after the pipeline completes. */} +
+ + + + what does the doc say about @kb-doc? + + + + +

+ [chunk 7 · vector + entity hit · 0.94]{" "} + The BM25 leg finds the exact term; the entity leg ties it to the canonical name. +

+ +
+
+ + {/* === Graph row === */} +
+
+ Entity graph + graph traversal +
+ +
+ + {/* ponytail: defs come first so the arrow marker is + available to every edge below. markerUnits="userSpaceOnUse" + keeps the arrowhead size constant regardless of the + line's stroke width — otherwise the head scales with + the edge and pulses look uneven. */} + + + + + + + + + {/* ponytail: edges render FIRST so nodes sit on top — the + "lines drawn over circles" bug is just SVG paint order + (later elements sit above earlier ones). */} + {EDGES.map(([fromId, toId], i) => { + const from = NODES.find((n) => n.id === fromId)!; + const to = NODES.find((n) => n.id === toId)!; + const fromIdx = NODES.findIndex((n) => n.id === fromId); + const toIdx = NODES.findIndex((n) => n.id === toId); + const lit = graphLit > Math.max(fromIdx, toIdx); + const isPulseEdge = graphPulse && PULSE_IDS.has(fromId) && PULSE_IDS.has(toId); + return ( + + ); + })} + {NODES.map((node, i) => { + const lit = i < graphLit; + const isPulse = graphPulse && PULSE_IDS.has(node.id); + // ponytail: z drives the SVG radius — front-side nodes + // (z=1) render larger, back-side (z=0) smaller. The + // 0.85..1.3 range keeps the smallest back-side node + // wide enough to fully cover the line behind it; the + // earlier 0.6 floor left a visible edge bleed. Query + // node sits at the centre so it gets the largest + // radius. Opacity stays 1.0 throughout — the + // perspective comes from size only, and a + // semi-transparent node lets the line behind show + // through, which reads as a bug. + const z = node.z ?? 0.5; + const rScale = 0.85 + z * 0.45; + const baseR = node.type === "query" ? NODE_R * 1.6 : NODE_R; + const r = isPulse ? baseR * rScale * 1.2 : baseR * rScale; + return ( + + + {/* ponytail: only the query node gets a label — 20 + nodes with text would be illegible at this + scale. The query reads as the "?" so the user + can identify the centre of the traversal. */} + {node.type === "query" && ( + + {node.label} + + )} + + ); + })} + +
+ +
+ + + Query ? lights up its 2-hop + neighborhood. The graph leg scores every entity it touches alongside BM25 + vector. + +
+
+
+ ); +}; + +const StageBox = ({ stage, lit }: { stage: Stage; lit: boolean }) => { + const Icon = stage.icon; + return ( + + + + {stage.label} + + + ); +}; diff --git a/components/landing/motion/memory-recall-demo.tsx b/components/landing/motion/memory-recall-demo.tsx index cf718c40..0e27bda1 100644 --- a/components/landing/motion/memory-recall-demo.tsx +++ b/components/landing/motion/memory-recall-demo.tsx @@ -27,7 +27,7 @@ export const MemoryRecallDemo = () => { const visible = reduced || inView; return ( -
+

Memory diff --git a/components/observability/panel.tsx b/components/observability/panel.tsx index 60dbc785..65977bd7 100644 --- a/components/observability/panel.tsx +++ b/components/observability/panel.tsx @@ -300,6 +300,26 @@ const TypedBadge: FC = () => { return ; }; +// ponytail: failed spans get a small AlertCircle next to the span name +// in the waterfall list. The timeline bar already draws a red border +// (WaterfallBar, lines 233-244), but a vertical list at-a-glance scan +// needs an inline cue too — clicking still opens the detail panel +// where the FAILED badge repeats. +const FailedIndicator: FC = () => { + const status = useAuiState( + (s) => (s as unknown as { span: SpanItemState }).span.status, + ) as SpanItemState["status"]; + if (status !== "failed") return null; + return ( + + ); +}; + const STATUS_STYLE: Record = { completed: { color: "hsl(142 71% 45%)", @@ -361,7 +381,7 @@ const WaterfallRow: FC = () => { { +

mounted at ThreadRoot. Subscribes to the // sheet-context to know which thread to load; holds the fetch lifecycle // (was previously inlined in button.tsx → per-message N-fold duplication). +// +// Decoupled from @assistant-ui/react's useAuiState — callers that DO have +// an AuiProvider pass `activeThreadId` so the sheet can show its +// "thread has changed" warning; callers without (e.g. +// /settings/knowledge-base) simply omit it and the warning never +// renders. This keeps the sheet reusable outside the chat runtime. import { useCallback, useEffect, useRef, useState } from "react"; import type { FC } from "react"; import { Sheet, SheetContent, SheetHeader, SheetTitle } from "@/components/ui/sheet"; @@ -14,6 +20,15 @@ import { useObservabilitySheetState } from "@/components/observability/sheet-con import type { AggregateDTO, InFlightRun, SpanDataDTO } from "@/lib/observability/validators"; const LOCAL_THREAD_PREFIX = "__LOCAL_"; +void LOCAL_THREAD_PREFIX; + +export type ObservabilitySheetProps = { + // ponytail: the chat runtime's currently-active thread id, when the + // caller is mounted under . Optional — when omitted, the + // "thread has changed" warning below is suppressed. Settings page + // passes nothing; the chat thread tree reads it via ChatObservabilitySheet. + activeThreadId?: string | null; +}; // ponytail: 10s poll while a bg agent is in flight. Once the API reports // in_flight_runs is empty, onRefresh resolves false and the polling + @@ -90,7 +105,21 @@ const RefreshCountdown: FC = ({ enabled, refreshIntervalM ); }; -export const ObservabilitySheet: FC = () => { +// ponytail: chat-side wrapper. Pulls the active thread id from the +// assistant-ui runtime via useAuiState and forwards it to the bare +// ObservabilitySheet so the "thread has changed" warning fires only +// inside the chat runtime. Settings page mounts +// directly (no provider) — no warning, no AuiProvider requirement. +export const ChatObservabilitySheet: FC = () => { + const auiThreadId = useAuiState((s) => { + const item = s.threads.threadItems.find((t) => t.id === s.threads.mainThreadId); + const candidate = item?.externalId ?? s.threads.mainThreadId; + return candidate && !candidate.startsWith(LOCAL_THREAD_PREFIX) ? candidate : null; + }); + return ; +}; + +export const ObservabilitySheet: FC = ({ activeThreadId = null }) => { const { open, threadId, parentMessageId, setOpen } = useObservabilitySheetState(); const [spans, setSpans] = useState([]); const [aggregate, setAggregate] = useState(null); @@ -100,15 +129,12 @@ export const ObservabilitySheet: FC = () => { const [loading, setLoading] = useState(false); const [error, setError] = useState(null); - // ponytail: when the sheet is closed we still track the latest threadId - // the user clicked on (saved in sheet-context state) so a re-open on the - // same thread shows cached data without a round trip. Reset load flags - // when the active threadId changes. - const auiThreadId = useAuiState((s) => { - const item = s.threads.threadItems.find((t) => t.id === s.threads.mainThreadId); - const candidate = item?.externalId ?? s.threads.mainThreadId; - return candidate && !candidate.startsWith(LOCAL_THREAD_PREFIX) ? candidate : null; - }); + // ponytail: `activeThreadId` is passed in by callers mounted under + // (ChatObservabilitySheet). Non-chat callers (settings + // page) omit it — the "thread has changed" warning below is + // suppressed, since there's no surrounding chat thread to compare + // against. The previous direct useAuiState call forced every mount + // site to live under AuiProvider, breaking settings usage. // ponytail: monotonic fetch id — every loadSpans() bumps it, and // older in-flight responses early-return when they no longer own @@ -198,7 +224,7 @@ export const ObservabilitySheet: FC = () => { - {auiThreadId && threadId && auiThreadId !== threadId ? ( + {activeThreadId && threadId && activeThreadId !== threadId ? (
Thread has changed — close and reopen to see the latest data.
diff --git a/components/settings/kb-view/dialogs.tsx b/components/settings/kb-view/dialogs.tsx new file mode 100644 index 00000000..5736f463 --- /dev/null +++ b/components/settings/kb-view/dialogs.tsx @@ -0,0 +1,564 @@ +import { useCallback, useEffect, useState } from "react"; +import { toast } from "sonner"; +import { Loader2 } from "lucide-react"; +import { Button } from "@/components/ui/button"; +import { + Dialog, + DialogContent, + DialogDescription, + DialogFooter, + DialogHeader, + DialogTitle, +} from "@/components/ui/dialog"; +import { Input } from "@/components/ui/input"; +import { Label } from "@/components/ui/label"; +import { cn } from "@/lib/utils"; +import { TOAST_DESCRIPTION_CLASS, getModeInfo, type ReprocessMode } from "./helpers"; +import { KbDocument, KbFolder } from "./types"; + +export function DocDeleteDialog({ + doc, + open, + onOpenChange, + onDeleted, +}: { + doc: KbDocument; + open: boolean; + onOpenChange: (open: boolean) => void; + onDeleted: () => void; +}) { + const [submitting, setSubmitting] = useState(false); + const [error, setError] = useState(null); + + const submit = useCallback(async () => { + setSubmitting(true); + setError(null); + try { + const res = await fetch(`/api/kb/documents/${doc.id}`, { method: "DELETE" }); + if (res.status === 204) { + toast.success("Document deleted", { + descriptionClassName: TOAST_DESCRIPTION_CLASS, + description: `「${doc.title}」was removed from this folder.`, + }); + onDeleted(); + return; + } + if (res.status === 404) { + setError("Already deleted"); + return; + } + setError(`Failed (${res.status})`); + } catch (e) { + setError(e instanceof Error ? e.message : String(e)); + } finally { + setSubmitting(false); + } + }, [doc.id, onDeleted]); + + return ( + { + onOpenChange(o); + if (!o) setError(null); + }} + > + + + Delete document? + + {/* ponytail: asChild swaps the default

for a

so + block-level children (ol, multi-paragraph copy) don't + trigger

-inside-

/

-inside-

hydration + errors. Slot merges text-sm + text-muted-foreground + from DialogDescription onto the div. */} +

+

+ This permanently removes {doc.title} and: +

+
    +
  1. The document row + all parsed chunks (embeddings, BM25 index, entity graph)
    2. +
  2. The observability run history for this doc
    3. +
  3. The standalone ingestion thread record (LangGraph metadata)
    4. +
+

+ Source PDF + rendered page PNGs stay in R2 until the v3 retention sweep. Raw + observability spans + LangGraph checkpoint state are kept by retention. +

+
+ + + {error &&

{error}

} + + + + + +
+ ); +} + +export function DocReprocessDialog({ + doc, + open, + onOpenChange, + onReprocessed, +}: { + doc: KbDocument; + open: boolean; + onOpenChange: (open: boolean) => void; + onReprocessed: () => void; +}) { + const [mode, setMode] = useState("full"); + const [submitting, setSubmitting] = useState(false); + const [error, setError] = useState(null); + + const hasPages = !!( + (doc.pages && doc.pages.length > 0) || + (doc.totalPages !== undefined && doc.totalPages > 0) + ); + + const totalPages = doc.pages ? doc.pages.length : (doc.totalPages ?? 0); + const failedPagesCount = doc.pages + ? doc.pages.filter((p) => !!p.errorMessage || !(p.markdown ?? "").trim()).length + : (doc.failedPages ?? 0); + + const hasUsableMarkdown = doc.pages + ? doc.pages.some((p) => (p.markdown ?? "").trim().length > 0) + : totalPages > failedPagesCount; + + const hasFailedPages = failedPagesCount > 0; + + const isChunksOnlyDisabled = !hasPages || !hasUsableMarkdown; + const isRetryFailedDisabled = !hasPages || !hasFailedPages; + + // ponytail: retryFailedChunks only makes sense once the doc + // reached a terminal indexing state. failedChunks comes from the + // doc-list endpoint aggregation, so we gate on it being + // populated. Disabled when 0 failed chunks — picking it would be + // a wasted API call. + const hasFailedChunks = (doc.failedChunks ?? 0) > 0; + const isRetryFailedChunksDisabled = doc.status !== "success" || !hasFailedChunks; + + useEffect(() => { + if (open) { + setMode("full"); + setError(null); + } + }, [open]); + + const submit = useCallback(async () => { + setSubmitting(true); + setError(null); + try { + const res = await fetch(`/api/kb/documents/${doc.id}/reprocess?mode=${mode}`, { + method: "POST", + }); + if (res.status === 202) { + let toastTitle = "Reprocess queued"; + let toastDesc = `「${doc.title}」is re-running OCR + chunking. Old chunks were cleared.`; + if (mode === "chunksOnly") { + toastTitle = "Rechunks queued"; + toastDesc = `「${doc.title}」- skipping OCR, rebuilding chunks from the cached pages. doc row stays Ready.`; + } else if (mode === "retryFailed") { + toastTitle = "Retry queued"; + toastDesc = `「${doc.title}」- retrying failed pages, then rebuilding chunks.`; + } else if (mode === "retryFailedChunks") { + toastTitle = "Chunk retry queued"; + toastDesc = `「${doc.title}」- re-running entity extraction on the failed chunks only. Successful chunks and the doc status stay untouched.`; + } + + toast.info(toastTitle, { + descriptionClassName: TOAST_DESCRIPTION_CLASS, + description: toastDesc, + }); + onReprocessed(); + return; + } + if (res.status === 409) { + const body = (await res.json().catch(() => ({}))) as { code?: string; reason?: string }; + if (body.code === "NOT_READY") { + setError( + "This option is not available — pages haven't been extracted yet. Pick 'Full re-run' first.", + ); + return; + } + const msg = + body.code === "ATTACHMENT_MISSING" + ? "Source attachment is missing — re-upload the file instead." + : body.code === "PROCESSING" + ? "Already processing — try again when the row settles." + : `Server rejected: ${body.code ?? res.status}`; + setError(msg); + return; + } + if (res.status === 404) { + setError("Document no longer exists."); + return; + } + setError(`Failed (${res.status})`); + } catch (e) { + setError(e instanceof Error ? e.message : String(e)); + } finally { + setSubmitting(false); + } + }, [doc.id, doc.title, mode, onReprocessed]); + + return ( + { + onOpenChange(o); + if (!o) setError(null); + }} + > + + + Reprocess document? + + {doc.title} - existing chunks are wiped before + re-running. Choose a reprocess mode. + + + +
+ {( + [ + { + value: "full", + disabled: false, + reason: "", + }, + { + value: "chunksOnly", + disabled: isChunksOnlyDisabled, + reason: "No pages cache", + }, + { + value: "retryFailed", + disabled: isRetryFailedDisabled, + reason: "No failed pages", + }, + { + value: "retryFailedChunks", + disabled: isRetryFailedChunksDisabled, + reason: doc.status !== "success" ? "Doc not indexed yet" : "No failed chunks", + }, + ] satisfies Array<{ value: ReprocessMode; disabled: boolean; reason: string }> + ).map(({ value, disabled, reason }) => { + const { title, description } = getModeInfo(value); + return ( + + ); + })} +
+ + {error &&

{error}

} + + + + + + + ); +} + +export function FolderDeleteDialog({ + folder, + onOpenChange, + onDeleted, +}: { + folder: KbFolder | null; + onOpenChange: (open: boolean) => void; + onDeleted: () => void; +}) { + const [submitting, setSubmitting] = useState(false); + const [error, setError] = useState(null); + + useEffect(() => { + setError(null); + }, [folder]); + + const submit = useCallback(async () => { + if (!folder) return; + setSubmitting(true); + setError(null); + try { + const res = await fetch(`/api/kb/folders/${folder.id}`, { method: "DELETE" }); + if (res.status === 204) { + toast.success("Folder deleted", { + descriptionClassName: TOAST_DESCRIPTION_CLASS, + description: `「${folder.name}」was removed. Documents inside were kept.`, + }); + onDeleted(); + return; + } + if (res.status === 404) { + toast.success("Folder deleted", { + descriptionClassName: TOAST_DESCRIPTION_CLASS, + description: `「${folder.name}」was already removed.`, + }); + onDeleted(); + return; + } + if (res.status === 409) { + const body = (await res.json().catch(() => ({}))) as { docCount?: number }; + setError( + `Folder still has ${body.docCount ?? "some"} document${body.docCount === 1 ? "" : "s"} — delete them first.`, + ); + return; + } + setError(`Failed (${res.status})`); + } catch (e) { + setError(e instanceof Error ? e.message : String(e)); + } finally { + setSubmitting(false); + } + }, [folder, onDeleted]); + + return ( + { + onOpenChange(o); + if (!o) setError(null); + }} + > + + + Delete folder? + + This permanently removes the {folder?.name} folder. + Folders that still contain documents can't be deleted — empty the folder first. + + + {error &&

{error}

} + + + + + + + ); +} + +export function FolderNameDialog({ + mode, + folder, + open, + onOpenChange, + onCreated, + onSaved, +}: { + mode: "create" | "edit"; + folder: KbFolder | null; + open: boolean; + onOpenChange: (open: boolean) => void; + onCreated?: (folder: KbFolder) => void; + onSaved?: (folder: KbFolder) => void; +}) { + const [name, setName] = useState(folder?.name ?? ""); + const [error, setError] = useState(null); + const [submitting, setSubmitting] = useState(false); + + useEffect(() => { + if (open) { + setName(folder?.name ?? ""); + setError(null); + } + }, [open, folder?.name, folder?.id]); + + const submit = useCallback(async () => { + const trimmed = name.trim(); + if (!trimmed) { + setError("Name is required"); + return; + } + setSubmitting(true); + setError(null); + try { + if (mode === "create") { + const res = await fetch("/api/kb/folders", { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify({ name: trimmed }), + }); + if (res.status === 201) { + const body = (await res.json()) as { folder: KbFolder }; + onCreated?.(body.folder); + setName(""); + onOpenChange(false); + return; + } + if (res.status === 409) { + setError("A folder with this name already exists"); + return; + } + setError(`Failed (${res.status})`); + } else { + if (!folder) { + setError("No folder to edit"); + return; + } + const res = await fetch(`/api/kb/folders/${folder.id}`, { + method: "PATCH", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify({ name: trimmed }), + }); + if (res.status === 200) { + const body = (await res.json()) as { folder: KbFolder }; + onSaved?.(body.folder); + onOpenChange(false); + return; + } + if (res.status === 409) { + setError("A folder with this name already exists"); + return; + } + setError(`Failed (${res.status})`); + } + } catch (e) { + setError(e instanceof Error ? e.message : String(e)); + } finally { + setSubmitting(false); + } + }, [mode, folder, name, onCreated, onSaved, onOpenChange]); + + return ( + { + onOpenChange(o); + if (!o) setError(null); + }} + > + + + {mode === "create" ? "New folder" : "Edit folder"} + + {mode === "create" + ? "Group your knowledge base documents by topic or project." + : "Rename this folder. Documents inside keep their content."} + + +
+ + setName(e.target.value)} + placeholder="e.g. Project Research" + maxLength={64} + autoFocus + onKeyDown={(e) => { + if (e.key === "Enter" && !submitting) void submit(); + }} + /> + {error &&

{error}

} +
+ + + + + + + ); +} diff --git a/components/settings/kb-view/doc-detail-dialog.tsx b/components/settings/kb-view/doc-detail-dialog.tsx new file mode 100644 index 00000000..cfb90eac --- /dev/null +++ b/components/settings/kb-view/doc-detail-dialog.tsx @@ -0,0 +1,862 @@ +import { useCallback, useEffect, useMemo, useState } from "react"; +import { + AlertCircle, + ArrowDown, + ArrowRight, + Blocks, + Check, + Copy, + ExternalLink, + FileImage, + FileText, + Loader2, + Network, + ScanText, +} from "lucide-react"; +import { Badge } from "@/components/ui/badge"; +import { Button } from "@/components/ui/button"; +import { + Dialog, + DialogContent, + DialogDescription, + DialogHeader, + DialogTitle, +} from "@/components/ui/dialog"; +import { Separator } from "@/components/ui/separator"; +import { Skeleton } from "@/components/ui/skeleton"; +import { cn } from "@/lib/utils"; +import { KB_POLL_INTERVAL_MS } from "@/lib/constants"; +import { KnowledgeGraph } from "./knowledge-graph"; +import { ChunkStatusBadge, DocStatusBadge, ChunksStatusBadge } from "./status-badge"; +import { entityColor, type EntityColor } from "@/lib/kb/entityColor"; +import { KbDocDetail } from "./types"; + +export function DocDetailDialog({ + docId, + open, + onOpenChange, +}: { + docId: string; + open: boolean; + onOpenChange: (open: boolean) => void; +}) { + const [detail, setDetail] = useState(null); + const [loading, setLoading] = useState(false); + const [activeTab, setActiveTab] = useState<"full_markdown" | "pages" | "chunks" | "graph">( + "full_markdown", + ); + const [copied, setCopied] = useState(false); + + useEffect(() => { + if (!open) return; + const controller = new AbortController(); + setLoading(true); + void fetch(`/api/kb/documents/${docId}`, { signal: controller.signal }) + .then((r) => (r.ok ? r.json() : Promise.reject(new Error(`${r.status}`)))) + .then((body: KbDocDetail) => { + if (!controller.signal.aborted) setDetail(body); + }) + .catch((err) => { + if (err?.name !== "AbortError") setDetail(null); + }) + .finally(() => { + if (!controller.signal.aborted) setLoading(false); + }); + return () => controller.abort(); + }, [open, docId]); + + useEffect(() => { + if (!open) return; + if (!detail) return; + const docInflight = detail.doc.status === "pending" || detail.doc.status === "parsing"; + const chunksInflight = detail.chunks.some( + (c) => c.status === "pending" || c.status === "parsing", + ); + if (!docInflight && !chunksInflight && detail.chunks.length > 0) return; + const controller = new AbortController(); + const t = setInterval(() => { + void fetch(`/api/kb/documents/${docId}`, { signal: controller.signal }) + .then((r) => (r.ok ? r.json() : null)) + .then((body: KbDocDetail | null) => { + if (!controller.signal.aborted && body) setDetail(body); + }) + .catch((err) => { + if (err?.name !== "AbortError") { + console.error("DocDetailDialog poll failed", err); + } + }); + }, KB_POLL_INTERVAL_MS); + return () => { + controller.abort(); + clearInterval(t); + }; + }, [open, docId, detail]); + + const handleCopy = useCallback((text: string) => { + void navigator.clipboard.writeText(text); + setCopied(true); + setTimeout(() => setCopied(false), 2000); + }, []); + + const fullMarkdown = useMemo(() => { + if (!detail) return ""; + if (detail.doc.pages && detail.doc.pages.length > 0) { + return detail.doc.pages + .map((p: any) => p.markdown) + .filter((m) => typeof m === "string" && m.length > 0) + .join("\n\n"); + } + return detail.chunks.map((c) => c.content).join("\n\n"); + }, [detail]); + + // ponytail: graphRAG-native entity color map for chunk-level + // entity badges. Same single-source-of-truth pattern as + // knowledge-graph.tsx — collect degree + neighbor signature from + // every chunk's relationships, then look up per entity. + // chunk-level entities reference the SAME entity across chunks + // (after alignment), so colors stay stable as the user scrolls. + const entityColorMap = useMemo(() => { + const out = new Map(); + if (!detail) return out; + const degreeByName = new Map(); + const neighborsByName = new Map>(); + for (const c of detail.chunks) { + for (const r of c.relationships ?? []) { + const src = r.source; + const tgt = r.target; + if (!src || !tgt) continue; + degreeByName.set(src, (degreeByName.get(src) ?? 0) + 1); + degreeByName.set(tgt, (degreeByName.get(tgt) ?? 0) + 1); + if (!neighborsByName.has(src)) neighborsByName.set(src, new Set()); + if (!neighborsByName.has(tgt)) neighborsByName.set(tgt, new Set()); + neighborsByName.get(src)!.add(tgt); + neighborsByName.get(tgt)!.add(src); + } + } + // ponytail: only color entities that actually appear in the + // current chunks. An entity not in neighborsByName (no outgoing + // or incoming relationship) gets the muted-slate fallback. + const seen = new Set(); + for (const c of detail.chunks) { + for (const e of c.entities ?? []) { + if (!e?.name || seen.has(e.name)) continue; + seen.add(e.name); + const degree = degreeByName.get(e.name) ?? 0; + out.set(e.name, entityColor(e.name, [...(neighborsByName.get(e.name) ?? [])], degree)); + } + } + return out; + }, [detail]); + + // ponytail: while the doc detail is loading, skip the entire + // ready-state render and emit a single dedicated skeleton dialog. + // Keeps the main render path shallow — one ternary per concept + // (detail / loading / error) instead of fighting the tabs row for + // space. + if (loading) { + return ( + { + onOpenChange(o); + if (!o) { + setDetail(null); + setLoading(false); + } + }} + > + + + {/* ponytail: skeleton DialogHeader mirrors the loaded + shape — title on the left, the Radix-managed X + close button stays in the top-right corner + (no skeleton needed), then the status-pill + 3 + meta-pill strip below. */} + + + + + {/* ponytail: meta strip mirrors the loaded layout — + status pill, dot separator, then 3 meta-text + placeholders (contentType / pages / chunks). One + separator per meta pair, matching the loaded + DialogDescription structure exactly. */} +
+ + {[0, 1, 2].map((i) => ( +
+ + +
+ ))} +
+ + + + {/* ponytail: 4 tab pill placeholders — same dimensions + (h-7 body, h-9 row) as the loaded Markdown/Pages/ + Chunks/Graph row so the body stays anchored. The + first pill is the active tab (rendered with the + loaded state styling — bg-background/text-foreground/ + shadow-sm) so the highlight isn't lost during loading; + the remaining three are bare skeleton bars + (rounded-md, semi-transparent) so they read as + inactive slots without competing for contrast. */} +
+ {/* active tab (Markdown) — match loaded style exactly. + Fixed width (no flex-1) so on mobile the white pill + doesn't expand to fill the whole row. */} + + + +
+ + {/* ponytail: Markdown body card — mirrors the + rendered markdown block (outer card + Copy pill + a + heading + alternating paragraph lines + a subheading + + another paragraph block). Sized to feel like 1-2 + screens of markdown content. */} +
+
+
+ + +
+
+ + + + + + + + + + + +
+
+
+ + + ); + } + + // ponytail: detail may still be null briefly between close + reopen. + // The loading skeleton path already returned, so this is the only + // possibility — narrow with a local copy so the JSX below reads + // `d.doc.title` instead of `detail?.doc.title` 15× in a row. + const d = detail; + if (!d) return null; + + const totalChunks = d.chunks.length; + const successChunks = d.chunks.filter((c) => c.status === "success").length; + const failedChunks = d.chunks.filter((c) => c.status === "failed").length; + const pendingChunks = d.chunks.filter((c) => c.status === "pending").length; + const parsingChunks = d.chunks.filter((c) => c.status === "parsing").length; + + // ponytail: derive page-level counts the same way chunksStatusBadge + // does. When pages carry an explicit `status` mirror it; legacy rows + // (status absent) fall back to the markdown/errorMessage heuristic. + const pagesTotal = d.doc.pages?.length ?? 0; + const inferPageStatus = ( + p: NonNullable[number], + ): "pending" | "parsing" | "success" | "failed" => { + if ( + p.status === "pending" || + p.status === "parsing" || + p.status === "success" || + p.status === "failed" + ) { + return p.status; + } + if (p.errorMessage) return "failed"; + if ((p.markdown ?? "").trim().length > 0) return "success"; + return "pending"; + }; + const pages = d.doc.pages ?? []; + const successPages = pages.filter((p) => inferPageStatus(p) === "success").length; + const failedPages = pages.filter((p) => inferPageStatus(p) === "failed").length; + const parsingPages = pages.filter((p) => inferPageStatus(p) === "parsing").length; + const pendingPages = pages.filter((p) => inferPageStatus(p) === "pending").length; + + return ( + { + onOpenChange(o); + if (!o) { + setDetail(null); + setActiveTab("full_markdown"); + } + }} + > + + +
+ {d.doc.title} + {d.doc.attachmentUrl && ( + + )} +
+ +
+ + + + + + {d.doc.contentType} + + + {d.doc.pages && + d.doc.pages.length > 0 && + (() => { + const totalPages = d.doc.pages.length; + const isReprocessing = d.doc.status === "pending" || d.doc.status === "parsing"; + const failedPagesCount = isReprocessing + ? 0 + : d.doc.pages.filter((p) => !!p.errorMessage || !(p.markdown ?? "").trim()) + .length; + return ( + <> + + + + {totalPages} pages + {failedPagesCount > 0 && ( + + ({failedPagesCount} failed) + + )} + + + + ); + })()} + + {d.chunks.length > 0 && ( + <> + + + {d.chunks.length} chunks + + + )} +
+ + + +
+ + + + +
+ +
+ {!detail ? ( +

Failed to load.

+ ) : activeTab === "full_markdown" ? ( +
+
+ + Markdown + + {fullMarkdown && ( + + )} +
+
+ {fullMarkdown || ( + No text extracted yet. + )} +
+
+ ) : activeTab === "pages" ? ( + // Tab 2: Pages — left (image + ref) → right (markdown) + d.doc.pages && d.doc.pages.length > 0 ? ( +
+ {d.doc.pages.map((p) => { + const page = p as { + pageIndex: number; + imageUrl: string; + markdown: string; + referenceText?: string; + errorMessage?: string; + status?: "pending" | "parsing" | "success" | "failed"; + }; + const hasRef = !!page.referenceText?.trim(); + return ( +
+ {/* Card Header */} +
+ + Page #{page.pageIndex + 1} + + {(() => { + // ponytail: prefer page.status when set (fresh + // ingest); fall back to errorMessage/markdown + // heuristic for legacy rows. + const s = + page.status ?? + (page.errorMessage + ? "failed" + : (page.markdown ?? "").trim().length > 0 + ? "success" + : "pending"); + const inFlight = d.doc.status === "pending" || d.doc.status === "parsing"; + if (s === "failed" && !inFlight) { + return ( + + Failed + + ); + } + if (s === "success") { + return ( + + Succeeded + + ); + } + if (s === "parsing" && inFlight) { + return ( + + Parsing… + + ); + } + return ( + + Pending + + ); + })()} +
+ + {/* Body: [Image + Ref] → [Markdown] */} +
+ {/* Left 1/3: Image stacked above "+" and Reference Text */} +
+ {/* Page Image */} +
+ + + Page Image + + + {`Page + +
+ + {/* + connector — matches ArrowRight size */} +
+ + + + +
+ + {/* Reference Text */} +
+ + + Reference Text + + {hasRef ? ( +
+ {page.referenceText} +
+ ) : ( +
+ + No text layer +
+ )} +
+
+ + {/* Center arrow (horizontal on desktop, vertical on mobile) */} +
+ +
+
+ +
+ + {/* Right 2/3: Markdown */} +
+ + + Markdown + + {page.markdown ? ( +
+ {page.markdown} +
+ ) : page.errorMessage ? ( +
+ + + Page OCR Failed + + + {page.errorMessage} + +
+ ) : ( +
+ + Pending… +
+ )} +
+
+
+ ); + })} +
+ ) : ( +

+ No page screenshots available for this document (e.g. legacy or non-PDF format). +

+ ) + ) : activeTab === "chunks" ? ( + // Tab 3: Embed Chunks +
+ {d.chunks.length === 0 ? ( + d.doc.status === "pending" || d.doc.status === "parsing" ? ( +
+
+ + Generating chunks… OCR finished, indexing in progress. +
+ + + +
+ ) : ( +

+ {d.doc.status === "success" + ? "Embedding chunks are still being calculated in the background. They will appear here in a few moments." + : d.doc.status === "failed" + ? "Ingestion failed — chunks not produced." + : "Ingestion in progress…"} +

+ ) + ) : ( + <> + {(() => { + const docInflight = d.doc.status === "pending" || d.doc.status === "parsing"; + const chunksInflight = d.chunks.some( + (c) => c.status === "pending" || c.status === "parsing", + ); + const showSpinner = docInflight || chunksInflight; + return ( +
+ + Indexed{" "} + + {d.chunks.filter((c) => c.status === "success").length} + {" "} + / {d.chunks.length} + {showSpinner && ( + + )} + + {d.chunks.some((c) => c.status === "failed") && ( + + {d.chunks.filter((c) => c.status === "failed").length} failed + + )} +
+ ); + })()} + {d.chunks.map((c) => ( +
+ {/* Card Header */} +
+
+ + Chunk #{c.ordinal + 1} + +
+ +
+ {/* Card Body */} +
+

+ {c.content} +

+ {c.status === "failed" && c.errorMessage && ( +

+ {c.errorMessage} +

+ )} + + {/* Metadata Box: Entities + Themes + Relationships */} + {((c.entities && c.entities.length > 0) || + (c.themes && c.themes.length > 0) || + (c.relationships && c.relationships.length > 0)) && ( +
+ {/* Entities in chunk */} + {c.entities && c.entities.length > 0 && ( +
+ + Entities + +
+ {c.entities.map((e, idx) => { + // ponytail: chunk-level entity badge + // color — graphRAG-native + // entityColor (hue = neighbor + // signature, sat/light = degree), + // single source of truth shared + // with knowledge-graph canvas + + // hover tooltip + uniqueEntities + // list. No string-type whitelist. + const entityColorFor = + entityColorMap.get(e.name) ?? entityColor(e.name, [], 0); + return ( + + {e.name} + + ); + })} +
+
+ )} + + {/* Separator if there are both entities and themes/relationships */} + {c.entities && + c.entities.length > 0 && + ((c.themes && c.themes.length > 0) || + (c.relationships && c.relationships.length > 0)) && ( + + )} + + {/* Themes in chunk */} + {c.themes && c.themes.length > 0 && ( +
+ + Themes + +
+ {c.themes.map((t, idx) => ( + + #{t} + + ))} +
+
+ )} + + {/* Separator if there are relationships and themes */} + {c.themes && + c.themes.length > 0 && + c.relationships && + c.relationships.length > 0 && ( + + )} + + {/* Relationships in chunk */} + {c.relationships && c.relationships.length > 0 && ( +
+ + Relationships + +
+ {c.relationships.map((r, idx) => ( +
+ + {r.source} + + + ({r.relation}) + + + {r.target} + + • + {r.description} +
+ ))} +
+
+ )} +
+ )} +
+
+ ))} + + )} +
+ ) : activeTab === "graph" ? ( + + ) : null} +
+ + + ); +} diff --git a/components/settings/kb-view/doc-table.tsx b/components/settings/kb-view/doc-table.tsx new file mode 100644 index 00000000..14f4fcdb --- /dev/null +++ b/components/settings/kb-view/doc-table.tsx @@ -0,0 +1,337 @@ +import { useEffect, useRef, useState } from "react"; +import { ExternalLink, Network, Plus, RefreshCw, Search, Trash2 } from "lucide-react"; +import { Card, CardContent } from "@/components/ui/card"; +import { Separator } from "@/components/ui/separator"; +import { Button } from "@/components/ui/button"; +import { Tooltip, TooltipContent, TooltipTrigger } from "@/components/ui/tooltip"; +import { cn } from "@/lib/utils"; +import { KB_POLL_INTERVAL_MS } from "@/lib/constants"; +import { KbResponse, KbDocument } from "./types"; +import { HeaderBar } from "./folder-sidebar"; +import { DocStatusBadge, ChunksStatusBadge, formatTimestamp } from "./status-badge"; +import { DocDetailDialog } from "./doc-detail-dialog"; +import { DocDeleteDialog, DocReprocessDialog } from "./dialogs"; +import { FolderGraphDialog } from "./folder-graph-dialog"; +import { LivePollIndicator } from "./live-poll-indicator"; +import { ObservabilityPopover } from "./observability-popover"; + +export function DocTable({ + group, + focusDocId, + onAddDoc, + onRefresh, + isLivePolling, +}: { + group: KbResponse["groups"][number] | null; + focusDocId: string | null; + onAddDoc: () => void; + onRefresh: () => Promise | void; + isLivePolling: boolean; +}) { + const [folderGraphOpen, setFolderGraphOpen] = useState(false); + + if (!group) { + return ( + + + Select a folder to view its documents. + + + ); + } + + return ( + + + + + + + + + Folder Graph + + + + + + + Add doc + +
+ } + /> + + {group.documents.length === 0 ? ( +
+ No documents in this folder yet. Click to upload a + PDF. +
+ ) : ( + + )} + + + + ); +} + +export function DocTableRows({ + docs, + focusDocId, + onRefresh, +}: { + docs: KbDocument[]; + focusDocId: string | null; + onRefresh: () => Promise | void; +}) { + return ( +
+ {docs.map((doc, i) => ( +
+ {i > 0 && } + +
+ ))} +
+ ); +} + +export function DocRow({ + doc, + isFocused, + onRefresh, +}: { + doc: KbDocument; + isFocused: boolean; + onRefresh: () => Promise | void; +}) { + const [previewOpen, setPreviewOpen] = useState(false); + const [deleteOpen, setDeleteOpen] = useState(false); + const [reprocessOpen, setReprocessOpen] = useState(false); + const type = doc.contentType.replace("application/", ""); + const rowRef = useRef(null); + + useEffect(() => { + if (!isFocused) return; + const t = setTimeout(() => { + rowRef.current?.scrollIntoView({ block: "center", behavior: "smooth" }); + }, 50); + return () => clearTimeout(t); + }, [isFocused]); + + const isInflight = doc.status === "pending" || doc.status === "parsing"; + const actions = ( + <> + + + + + + {isInflight ? "Already processing" : "Reprocess"} + + + + + + + + Preview + + {doc.attachmentUrl && ( + + + + + Open source + + )} + + + + + Delete + + + ); + const meta = ( +
+
+ {type} + · + +
+
+ + +
+
+ ); + return ( +
+
+
+
+ {doc.title} +
+
{actions}
+
+
{meta}
+
+
+
+ {doc.title} +
+
{type}
+ +
+ + +
+
{actions}
+
+ + {previewOpen && ( + + )} + {deleteOpen && ( + { + setDeleteOpen(false); + void onRefresh(); + }} + /> + )} + {reprocessOpen && ( + { + setReprocessOpen(false); + void onRefresh(); + }} + /> + )} +
+ ); +} diff --git a/components/settings/kb-view/folder-graph-dialog.tsx b/components/settings/kb-view/folder-graph-dialog.tsx new file mode 100644 index 00000000..ca24ca72 --- /dev/null +++ b/components/settings/kb-view/folder-graph-dialog.tsx @@ -0,0 +1,160 @@ +"use client"; + +import { useEffect, useState } from "react"; +import { Badge } from "@/components/ui/badge"; +import { + Dialog, + DialogContent, + DialogDescription, + DialogHeader, + DialogTitle, +} from "@/components/ui/dialog"; +import { Skeleton } from "@/components/ui/skeleton"; +import { KnowledgeGraph, type KnowledgeGraphChunk } from "./knowledge-graph"; + +type FolderDetail = { + folder: { + id: string; + name: string; + createdAt: string; + updatedAt: string; + }; + chunks: (KnowledgeGraphChunk & { errorMessage: string | null })[]; +}; + +// ponytail: thin dialog shell over KnowledgeGraph. Header summarises +// folder + chunk count; body delegates to the shared component with +// `skipFailedChunks` so the cross-folder rollup ignores pending / +// parsing / failed chunks (per-doc view folds them in by contrast). +export function FolderGraphDialog({ + folderId, + folderName, + open, + onOpenChange, +}: { + folderId: string; + folderName: string; + open: boolean; + onOpenChange: (open: boolean) => void; +}) { + const [detail, setDetail] = useState(null); + const [loading, setLoading] = useState(false); + + useEffect(() => { + if (!open) return; + setLoading(true); + const controller = new AbortController(); + + fetch(`/api/kb/folders/${folderId}`, { signal: controller.signal }) + .then((res) => { + if (!res.ok) throw new Error("Failed to load"); + return res.json(); + }) + .then((data) => { + setDetail(data); + setLoading(false); + }) + .catch((err) => { + if (err.name !== "AbortError") { + setLoading(false); + } + }); + + return () => { + controller.abort(); + }; + }, [open, folderId]); + + return ( + { + onOpenChange(o); + if (!o) setDetail(null); + }} + > + + +
+ + Knowledge Graph: {folderName} + +
+ {/* ponytail: DialogDescription renders as a

by default + which can't contain a

. asChild swaps it to a +
so the badge + dot + folder name stack renders + semantically. */} + +
+ + Folder + + • + + {folderName} + + {detail && ( + <> + • + + {detail.chunks.length} Chunks + + + )} +
+ + + +
+ {loading ? ( +
+ {/* ponytail: skeleton mirrors the KnowledgeGraph layout + — 4 tab pills at the top, then a node cloud framed + by the same h-[500px] surface. We don't try to fake + edges; the real graph covers that on mount. */} +
+ {[0, 1, 2, 3].map((i) => ( + + ))} +
+
+
+ {[ + "size-3", + "size-4", + "size-5", + "size-6", + "size-4", + "size-3", + "size-5", + "size-4", + "size-3", + "size-5", + "size-6", + "size-4", + "size-3", + "size-5", + "size-4", + ].map((c, i) => ( + + ))} +
+
+
+ ) : !detail ? ( +

Failed to load.

+ ) : ( + + )} +
+ +
+ ); +} diff --git a/components/settings/kb-view/folder-sidebar.tsx b/components/settings/kb-view/folder-sidebar.tsx new file mode 100644 index 00000000..52effcac --- /dev/null +++ b/components/settings/kb-view/folder-sidebar.tsx @@ -0,0 +1,187 @@ +import { useState } from "react"; +import { Folder, MoreHorizontal, Pencil, Plus, Trash2 } from "lucide-react"; +import { Card, CardContent } from "@/components/ui/card"; +import { Separator } from "@/components/ui/separator"; +import { Button } from "@/components/ui/button"; +import { Tooltip, TooltipContent, TooltipTrigger } from "@/components/ui/tooltip"; +import { + DropdownMenu, + DropdownMenuContent, + DropdownMenuItem, + DropdownMenuTrigger, +} from "@/components/ui/dropdown-menu"; +import { cn } from "@/lib/utils"; +import { KbResponse, KbFolder } from "./types"; +import { FolderDeleteDialog, FolderNameDialog } from "./dialogs"; + +// ponytail: shared row chrome for the Folders sidebar header and the +// doc-table header — same padding (px-4) so the right-edge action +// (folder `+`, doc `+`) lines up with the action buttons in each +// data row. h-9 keeps the heights identical. +export function HeaderBar({ label, action }: { label: string; action: React.ReactNode }) { + return ( +
+ + {label} + + {action} +
+ ); +} + +export function FolderSidebar({ + groups, + selectedId, + onSelect, + onNewFolder, + onRefresh, +}: { + groups: KbResponse["groups"]; + selectedId: string | null; + onSelect: (id: string) => void; + onNewFolder: () => void; + onRefresh: () => Promise | void; +}) { + const [deleteTarget, setDeleteTarget] = useState(null); + const [editTarget, setEditTarget] = useState(null); + const [openFolderId, setOpenFolderId] = useState(null); + + // ponytail: mirror folder selection onto ?folder= so refreshing + // the page (or sharing the URL) lands back on the same folder. + // replaceState (not pushState) — folder switching is a sub-state of + // the page, not a history step worth a back-button entry. + const handleSelect = (id: string) => { + onSelect(id); + if (typeof window !== "undefined") { + const url = new URL(window.location.href); + if (url.searchParams.get("folder") === id) return; + url.searchParams.set("folder", id); + window.history.replaceState(window.history.state, "", url.toString()); + } + }; + + return ( + <> + + + + + + + New folder + + } + /> + +
    + {groups.length === 0 ? ( +
    + + No folders yet + +
    + ) : ( + groups.map((g) => { + const active = g.folder.id === selectedId; + const menuOpen = openFolderId === g.folder.id; + return ( +
  • + + setOpenFolderId(o ? g.folder.id : null)} + > + + + + + setEditTarget(g.folder)} + className="hover:bg-muted focus:bg-muted flex cursor-pointer items-center gap-2 rounded-md px-2 py-1.5 text-sm outline-none select-none" + > + + Edit + + setDeleteTarget(g.folder)} + className="text-destructive hover:bg-destructive/10 hover:text-destructive focus:bg-destructive/10 focus:text-destructive flex cursor-pointer items-center gap-2 rounded-md px-2 py-1.5 text-sm outline-none select-none" + > + + Delete + + + +
    • + ); + }) + )} +
+ + + + !o && setDeleteTarget(null)} + onDeleted={() => { + setDeleteTarget(null); + void onRefresh(); + }} + /> + !o && setEditTarget(null)} + onSaved={() => { + setEditTarget(null); + void onRefresh(); + }} + /> + + ); +} diff --git a/components/settings/kb-view/helpers.ts b/components/settings/kb-view/helpers.ts new file mode 100644 index 00000000..0c6c5ba2 --- /dev/null +++ b/components/settings/kb-view/helpers.ts @@ -0,0 +1,114 @@ +import { toast } from "sonner"; + +export const TOAST_DESCRIPTION_CLASS = "!text-foreground"; + +// ponytail: shared reprocess mode vocabulary. The reprocess dialog +// shows full {title, description} per option; the observability popover +// shows just title alongside the source label. Keeping a single source +// of truth means the two surfaces can never disagree on what a mode +// means. +export type ReprocessMode = "full" | "chunksOnly" | "retryFailed" | "retryFailedChunks"; + +export type ModeInfo = { title: string; description: string }; + +const MODE_INFO: Record = { + full: { + title: "Full run", + description: "Re-render the PDF, re-run OCR, then re-chunk and re-embed.", + }, + chunksOnly: { + title: "Chunks only", + description: "Skip OCR — reuse the cached pages markdown to rebuild chunks + entities.", + }, + retryFailed: { + title: "Retry failed OCR", + description: "Re-OCR failed pages only, keep successful pages, then rebuild chunks.", + }, + retryFailedChunks: { + title: "Retry failed chunks", + description: "Keep successful chunks, only re-embed + re-extract entities for failed ones.", + }, +}; + +export function getModeInfo(mode: ReprocessMode): ModeInfo { + return MODE_INFO[mode]; +} + +export async function sha256Hex(file: File): Promise { + try { + const buf = await file.arrayBuffer(); + const digest = await crypto.subtle.digest("SHA-256", buf); + return Array.from(new Uint8Array(digest)) + .map((b) => b.toString(16).padStart(2, "0")) + .join(""); + } catch { + return undefined; + } +} + +export async function handleAddDoc( + file: File, + folderId: string, + onRefresh: () => Promise | void, +) { + try { + const sha = await sha256Hex(file); + const presignRes = await fetch("/api/attachments/presign", { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify({ + name: file.name, + contentType: file.type, + sizeBytes: file.size, + ...(sha ? { sha256: sha } : {}), + }), + }); + if (!presignRes.ok) throw new Error(`presign failed: ${presignRes.status}`); + const presign = (await presignRes.json()) as { + id: string; + uploadUrl: string; + uploadHeaders: Record; + publicUrl: string; + skipUpload?: boolean; + }; + + if (!presign.skipUpload) { + const putRes = await fetch(presign.uploadUrl, { + method: "PUT", + headers: presign.uploadHeaders, + body: file, + }); + if (!putRes.ok) throw new Error(`upload failed: ${putRes.status}`); + + const confirmRes = await fetch(`/api/attachments/${presign.id}/confirm`, { method: "POST" }); + if (!confirmRes.ok) throw new Error(`confirm failed: ${confirmRes.status}`); + } + + const uploadRes = await fetch("/api/kb/upload", { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify({ folderId, attachmentId: presign.id, title: file.name }), + }); + if (!uploadRes.ok && uploadRes.status !== 202) { + throw new Error(`kb upload failed: ${uploadRes.status}`); + } + if (uploadRes.status === 200) { + const body = (await uploadRes.json()) as { deduped?: boolean; doc?: { title?: string } }; + if (body.deduped) { + toast.info("Already in knowledge base", { + descriptionClassName: TOAST_DESCRIPTION_CLASS, + description: `「${body.doc?.title ?? file.name}」was previously uploaded — skipped duplicate.`, + }); + } + } else if (uploadRes.status === 202) { + const body = (await uploadRes.json()) as { doc?: { title?: string } }; + toast.success("Upload queued", { + descriptionClassName: TOAST_DESCRIPTION_CLASS, + description: `「${body.doc?.title ?? file.name}」is being ingested. Status will flip Pending → Parsing → Ready.`, + }); + } + void onRefresh(); + } catch (err) { + console.error("Add Doc failed", err); + } +} diff --git a/components/settings/kb-view/index.tsx b/components/settings/kb-view/index.tsx new file mode 100644 index 00000000..92abf002 --- /dev/null +++ b/components/settings/kb-view/index.tsx @@ -0,0 +1 @@ +export { KbView } from "./kb-view"; diff --git a/components/settings/kb-view/kb-view.tsx b/components/settings/kb-view/kb-view.tsx new file mode 100644 index 00000000..ef6c07aa --- /dev/null +++ b/components/settings/kb-view/kb-view.tsx @@ -0,0 +1,285 @@ +"use client"; + +import { Suspense, useCallback, useEffect, useMemo, useRef, useState } from "react"; +import { useSearchParams } from "next/navigation"; +import { Card, CardContent } from "@/components/ui/card"; +import { Separator } from "@/components/ui/separator"; +import { Skeleton } from "@/components/ui/skeleton"; +import { TooltipProvider } from "@/components/ui/tooltip"; +import { cn } from "@/lib/utils"; +import { KB_POLL_INTERVAL_MS } from "@/lib/constants"; +import { ObservabilitySheet } from "@/components/observability/sheet"; +import { ObservabilitySheetProvider } from "@/components/observability/sheet-context"; +import { KbResponse } from "./types"; +import { FolderSidebar } from "./folder-sidebar"; +import { DocTable } from "./doc-table"; +import { FolderNameDialog } from "./dialogs"; +import { handleAddDoc } from "./helpers"; + +export function KbView({ className }: { className?: string }) { + // ponytail: KB doc rows open the singleton ObservabilitySheet via the + // same context the chat thread uses. Provider + sheet are mounted at + // the KbView root so any DocRow (or future descendants) can call + // useOpenObservabilitySheet() without each subtree wiring its own. + return ( + + }> + + + + + ); +} + +function KbViewContent({ className }: { className?: string }) { + const searchParams = useSearchParams(); + const focusDocId = searchParams.get("doc"); + const initialFolderId = searchParams.get("folder"); + const [data, setData] = useState(null); + const [error, setError] = useState(null); + const [selectedFolderId, setSelectedFolderId] = useState(initialFolderId); + const [newFolderOpen, setNewFolderOpen] = useState(false); + const [isLivePolling, setIsLivePolling] = useState(false); + const fileInputRef = useRef(null); + + const load = useCallback(async () => { + try { + // ponytail: scope the payload to the currently-selected folder — + // the sidebar still gets the full folder list, but only the + // selected folder's documents are populated (other folders + // return `documents: []`). Cuts the JOIN cost on the KB-doc + // list query, and lets `anyInflight` stay scoped to the + // folder the user is actually looking at. + const qs = selectedFolderId ? `?folderId=${encodeURIComponent(selectedFolderId)}` : ""; + const res = await fetch(`/api/kb/documents${qs}`); + if (!res.ok) { + setError(`failed to load (${res.status})`); + return; + } + const body = (await res.json()) as KbResponse; + setData(body); + setSelectedFolderId((prev) => { + if (prev && body.groups.some((g) => g.folder.id === prev)) return prev; + // ponytail: ?folder= from the URL outranks the + // doc-derived heuristic on cold-load — refreshing a deep + // link lands on the requested folder, not on the folder + // owning the focus doc. Falls through to focusDocId only + // when the URL didn't pin a folder. + if (initialFolderId && body.groups.some((g) => g.folder.id === initialFolderId)) { + return initialFolderId; + } + if (focusDocId) { + const owning = body.groups.find((g) => g.documents.some((d) => d.id === focusDocId)); + if (owning) return owning.folder.id; + } + const firstWithDocs = body.groups.find((g) => g.documents.length > 0); + return firstWithDocs?.folder.id ?? body.groups[0]?.folder.id ?? null; + }); + } catch (e) { + setError(e instanceof Error ? e.message : String(e)); + } + }, [focusDocId, initialFolderId, selectedFolderId]); + + // ponytail: after any "reprocess" / "delete" / "upload" action we + // brute-force polling for a short window (~12s — covers the worst + // case of "kbAgent gets dispatched → wipe commits → fire-and-forget + // chunk INSERT lands 1-5s later"). Without this window the table's + // `anyInflight` heuristic can race: the wipe returns success/failed + // with totalChunks=0 BEFORE kbAgent has had a chance to INSERT the + // new chunks, so 0 > 0+0 is false and the polling timer stops — + // table then stays on "No Chunks" forever (until the user clicks + // Preview or refresh). + const recentlyDispatchedUntilRef = useRef(0); + const markRecentlyDispatched = useCallback(() => { + recentlyDispatchedUntilRef.current = Date.now() + 12_000; + }, []); + // wraps `load` so any caller using it as `onRefresh` automatically + // primes the post-dispatch polling window. DocDetailDialog / + // DocTable / FolderSidebar all keep their existing onRefresh wiring. + const loadWithHeartbeat = useCallback(async () => { + markRecentlyDispatched(); + await load(); + }, [load, markRecentlyDispatched]); + + useEffect(() => { + void load(); + // ponytail: re-fetch when the user switches folders. The API + // payload is scoped via `?folderId=`, so the doc table on + // the right side lands on the new folder's data without us + // having to do any client-side filtering. + }, [load]); + + useEffect(() => { + if (!data) return; + // ponytail: keep polling while EITHER the doc-row status is in + // flight OR any chunk is still pending/parsing inside an + // otherwise-success doc. OCR finalises (status flips to "success") + // well before chunks finish their embedding + entity-extract pass, + // so the original "doc.status === pending|parsing" check froze + // the badge on "Indexing" until the user opened the detail dialog + // (which has its own /api/kb/documents/[id] polling) and missed + // the table-level refresh entirely. + const anyInflight = data.groups.some((g) => + g.documents.some( + (d) => + d.status === "pending" || + d.status === "parsing" || + (d.totalChunks ?? 0) > (d.successChunks ?? 0) + (d.failedChunks ?? 0) || + ((d.totalChunks ?? 0) == 0 && (d.totalPages ?? 0) > 0), + ), + ); + + // brute-force window after a Reprocess/Upload/Delete dispatch so + // the wipe→INSERT race doesn't strand the table on stale counts. + const inDispatchWindow = Date.now() < recentlyDispatchedUntilRef.current; + if (!anyInflight && !inDispatchWindow) { + setIsLivePolling(false); + return; + } + setIsLivePolling(true); + // ponytail: if polling is sustained only by the post-dispatch + // window (no in-flight docs to drive it), schedule a state flip + // so the live indicator turns off when the window expires. + if (!anyInflight && inDispatchWindow) { + const ms = recentlyDispatchedUntilRef.current - Date.now(); + const t = setTimeout(() => setIsLivePolling(false), ms); + return () => clearTimeout(t); + } + const t = setInterval(() => void load(), KB_POLL_INTERVAL_MS); + return () => clearInterval(t); + }, [data, load]); + + const selectedGroup = useMemo( + () => data?.groups.find((g) => g.folder.id === selectedFolderId) ?? null, + [data, selectedFolderId], + ); + + if (error) { + return ( +
+ {error} +
+ ); + } + + if (!data) { + return ; + } + + return ( + +
+
+

Knowledge base

+

+ PDFs and other attachments you upload in chat land here as searchable documents. Drop a + file into the composer and the assistant will use it to ground its replies. +

+
+ +
+ setNewFolderOpen(true)} + onRefresh={loadWithHeartbeat} + /> + fileInputRef.current?.click()} + onRefresh={loadWithHeartbeat} + isLivePolling={isLivePolling} + /> +
+ + { + const file = e.target.files?.[0]; + if (file && selectedGroup) { + void handleAddDoc(file, selectedGroup.folder.id, loadWithHeartbeat); + } + e.target.value = ""; // allow re-pick same file + }} + /> + + { + setSelectedFolderId(folder.id); + void loadWithHeartbeat(); + }} + /> +
+ + ); +} + +function KbViewSkeleton({ className }: { className?: string }) { + return ( +
+
+ + +
+
+ + +
+ + +
+ +
+ {[0, 1].map((i) => ( + + ))} +
+ + + + +
+ + +
+ + {[0, 1, 2].map((i) => ( +
+ {i > 0 && } +
+
+ + + + +
+ +
+
+ + + + +
+ + + +
+
+
+ ))} + + +
+
+ ); +} diff --git a/components/settings/kb-view/knowledge-graph.tsx b/components/settings/kb-view/knowledge-graph.tsx new file mode 100644 index 00000000..81e1ef1f --- /dev/null +++ b/components/settings/kb-view/knowledge-graph.tsx @@ -0,0 +1,818 @@ +"use client"; + +import { useCallback, useEffect, useMemo, useRef, useState } from "react"; +import dynamic from "next/dynamic"; +import { ArrowRight, Eye, Hash, Link2, Tags } from "lucide-react"; +import { Badge } from "@/components/ui/badge"; +import { cn } from "@/lib/utils"; +import { entityColor, type EntityColor } from "@/lib/kb/entityColor"; + +// ponytail: shared KnowledgeGraph card used by both the per-document +// preview dialog (doc-detail-dialog) and the per-folder graph dialog +// (folder-graph-dialog). Caller hands us a chunk list and decides +// whether non-success chunks should be folded into the dedup pass +// (folder-graph view skips them; per-doc view historically kept +// them — same wire shape either way so the children can render +// anything the dedup drops). + +const ForceGraph2D = dynamic(() => import("react-force-graph-2d"), { + ssr: false, + loading: () => ( +
+ Loading visual map... +
+ ), +}); + +type KgEntity = { name: string; type: string; description: string }; +type KgRelationship = { + source: string; + target: string; + relation: string; + description: string; +}; + +export type KnowledgeGraphChunk = { + entities: KgEntity[]; + relationships: KgRelationship[]; + themes: string[]; + status: "pending" | "parsing" | "success" | "failed"; +}; + +export type KnowledgeGraphProps = { + chunks: KnowledgeGraphChunk[]; + /** + * ponytail: dedup entities by `name::type` (case-folded) so two + * chunks tagging "Acme" with different Types both surface on the + * graph. Relationship key is `source::target::relation`. Themes + * dedup by exact string. Only chunks with status='success' land + * in the rollup when this is on. + */ + skipFailedChunks?: boolean; + /** + * Copy shown when chunks.length === 0. Defaults to "Upload a + * document into this folder to extract entities and relationships." + */ + emptyMessage?: string; +}; + +export function KnowledgeGraph({ + chunks, + skipFailedChunks = false, + emptyMessage = "Upload a document into this folder to extract entities and relationships.", +}: KnowledgeGraphProps) { + const [graphView, setGraphView] = useState<"visual" | "themes" | "entities" | "relationships">( + "visual", + ); + + const containerRef = useRef(null); + const [dimensions, setDimensions] = useState({ width: 600, height: 500 }); + const [hoveredNode, setHoveredNode] = useState(null); + const [hoveredLink, setHoveredLink] = useState(null); + const [tooltipPos, setTooltipPos] = useState({ x: 0, y: 0 }); + + const handleMouseMove = useCallback((e: React.MouseEvent) => { + if (!containerRef.current) return; + const rect = containerRef.current.getBoundingClientRect(); + setTooltipPos({ + x: e.clientX - rect.left, + y: e.clientY - rect.top, + }); + }, []); + + const handleMouseLeave = useCallback(() => { + setHoveredNode(null); + setHoveredLink(null); + }, []); + + // ponytail: stable callbacks for ForceGraph2D so the lib doesn't see + // a new function identity every poll and re-stage its render loop. + // nodeVal / nodeLabel / linkLabel are stateless — empty fns are fine. + // ponytail: nodeVal formula mirrors nodeCanvasObject's r formula + // (2.5 + degree*0.6, max 8) so react-force-graph-2d's internal + // collision / label-clearance uses the SAME circle we paint. + // nodeRelSize is the per-unit diameter, so nodeRelSize * sqrt(nodeVal) + // = the visible radius on screen. + const nodeVal = useCallback((node: any) => { + const degree = node.degree || 1; + const r = Math.min(8, Math.max(2.5, 2.5 + degree * 0.6)); + return Math.pow(r / 6, 2); + }, []); + const emptyLabel = useCallback(() => "", []); + // ponytail: thinner, more transparent links + arrows after the + // 321-entity folder view flooded the canvas with overlapping + // geometry. 0.2-0.6px line, 0.1-0.32 alpha — halve the second + // pass too. Structural meaning (hub links slightly thicker) is + // still there but links no longer dominate the canvas. + const linkWidth = useCallback((link: any) => { + const degree = Math.max(link.source?.degree ?? 1, link.target?.degree ?? 1); + return Math.min(0.6, 0.2 + Math.log2(1 + degree) * 0.1); + }, []); + const linkColor = useCallback((link: any) => { + const degree = Math.max(link.source?.degree ?? 1, link.target?.degree ?? 1); + const alpha = Math.min(0.32, 0.1 + Math.log2(1 + degree) * 0.06); + return `rgba(100, 116, 139, ${alpha.toFixed(3)})`; + }, []); + // ponytail: matches nodeCanvasObject's r formula so the link's + // endpoint offset lands on the visible circle, not the internal + // collision radius that react-force-graph-2d uses for layout. + const nodeRadius = useCallback((node: any) => { + if (!node) return 4; + const degree = node.degree || 1; + return Math.min(8, Math.max(2.5, 2.5 + degree * 0.6)); + }, []); + // ponytail: custom link render so the directional arrow lands on + // the target node's circle edge instead of its center. Default + // react-force-graph-2d arrow drawing uses the raw line endpoints + // (node.x/y), so the arrowhead is buried under the node. We + // shorten both ends by each node's radius along the link + // direction, then draw line + filled triangle arrow. + const linkCanvasObject = useCallback( + (link: any, ctx: CanvasRenderingContext2D) => { + const s = link.source as { x: number; y: number; degree?: number }; + const t = link.target as { x: number; y: number; degree?: number }; + if ( + typeof s?.x !== "number" || + typeof s?.y !== "number" || + typeof t?.x !== "number" || + typeof t?.y !== "number" + ) { + return; + } + const dx = t.x - s.x; + const dy = t.y - s.y; + const dist = Math.hypot(dx, dy); + if (dist < 0.001) return; + const ux = dx / dist; + const uy = dy / dist; + const sr = nodeRadius(s); + const tr = nodeRadius(t); + // Shorten by source radius on the source end, by (target + // radius + arrow length) on the target end so the arrowhead + // tip lands exactly on the target circle edge. Arrow shrunk + // from 5px / 2.5 half-width to 3.5 / 1.8 to match the new + // thinner lines (was 5x5 triangle on a 0.4-1.2px line — + // looked heavier than the line itself). + const arrowLen = 2.5; + const halfWidth = 1.2; + const sx = s.x + ux * sr; + const sy = s.y + uy * sr; + const tx = t.x - ux * (tr + arrowLen); + const ty = t.y - uy * (tr + arrowLen); + + // Line: matches linkWidth / linkColor formulas above (halved + // the second pass: 0.2-0.6px line, 0.1-0.32 alpha). + const degree = Math.max(s.degree ?? 1, t.degree ?? 1); + const alpha = Math.min(0.32, 0.1 + Math.log2(1 + degree) * 0.06); + ctx.strokeStyle = `rgba(100, 116, 139, ${alpha.toFixed(3)})`; + ctx.lineWidth = Math.min(0.6, 0.2 + Math.log2(1 + degree) * 0.1); + ctx.beginPath(); + ctx.moveTo(sx, sy); + ctx.lineTo(tx, ty); + ctx.stroke(); + + // Filled triangle arrowhead pointing along the link direction + // at the (tx, ty) tip. + const nx = -uy; + const ny = ux; + ctx.fillStyle = ctx.strokeStyle; + ctx.beginPath(); + ctx.moveTo(tx + ux * arrowLen, ty + uy * arrowLen); + ctx.lineTo(tx + nx * halfWidth, ty + ny * halfWidth); + ctx.lineTo(tx - nx * halfWidth, ty - ny * halfWidth); + ctx.closePath(); + ctx.fill(); + }, + [nodeRadius], + ); + const onNodeHover = useCallback((node: any) => { + setHoveredNode(node || null); + if (node) setHoveredLink(null); + }, []); + const onLinkHover = useCallback((link: any) => { + setHoveredLink(link || null); + if (link) setHoveredNode(null); + }, []); + + // ponytail: react-force-graph-2d instance ref so we can call + // zoomToFit() once after first layout, animating into the cluster + // instead of forcing the user to scroll-zoom in. We re-fit whenever + // graph data shape materially changes (entity / relationship count + // jumps by > 5) so a re-render lands on a useful framing. + const graphRef = useRef(null); + const lastFittedKey = useRef(""); + + useEffect(() => { + if (graphView !== "visual" || !containerRef.current) return; + const resizeObserver = new ResizeObserver((entries) => { + for (const entry of entries) { + const { width } = entry.contentRect; + setDimensions({ width: width || 600, height: 500 }); + } + }); + resizeObserver.observe(containerRef.current); + return () => resizeObserver.disconnect(); + }, [graphView]); + + // ponytail: graphRAG-native node color map. Built once per graphData + // shape. Each entry: name → {h,s,l,bg,fg,border}. Consumers: + // nodeCanvasObject (canvas fill), hover tooltip badge, uniqueEntities + // list badge, doc-detail-dialog entity badge. All read from this + // single source. Neighbor signature drives hue (catches "two + // entities in the same neighborhood"); degree drives saturation + + // lightness (catches hubs vs leaves). Replaces the person / + // organization / concept string whitelist that left ~40 LLM-extracted + // types grey. + // ponytail: doc-detail-dialog polls /api/kb/documents/[id] every 2s. + // Each poll lands a fresh `chunks` array reference, which would force a + // re-dedup + new graphData object, which would re-trigger + // react-force-graph-2d's D3 simulation from scratch (visual flicker + + // layout thrash). We pin the dedup OUTPUT to a fingerprint: if the + // rolled-up (entity name + description) and (rel triple + description) + // bag is identical across polls, we keep the previous graphData + // reference. The fingerprint is a sorted, name-folded string concat — + // cheap (≤ a few KB for thousands of entities) and stable across polls + // when content hasn't changed (e.g. one chunk flips parsing → success). + const dedupFingerprint = useMemo(() => { + const parts: string[] = []; + for (const c of chunks ?? []) { + if (skipFailedChunks && c.status !== "success") continue; + for (const e of c.entities ?? []) { + parts.push(`E|${e.name.toLowerCase()}|${e.type}|${e.description}`); + } + for (const r of c.relationships ?? []) { + parts.push( + `R|${r.source.toLowerCase()}|${r.target.toLowerCase()}|${r.relation.toLowerCase()}|${r.description}`, + ); + } + for (const t of c.themes ?? []) { + parts.push(`T|${t}`); + } + } + parts.sort(); + return parts.join("\n"); + }, [chunks, skipFailedChunks]); + + const { uniqueThemes, uniqueEntities, uniqueRelationships, graphData } = useMemo(() => { + if (!chunks || chunks.length === 0) { + return { + uniqueThemes: [], + uniqueEntities: [], + uniqueRelationships: [], + graphData: { nodes: [], links: [] }, + }; + } + + const themesSet = new Set(); + // ponytail: entity dedup is name-only (case-folded). Whether LLM + // tagged "Acme" as `Organization` or `Tool`, we collapse to one + // node — keeps the graph dense for per-doc preview where 17 + // chunks mentioning the same person shouldn't render as 17 islands. + // Longest description wins for that bucket. + const entityMap = new Map(); + const relMap = new Map(); + + for (const c of chunks) { + if (skipFailedChunks && c.status !== "success") continue; + for (const t of c.themes ?? []) { + themesSet.add(t); + } + for (const e of c.entities ?? []) { + const key = e.name.toLowerCase(); + const existing = entityMap.get(key); + if (!existing || e.description.length > existing.description.length) { + entityMap.set(key, e); + } + } + for (const r of c.relationships ?? []) { + // ponytail: relationships fold on (source, target, relation). + // Direction matters — A->B and B->A stay separate rows. + const key = `${r.source.toLowerCase()}::${r.target.toLowerCase()}::${r.relation.toLowerCase()}`; + const existing = relMap.get(key); + if (!existing || r.description.length > existing.description.length) { + relMap.set(key, r); + } + } + } + + const uniqueThemes = Array.from(themesSet).sort(); + const uniqueEntities = Array.from(entityMap.values()).sort((a, b) => + a.name.localeCompare(b.name), + ); + const uniqueRelationships = Array.from(relMap.values()).sort((a, b) => + a.source.localeCompare(b.source), + ); + + // ponytail: degree for each entity — incoming or outgoing edge + // count under case-folded names so a rel like "Acme uses Beta" + // contributes to `Acme.degree` and `Beta.degree` symmetrically. + const degrees = new Map(); + for (const r of uniqueRelationships) { + const src = r.source.toLowerCase(); + const tgt = r.target.toLowerCase(); + degrees.set(src, (degrees.get(src) || 0) + 1); + degrees.set(tgt, (degrees.get(tgt) || 0) + 1); + } + + const nodes = uniqueEntities.map((e) => ({ + id: e.name.toLowerCase(), + name: e.name, + type: e.type, + description: e.description, + degree: degrees.get(e.name.toLowerCase()) || 0, + })); + + const links = uniqueRelationships.map((r) => ({ + source: r.source.toLowerCase(), + target: r.target.toLowerCase(), + relation: r.relation, + description: r.description, + })); + + // Filter out links pointing to non-existent nodes to prevent D3-force from crashing + const validNodeIds = new Set(nodes.map((n) => n.id)); + const filteredLinks = links.filter( + (l) => validNodeIds.has(l.source) && validNodeIds.has(l.target), + ); + + return { + uniqueThemes, + uniqueEntities, + uniqueRelationships, + graphData: { nodes, links: filteredLinks }, + }; + // ponytail: depend on the fingerprint (stable across same-content polls), + // not the raw `chunks` array (which flips reference every 2s poll). + }, [dedupFingerprint, chunks?.length, skipFailedChunks]); + + // ponytail: graphRAG-native node color map. Built once per graphData + // shape (graphData refs are pinned by the dedup fingerprint above so + // polls with unchanged content don't churn this). Each entry: + // name → {h,s,l,bg,fg,border}. Consumers: nodeCanvasObject (canvas + // fill), hover tooltip badge, uniqueEntities list badge, and + // doc-detail-dialog entity badge (all import entityColor from + // lib/kb/entityColor — single source of truth). + // ponytail: same neighbor signature = same hue (catches "two + // entities in the same neighborhood"); degree drives saturation + + // lightness (catches hubs vs leaves). Replaces the person / + // organization / concept string whitelist that left ~40 LLM- + // extracted types grey. + const nodeColors = useMemo(() => { + const degreeByName = new Map(); + const neighborsByName = new Map>(); + for (const link of graphData.links) { + const s = link.source as unknown as { name?: string; id?: string }; + const t = link.target as unknown as { name?: string; id?: string }; + const src = (s?.name ?? s?.id ?? String(link.source)) as string; + const tgt = (t?.name ?? t?.id ?? String(link.target)) as string; + degreeByName.set(src, (degreeByName.get(src) ?? 0) + 1); + degreeByName.set(tgt, (degreeByName.get(tgt) ?? 0) + 1); + if (!neighborsByName.has(src)) neighborsByName.set(src, new Set()); + if (!neighborsByName.has(tgt)) neighborsByName.set(tgt, new Set()); + neighborsByName.get(src)!.add(tgt); + neighborsByName.get(tgt)!.add(src); + } + const out = new Map(); + for (const n of graphData.nodes) { + const name = (n as { name?: string; id?: string }).name ?? (n as { id: string }).id; + const degree = degreeByName.get(name) ?? (n as { degree?: number }).degree ?? 0; + out.set(name, entityColor(name, [...(neighborsByName.get(name) ?? [])], degree)); + } + return out; + }, [graphData.nodes, graphData.links]); + + // ponytail: lookup helper with muted-slate fallback for any entity + // name not in the color map (e.g. a row that survived dedup + // filtering but lost its neighbors). + const colorFor = useCallback( + (name: string): EntityColor => nodeColors.get(name) ?? entityColor(name, [], 0), + [nodeColors], + ); + + // ponytail: canvas node paint. Reads bg/fg/border from the + // graphRAG-native color map instead of the person/organization/ + // concept whitelist (which left 39/40 of this doc's entity types + // grey). Pulled after nodeColors so the closure can read it. + const nodeCanvasObject = useCallback( + (node: any, ctx: CanvasRenderingContext2D, globalScale: number) => { + if ( + typeof node.x !== "number" || + typeof node.y !== "number" || + !isFinite(node.x) || + !isFinite(node.y) + ) { + return; + } + + const label = node.name; + const degree = node.degree || 1; + const zoomBoost = Math.min(0.6, Math.max(-0.5, (globalScale - 1) * -0.55)); + // ponytail: smaller radius range after the redesign. degree + // 1-2 sits at 2.5 (tiny), 8+ at 8 (visible hub). Avoids the + // "1.5x the size of everything else" effect we saw in the first + // visual pass. + const r = Math.min(8, Math.max(2.5, 2.5 + degree * 0.6 + zoomBoost)); + + // ponytail: flat fill + 1px stroke. Replaces the radial gradient + // (which gave every node a 3D ball look) — now matches the + // outline style of the badges: light bg fill, same-hue border. + // Hue from entityColor (neighbor signature); saturation pulled + // down so dots don't compete with each other on the slate bg. + let fillColor = "#f1f5f9"; + let strokeColor = "#475569"; + const nodeColor = nodeColors.get(label); + if (nodeColor) { + fillColor = nodeColor.bg; + strokeColor = nodeColor.border; + } + + ctx.beginPath(); + ctx.arc(node.x, node.y, r, 0, 2 * Math.PI, false); + ctx.fillStyle = fillColor; + ctx.fill(); + ctx.lineWidth = 1; + ctx.strokeStyle = strokeColor; + ctx.stroke(); + + // ponytail: label visibility tightened — only hubs (degree >= 6) + // and nodes the user explicitly zoomed into (globalScale > 1.6) + // get a label. The old "globalScale > 0.6 + degree 3" rule + // flooded the canvas with overlapping text on the 321-entity + // folder view. + const showLabel = degree >= 6 || globalScale > 1.6; + if (showLabel) { + const fontSize = Math.max(2.6, 3.4 - Math.min(0.8, degree * 0.06)); + ctx.font = `${fontSize}px Inter, system-ui, -apple-system, sans-serif`; + ctx.textAlign = "center"; + ctx.textBaseline = "top"; + + const maxLen = globalScale > 1.5 ? 18 : globalScale > 0.9 ? 12 : 8; + let displayLabel = label; + if (label.length > maxLen) { + displayLabel = label.substring(0, maxLen - 2) + "..."; + } + + // ponytail: thin white halo around the label. 3px was + // overpowering the small text (2.6-3.4px fontSize). 1.5px + // + low alpha keeps overlapping geometry from breaking + // legibility without making each label look like it's + // outlined in chalk. + ctx.lineWidth = 1.5; + ctx.strokeStyle = "rgba(248, 250, 252, 0.55)"; + ctx.lineJoin = "round"; + ctx.strokeText(displayLabel, node.x, node.y + r + 3.5); + ctx.fillStyle = "#475569"; + ctx.fillText(displayLabel, node.x, node.y + r + 3.5); + } + }, + [nodeColors], + ); + + // ponytail: re-fit whenever the node/edge count materially changes + // so the graph lands on a useful framing for the user. Defers a + // tick so the simulation has its first coords before we ask for a + // fit. Idempotent across renders — lastFittedKey dedupes. + useEffect(() => { + if (graphView !== "visual") return; + const key = `${graphData.nodes.length}-${graphData.links.length}`; + if (key === lastFittedKey.current) return; + lastFittedKey.current = key; + const raf = requestAnimationFrame(() => graphRef.current?.zoomToFit(400, 60)); + const id = setTimeout(() => graphRef.current?.zoomToFit(400, 60), 300); + return () => { + cancelAnimationFrame(raf); + clearTimeout(id); + }; + }, [graphData.nodes.length, graphData.links.length, graphView]); + + if (chunks.length === 0) { + return ( +

+ {emptyMessage} +

+ ); + } + + return ( +
+
+ + + + +
+ +
+ {graphView === "visual" && ( +
+ {graphData.nodes.length === 0 ? ( + + No entities to visualize yet. + + ) : ( + <> + "replace"} + linkCanvasObject={linkCanvasObject} + linkWidth={linkWidth} + linkColor={linkColor} + cooldownTicks={80} + nodeCanvasObject={nodeCanvasObject} + onNodeHover={onNodeHover} + onLinkHover={onLinkHover} + /> + + {(() => { + const isNode = !!hoveredNode; + const w = isNode ? 240 : 260; + const h = isNode ? 120 : 100; // estimated height offset + + const rawX = + tooltipPos.x + w > dimensions.width ? tooltipPos.x - w - 6 : tooltipPos.x + 6; + const rawY = + tooltipPos.y + h > dimensions.height ? tooltipPos.y - h - 6 : tooltipPos.y + 6; + const left = Math.max(8, rawX); + const top = Math.max(8, rawY); + + if (hoveredNode) { + // ponytail: hover tooltip type badge — color from + // graphRAG-native entityColor (hue = neighbor + // signature, sat/light = degree). Replaces the + // person/organization/concept/team/job/motto + // string whitelist. + const hoveredColor = colorFor(hoveredNode.name); + + return ( +
+
+
+ Entity Info +
+
+ + {hoveredNode.name} + + + {hoveredNode.type || "Other"} + +
+
+ + {hoveredNode.description && ( +
+ + Description + +

+ {hoveredNode.description} +

+
+ )} +
+ ); + } + + if (hoveredLink) { + return ( +
+
+
+ Relationship Path +
+
+ + {hoveredLink.source.name || hoveredLink.source} + + → + + {hoveredLink.target.name || hoveredLink.target} + +
+
+ +
+
+ Relation Type +
+
+ + {hoveredLink.relation} + +
+
+ + {hoveredLink.description && ( +
+ + Description + +

+ {hoveredLink.description} +

+
+ )} +
+ ); + } + + return null; + })()} + + )} +
+ )} + + {graphView === "themes" && uniqueThemes.length > 0 && ( +
+
+ {uniqueThemes.map((theme, i) => ( + + #{theme} + + ))} +
+
+ )} + + {graphView === "entities" && uniqueEntities.length > 0 && ( +
+
+
+
Name
+
Type
+
Description
+
+ {uniqueEntities.map((e, idx) => { + // ponytail: list-view type badge — color from + // graphRAG-native entityColor (same source as canvas + // nodes + hover tooltip + doc-detail-dialog). No + // string-type whitelist; hue = neighbor signature, + // saturation/lightness = degree. + const listColor = colorFor(e.name); + return ( +
+ {e.name} +
+ + Type: + + + {e.type} + +
+
+ + Description: + + {e.description} +
+
+ ); + })} +
+
+ )} + + {graphView === "relationships" && uniqueRelationships.length > 0 && ( +
+
+
+
Connection
+
Relationship
+
Context / Description
+
+ {uniqueRelationships.map((r, idx) => ( +
+
+ + Connection: + + {r.source} + + + {r.target} + +
+
+ + Relationship: + + + {r.relation} + +
+
+ + Context / Description: + + {r.description} +
+
+ ))} +
+
+ )} +
+
+ ); +} diff --git a/components/settings/kb-view/live-poll-indicator.tsx b/components/settings/kb-view/live-poll-indicator.tsx new file mode 100644 index 00000000..71e464b0 --- /dev/null +++ b/components/settings/kb-view/live-poll-indicator.tsx @@ -0,0 +1,66 @@ +"use client"; + +import { useEffect, useState } from "react"; +import { Tooltip, TooltipContent, TooltipTrigger } from "@/components/ui/tooltip"; +import { cn } from "@/lib/utils"; + +// ponytail: small pulsing dot in the table header that lights up +// when the auto-refresh loop is active. The tooltip counts down to +// the next refresh so the user has a sense of when the table will +// land on fresh data — useful during reprocess / upload windows when +// the doc rows look stuck on a stale status. +// +// `active` is the parent's view of "polling is running right now"; +// when it flips false, the indicator unmounts. The internal +// countdown is best-effort (resets every `intervalMs`) — the actual +// refresh fires a few hundred ms after the timer hits 0, which is +// close enough that the user never notices the drift. +export function LivePollIndicator({ + active, + intervalMs, + className, +}: { + active: boolean; + intervalMs: number; + className?: string; +}) { + const [secondsLeft, setSecondsLeft] = useState(Math.ceil(intervalMs / 1000)); + + useEffect(() => { + if (!active) return; + setSecondsLeft(Math.ceil(intervalMs / 1000)); + const t = setInterval(() => { + setSecondsLeft((s) => (s <= 1 ? Math.ceil(intervalMs / 1000) : s - 1)); + }, 1000); + return () => clearInterval(t); + }, [active, intervalMs]); + + if (!active) return null; + + return ( + + + + + Auto-refresh in {secondsLeft}s + + ); +} diff --git a/components/settings/kb-view/observability-popover.tsx b/components/settings/kb-view/observability-popover.tsx new file mode 100644 index 00000000..27076767 --- /dev/null +++ b/components/settings/kb-view/observability-popover.tsx @@ -0,0 +1,223 @@ +"use client"; + +// ponytail: Activity icon → Popover listing every kbAgent invocation +// for this doc. Click a run → opens the singleton ObservabilitySheet +// for that (threadId, parentMessageId) pair. Each row carries its own +// threadId — chat-uploaded docs have rows on the chat thread, while +// Settings standalone rows live on the docId-derived thread. The API +// stitches them together via kb_observability.docId, so the popover +// shows the full per-doc history in one place. +// +// Why a Popover (not a Dialog / Tab inside DocDetailDialog): the +// affordance is "look up runs for this one doc" — a transient query, +// not a destination. Putting it next to RefreshCw / Search / Delete +// keeps related actions grouped. The Popover reuses shadcn Popover +// primitives so styling matches the rest of the row. +import { useEffect, useState } from "react"; +import { Activity, MessageSquare, RefreshCw, Upload } from "lucide-react"; + +import { Button } from "@/components/ui/button"; +import { + Popover, + PopoverContent, + PopoverHeader, + PopoverTitle, + PopoverTrigger, +} from "@/components/ui/popover"; +import { Tooltip, TooltipContent, TooltipTrigger } from "@/components/ui/tooltip"; +import { useOpenObservabilitySheet } from "@/components/observability/sheet-context"; +import { getModeInfo, type ReprocessMode } from "./helpers"; + +type ObservabilityRun = { + runId: string | null; + threadId: string; + parentMessageId: string; + source: string; + mode: ReprocessMode; + createdAt: string; +}; + +type ObservabilityResponse = { + doc_id: string; + runs: ObservabilityRun[]; +}; + +export function ObservabilityPopover({ docId }: { docId: string }) { + const [open, setOpen] = useState(false); + const [loading, setLoading] = useState(false); + const [error, setError] = useState(null); + const [data, setData] = useState(null); + const openSheet = useOpenObservabilitySheet(); + + // ponytail: lazy-fetch on each open — close resets state so the + // next click re-requests. Matters after a reprocess or chat + // upload: the user reopens the popover and expects the new run + // to show up. Caching across opens would hide it. + useEffect(() => { + if (!open || data || loading) return; + setLoading(true); + setError(null); + fetch(`/api/kb/documents/${docId}/observability`, { credentials: "include" }) + .then(async (res) => { + if (!res.ok) throw new Error(`HTTP ${res.status}`); + return (await res.json()) as ObservabilityResponse; + }) + .then((body) => { + setData(body); + }) + .catch((e: unknown) => { + setError(e instanceof Error ? e.message : "Failed to load observability"); + }) + .finally(() => { + setLoading(false); + }); + }, [open, data, loading, docId]); + + const handleRunClick = (run: ObservabilityRun) => { + setOpen(false); + // ponytail: each row carries its own threadId — chat uploads + // open against the chat thread, standalone against the + // docId-derived thread. parentMessageId is the synthetic + // HumanMessage id (standalone) or the user's chat msg id (chat). + openSheet({ threadId: run.threadId, parentMessageId: run.parentMessageId }); + }; + + // ponytail: reset cached state on close so the next click fires a + // fresh /observability request — otherwise the useEffect's `data` + // guard short-circuits and the popover reopens with stale rows. + const handleOpenChange = (nextOpen: boolean) => { + if (!nextOpen) { + setData(null); + setError(null); + setLoading(false); + } + setOpen(nextOpen); + }; + + return ( + + + + + + + + View observability + + + + Observability List + +
+ {loading ? ( + // ponytail: skeleton mirrors the real row layout so the + // popover doesn't reflow when data arrives — same icon + // column, same source/badge line, same time line. +
+
+
+
+
+
+
+
+
+
+
+
+ ) : error ? ( +
+ {error} +
+ ) : !data || data.runs.length === 0 ? ( +
+ No re-runs yet. Use the reprocess button to record one. +
+ ) : ( +
    + {data.runs.map((run, i) => { + const Icon = sourceIcon(run.source); + return ( +
  • + +
    • + ); + })} +
+ )} +
+ + + ); +} + +// ponytail: friendly source labels. The full path doesn't write +// observability rows anymore (the kb_document row IS the event for +// initial uploads); only chunksOnly / retryFailed / retryFailedChunks +// runs land here, so source is effectively always "kb-reprocess" in +// practice — but the labels handle the other values defensively in +// case chat-path chunksOnly shows up later. +function sourceLabel(source: ObservabilityRun["source"]): string { + switch (source) { + case "kb-upload": + return "Upload"; + case "kb-reprocess": + return "Reprocess"; + case "chat": + return "Chat upload"; + } + return source; +} + +function sourceIcon(source: ObservabilityRun["source"]) { + switch (source) { + case "kb-upload": + return Upload; + case "kb-reprocess": + return RefreshCw; + case "chat": + return MessageSquare; + } + return Activity; +} + +// ponytail: mode labels match the reprocess dialog copy so the popover +// reads as a continuation of the same vocabulary. + +function formatTimestamp(iso: string): string { + const t = new Date(iso); + if (Number.isNaN(t.getTime())) return iso; + const pad = (n: number) => String(n).padStart(2, "0"); + return `${t.getFullYear()}/${pad(t.getMonth() + 1)}/${pad(t.getDate())} ${pad(t.getHours())}:${pad(t.getMinutes())}:${pad(t.getSeconds())}`; +} diff --git a/components/settings/kb-view/status-badge.tsx b/components/settings/kb-view/status-badge.tsx new file mode 100644 index 00000000..95b2e820 --- /dev/null +++ b/components/settings/kb-view/status-badge.tsx @@ -0,0 +1,75 @@ +import { AlertCircle, CheckCircle2, FileText, Loader2 } from "lucide-react"; +import { Badge } from "@/components/ui/badge"; +import { Tooltip, TooltipContent, TooltipTrigger } from "@/components/ui/tooltip"; +import { cn } from "@/lib/utils"; + +import { + DocStatusBadge, + ChunksStatusBadge, + type KbStatus, +} from "@/components/tool-ui/kb/status-badge"; + +export { DocStatusBadge, ChunksStatusBadge, type KbStatus }; + +// ponytail: ChunkStatusBadge is per-chunk and only renders in the +// doc-detail dialog's Chunks tab. It stayed in the settings view +// because the chat-side list_documents card only needs per-doc +// status, not per-chunk. +export function ChunkStatusBadge({ + status, + errorMessage, +}: { + status: KbStatus; + errorMessage: string | null; +}) { + // ponytail: match the Pages tab's "Succeeded"/"Failed"/"Parsing…"/"Pending" + // badge shape exactly — same variant, same dimensions + // (text-[9px] py-0 px-1.5), same label text. Without this, the + // Pages "SUCCEEDED" pill and the Chunks "SUCCESS" pill looked like + // two different status indicators even though they meant the same + // thing. + const variant: "success" | "destructive" | "muted" = + status === "success" ? "success" : status === "failed" ? "destructive" : "muted"; + const label = + status === "success" + ? "Succeeded" + : status === "failed" + ? "Failed" + : status === "parsing" + ? "Parsing…" + : "Pending"; + const content = ( + + {label} + + ); + if (status === "failed" && errorMessage) { + return ( + + {content} + {errorMessage} + + ); + } + return content; +} + +export function StatusIcon({ status }: { status: KbStatus }) { + switch (status) { + case "success": + return ; + case "failed": + return ; + case "parsing": + case "pending": + return ; + default: + return ; + } +} + +export function formatTimestamp(iso: string): string { + const d = new Date(iso); + if (Number.isNaN(d.getTime())) return iso; + return d.toLocaleString(); +} diff --git a/components/settings/kb-view/types.ts b/components/settings/kb-view/types.ts new file mode 100644 index 00000000..0b138236 --- /dev/null +++ b/components/settings/kb-view/types.ts @@ -0,0 +1,51 @@ +export type KbStatus = "pending" | "parsing" | "success" | "failed"; + +export type KbDocument = { + id: string; + title: string; + status: KbStatus; + errorMessage: string | null; + contentType: string; + attachmentId: string | null; + attachmentUrl: string | null; + pages?: Array<{ + pageIndex: number; + imageUrl: string; + markdown: string; + referenceText?: string; + errorMessage?: string; + status?: "pending" | "parsing" | "success" | "failed"; + }>; + createdAt: string; + updatedAt: string; + totalChunks?: number; + successChunks?: number; + failedChunks?: number; + pendingChunks?: number; + parsingChunks?: number; + totalPages?: number; + failedPages?: number; + pendingPages?: number; + parsingPages?: number; +}; + +export type KbFolder = { id: string; name: string }; + +export type KbResponse = { + groups: Array<{ folder: KbFolder; documents: KbDocument[] }>; +}; + +export type KbChunkPreviewLocal = { + ordinal: number; + content: string; + entities: Array<{ name: string; type: string; description: string }>; + relationships: Array<{ source: string; target: string; relation: string; description: string }>; + themes: string[]; + status: "pending" | "parsing" | "success" | "failed"; + errorMessage: string | null; +}; + +export type KbDocDetail = { + doc: KbDocument & { folderId: string; contentHash: string }; + chunks: KbChunkPreviewLocal[]; +}; diff --git a/components/tool-ui/kb/chunk-list.tsx b/components/tool-ui/kb/chunk-list.tsx new file mode 100644 index 00000000..8a589509 --- /dev/null +++ b/components/tool-ui/kb/chunk-list.tsx @@ -0,0 +1,81 @@ +"use client"; + +import { useState } from "react"; +import { FileTextIcon } from "lucide-react"; + +import { chunkPreview, legBadges } from "./parser"; +import type { KbDocument } from "./types"; + +// ponytail: shared chunk list — search_kb renders +// the same { [n] · title · leg badges · preview } rows. Order is locked +// to the backend's RRF ranking (.claude/13-kb-v3.md "顺序锁定"). +// Now collapsible (default 3) to prevent large context lists from +// cluttering the chat view. + +export function KbChunkList({ docs, slot }: { docs: KbDocument[]; slot: string }) { + const [expanded, setExpanded] = useState(false); + + if (docs.length === 0) return null; + + const isRerank = docs.some((d) => d.rrfScore > 0.05); + const visibleDocs = expanded ? docs : docs.slice(0, 3); + const hasMore = docs.length > 3; + + return ( +
+
    + {visibleDocs.map((doc, i) => { + const badges = legBadges(doc.legsHit); + return ( +
  1. +
    + + [{i + 1}] + + + + {doc.docTitle} + +
    + {badges.length > 0 && ( +
    + {badges.map((b) => ( + + {b} + + ))} +
    + )} + {typeof doc.rrfScore === "number" && doc.rrfScore > 0 && ( + + Score:{" "} + {isRerank ? `${Math.round(doc.rrfScore * 100)}%` : doc.rrfScore.toFixed(3)} + + )} +
    +
    +

    + {chunkPreview(doc.content)} +

    +
    2. + ); + })} +
+ {hasMore && ( + + )} +
+ ); +} diff --git a/components/tool-ui/kb/index.ts b/components/tool-ui/kb/index.ts new file mode 100644 index 00000000..9d74a781 --- /dev/null +++ b/components/tool-ui/kb/index.ts @@ -0,0 +1,2 @@ +export { KbSearchToolUI } from "./search-kb-card"; +export { KbListDocumentsToolUI } from "./list-documents-card"; diff --git a/components/tool-ui/kb/list-documents-card.tsx b/components/tool-ui/kb/list-documents-card.tsx new file mode 100644 index 00000000..bd07a1b4 --- /dev/null +++ b/components/tool-ui/kb/list-documents-card.tsx @@ -0,0 +1,190 @@ +"use client"; + +import { useState } from "react"; +import { ChevronRightIcon, FileTextIcon, FolderIcon, LoaderIcon } from "lucide-react"; +import type { ToolCallMessagePartComponent } from "@assistant-ui/react"; + +import { CardShell, CardHeader } from "@/components/tool-ui/primitives/card"; +import { Collapsible, CollapsibleContent, CollapsibleTrigger } from "@/components/ui/collapsible"; +import { unwrapToolResult } from "@/components/tool-ui/tool-result"; +import { cn } from "@/lib/utils"; + +import { ChunksStatusBadge, DocStatusBadge } from "./status-badge"; +import type { ListDocResult, ListDocumentsFolder } from "./types"; + +type ListArgs = { + folderId?: string; + status?: string; + titleQuery?: string; + page?: number; + pageSize?: number; +}; + +// ponytail: visible cap before the "show more" button. Matches the +// per-folder default in the search_kb chunk list so the two cards +// feel consistent in the chat thread. +const VISIBLE_DOCS_PER_FOLDER = 3; + +// ponytail: two independent collapses per folder — clicking the +// header chevron toggles the whole section, "Show more" toggles +// whether the doc list shows 3 docs or the full set. Radix +// Collapsible drives both so the open/close animates (height + +// chevron rotation) instead of jumping. +function FolderSection({ folder }: { folder: ListDocumentsFolder }) { + const [folderOpen, setFolderOpen] = useState(true); + const [docsExpanded, setDocsExpanded] = useState(false); + if (folder.documents.length === 0) return null; + + // ponytail: the first 3 always render in the main list; the + // rest go into the second Collapsible so their height animates + // in/out (rather than the main list re-rendering with a + // different slice and the new docs popping in un-animated). + const head = folder.documents.slice(0, VISIBLE_DOCS_PER_FOLDER); + const tail = folder.documents.slice(VISIBLE_DOCS_PER_FOLDER); + const hidden = tail.length; + + return ( + + + + + {folder.name} + + · {folder.documents.length} + + + +
    + {head.map((d) => ( + + ))} +
+ {hidden > 0 && ( + // ponytail: button lives OUTSIDE the inner Collapsible so + // it always sits at the bottom of the folder content (not + // jumps to the middle of the row when the tail expands). + // The inner Collapsible only owns the tail list's height + // animation. + + +
    + {tail.map((d) => ( + + ))} +
+ + + )} + {hidden > 0 && ( + + )} + + + ); +} + +// ponytail: shared per-doc row. Same chrome in the head list and +// the (collapsed) tail list, so a CSS tweak only lands in one place. +function DocRow({ d }: { d: ListDocumentsFolder["documents"][number] }) { + return ( +
  • +
    + + + {d.title} + + +
    +
    + + +
    +
    • + ); +} + +// ponytail: per-doc date formatter. Short locale string + the +//