diff --git a/apps/cli/src/ai/chat-command.test.ts b/apps/cli/src/ai/chat-command.test.ts new file mode 100644 index 00000000..ca6b1a5b --- /dev/null +++ b/apps/cli/src/ai/chat-command.test.ts @@ -0,0 +1,120 @@ +import { describe, it } from 'node:test' +import assert from 'node:assert/strict' +import { Writable } from 'node:stream' +import type { AIChatManager, DatabaseManager } from '@openchatlab/node-runtime' +import { resolveAIChatTarget, runChatTurn } from './chat-command' + +function createDbManager(sessionIds: string[]): DatabaseManager { + return { + listSessionIds: () => sessionIds, + open: (sessionId: string) => (sessionIds.includes(sessionId) ? {} : null), + } as unknown as DatabaseManager +} + +function createAIChatManager(existing: Array<{ id: string; sessionId: string }> = []): AIChatManager { + const chats = new Map( + existing.map((chat) => [chat.id, { ...chat, title: null, assistantId: 'general_cn' }]) + ) + const messages: Array<{ aiChatId: string; role: string; content: string }> = [] + + return { + getAIChat: (aiChatId: string) => chats.get(aiChatId) ?? null, + createAIChat: (sessionId: string, title: string | undefined, assistantId: string) => { + const id = `ai_chat_${chats.size + 1}` + const chat = { id, sessionId, title: title ?? null, assistantId } + chats.set(id, chat) + return chat + }, + addMessage: (aiChatId: string, role: string, content: string) => { + messages.push({ aiChatId, role, content }) + return { id: `msg_${messages.length}`, aiChatId, role, content, timestamp: 1 } + }, + __messages: messages, + } as unknown as AIChatManager +} + +class MemoryWritable extends Writable { + chunks: string[] = [] + + _write(chunk: Buffer | string, _encoding: BufferEncoding, callback: (error?: Error | null) => void): void { + this.chunks.push(String(chunk)) + callback() + } + + text(): string { + return this.chunks.join('') + } +} + +describe('resolveAIChatTarget', () => { + it('creates a new AI chat for an explicit session id', () => { + const target = resolveAIChatTarget( + { sessionId: 'session-1', question: 'hello' }, + { dbManager: createDbManager(['session-1']), aiChatManager: createAIChatManager() } + ) + + assert.equal(target.sessionId, 'session-1') + assert.equal(target.aiChatId, 'ai_chat_1') + assert.equal(target.created, true) + }) + + it('recovers session id from a globally unique aiChatId', () => { + const target = resolveAIChatTarget( + { aiChatId: 'ai-chat-1' }, + { + dbManager: createDbManager(['session-1']), + aiChatManager: createAIChatManager([{ id: 'ai-chat-1', sessionId: 'session-1' }]), + } + ) + + assert.deepEqual(target, { sessionId: 'session-1', aiChatId: 'ai-chat-1', created: false }) + }) + + it('rejects mismatched explicit session id and aiChatId', () => { + assert.throws( + () => + resolveAIChatTarget( + { sessionId: 'session-2', aiChatId: 'ai-chat-1' }, + { + dbManager: createDbManager(['session-1', 'session-2']), + aiChatManager: createAIChatManager([{ id: 'ai-chat-1', sessionId: 'session-1' }]), + } + ), + /belongs to session session-1/ + ) + }) +}) + +describe('runChatTurn', () => { + it('collects streamed answer and persists user and assistant messages', async () => { + const stdout = new MemoryWritable() + const aiChatManager = createAIChatManager() + const result = await runChatTurn( + { sessionId: 'session-1', question: 'hello', json: true }, + { + dbManager: createDbManager(['session-1']), + pathProvider: {} as never, + aiChatManager, + stdout, + createRunAgentStream: () => async (_params, onEvent) => { + onEvent({ type: 'content', content: 'hi' }) + onEvent({ + type: 'done', + isFinished: true, + usage: { promptTokens: 1, completionTokens: 1, totalTokens: 2 }, + }) + }, + } + ) + + assert.equal(result.sessionId, 'session-1') + assert.equal(result.aiChatId, 'ai_chat_1') + assert.equal(result.answer, 'hi') + assert.equal(result.usage.tokenUsage?.totalTokens, 2) + assert.equal(stdout.text(), '') + assert.deepEqual((aiChatManager as unknown as { __messages: unknown[] }).__messages, [ + { aiChatId: 'ai_chat_1', role: 'user', content: 'hello' }, + { aiChatId: 'ai_chat_1', role: 'assistant', content: 'hi' }, + ]) + }) +}) diff --git a/apps/cli/src/ai/chat-command.ts b/apps/cli/src/ai/chat-command.ts new file mode 100644 index 00000000..750fb917 --- /dev/null +++ b/apps/cli/src/ai/chat-command.ts @@ -0,0 +1,184 @@ +import { createInterface } from 'node:readline/promises' +import type { Readable, Writable } from 'node:stream' +import type { DatabaseManager, AIChatManager, AgentStreamChunk, TokenUsageData } from '@openchatlab/node-runtime' +import type { PathProvider } from '@openchatlab/core' +import { createCliRunAgentStream } from './agent-stream-runner' + +export interface ChatCommandOptions { + sessionId?: string + aiChatId?: string + question?: string + json?: boolean + stream?: boolean + locale?: string +} + +export interface ChatCommandDeps { + dbManager: DatabaseManager + pathProvider: PathProvider + aiChatManager: AIChatManager + stdout?: Pick + stderr?: Pick + stdin?: Readable + createRunAgentStream?: typeof createCliRunAgentStream +} + +export interface ResolvedAIChatTarget { + sessionId: string + aiChatId: string + created: boolean +} + +export interface ChatTurnResult { + sessionId: string + aiChatId: string + question: string + answer: string + usage: { + durationMs: number + tokenUsage: TokenUsageData | null + } +} + +function write(stream: Pick, text: string): void { + stream.write(text) +} + +function buildTitle(question?: string): string { + const trimmed = question?.trim() + if (!trimmed) return 'CLI Chat' + return trimmed.length > 50 ? `${trimmed.slice(0, 50)}...` : trimmed +} + +function resolveSingleSessionId(dbManager: DatabaseManager): string { + const sessionIds = dbManager.listSessionIds() + if (sessionIds.length === 1) return sessionIds[0]! + if (sessionIds.length === 0) { + throw new Error('No chat sessions found. Import data first or run `chatlab sessions` to verify available sessions.') + } + throw new Error('Missing --session-id. Run `chatlab sessions` to choose a session ID.') +} + +function assertSessionExists(dbManager: DatabaseManager, sessionId: string): void { + const db = dbManager.open(sessionId) + if (!db) { + throw new Error(`Session ${sessionId} not found`) + } +} + +export function resolveAIChatTarget( + options: Pick, + deps: Pick +): ResolvedAIChatTarget { + if (options.aiChatId) { + const aiChat = deps.aiChatManager.getAIChat(options.aiChatId) + if (!aiChat) { + throw new Error(`AI chat ${options.aiChatId} not found`) + } + if (options.sessionId && options.sessionId !== aiChat.sessionId) { + throw new Error(`AI chat ${options.aiChatId} belongs to session ${aiChat.sessionId}, not ${options.sessionId}`) + } + assertSessionExists(deps.dbManager, aiChat.sessionId) + return { sessionId: aiChat.sessionId, aiChatId: aiChat.id, created: false } + } + + const sessionId = options.sessionId ?? resolveSingleSessionId(deps.dbManager) + assertSessionExists(deps.dbManager, sessionId) + const aiChat = deps.aiChatManager.createAIChat(sessionId, buildTitle(options.question), 'general_cn') + return { sessionId, aiChatId: aiChat.id, created: true } +} + +export async function runChatTurn( + options: Required> & + Pick, + deps: ChatCommandDeps +): Promise { + const stdout = deps.stdout ?? process.stdout + const target = resolveAIChatTarget(options, deps) + const startedAt = Date.now() + let answer = '' + let tokenUsage: TokenUsageData | null = null + + const runAgentStream = (deps.createRunAgentStream ?? createCliRunAgentStream)(deps.dbManager, deps.aiChatManager) + + await runAgentStream( + { + userMessage: options.question, + sessionId: target.sessionId, + aiChatId: target.aiChatId, + chatType: 'group', + locale: options.locale ?? 'zh-CN', + }, + (chunk: AgentStreamChunk) => { + if (chunk.type === 'content' && chunk.content) { + answer += chunk.content + if (!options.json && options.stream !== false) write(stdout, chunk.content) + } + if (chunk.type === 'done' && chunk.usage) { + tokenUsage = chunk.usage + } + }, + new AbortController().signal + ) + + if (!options.json && options.stream !== false && answer && !answer.endsWith('\n')) { + write(stdout, '\n') + } + + deps.aiChatManager.addMessage(target.aiChatId, 'user', options.question) + deps.aiChatManager.addMessage( + target.aiChatId, + 'assistant', + answer, + undefined, + undefined, + undefined, + tokenUsage ?? undefined + ) + + return { + sessionId: target.sessionId, + aiChatId: target.aiChatId, + question: options.question, + answer, + usage: { + durationMs: Date.now() - startedAt, + tokenUsage, + }, + } +} + +export async function runChatCommand(options: ChatCommandOptions, deps: ChatCommandDeps): Promise { + const stdout = deps.stdout ?? process.stdout + const stderr = deps.stderr ?? process.stderr + + if (options.question?.trim()) { + const result = await runChatTurn({ ...options, question: options.question.trim() }, deps) + if (options.json) { + write(stdout, `${JSON.stringify(result, null, 2)}\n`) + } else if (options.stream === false) { + write(stdout, `${result.answer}\n`) + } + return + } + + const rl = createInterface({ + input: deps.stdin ?? process.stdin, + output: stdout as Writable, + }) + + try { + let aiChatId = options.aiChatId + while (true) { + const question = (await rl.question('chatlab> ')).trim() + if (!question || question === 'exit' || question === 'quit') break + const result = await runChatTurn({ ...options, aiChatId, question, json: false, stream: true }, deps) + aiChatId = result.aiChatId + } + } catch (error) { + write(stderr, `${error instanceof Error ? error.message : String(error)}\n`) + throw error + } finally { + rl.close() + } +} diff --git a/apps/cli/src/cli.ts b/apps/cli/src/cli.ts index 62f1b568..e9153cad 100644 --- a/apps/cli/src/cli.ts +++ b/apps/cli/src/cli.ts @@ -11,6 +11,7 @@ import { DEFAULT_API_PORT, loadConfig, getConfigPath } from '@openchatlab/config import { NodePathProvider, DatabaseManager, + AIChatManager, applyPendingNodeDataDirMigrationIfNeeded, hasPendingElectronDataWarning, verifyCliDataPath, @@ -299,6 +300,41 @@ program await startCliMcpServer() }) +program + .command('chat') + .description('Ask ChatLab AI about an imported chat session') + .option('--session-id ', 'Source chat session ID') + .option('--ai-chat-id ', 'Existing AI chat ID to continue') + .option('-q, --question ', 'Question to ask') + .option('--json', 'Output structured JSON') + .option('--no-stream', 'Disable streaming output') + .option('--locale ', 'AI response locale', 'zh-CN') + .action(async (options) => { + const { runChatCommand } = await import('./ai/chat-command') + const { dbManager, pathProvider } = initRuntime() + const aiChatManager = new AIChatManager(pathProvider.getAiDataDir(), { nativeBinding: resolveNativeBinding() }) + + try { + await runChatCommand( + { + sessionId: options.sessionId, + aiChatId: options.aiChatId, + question: options.question, + json: !!options.json, + stream: options.stream, + locale: options.locale, + }, + { dbManager, pathProvider, aiChatManager } + ) + } catch (error) { + console.error(error instanceof Error ? error.message : String(error)) + process.exitCode = 1 + } finally { + aiChatManager.close() + dbManager.closeAll() + } + }) + program .command('start') .description('Start ChatLab (HTTP API + Web UI)')