diff --git a/CHANGELOG.md b/CHANGELOG.md index c2368d5..91ada5f 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,29 @@ All notable changes to DocMind are logged here, phase by phase. This is the publ Format loosely follows [Keep a Changelog](https://keepachangelog.com/en/1.1.0/). +## Agent Citations + Multi-Turn Conversations — 2026-07-26 + +### Added + +#### Multi-Turn Agent Conversations +- **`Conversation` + `Message` Prisma models** — new `MessageRole` enum (`user`/`assistant`); `Message.citations` stores citation data for assistant turns produced via the short-circuit path below. +- **`ConversationsModule`** (`ConversationsService`, `ConversationsController`) — `GET /v1/conversations` (list current user's conversations), `GET /v1/conversations/:id/messages` (full history), ownership-checked via `assertOwnership()`. +- **`trimHistory()` util** — hybrid history-window strategy: hard ceiling of the last 20 messages, then a further ~3000-token budget trim within that window, always preserving at least the most recent turn even if it alone exceeds the budget. +- **`AgentChatDto.conversationId`** (optional) — omit to start a new conversation; the new id is returned via a `conversation_started` SSE event before the answer stream begins. +- **`AgentService.run()`** now loads and seeds trimmed prior history into the graph's initial `messages` state, and persists the user query immediately plus the final assistant answer once produced. Nothing is persisted for turns that pause on an `external_write` confirmation (no final answer yet). +- Scope: multi-turn ships on `/v1/agent/chat` only; `/v1/chat/stream` remains single-turn. + +#### Agent Citations +- **Citation short-circuit** — when a dispatched tool (`query_documents`) returns `{answer, citations}`, `AgentService` now emits a `citations` SSE event and streams that answer directly instead of routing back through `modelTurn` for re-synthesis. This guarantees `[N]` markers in the answer text stay aligned with the emitted citation array, and closes a previously dead `citations` event type in `agent-sse.types.ts` that the frontend already knew how to render. + +#### Frontend +- **`useChatStream`** rewritten around a `messages: ChatMessage[]` thread and a `conversationId` ref, replacing the old single-answer `content`/`citations` state. +- **`/chat` page** rewritten as a scrollable message thread (`MessageBubble` per turn) with a "New conversation" control (`startNewConversation()`), instead of rendering only the latest Q&A pair. + +### Tests +- `agent.service.spec.ts` — citation short-circuit (single `generate()` call, correct citation payload), history seeding into the first `generate()` call, message persistence across the cited-answer/normal/proposal-pause paths. +- `history.util.spec.ts` — hard-ceiling cap, token-budget trim, and the always-keep-latest-turn safeguard. + ## JWT Auth — 2026-07-24 ### Added @@ -337,4 +360,4 @@ Format loosely follows [Keep a Changelog](https://keepachangelog.com/en/1.1.0/). - Starter branding renamed `jsstack` → `docmind` across package names, Docker Compose services, env files, and CI registry paths. - Postgres image `postgres:16-alpine` → `pgvector/pgvector:pg16`; added missing `migrate` service. - Port assignments set to non-default values (backend 4500, frontend 3400, Postgres 5349, Redis 6399) to avoid VPS conflicts. -- CI/CD split into `ci` / `build` / `deploy` workflows with EC2 deployment via WireGuard SSH. +- CI/CD split into `ci` / `build` / `deploy` workflows with EC2 deployment via WireGuard SSH. \ No newline at end of file diff --git a/backend/Dockerfile b/backend/Dockerfile index 64c8aba..6bced1b 100644 --- a/backend/Dockerfile +++ b/backend/Dockerfile @@ -53,7 +53,7 @@ USER nestjs EXPOSE 4500 HEALTHCHECK --interval=10s --timeout=5s --start-period=30s --retries=5 \ - CMD node -e "require('http').get('http://localhost:4500/health', r => process.exit(r.statusCode === 200 ? 0 : 1)).on('error', () => process.exit(1))" + CMD node -e "require('http').get('http://127.0.0.1:4500/health', r => process.exit(r.statusCode === 200 ? 0 : 1)).on('error', () => process.exit(1))" ENTRYPOINT ["dumb-init", "--"] # Running from the workspace directory structure diff --git a/backend/prisma.config.ts b/backend/prisma.config.ts index f4b9849..cfc775d 100644 --- a/backend/prisma.config.ts +++ b/backend/prisma.config.ts @@ -1,5 +1,7 @@ import { defineConfig } from 'prisma/config'; +import dotenv from 'dotenv'; +dotenv.config(); export default defineConfig({ schema: 'prisma/schema.prisma', migrations: { diff --git a/backend/prisma/migrations/20260727045518_add_conversations/migration.sql b/backend/prisma/migrations/20260727045518_add_conversations/migration.sql new file mode 100644 index 0000000..e3a42b5 --- /dev/null +++ b/backend/prisma/migrations/20260727045518_add_conversations/migration.sql @@ -0,0 +1,108 @@ +/* + Warnings: + + - You are about to drop the column `content_tsv` on the `chunks` table. All the data in the column will be lost. + +*/ +-- CreateEnum +CREATE TYPE "MessageRole" AS ENUM ('user', 'assistant'); + +-- DropForeignKey +ALTER TABLE "chunks" DROP CONSTRAINT "chunks_documentId_fkey"; + +-- DropForeignKey +ALTER TABLE "notes" DROP CONSTRAINT "notes_userId_fkey"; + +-- DropForeignKey +ALTER TABLE "query_traces" DROP CONSTRAINT "query_traces_userId_fkey"; + +-- DropForeignKey +ALTER TABLE "tasks" DROP CONSTRAINT "tasks_userId_fkey"; + +-- DropForeignKey +ALTER TABLE "tool_call_audits" DROP CONSTRAINT "tool_call_audits_userId_fkey"; + +-- DropIndex +DROP INDEX "chunks_content_tsv_idx"; + +-- DropIndex +DROP INDEX "idx_chunks_embedding_hnsw"; + +-- AlterTable +ALTER TABLE "chunks" DROP COLUMN "content_tsv", +ALTER COLUMN "id" DROP DEFAULT; + +-- AlterTable +ALTER TABLE "documents" ALTER COLUMN "id" DROP DEFAULT, +ALTER COLUMN "updatedAt" DROP DEFAULT; + +-- AlterTable +ALTER TABLE "notes" ALTER COLUMN "id" DROP DEFAULT, +ALTER COLUMN "updatedAt" DROP DEFAULT; + +-- AlterTable +ALTER TABLE "query_traces" ALTER COLUMN "id" DROP DEFAULT, +ALTER COLUMN "toolCallAuditIds" DROP DEFAULT; + +-- AlterTable +ALTER TABLE "tasks" ALTER COLUMN "id" DROP DEFAULT, +ALTER COLUMN "updatedAt" DROP DEFAULT; + +-- AlterTable +ALTER TABLE "tool_call_audits" ALTER COLUMN "id" DROP DEFAULT; + +-- AlterTable +ALTER TABLE "users" ALTER COLUMN "id" DROP DEFAULT, +ALTER COLUMN "updatedAt" DROP DEFAULT; + +-- CreateTable +CREATE TABLE "conversations" ( + "id" TEXT NOT NULL, + "userId" TEXT NOT NULL, + "createdAt" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP, + "updatedAt" TIMESTAMP(3) NOT NULL, + + CONSTRAINT "conversations_pkey" PRIMARY KEY ("id") +); + +-- CreateTable +CREATE TABLE "messages" ( + "id" TEXT NOT NULL, + "conversationId" TEXT NOT NULL, + "role" "MessageRole" NOT NULL, + "content" TEXT NOT NULL, + "citations" JSONB, + "createdAt" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP, + + CONSTRAINT "messages_pkey" PRIMARY KEY ("id") +); + +-- CreateIndex +CREATE INDEX "conversations_userId_updatedAt_idx" ON "conversations"("userId", "updatedAt"); + +-- CreateIndex +CREATE INDEX "messages_conversationId_createdAt_idx" ON "messages"("conversationId", "createdAt"); + +-- AddForeignKey +ALTER TABLE "chunks" ADD CONSTRAINT "chunks_documentId_fkey" FOREIGN KEY ("documentId") REFERENCES "documents"("id") ON DELETE RESTRICT ON UPDATE CASCADE; + +-- AddForeignKey +ALTER TABLE "notes" ADD CONSTRAINT "notes_userId_fkey" FOREIGN KEY ("userId") REFERENCES "users"("id") ON DELETE CASCADE ON UPDATE CASCADE; + +-- AddForeignKey +ALTER TABLE "tasks" ADD CONSTRAINT "tasks_userId_fkey" FOREIGN KEY ("userId") REFERENCES "users"("id") ON DELETE CASCADE ON UPDATE CASCADE; + +-- AddForeignKey +ALTER TABLE "tool_call_audits" ADD CONSTRAINT "tool_call_audits_userId_fkey" FOREIGN KEY ("userId") REFERENCES "users"("id") ON DELETE CASCADE ON UPDATE CASCADE; + +-- AddForeignKey +ALTER TABLE "query_traces" ADD CONSTRAINT "query_traces_userId_fkey" FOREIGN KEY ("userId") REFERENCES "users"("id") ON DELETE CASCADE ON UPDATE CASCADE; + +-- AddForeignKey +ALTER TABLE "conversations" ADD CONSTRAINT "conversations_userId_fkey" FOREIGN KEY ("userId") REFERENCES "users"("id") ON DELETE CASCADE ON UPDATE CASCADE; + +-- AddForeignKey +ALTER TABLE "messages" ADD CONSTRAINT "messages_conversationId_fkey" FOREIGN KEY ("conversationId") REFERENCES "conversations"("id") ON DELETE CASCADE ON UPDATE CASCADE; + +-- RenameIndex +ALTER INDEX "idx_chunks_documentId" RENAME TO "chunks_documentId_idx"; diff --git a/backend/prisma/migrations/20260727060000_restore_keyword_and_vector_indexes/migration.sql b/backend/prisma/migrations/20260727060000_restore_keyword_and_vector_indexes/migration.sql new file mode 100644 index 0000000..57e44f4 --- /dev/null +++ b/backend/prisma/migrations/20260727060000_restore_keyword_and_vector_indexes/migration.sql @@ -0,0 +1,29 @@ +-- Defensive restore migration. +-- +-- Root cause: content_tsv (0007_keyword_search) and idx_chunks_embedding_hnsw +-- (0003_chunk_schema) are hand-written SQL objects that Prisma cannot express +-- in schema.prisma (a GENERATED ALWAYS AS STORED column, and a vector HNSW +-- index respectively). Neither was declared as an Unsupported() field/index, +-- so a subsequent `prisma migrate dev` run diffed the live DB against +-- schema.prisma, saw both as "not in schema", and generated DROP statements +-- for them as part of an unrelated migration (add_conversations). +-- +-- This migration is fully idempotent (IF NOT EXISTS / IF EXISTS guards) so +-- it's safe to run whether or not the drop actually happened on a given +-- database. schema.prisma now declares content_tsv as Unsupported("tsvector") +-- to prevent this from recurring for the column; the two indexes below have +-- no equivalent protection in Prisma and must stay hand-maintained — see the +-- comments in schema.prisma next to the Chunk model. + +-- Restore the generated tsvector column for full-text search on chunks.content +ALTER TABLE "chunks" + ADD COLUMN IF NOT EXISTS "content_tsv" tsvector + GENERATED ALWAYS AS (to_tsvector('english', content)) STORED; + +-- Restore the GIN index backing keyword search +CREATE INDEX IF NOT EXISTS "chunks_content_tsv_idx" + ON "chunks" USING GIN ("content_tsv"); + +-- Restore the HNSW index backing vector similarity search +CREATE INDEX IF NOT EXISTS "idx_chunks_embedding_hnsw" ON "chunks" + USING hnsw ("embedding" vector_cosine_ops); diff --git a/backend/prisma/migrations/migration_lock.toml b/backend/prisma/migrations/migration_lock.toml index 99e4f20..044d57c 100644 --- a/backend/prisma/migrations/migration_lock.toml +++ b/backend/prisma/migrations/migration_lock.toml @@ -1,3 +1,3 @@ # Please do not edit this file manually -# It should be added in your version-control system (i.e. Git) +# It should be added in your version-control system (e.g., Git) provider = "postgresql" diff --git a/backend/prisma/schema.prisma b/backend/prisma/schema.prisma index 52bc0d0..275e1e9 100644 --- a/backend/prisma/schema.prisma +++ b/backend/prisma/schema.prisma @@ -60,6 +60,16 @@ model Chunk { chunkIndex Int embeddingProvider String @default("gemini-embedding-001") embedding Unsupported("vector(768)") + // Generated column created by hand-written SQL (0007_keyword_search) — + // Prisma can't create GENERATED ALWAYS AS columns itself. Declaring it + // here as Unsupported only tells Prisma "this column exists, don't try + // to drop it" on future `migrate dev` diffs. It does NOT protect the + // chunks_content_tsv_idx GIN index, or idx_chunks_embedding_hnsw above — + // Prisma cannot represent indexes on Unsupported-typed columns at all. + // Always run `prisma migrate dev --create-only` for schema changes and + // inspect the generated SQL for an unexpected DROP INDEX/DROP COLUMN on + // either of those two objects before applying. + contentTsv Unsupported("tsvector")? @map("content_tsv") createdAt DateTime @default(now()) document Document @relation(fields: [documentId], references: [id]) @@ -86,6 +96,7 @@ model User { tasks Task[] toolCallAudits ToolCallAudit[] queryTraces QueryTrace[] + conversations Conversation[] @@map("users") } @@ -152,3 +163,43 @@ model QueryTrace { @@map("query_traces") } + +// ── Multi-turn agent chat history ─────────────────────────────────── +// Scoped to /v1/agent/chat only (see plan doc). Tool-call/tool-result +// scaffolding is intentionally NOT persisted here — only user turns and +// the final assistant answer, so replayed history stays clean. Per-turn +// tool-call detail is already captured separately in ToolCallAudit. + +enum MessageRole { + user + assistant +} + +model Conversation { + id String @id @default(uuid()) + userId String + createdAt DateTime @default(now()) + updatedAt DateTime @updatedAt + + user User @relation(fields: [userId], references: [id], onDelete: Cascade) + messages Message[] + + @@index([userId, updatedAt]) + @@map("conversations") +} + +model Message { + id String @id @default(uuid()) + conversationId String + role MessageRole + content String + // Present only on assistant messages produced by a citation-bearing + // tool result (see agent.service.ts finalAnswerOverride short-circuit). + citations Json? + createdAt DateTime @default(now()) + + conversation Conversation @relation(fields: [conversationId], references: [id], onDelete: Cascade) + + @@index([conversationId, createdAt]) + @@map("messages") +} diff --git a/backend/src/app.module.ts b/backend/src/app.module.ts index 9dcae47..335a406 100644 --- a/backend/src/app.module.ts +++ b/backend/src/app.module.ts @@ -14,6 +14,7 @@ import { RetrievalModule } from './modules/retrieval/retrieval.module'; import { QueryModule } from './modules/query/query.module'; import { ToolsModule } from './modules/tools/tools.module'; import { AgentModule } from './modules/agent/agent.module'; +import { ConversationsModule } from './modules/conversations/conversations.module'; import { NotesModule } from './modules/notes/notes.module'; import { TasksModule } from './modules/tasks/tasks.module'; import { TraceModule } from './modules/trace/trace.module'; @@ -55,6 +56,7 @@ const configValidationSchema = Joi.object({ QueryModule, ToolsModule, AgentModule, + ConversationsModule, NotesModule, TasksModule, TraceModule, diff --git a/backend/src/modules/agent/agent-sse.types.ts b/backend/src/modules/agent/agent-sse.types.ts index f863c70..36c1986 100644 --- a/backend/src/modules/agent/agent-sse.types.ts +++ b/backend/src/modules/agent/agent-sse.types.ts @@ -1,6 +1,7 @@ import type { ToolProposal } from '../tools/tool-proposal.type'; export type AgentSseEvent = + | { type: 'conversation_started'; data: { conversationId: string } } | { type: 'token'; data: string } | { type: 'citations'; data: unknown[] } | { type: 'done'; data: string } diff --git a/backend/src/modules/agent/agent.controller.ts b/backend/src/modules/agent/agent.controller.ts index e42857c..0db3958 100644 --- a/backend/src/modules/agent/agent.controller.ts +++ b/backend/src/modules/agent/agent.controller.ts @@ -27,6 +27,7 @@ import { Observable, Subject } from 'rxjs'; import type Redis from 'ioredis'; import { REDIS_CLIENT } from '../../redis/redis.module'; import { ToolRegistryService } from '../tools/tool-registry.service'; +import { ConversationsService } from '../conversations/conversations.service'; import { AgentService } from './agent.service'; import { CurrentUser, @@ -44,6 +45,16 @@ export class AgentChatDto { @MinLength(1) @MaxLength(5000) query!: string; + + @ApiProperty({ + required: false, + description: + 'Existing conversation to continue. Omit to start a new conversation ' + + '— the new id is returned via a conversation_started SSE event.', + }) + @IsOptional() + @IsUUID() + conversationId?: string; } export class ConfirmDto { @@ -71,6 +82,7 @@ export class AgentController { constructor( private readonly agentService: AgentService, private readonly toolRegistry: ToolRegistryService, + private readonly conversations: ConversationsService, @Optional() @Inject(REDIS_CLIENT) private readonly redis: Redis | null, ) {} @@ -86,18 +98,40 @@ export class AgentController { ): Observable { const subject = new Subject(); - void this.agentService - .run(dto.query, user.sub, (event: AgentSseEvent) => { - subject.next({ data: JSON.stringify(event) }); - if (event.type === 'done' || event.type === 'error') { - subject.complete(); + void (async () => { + try { + let conversationId = dto.conversationId; + + if (conversationId) { + await this.conversations.assertOwnership(user.sub, conversationId); + } else { + const conversation = await this.conversations.create(user.sub); + conversationId = conversation.id; + subject.next({ + data: JSON.stringify({ + type: 'conversation_started', + data: { conversationId }, + }), + }); } - }) - .catch((err: unknown) => { + + await this.agentService.run( + dto.query, + user.sub, + conversationId, + (event: AgentSseEvent) => { + subject.next({ data: JSON.stringify(event) }); + if (event.type === 'done' || event.type === 'error') { + subject.complete(); + } + }, + ); + } catch (err) { const msg = err instanceof Error ? err.message : 'Agent error'; subject.next({ data: JSON.stringify({ type: 'error', data: msg }) }); subject.complete(); - }); + } + })(); return subject.asObservable(); } diff --git a/backend/src/modules/agent/agent.module.ts b/backend/src/modules/agent/agent.module.ts index a960e42..b88c8b0 100644 --- a/backend/src/modules/agent/agent.module.ts +++ b/backend/src/modules/agent/agent.module.ts @@ -2,11 +2,12 @@ import { Module } from '@nestjs/common'; import { PrismaModule } from '../../prisma/prisma.module'; import { ProvidersModule } from '../providers/providers.module'; import { ToolsModule } from '../tools/tools.module'; +import { ConversationsModule } from '../conversations/conversations.module'; import { AgentController } from './agent.controller'; import { AgentService } from './agent.service'; @Module({ - imports: [ToolsModule, ProvidersModule, PrismaModule], + imports: [ToolsModule, ProvidersModule, PrismaModule, ConversationsModule], controllers: [AgentController], providers: [AgentService], exports: [AgentService], diff --git a/backend/src/modules/agent/agent.service.spec.ts b/backend/src/modules/agent/agent.service.spec.ts index 8a3a7f5..327846d 100644 --- a/backend/src/modules/agent/agent.service.spec.ts +++ b/backend/src/modules/agent/agent.service.spec.ts @@ -3,6 +3,7 @@ import { EventEmitter2 } from '@nestjs/event-emitter'; import { ConfigService } from '@nestjs/config'; import { AgentService } from './agent.service'; import { ToolRegistryService } from '../tools/tool-registry.service'; +import { ConversationsService } from '../conversations/conversations.service'; import { GENERATION_PROVIDER } from '../providers/generation.provider'; import { RiskTier } from '../../common/constants'; import type { AgentSseEvent } from './agent-sse.types'; @@ -12,10 +13,15 @@ function makeProvider(responses: string[]) { let callCount = 0; return { model: 'mock-model', - generate: jest.fn((): Promise<{ content: string }> => - Promise.resolve({ - content: responses[Math.min(callCount++, responses.length - 1)], - }), + generate: jest.fn( + (_options: { messages: unknown[]; systemPrompt: string }): Promise<{ + content: string; + }> => { + void _options; // kept only so Parameters<> isn't inferred as [] — see mock.calls[0][0] usage below + return Promise.resolve({ + content: responses[Math.min(callCount++, responses.length - 1)], + }); + }, ), generateStream: jest.fn(), }; @@ -28,12 +34,25 @@ function makeRegistry() { }; } -async function buildService(provider: unknown, registry: unknown) { +function makeConversations() { + return { + loadHistory: jest.fn().mockResolvedValue([]), + appendUserMessage: jest.fn().mockResolvedValue(undefined), + appendAssistantMessage: jest.fn().mockResolvedValue(undefined), + }; +} + +async function buildService( + provider: unknown, + registry: unknown, + conversations: unknown = makeConversations(), +) { const mod = await Test.createTestingModule({ providers: [ AgentService, { provide: GENERATION_PROVIDER, useValue: provider }, { provide: ToolRegistryService, useValue: registry }, + { provide: ConversationsService, useValue: conversations }, { provide: ConfigService, useValue: { get: jest.fn().mockReturnValue(undefined) }, @@ -47,9 +66,12 @@ async function buildService(provider: unknown, registry: unknown) { function collectEvents( service: AgentService, query: string, + conversationId = 'conv-1', ): Promise { const events: AgentSseEvent[] = []; - return service.run(query, 'user-1', (e) => events.push(e)).then(() => events); + return service + .run(query, 'user-1', conversationId, (e) => events.push(e)) + .then(() => events); } describe('AgentService', () => { @@ -132,6 +154,100 @@ describe('AgentService', () => { expect(events.some((e) => e.type === 'done')).toBe(true); }); + it('short-circuits on a cited answer: emits citations, no second generate call', async () => { + const provider = makeProvider([ + JSON.stringify({ tool: 'query_documents', params: { query: 'MVCC' } }), + ]); + const registry = makeRegistry(); + const citations = [ + { marker: '[1]', chunkId: 'chunk-1', documentTitle: 'Doc A', snippet: '...' }, + ]; + registry.dispatch.mockResolvedValue({ + answer: 'MVCC avoids locking [1].', + citations, + }); + registry.listTools.mockReturnValue([ + { + name: 'query_documents', + description: 'Search documents and return an answer with citations', + riskTier: RiskTier.read, + }, + ]); + const service = await buildService(provider, registry); + + const events = await collectEvents(service, 'How does MVCC work?'); + + const citationEvent = events.find((e) => e.type === 'citations'); + expect(citationEvent).toBeDefined(); + if (citationEvent?.type === 'citations') { + expect(citationEvent.data).toEqual(citations); + } + + const tokenContent = events + .filter((e) => e.type === 'token') + .map((e) => e.data) + .join(''); + expect(tokenContent.trim()).toBe('MVCC avoids locking [1].'); + + expect(events.some((e) => e.type === 'done')).toBe(true); + // Only one generate() call — the outer model never re-synthesizes + // the tool's own answer. + expect(provider.generate).toHaveBeenCalledTimes(1); + }); + + it('seeds prior history into the first generate() call and persists both turns', async () => { + const provider = makeProvider(['Sure, following up on that.']); + const registry = makeRegistry(); + const conversations = makeConversations(); + const priorHistory = [ + { role: 'user' as const, content: 'What is MVCC?' }, + { role: 'assistant' as const, content: 'MVCC avoids locking.' }, + ]; + conversations.loadHistory.mockResolvedValue(priorHistory); + const service = await buildService(provider, registry, conversations); + + await collectEvents(service, 'Tell me more', 'conv-42'); + + expect(conversations.loadHistory).toHaveBeenCalledWith('conv-42'); + // First generate() call must include prior history ahead of the new query + const firstCallArgs = provider.generate.mock.calls[0][0]; + expect(firstCallArgs.messages).toEqual([ + ...priorHistory, + { role: 'user', content: 'Tell me more' }, + ]); + + // User query persisted immediately, assistant answer persisted at the end + expect(conversations.appendUserMessage).toHaveBeenCalledWith( + 'conv-42', + 'Tell me more', + ); + expect(conversations.appendAssistantMessage).toHaveBeenCalledWith( + 'conv-42', + 'Sure, following up on that.', + ); + }); + + it('does not persist an assistant message when the turn pauses on confirmation', async () => { + const proposal = { + type: 'proposal' as const, + toolName: 'send_email', + preview: 'Send digest to user@example.com', + confirmationToken: 'tok-abc', + }; + const provider = makeProvider([ + JSON.stringify({ tool: 'send_email', params: { subject: 'Digest' } }), + ]); + const registry = makeRegistry(); + registry.dispatch.mockResolvedValue(proposal); + const conversations = makeConversations(); + const service = await buildService(provider, registry, conversations); + + await collectEvents(service, 'Send email digest'); + + expect(conversations.appendUserMessage).toHaveBeenCalled(); + expect(conversations.appendAssistantMessage).not.toHaveBeenCalled(); + }); + it('respects max iterations guard', async () => { const provider = makeProvider([ JSON.stringify({ tool: 'loop_tool', params: {} }), @@ -144,6 +260,7 @@ describe('AgentService', () => { AgentService, { provide: GENERATION_PROVIDER, useValue: provider }, { provide: ToolRegistryService, useValue: registry }, + { provide: ConversationsService, useValue: makeConversations() }, { provide: ConfigService, useValue: { get: jest.fn().mockReturnValue(2) }, @@ -252,4 +369,4 @@ describe('AgentService — parseModelOutput fence stripping', () => { // Embedded raw JSON object must not appear verbatim in streamed tokens expect(tokenContent).not.toContain('{"tool":"search","params":{}}'); }); -}); +}); \ No newline at end of file diff --git a/backend/src/modules/agent/agent.service.ts b/backend/src/modules/agent/agent.service.ts index fc93a77..25d284d 100644 --- a/backend/src/modules/agent/agent.service.ts +++ b/backend/src/modules/agent/agent.service.ts @@ -9,6 +9,7 @@ import { } from '../providers/generation.provider'; import { ToolRegistryService } from '../tools/tool-registry.service'; import { isToolProposal, ToolProposal } from '../tools/tool-proposal.type'; +import { ConversationsService } from '../conversations/conversations.service'; import type { AgentSseEvent } from './agent-sse.types'; const SYSTEM_PROMPT = `You are DocMind, an AI assistant that can use tools to help users explore their documents. @@ -37,6 +38,14 @@ const AgentState = Annotation.Root({ error?: string; } | null>(), proposal: Annotation(), + // Set when a tool result already contains a final answer + citations + // (e.g. query_documents). Short-circuits the loop instead of routing + // back through modelTurn, so the [N] markers in `answer` stay aligned + // with the `citations` array we emit alongside it. + finalAnswerOverride: Annotation({ + value: (_: CitedAnswer | null, b: CitedAnswer | null) => b, + default: () => null, + }), iterationCount: Annotation({ value: (_: number, b: number) => b, default: () => 0, @@ -49,6 +58,21 @@ interface ParsedAction { toolCall: { name: string; params: unknown } | null; } +/** Shape returned by citation-bearing tools (e.g. query_documents). */ +interface CitedAnswer { + answer: string; + citations: unknown[]; +} + +function isCitedAnswer(value: unknown): value is CitedAnswer { + return ( + typeof value === 'object' && + value !== null && + typeof (value as Record).answer === 'string' && + Array.isArray((value as Record).citations) + ); +} + @Injectable() export class AgentService { private readonly logger = new Logger(AgentService.name); @@ -59,6 +83,7 @@ export class AgentService { @Inject(GENERATION_PROVIDER) private readonly provider: GenerationProvider, private readonly config: ConfigService, private readonly eventEmitter: EventEmitter2, + private readonly conversations: ConversationsService, ) { this.maxIterations = this.config.get('AGENT_MAX_ITERATIONS') ?? 10; } @@ -72,13 +97,26 @@ export class AgentService { * the dispatched tool is `external_write`; the stream consumer emits * `confirmation_required` and stops. Resume happens via POST /agent/confirm * which calls ToolRegistryService.executeConfirmed() independently. + * + * Multi-turn: prior messages for `conversationId` are loaded (trimmed via + * ConversationsService/history.util) and seeded into the graph's initial + * `messages` state ahead of the new query. The user's query is persisted + * immediately; the final assistant answer is persisted once produced — + * either via the normal final-answer path or the citation short-circuit + * below. Tool-call/tool-result scaffolding is never persisted (see + * ConversationsService for rationale), and nothing is persisted if the + * turn pauses on an external_write confirmation (no final answer yet). */ async run( query: string, userId: string, + conversationId: string, emit: (event: AgentSseEvent) => void, queryId?: string, ): Promise { + const history = await this.conversations.loadHistory(conversationId); + await this.conversations.appendUserMessage(conversationId, query); + const toolList = this.toolRegistry .listTools() .map((t) => `- ${t.name}: ${t.description}`) @@ -142,6 +180,18 @@ export class AgentService { }; } + // Citation-bearing tool result: short-circuit to a final answer + // instead of looping back through modelTurn (see finalAnswerOverride + // comment above for why we don't let the outer model re-synthesize). + if (!error && isCitedAnswer(dispatchResult)) { + return { + finalAnswerOverride: dispatchResult, + toolResult: { toolName: name, result: dispatchResult }, + proposal: null, + pendingToolCall: null, + }; + } + const toolResultMsg: ChatMessage = error ? { role: 'user', @@ -173,6 +223,7 @@ export class AgentService { const routeAfterToolDispatch = (state: AgentStateType): string => { if (state.proposal) return 'proposalPending'; + if (state.finalAnswerOverride) return 'citedAnswer'; if (state.iterationCount >= maxIter) return 'maxReached'; return 'loop'; }; @@ -188,6 +239,7 @@ export class AgentService { .addConditionalEdges('toolDispatch', routeAfterToolDispatch, { loop: 'modelTurn', proposalPending: END, + citedAnswer: END, maxReached: END, }) .addEdge(START, 'modelTurn') @@ -198,10 +250,13 @@ export class AgentService { let currentIterationCount = 0; let lastNode = ''; let proposalEmitted = false; + let citedAnswerEmitted = false; + let finalAnswerText: string | null = null; + let finalCitations: unknown[] | undefined; for await (const stepOutput of await graph.stream( { - messages: [{ role: 'user', content: query }], + messages: [...history, { role: 'user', content: query }], systemPrompt, }, { streamMode: 'updates' }, @@ -233,6 +288,7 @@ export class AgentService { EMBEDDED_TOOL_JSON_RE.test(answer) ? "I couldn't complete that request." : answer; + finalAnswerText = safeAnswer; for (const token of safeAnswer.split(' ')) { emit({ type: 'token', data: token + ' ' }); } @@ -247,6 +303,18 @@ export class AgentService { break; } + if (update.finalAnswerOverride) { + const { answer, citations } = update.finalAnswerOverride; + finalAnswerText = answer; + finalCitations = citations; + emit({ type: 'citations', data: citations }); + for (const token of answer.split(' ')) { + emit({ type: 'token', data: token + ' ' }); + } + citedAnswerEmitted = true; + continue; // graph already routed to END; let the stream finish naturally + } + const tr = update.toolResult; if (tr?.error) { emit({ @@ -265,7 +333,7 @@ export class AgentService { } } - if (proposalEmitted) { + const emitTurnCompleted = (): void => { this.eventEmitter.emit('TurnCompleted', { userId, queryId, @@ -276,27 +344,44 @@ export class AgentService { cacheFlags: { embeddingHit: false, answerHit: false }, toolCallAuditIds: [], }); + }; + + if (proposalEmitted) { + emitTurnCompleted(); + return; + } + + if (citedAnswerEmitted) { + if (finalAnswerText !== null) { + await this.conversations.appendAssistantMessage( + conversationId, + finalAnswerText, + finalCitations, + ); + } + emit({ type: 'done', data: '' }); + emitTurnCompleted(); return; } if (lastNode === 'toolDispatch' && currentIterationCount >= maxIter) { + await this.conversations.appendAssistantMessage( + conversationId, + '(reached max tool iterations without a final answer)', + ); emit({ type: 'done', data: 'max_iterations_reached' }); return; } // Final answer path + if (finalAnswerText !== null) { + await this.conversations.appendAssistantMessage( + conversationId, + finalAnswerText, + ); + } emit({ type: 'done', data: '' }); - - this.eventEmitter.emit('TurnCompleted', { - userId, - queryId, - query, - provider: this.provider.model, - model: this.provider.model, - latencyBreakdown: { total: Date.now() - startMs }, - cacheFlags: { embeddingHit: false, answerHit: false }, - toolCallAuditIds: [], - }); + emitTurnCompleted(); } // ── Private helpers ───────────────────────────────────────────── diff --git a/backend/src/modules/conversations/conversations.controller.ts b/backend/src/modules/conversations/conversations.controller.ts new file mode 100644 index 0000000..4d837ae --- /dev/null +++ b/backend/src/modules/conversations/conversations.controller.ts @@ -0,0 +1,26 @@ +import { Controller, Get, Param } from '@nestjs/common'; +import { ApiOperation, ApiParam, ApiTags } from '@nestjs/swagger'; +import { + CurrentUser, + JwtPayload, +} from '../../common/decorators/current-user.decorator'; +import { ConversationsService } from './conversations.service'; + +@ApiTags('conversations') +@Controller('v1/conversations') +export class ConversationsController { + constructor(private readonly conversations: ConversationsService) {} + + @Get() + @ApiOperation({ summary: "List the current user's conversations" }) + list(@CurrentUser() user: JwtPayload) { + return this.conversations.listForUser(user.sub); + } + + @Get(':id/messages') + @ApiOperation({ summary: 'Get full message history for a conversation' }) + @ApiParam({ name: 'id' }) + getMessages(@CurrentUser() user: JwtPayload, @Param('id') id: string) { + return this.conversations.getMessages(user.sub, id); + } +} diff --git a/backend/src/modules/conversations/conversations.module.ts b/backend/src/modules/conversations/conversations.module.ts new file mode 100644 index 0000000..9ccf166 --- /dev/null +++ b/backend/src/modules/conversations/conversations.module.ts @@ -0,0 +1,12 @@ +import { Module } from '@nestjs/common'; +import { PrismaModule } from '../../prisma/prisma.module'; +import { ConversationsController } from './conversations.controller'; +import { ConversationsService } from './conversations.service'; + +@Module({ + imports: [PrismaModule], + controllers: [ConversationsController], + providers: [ConversationsService], + exports: [ConversationsService], +}) +export class ConversationsModule {} diff --git a/backend/src/modules/conversations/conversations.service.ts b/backend/src/modules/conversations/conversations.service.ts new file mode 100644 index 0000000..2c88ffc --- /dev/null +++ b/backend/src/modules/conversations/conversations.service.ts @@ -0,0 +1,100 @@ +import { Injectable, NotFoundException } from '@nestjs/common'; +import { Prisma } from '../../../generated/prisma/client'; +import { PrismaService } from '../../prisma/prisma.service'; +import type { ChatMessage } from '../providers/generation.provider'; +import { trimHistory } from './history.util'; + +@Injectable() +export class ConversationsService { + constructor(private readonly prisma: PrismaService) {} + + async create(userId: string) { + return this.prisma.conversation.create({ data: { userId } }); + } + + async listForUser(userId: string) { + return this.prisma.conversation.findMany({ + where: { userId }, + orderBy: { updatedAt: 'desc' }, + }); + } + + /** Throws NotFoundException if the conversation doesn't exist or isn't owned by userId. */ + async assertOwnership(userId: string, conversationId: string): Promise { + const conversation = await this.prisma.conversation.findFirst({ + where: { id: conversationId, userId }, + select: { id: true }, + }); + if (!conversation) { + throw new NotFoundException(`Conversation ${conversationId} not found`); + } + } + + /** Full raw message history for a conversation (for reload/display), ownership-checked. */ + async getMessages(userId: string, conversationId: string) { + await this.assertOwnership(userId, conversationId); + return this.prisma.message.findMany({ + where: { conversationId }, + orderBy: { createdAt: 'asc' }, + }); + } + + /** + * Loads and trims history for replay into the agent's `messages` state. + * Caller is responsible for ownership checks (agent.chat already resolves + * conversationId via the controller, which validates or creates it). + */ + async loadHistory(conversationId: string): Promise { + const rows = await this.prisma.message.findMany({ + where: { conversationId }, + orderBy: { createdAt: 'asc' }, + select: { role: true, content: true }, + }); + const messages: ChatMessage[] = rows.map((r) => ({ + role: r.role, + content: r.content, + })); + return trimHistory(messages); + } + + async appendUserMessage( + conversationId: string, + content: string, + ): Promise { + await this.appendMessage(conversationId, 'user', content); + } + + async appendAssistantMessage( + conversationId: string, + content: string, + citations?: unknown[], + ): Promise { + await this.appendMessage(conversationId, 'assistant', content, citations); + } + + private async appendMessage( + conversationId: string, + role: 'user' | 'assistant', + content: string, + citations?: unknown[], + ): Promise { + await this.prisma.$transaction([ + this.prisma.message.create({ + data: { + conversationId, + role, + content, + citations: citations + ? (citations as Prisma.InputJsonValue) + : undefined, + }, + }), + // Bump updatedAt so conversations sort by recent activity (used by + // the future conversation-list sidebar). + this.prisma.conversation.update({ + where: { id: conversationId }, + data: { updatedAt: new Date() }, + }), + ]); + } +} diff --git a/backend/src/modules/conversations/history.util.spec.ts b/backend/src/modules/conversations/history.util.spec.ts new file mode 100644 index 0000000..bbb2e83 --- /dev/null +++ b/backend/src/modules/conversations/history.util.spec.ts @@ -0,0 +1,64 @@ +import { trimHistory } from './history.util'; +import type { ChatMessage } from '../providers/generation.provider'; + +function msg(role: 'user' | 'assistant', content: string): ChatMessage { + return { role, content }; +} + +describe('trimHistory', () => { + it('returns everything when under both caps', () => { + const messages = [msg('user', 'hi'), msg('assistant', 'hello')]; + expect(trimHistory(messages)).toEqual(messages); + }); + + it('caps at maxMessages, keeping the most recent ones', () => { + const messages = Array.from({ length: 30 }, (_, i) => + msg(i % 2 === 0 ? 'user' : 'assistant', `msg-${i}`), + ); + const result = trimHistory(messages, { maxMessages: 10 }); + expect(result).toHaveLength(10); + expect(result[0].content).toBe('msg-20'); + expect(result[result.length - 1].content).toBe('msg-29'); + }); + + it('trims older messages once the token budget is exceeded', () => { + // Each message ~250 chars ≈ 63 tokens. Budget 100 tokens → keep ~1-2. + const long = 'x'.repeat(250); + const messages = [ + msg('user', long), + msg('assistant', long), + msg('user', long), + msg('assistant', long), + ]; + const result = trimHistory(messages, { + maxMessages: 20, + tokenBudget: 100, + }); + expect(result.length).toBeLessThan(messages.length); + // Must keep the newest message + expect(result[result.length - 1]).toEqual(messages[messages.length - 1]); + }); + + it('always keeps at least the latest turn even if it alone exceeds the token budget', () => { + const huge = 'x'.repeat(20000); // ~5000 tokens, way over any small budget + const messages = [ + msg('user', 'earlier question'), + msg('assistant', 'earlier answer'), + msg('user', huge), + ]; + const result = trimHistory(messages, { maxMessages: 20, tokenBudget: 50 }); + // minKeep = 2 → keeps at least the last 2 messages regardless of size + expect(result).toHaveLength(2); + expect(result[result.length - 1].content).toBe(huge); + }); + + it('applies the hard ceiling before the token budget pass', () => { + const messages = Array.from({ length: 25 }, (_, i) => msg('user', `m${i}`)); + const result = trimHistory(messages, { + maxMessages: 5, + tokenBudget: 1_000_000, // budget effectively irrelevant here + }); + expect(result).toHaveLength(5); + expect(result[0].content).toBe('m20'); + }); +}); diff --git a/backend/src/modules/conversations/history.util.ts b/backend/src/modules/conversations/history.util.ts new file mode 100644 index 0000000..a15d4ea --- /dev/null +++ b/backend/src/modules/conversations/history.util.ts @@ -0,0 +1,60 @@ +import type { ChatMessage } from '../providers/generation.provider'; + +/** Hard ceiling on the number of prior messages loaded into agent history. */ +export const HISTORY_MAX_MESSAGES = 20; + +/** Approximate token budget for the trimmed history window. */ +export const HISTORY_TOKEN_BUDGET = 3000; + +/** Rough chars-per-token heuristic — no tokenizer dependency needed for v1. */ +const CHARS_PER_TOKEN = 4; + +export interface TrimHistoryOptions { + maxMessages?: number; + tokenBudget?: number; +} + +/** + * Trims conversation history for replay into the agent's `messages` state. + * + * Two passes, applied in order: + * 1. Hard ceiling — keep at most the last `maxMessages` rows. Bounds the + * DB read and gives a cheap worst case regardless of message length. + * 2. Token budget — walk the remaining window newest-to-oldest, dropping + * older messages once the approximate token budget is exceeded. + * + * Always keeps at least the latest turn (last 2 messages, or fewer if the + * conversation doesn't have that many yet), even if it alone exceeds the + * token budget — a single long turn should never be silently dropped. + */ +export function trimHistory( + messages: ChatMessage[], + options: TrimHistoryOptions = {}, +): ChatMessage[] { + const maxMessages = options.maxMessages ?? HISTORY_MAX_MESSAGES; + const tokenBudget = options.tokenBudget ?? HISTORY_TOKEN_BUDGET; + + const capped = + messages.length > maxMessages ? messages.slice(-maxMessages) : messages; + + const minKeep = Math.min(2, capped.length); + + let tokens = 0; + let keepFrom = 0; + + for (let i = capped.length - 1; i >= 0; i--) { + const approxTokens = Math.ceil(capped[i].content.length / CHARS_PER_TOKEN); + const keptSoFar = capped.length - i; + const wouldExceed = tokens + approxTokens > tokenBudget; + + if (wouldExceed && keptSoFar > minKeep) { + keepFrom = i + 1; + break; + } + + tokens += approxTokens; + keepFrom = i; + } + + return capped.slice(keepFrom); +} diff --git a/backend/test/ownership.integration.spec.ts b/backend/test/ownership.integration.spec.ts index a59c65a..2a670af 100644 --- a/backend/test/ownership.integration.spec.ts +++ b/backend/test/ownership.integration.spec.ts @@ -16,7 +16,8 @@ // is imported below), so these must be set BEFORE the import to pass // ConfigModule validation. Actual values (testcontainer port, etc.) are // overwritten in beforeAll. -process.env['DATABASE_URL'] = 'postgresql://placeholder:placeholder@localhost:9999/placeholder'; +process.env['DATABASE_URL'] = + 'postgresql://placeholder:placeholder@localhost:9999/placeholder'; process.env['REDIS_HOST'] = 'localhost'; process.env['REDIS_PORT'] = '6399'; process.env['REDIS_URL'] = 'redis://localhost:6399'; diff --git a/docker-compose.yml b/docker-compose.yml index 61a2c39..88717df 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -87,7 +87,7 @@ services: migrate: condition: service_completed_successfully healthcheck: - test: [ "CMD", "wget", "-qO-", "http://localhost:4500/health" ] + test: [ "CMD", "wget", "-qO-", "http://127.0.0.1:4500/health" ] interval: 15s timeout: 5s retries: 5 diff --git a/frontend/20260727052929_add_conversations/migration.sql b/frontend/20260727052929_add_conversations/migration.sql new file mode 100644 index 0000000..fa39fdc --- /dev/null +++ b/frontend/20260727052929_add_conversations/migration.sql @@ -0,0 +1,8 @@ +-- DropIndex +DROP INDEX "chunks_content_tsv_idx"; + +-- DropIndex +DROP INDEX "idx_chunks_embedding_hnsw"; + +-- AlterTable +ALTER TABLE "chunks" ALTER COLUMN "content_tsv" DROP DEFAULT; diff --git a/frontend/src/app/chat/page.tsx b/frontend/src/app/chat/page.tsx index 0734a5d..92f67f3 100644 --- a/frontend/src/app/chat/page.tsx +++ b/frontend/src/app/chat/page.tsx @@ -1,7 +1,7 @@ 'use client'; import { useState } from 'react'; -import { useChatStream } from '@/hooks/useChatStream'; +import { useChatStream, type ChatMessage } from '@/hooks/useChatStream'; import { ConfirmationCard } from '@/components/ConfirmationCard'; import type { Citation } from '@/types/api'; @@ -59,31 +59,129 @@ function AnswerWithCitations({ ); } +function MessageBubble({ message }: { message: ChatMessage }) { + const isUser = message.role === 'user'; + + return ( +
+
+ {isUser ? ( +

{message.content}

+ ) : ( + <> + {message.content.length > 0 ? ( + + ) : message.streaming ? ( + + ) : null} + {message.streaming && message.content.length > 0 && ( + + )} + + {message.citations.length > 0 && ( +
+

+ Sources ({message.citations.length}) +

+
+ {message.citations.map((citation) => ( +
+ + {citation.marker} · {citation.documentTitle} + +
+ {citation.snippet} +
+
+ ))} +
+
+ )} + + )} +
+
+ ); +} + export default function ChatPage() { const [query, setQuery] = useState(''); const { - content, - citations, + messages, loading, error, pendingConfirmation, ask, abort, clearConfirmation, + startNewConversation, } = useChatStream(); const handleSubmit = (e: React.FormEvent) => { e.preventDefault(); if (!query.trim()) return; ask(query.trim()); + setQuery(''); }; - const hasAnswer = content.length > 0; - return ( -
-

Chat

-

Ask questions about your ingested documents.

+
+
+
+

Chat

+

+ Ask questions about your ingested documents. Follow-up questions keep the conversation's context. +

+
+ {messages.length > 0 && ( + + )} +
+ +
+ {messages.map((message) => ( + + ))} + + {messages.length === 0 && !loading && !error && ( +
+ No messages yet. Ask a question below. +
+ )} +
+ + {error && ( +
+ {error} +
+ )} + + {pendingConfirmation && ( + { + // After confirmation, the tool executes server-side. + // The SSE stream has ended, so we just clean up. + clearConfirmation(); + }} + onCancel={clearConfirmation} + /> + )}
)}
- - {error && ( -
- {error} -
- )} - - {pendingConfirmation && ( - { - // After confirmation, the tool executes server-side. - // The SSE stream has ended, so we just clean up. - clearConfirmation(); - }} - onCancel={clearConfirmation} - /> - )} - - {hasAnswer && ( -
-
- - {loading && ( - - )} -
- - {citations.length > 0 && ( -
-

- Sources ({citations.length}) -

-
- {citations.map((citation) => ( -
- - {citation.marker} · {citation.documentTitle} - -
- {citation.snippet} -
-
- ))} -
-
- )} -
- )} - - {!hasAnswer && !loading && !error && ( -
- No answer yet. Ask a question above. -
- )}
); -} +} \ No newline at end of file diff --git a/frontend/src/hooks/useChatStream.ts b/frontend/src/hooks/useChatStream.ts index dd04477..e3771d9 100644 --- a/frontend/src/hooks/useChatStream.ts +++ b/frontend/src/hooks/useChatStream.ts @@ -11,7 +11,17 @@ export interface ToolProposal { confirmationToken: string; } +export interface ChatMessage { + id: string; + role: 'user' | 'assistant'; + content: string; + citations: Citation[]; + /** True while this assistant message's tokens are still streaming in. */ + streaming?: boolean; +} + type StreamEvent = + | { type: 'conversation_started'; data: { conversationId: string } } | { type: 'citations'; data: Citation[] } | { type: 'token'; data: string } | { type: 'done'; data: string } @@ -21,29 +31,36 @@ type StreamEvent = | { type: 'confirmation_required'; data: ToolProposal }; interface ChatStreamState { - content: string; - citations: Citation[]; + conversationId: string | null; + messages: ChatMessage[]; loading: boolean; error: string | null; pendingConfirmation: ToolProposal | null; } interface UseChatStreamReturn extends ChatStreamState { - ask: (query: string, topK?: number) => void; + ask: (query: string) => void; abort: () => void; clearConfirmation: () => void; + /** Starts a brand-new conversation on the next ask() instead of continuing the current one. */ + startNewConversation: () => void; +} + +function makeId(): string { + return `${Date.now()}-${Math.random().toString(36).slice(2)}`; } export function useChatStream(): UseChatStreamReturn { const [state, setState] = useState({ - content: '', - citations: [], + conversationId: null, + messages: [], loading: false, error: null, pendingConfirmation: null, }); const abortRef = useRef(null); + const conversationIdRef = useRef(null); const abort = useCallback(() => { abortRef.current?.abort(); @@ -55,20 +72,64 @@ export function useChatStream(): UseChatStreamReturn { setState((prev) => ({ ...prev, pendingConfirmation: null })); }, []); - const ask = useCallback((query: string, topK?: number) => { + const startNewConversation = useCallback(() => { + conversationIdRef.current = null; + setState({ + conversationId: null, + messages: [], + loading: false, + error: null, + pendingConfirmation: null, + }); + }, []); + + const ask = useCallback((query: string) => { // Cancel any in-flight request abortRef.current?.abort(); const controller = new AbortController(); abortRef.current = controller; - setState({ content: '', citations: [], loading: true, error: null, pendingConfirmation: null }); + const userMessage: ChatMessage = { + id: makeId(), + role: 'user', + content: query, + citations: [], + }; + const assistantId = makeId(); + const assistantMessage: ChatMessage = { + id: assistantId, + role: 'assistant', + content: '', + citations: [], + streaming: true, + }; + + setState((prev) => ({ + ...prev, + messages: [...prev.messages, userMessage, assistantMessage], + loading: true, + error: null, + pendingConfirmation: null, + })); + + const updateAssistant = (patch: Partial) => { + setState((prev) => ({ + ...prev, + messages: prev.messages.map((m) => + m.id === assistantId ? { ...m, ...patch } : m, + ), + })); + }; void (async () => { try { const res = await fetch(`${API_BASE_URL}/v1/agent/chat`, { method: 'POST', headers: { 'Content-Type': 'application/json', ...(await getAuthHeaders()) }, - body: JSON.stringify({ query, topK }), + body: JSON.stringify({ + query, + conversationId: conversationIdRef.current ?? undefined, + }), signal: controller.signal, }); @@ -100,19 +161,35 @@ export function useChatStream(): UseChatStreamReturn { continue; } - if (event.type === 'citations') { - setState((prev) => ({ ...prev, citations: event.data as Citation[] })); + if (event.type === 'conversation_started') { + conversationIdRef.current = event.data.conversationId; + setState((prev) => ({ + ...prev, + conversationId: event.data.conversationId, + })); + } else if (event.type === 'citations') { + updateAssistant({ citations: event.data as Citation[] }); } else if (event.type === 'token') { - setState((prev) => ({ ...prev, content: prev.content + (event.data as string) })); + setState((prev) => ({ + ...prev, + messages: prev.messages.map((m) => + m.id === assistantId + ? { ...m, content: m.content + (event.data as string) } + : m, + ), + })); } else if (event.type === 'done') { + updateAssistant({ streaming: false }); setState((prev) => ({ ...prev, loading: false })); } else if (event.type === 'error') { + updateAssistant({ streaming: false }); setState((prev) => ({ ...prev, error: event.data as string, loading: false, })); } else if (event.type === 'confirmation_required') { + updateAssistant({ streaming: false }); setState((prev) => ({ ...prev, loading: false, @@ -123,6 +200,7 @@ export function useChatStream(): UseChatStreamReturn { } } catch (err: unknown) { if (err instanceof Error && err.name === 'AbortError') return; + updateAssistant({ streaming: false }); setState((prev) => ({ ...prev, error: err instanceof Error ? err.message : 'Stream failed', @@ -132,5 +210,5 @@ export function useChatStream(): UseChatStreamReturn { })(); }, []); - return { ...state, ask, abort, clearConfirmation }; + return { ...state, ask, abort, clearConfirmation, startNewConversation }; }