feat: add cli ai chat command

This commit is contained in:
digua committed 2026-06-07 20:27:27 +08:00
1 parent c2295e317f
commit 9b17b09354
3 files changed
+340

No files matched your search

+120
View File
@@ -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<string, { id: string; sessionId: string; title: string | null; assistantId: string }>(
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' },
])
})
})
+184
View File
@@ -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<Writable, 'write'>
stderr?: Pick<Writable, 'write'>
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<Writable, 'write'>, 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<ChatCommandOptions, 'sessionId' | 'aiChatId' | 'question'>,
deps: Pick<ChatCommandDeps, 'dbManager' | 'aiChatManager'>
): 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<ChatCommandOptions, 'question'>> &
Pick<ChatCommandOptions, 'sessionId' | 'aiChatId' | 'json' | 'stream' | 'locale'>,
deps: ChatCommandDeps
): Promise<ChatTurnResult> {
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<void> {
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()
}
}
+36
View File
@@ -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 <id>', 'Source chat session ID')
.option('--ai-chat-id <id>', 'Existing AI chat ID to continue')
.option('-q, --question <text>', 'Question to ask')
.option('--json', 'Output structured JSON')
.option('--no-stream', 'Disable streaming output')
.option('--locale <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)')