mirror of
https://github.com/lobehub/lobe-chat.git
synced 2026-06-14 03:30:19 +00:00
✨ feat: support history context auto compress (#11790)
* add compress implement * push * update * update * update * fix auto scroll * db schema update * fix compress * fix types * fix lint * update get compressedMessages * update content * update * fix tests * fix tests * fix tests
This commit is contained in:
Arvin Xu
@@ -249,13 +249,6 @@ jobs:
|
||||
- name: Lint
|
||||
run: npm run lint
|
||||
|
||||
- name: Test Client DB
|
||||
run: pnpm --filter @lobechat/database test:client-db
|
||||
env:
|
||||
KEY_VAULTS_SECRET: Kix2wcUONd4CX51E/ZPAd36BqM4wzJgKjPtz2sGztqQ=
|
||||
S3_PUBLIC_DOMAIN: https://example.com
|
||||
APP_URL: https://home.com
|
||||
|
||||
- name: Test Coverage
|
||||
run: pnpm --filter @lobechat/database test:coverage
|
||||
env:
|
||||
|
||||
@@ -492,6 +492,7 @@ table message_groups {
|
||||
type text
|
||||
content text
|
||||
editor_data jsonb
|
||||
metadata jsonb
|
||||
client_id varchar(255)
|
||||
accessed_at "timestamp with time zone" [not null, default: `now()`]
|
||||
created_at "timestamp with time zone" [not null, default: `now()`]
|
||||
|
||||
@@ -4,6 +4,7 @@ import { DEFAULT_SECURITY_BLACKLIST, InterventionChecker } from '../core';
|
||||
import {
|
||||
Agent,
|
||||
AgentInstruction,
|
||||
AgentInstructionCompressContext,
|
||||
AgentRuntimeContext,
|
||||
AgentState,
|
||||
GeneralAgentCallLLMInstructionPayload,
|
||||
@@ -11,11 +12,13 @@ import {
|
||||
GeneralAgentCallToolResultPayload,
|
||||
GeneralAgentCallToolsBatchInstructionPayload,
|
||||
GeneralAgentCallingToolInstructionPayload,
|
||||
GeneralAgentCompressionResultPayload,
|
||||
GeneralAgentConfig,
|
||||
HumanAbortPayload,
|
||||
TaskResultPayload,
|
||||
TasksBatchResultPayload,
|
||||
} from '../types';
|
||||
import { shouldCompress } from '../utils/tokenCounter';
|
||||
|
||||
/**
|
||||
* ChatAgent - The "Brain" of the chat agent
|
||||
@@ -219,6 +222,26 @@ export class GeneralChatAgent implements Agent {
|
||||
return { hasToolsCalling, parentMessageId, toolsCalling };
|
||||
}
|
||||
|
||||
/**
|
||||
* Find existing compression summary from messages
|
||||
* Looks for MessageGroup with type 'compression' and extracts its content
|
||||
*/
|
||||
private findExistingSummary(messages: any[]): string | undefined {
|
||||
// Look for compression group summary in messages
|
||||
// The summary is typically stored as a system message with compression metadata
|
||||
// or as a MessageGroup content field
|
||||
for (const msg of messages) {
|
||||
if (msg.role === 'system' && msg.metadata?.compressionSummary) {
|
||||
return msg.content;
|
||||
}
|
||||
// Check for MessageGroup type compression
|
||||
if (msg.messageGroupType === 'compression' && msg.content) {
|
||||
return msg.content;
|
||||
}
|
||||
}
|
||||
return undefined;
|
||||
}
|
||||
|
||||
/**
|
||||
* Handle abort scenario - unified abort handling logic
|
||||
*/
|
||||
@@ -260,6 +283,23 @@ export class GeneralChatAgent implements Agent {
|
||||
switch (context.phase) {
|
||||
case 'init':
|
||||
case 'user_input': {
|
||||
// Check if context compression is needed before calling LLM
|
||||
const compressionCheck = shouldCompress(state.messages, {
|
||||
maxWindowToken: this.config.compressionConfig?.maxWindowToken,
|
||||
});
|
||||
|
||||
if (compressionCheck.needsCompression) {
|
||||
// Context exceeds threshold, compress ALL messages into a single summary
|
||||
return {
|
||||
payload: {
|
||||
currentTokenCount: compressionCheck.currentTokenCount,
|
||||
existingSummary: this.findExistingSummary(state.messages),
|
||||
messages: state.messages,
|
||||
},
|
||||
type: 'compress_context',
|
||||
} as AgentInstructionCompressContext;
|
||||
}
|
||||
|
||||
// User input received, call LLM to generate response
|
||||
// At this point, messages may have been preprocessed with RAG/Search
|
||||
return {
|
||||
@@ -494,6 +534,27 @@ export class GeneralChatAgent implements Agent {
|
||||
};
|
||||
}
|
||||
|
||||
case 'compression_result': {
|
||||
// Context compression completed, continue to call LLM
|
||||
const compressionPayload = context.payload as GeneralAgentCompressionResultPayload;
|
||||
|
||||
// If compression was skipped (no messages to compress), just call LLM
|
||||
// Otherwise, messages have been updated with compressed content
|
||||
// Pass parentMessageId and createAssistantMessage=true to force new message creation
|
||||
return {
|
||||
payload: {
|
||||
// Force create new assistant message after compression
|
||||
createAssistantMessage: true,
|
||||
messages: compressionPayload.compressedMessages,
|
||||
model: this.config.modelRuntimeConfig?.model,
|
||||
parentMessageId: compressionPayload.parentMessageId,
|
||||
provider: this.config.modelRuntimeConfig?.provider,
|
||||
tools: state.tools,
|
||||
} as GeneralAgentCallLLMInstructionPayload,
|
||||
type: 'call_llm',
|
||||
};
|
||||
}
|
||||
|
||||
case 'human_abort': {
|
||||
// User aborted the operation
|
||||
const { hasToolsCalling, parentMessageId, toolsCalling, reason } =
|
||||
|
||||
@@ -101,6 +101,7 @@ export interface AgentEventResumed {
|
||||
export interface AgentEventCompressionComplete {
|
||||
type: 'compression_complete';
|
||||
groupId: string;
|
||||
parentMessageId?: string;
|
||||
}
|
||||
|
||||
export interface AgentEventCompressionError {
|
||||
|
||||
@@ -1,6 +1,8 @@
|
||||
import { ChatToolPayload, MessageToolCall } from '@lobechat/types';
|
||||
|
||||
export interface GeneralAgentCallLLMInstructionPayload {
|
||||
/** Force create a new assistant message (e.g., after compression) */
|
||||
createAssistantMessage?: boolean;
|
||||
isFirstMessage?: boolean;
|
||||
messages: any[];
|
||||
model: string;
|
||||
@@ -68,17 +70,23 @@ export interface GeneralAgentConfig {
|
||||
/**
|
||||
* Context compression configuration
|
||||
* Note: Compression checking is always enabled to prevent context overflow.
|
||||
* This config only controls the compression parameters.
|
||||
* When triggered, ALL messages are compressed into a single MessageGroup summary.
|
||||
*/
|
||||
compressionConfig?: {
|
||||
/** Number of recent messages to keep uncompressed (default: 10) */
|
||||
keepRecentCount?: number;
|
||||
/** Model's max context window token count (default: 128k) */
|
||||
maxWindowToken?: number;
|
||||
};
|
||||
modelRuntimeConfig?: {
|
||||
model: string;
|
||||
provider: string;
|
||||
/**
|
||||
* Compression model configuration
|
||||
* Used for context compression tasks
|
||||
*/
|
||||
compressionModel?: {
|
||||
model: string;
|
||||
provider: string;
|
||||
};
|
||||
};
|
||||
operationId: string;
|
||||
userId?: string;
|
||||
@@ -90,12 +98,10 @@ export interface GeneralAgentConfig {
|
||||
export interface GeneralAgentCompressionResultPayload {
|
||||
/** Compressed messages (summary + pinned + recent) */
|
||||
compressedMessages: any[];
|
||||
/** Token count after compression */
|
||||
compressedTokenCount: number;
|
||||
/** Compression group ID in database */
|
||||
groupId: string;
|
||||
/** Token count before compression */
|
||||
originalTokenCount: number;
|
||||
/** Parent message ID for subsequent LLM call (last assistant message before compression) */
|
||||
parentMessageId?: string;
|
||||
/** Whether compression was skipped (no messages to compress) */
|
||||
skipped?: boolean;
|
||||
}
|
||||
|
||||
@@ -209,6 +209,7 @@ export interface AgentInstructionResolveAbortedTools {
|
||||
|
||||
/**
|
||||
* Instruction to execute context compression
|
||||
* When triggered, compresses ALL messages into a single MessageGroup summary
|
||||
*/
|
||||
export interface AgentInstructionCompressContext {
|
||||
payload: {
|
||||
@@ -216,12 +217,8 @@ export interface AgentInstructionCompressContext {
|
||||
currentTokenCount: number;
|
||||
/** Existing summary to incorporate (for incremental compression) */
|
||||
existingSummary?: string;
|
||||
/** Number of recent messages to keep uncompressed */
|
||||
keepRecentCount: number;
|
||||
/** Messages to compress */
|
||||
messages: any[];
|
||||
/** Topic ID for the conversation */
|
||||
topicId: string;
|
||||
};
|
||||
type: 'compress_context';
|
||||
}
|
||||
|
||||
@@ -27,6 +27,14 @@ export interface AgentState {
|
||||
modelRuntimeConfig?: {
|
||||
model: string;
|
||||
provider: string;
|
||||
/**
|
||||
* Compression model configuration
|
||||
* Used for context compression tasks
|
||||
*/
|
||||
compressionModel?: {
|
||||
model: string;
|
||||
provider: string;
|
||||
};
|
||||
};
|
||||
|
||||
/**
|
||||
|
||||
@@ -109,12 +109,12 @@ describe('tokenCounter', () => {
|
||||
it('should use default values', () => {
|
||||
const threshold = getCompressionThreshold();
|
||||
expect(threshold).toBe(Math.floor(DEFAULT_MAX_CONTEXT * DEFAULT_THRESHOLD_RATIO));
|
||||
expect(threshold).toBe(96_000); // 128k * 0.75
|
||||
expect(threshold).toBe(64_000); // 128k * 0.5
|
||||
});
|
||||
|
||||
it('should use custom maxWindowToken', () => {
|
||||
const threshold = getCompressionThreshold({ maxWindowToken: 200_000 });
|
||||
expect(threshold).toBe(150_000); // 200k * 0.75
|
||||
expect(threshold).toBe(100_000); // 200k * 0.5
|
||||
});
|
||||
|
||||
it('should use custom thresholdRatio', () => {
|
||||
@@ -146,7 +146,7 @@ describe('tokenCounter', () => {
|
||||
|
||||
expect(result.needsCompression).toBe(false);
|
||||
expect(result.currentTokenCount).toBeGreaterThan(0);
|
||||
expect(result.threshold).toBe(96_000);
|
||||
expect(result.threshold).toBe(64_000); // 128k * 0.5
|
||||
});
|
||||
|
||||
it('should return needsCompression=true when over threshold', () => {
|
||||
@@ -154,22 +154,22 @@ describe('tokenCounter', () => {
|
||||
const messages = [
|
||||
{
|
||||
content: '',
|
||||
metadata: { usage: { totalOutputTokens: 100_000 } },
|
||||
metadata: { usage: { totalOutputTokens: 70_000 } },
|
||||
role: 'assistant',
|
||||
},
|
||||
];
|
||||
const result = shouldCompress(messages);
|
||||
|
||||
expect(result.needsCompression).toBe(true);
|
||||
expect(result.currentTokenCount).toBe(100_000);
|
||||
expect(result.threshold).toBe(96_000);
|
||||
expect(result.currentTokenCount).toBe(70_000);
|
||||
expect(result.threshold).toBe(64_000); // 128k * 0.5
|
||||
});
|
||||
|
||||
it('should return needsCompression=false when exactly at threshold', () => {
|
||||
const messages = [
|
||||
{
|
||||
content: '',
|
||||
metadata: { usage: { totalOutputTokens: 96_000 } },
|
||||
metadata: { usage: { totalOutputTokens: 64_000 } },
|
||||
role: 'assistant',
|
||||
},
|
||||
];
|
||||
@@ -177,7 +177,7 @@ describe('tokenCounter', () => {
|
||||
|
||||
// Exactly at threshold should not trigger compression
|
||||
expect(result.needsCompression).toBe(false);
|
||||
expect(result.currentTokenCount).toBe(96_000);
|
||||
expect(result.currentTokenCount).toBe(64_000);
|
||||
});
|
||||
|
||||
it('should use custom options', () => {
|
||||
|
||||
@@ -13,8 +13,8 @@ export interface TokenCountOptions {
|
||||
/** Default max context window (128k tokens) */
|
||||
export const DEFAULT_MAX_CONTEXT = 128_000;
|
||||
|
||||
/** Default threshold ratio (75% of max context) */
|
||||
export const DEFAULT_THRESHOLD_RATIO = 0.75;
|
||||
/** Default threshold ratio (50% of max context) */
|
||||
export const DEFAULT_THRESHOLD_RATIO = 0.5;
|
||||
|
||||
/**
|
||||
* Message interface for token counting
|
||||
|
||||
@@ -27,6 +27,7 @@ export const DEFAULT_AGENT_CHAT_CONFIG: LobeAgentChatConfig = {
|
||||
autoCreateTopicThreshold: 2,
|
||||
enableAutoCreateTopic: true,
|
||||
enableCompressHistory: true,
|
||||
enableContextCompression: true,
|
||||
enableHistoryCount: true,
|
||||
enableReasoning: false,
|
||||
enableStreaming: true,
|
||||
|
||||
@@ -6,6 +6,7 @@ import type { OpenAIChatMessage } from '@/types/index';
|
||||
import { ContextEngine } from '../../pipeline';
|
||||
import {
|
||||
AgentCouncilFlattenProcessor,
|
||||
CompressedGroupRoleTransformProcessor,
|
||||
GroupMessageFlattenProcessor,
|
||||
GroupOrchestrationFilterProcessor,
|
||||
GroupRoleTransformProcessor,
|
||||
@@ -266,6 +267,9 @@ export class MessagesEngine {
|
||||
// 19. Supervisor role restore (convert role=supervisor back to role=assistant for model)
|
||||
new SupervisorRoleRestoreProcessor(),
|
||||
|
||||
// 19b. Compressed group role transform (convert role=compressedGroup to role=user for model)
|
||||
new CompressedGroupRoleTransformProcessor(),
|
||||
|
||||
// 20. Group orchestration filter (remove supervisor's orchestration messages like broadcast/speak)
|
||||
// This must be BEFORE GroupRoleTransformProcessor so we filter based on original agentId/tools
|
||||
...(isAgentGroupEnabled && agentGroup.agentMap && agentGroup.currentAgentId
|
||||
|
||||
@@ -0,0 +1,70 @@
|
||||
import debug from 'debug';
|
||||
|
||||
import { BaseProcessor } from '../base/BaseProcessor';
|
||||
import type { Message, PipelineContext, ProcessorOptions } from '../types';
|
||||
|
||||
const log = debug('context-engine:processor:CompressedGroupRoleTransformProcessor');
|
||||
|
||||
/**
|
||||
* Compressed Group Role Transform Processor
|
||||
*
|
||||
* Transforms messages with role='compressedGroup' to role='user' before
|
||||
* sending to the model. The 'compressedGroup' role is used for UI rendering
|
||||
* to display compressed/summarized conversation history, but models don't
|
||||
* understand this role.
|
||||
*
|
||||
* The compressed summary content is wrapped in a system context block to
|
||||
* provide historical context to the model.
|
||||
*
|
||||
* Flow:
|
||||
* 1. DB stores compression groups with role='compressedGroup'
|
||||
* 2. conversation-flow passes them through for UI rendering
|
||||
* 3. This processor transforms to role='user' with wrapped content before model API call
|
||||
*
|
||||
* @example
|
||||
* ```typescript
|
||||
* const processor = new CompressedGroupRoleTransformProcessor();
|
||||
* const result = await processor.process(context);
|
||||
* // All compressedGroup messages are now user messages with wrapped content
|
||||
* ```
|
||||
*/
|
||||
export class CompressedGroupRoleTransformProcessor extends BaseProcessor {
|
||||
readonly name = 'CompressedGroupRoleTransformProcessor';
|
||||
|
||||
constructor(options: ProcessorOptions = {}) {
|
||||
super(options);
|
||||
}
|
||||
|
||||
protected async doProcess(context: PipelineContext): Promise<PipelineContext> {
|
||||
const clonedContext = this.cloneContext(context);
|
||||
|
||||
let processedCount = 0;
|
||||
|
||||
clonedContext.messages = clonedContext.messages.map((msg: Message) => {
|
||||
if (msg.role === 'compressedGroup') {
|
||||
processedCount++;
|
||||
log(`Transforming compressedGroup message to user role`);
|
||||
|
||||
// Wrap the compressed summary content in a context block
|
||||
const wrappedContent = msg.content
|
||||
? `<compressed_history_summary>\n${msg.content}\n</compressed_history_summary>`
|
||||
: '';
|
||||
|
||||
return {
|
||||
...msg,
|
||||
content: wrappedContent,
|
||||
role: 'user',
|
||||
};
|
||||
}
|
||||
|
||||
return msg;
|
||||
});
|
||||
|
||||
// Update metadata
|
||||
clonedContext.metadata.compressedGroupRoleTransformProcessed = processedCount;
|
||||
|
||||
log(`Compressed group role transform completed: ${processedCount} messages processed`);
|
||||
|
||||
return this.markAsExecuted(clonedContext);
|
||||
}
|
||||
}
|
||||
@@ -1,5 +1,6 @@
|
||||
// Transformer processors
|
||||
export { AgentCouncilFlattenProcessor } from './AgentCouncilFlatten';
|
||||
export { CompressedGroupRoleTransformProcessor } from './CompressedGroupRoleTransform';
|
||||
export { GroupMessageFlattenProcessor } from './GroupMessageFlatten';
|
||||
export {
|
||||
type GroupOrchestrationFilterConfig,
|
||||
|
||||
@@ -0,0 +1,11 @@
|
||||
import mixedGroups from './mixed-groups.json';
|
||||
import multipleCompressions from './multiple-compressions.json';
|
||||
import parallelGroup from './parallel-group.json';
|
||||
import simpleCompression from './simple-compression.json';
|
||||
|
||||
export const compression = {
|
||||
mixedGroups,
|
||||
multipleCompressions,
|
||||
parallelGroup,
|
||||
simpleCompression,
|
||||
};
|
||||
+80
@@ -0,0 +1,80 @@
|
||||
[
|
||||
{
|
||||
"id": "comp-group-1",
|
||||
"role": "compressedGroup",
|
||||
"content": "Summary: Initial project discussion about building a chat application.",
|
||||
"parentId": null,
|
||||
"createdAt": 1704060000000,
|
||||
"updatedAt": 1704060000000,
|
||||
"meta": {},
|
||||
"pinnedMessages": [
|
||||
{
|
||||
"id": "pinned-architecture",
|
||||
"role": "assistant",
|
||||
"content": "Key architecture decision: React + tRPC + PostgreSQL",
|
||||
"createdAt": "2024-01-01T08:00:00.000Z",
|
||||
"model": "gpt-4",
|
||||
"provider": "openai"
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"id": "msg-compare-trigger",
|
||||
"role": "user",
|
||||
"content": "Which state management should I use: Redux, Zustand, or Jotai?",
|
||||
"parentId": null,
|
||||
"createdAt": 1704067200000,
|
||||
"updatedAt": 1704067200000,
|
||||
"meta": {}
|
||||
},
|
||||
{
|
||||
"id": "parallel-group-state",
|
||||
"role": "compareGroup",
|
||||
"content": null,
|
||||
"parentId": null,
|
||||
"createdAt": 1704067205000,
|
||||
"updatedAt": 1704067205000,
|
||||
"meta": {},
|
||||
"children": [
|
||||
{
|
||||
"id": "resp-redux",
|
||||
"role": "assistant",
|
||||
"content": "Redux: Best for large-scale apps with complex state logic...",
|
||||
"createdAt": "2024-01-01T10:00:05.000Z",
|
||||
"model": "gpt-4",
|
||||
"provider": "openai"
|
||||
},
|
||||
{
|
||||
"id": "resp-zustand",
|
||||
"role": "assistant",
|
||||
"content": "Zustand: Lightweight and simple, great for most React apps...",
|
||||
"createdAt": "2024-01-01T10:00:06.000Z",
|
||||
"model": "claude-3-5-sonnet",
|
||||
"provider": "anthropic"
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"id": "msg-decision",
|
||||
"role": "user",
|
||||
"content": "I'll go with Zustand then. Let's start implementation.",
|
||||
"parentId": null,
|
||||
"createdAt": 1704067210000,
|
||||
"updatedAt": 1704067210000,
|
||||
"meta": {}
|
||||
},
|
||||
{
|
||||
"id": "msg-implementation",
|
||||
"role": "assistant",
|
||||
"content": "Great choice! Here's how to set up Zustand with TypeScript...",
|
||||
"parentId": "msg-decision",
|
||||
"createdAt": 1704067215000,
|
||||
"updatedAt": 1704067215000,
|
||||
"meta": {},
|
||||
"metadata": {
|
||||
"totalInputTokens": 300,
|
||||
"totalOutputTokens": 450,
|
||||
"totalTokens": 750
|
||||
}
|
||||
}
|
||||
]
|
||||
+67
@@ -0,0 +1,67 @@
|
||||
[
|
||||
{
|
||||
"id": "comp-group-1",
|
||||
"role": "compressedGroup",
|
||||
"content": "Summary: Early conversation about project setup and initial requirements.",
|
||||
"parentId": null,
|
||||
"createdAt": 1704060000000,
|
||||
"updatedAt": 1704060000000,
|
||||
"meta": {},
|
||||
"pinnedMessages": [
|
||||
{
|
||||
"id": "pinned-1-1",
|
||||
"role": "assistant",
|
||||
"content": "Key decision: Use TypeScript with strict mode",
|
||||
"createdAt": "2024-01-01T08:00:00.000Z",
|
||||
"model": "gpt-4",
|
||||
"provider": "openai"
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"id": "msg-middle-001",
|
||||
"role": "user",
|
||||
"content": "Let's discuss the database schema",
|
||||
"parentId": null,
|
||||
"createdAt": 1704063600000,
|
||||
"updatedAt": 1704063600000,
|
||||
"meta": {}
|
||||
},
|
||||
{
|
||||
"id": "msg-middle-002",
|
||||
"role": "assistant",
|
||||
"content": "Here's my suggested schema design...",
|
||||
"parentId": "msg-middle-001",
|
||||
"createdAt": 1704063605000,
|
||||
"updatedAt": 1704063605000,
|
||||
"meta": {}
|
||||
},
|
||||
{
|
||||
"id": "comp-group-2",
|
||||
"role": "compressedGroup",
|
||||
"content": "Summary: Database schema discussions concluded with PostgreSQL + Drizzle ORM decision.",
|
||||
"parentId": null,
|
||||
"createdAt": 1704067200000,
|
||||
"updatedAt": 1704067200000,
|
||||
"meta": {},
|
||||
"pinnedMessages": []
|
||||
},
|
||||
{
|
||||
"id": "msg-recent-001",
|
||||
"role": "user",
|
||||
"content": "Now let's implement the API endpoints",
|
||||
"parentId": null,
|
||||
"createdAt": 1704070800000,
|
||||
"updatedAt": 1704070800000,
|
||||
"meta": {}
|
||||
},
|
||||
{
|
||||
"id": "msg-recent-002",
|
||||
"role": "assistant",
|
||||
"content": "I'll create RESTful endpoints for CRUD operations...",
|
||||
"parentId": "msg-recent-001",
|
||||
"createdAt": 1704070805000,
|
||||
"updatedAt": 1704070805000,
|
||||
"meta": {}
|
||||
}
|
||||
]
|
||||
+47
@@ -0,0 +1,47 @@
|
||||
[
|
||||
{
|
||||
"id": "msg-user-001",
|
||||
"role": "user",
|
||||
"content": "Compare REST vs GraphQL for my use case",
|
||||
"parentId": null,
|
||||
"createdAt": 1704067200000,
|
||||
"updatedAt": 1704067200000,
|
||||
"meta": {}
|
||||
},
|
||||
{
|
||||
"id": "parallel-group-1",
|
||||
"role": "compareGroup",
|
||||
"content": null,
|
||||
"parentId": null,
|
||||
"createdAt": 1704067205000,
|
||||
"updatedAt": 1704067205000,
|
||||
"meta": {},
|
||||
"children": [
|
||||
{
|
||||
"id": "parallel-resp-gpt4",
|
||||
"role": "assistant",
|
||||
"content": "From GPT-4's perspective: REST is simpler and well-suited for CRUD operations...",
|
||||
"createdAt": "2024-01-01T10:00:05.000Z",
|
||||
"model": "gpt-4",
|
||||
"provider": "openai"
|
||||
},
|
||||
{
|
||||
"id": "parallel-resp-claude",
|
||||
"role": "assistant",
|
||||
"content": "From Claude's perspective: GraphQL offers more flexibility with its query language...",
|
||||
"createdAt": "2024-01-01T10:00:06.000Z",
|
||||
"model": "claude-3-5-sonnet",
|
||||
"provider": "anthropic"
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"id": "msg-user-002",
|
||||
"role": "user",
|
||||
"content": "Thanks for the comparison!",
|
||||
"parentId": null,
|
||||
"createdAt": 1704067210000,
|
||||
"updatedAt": 1704067210000,
|
||||
"meta": {}
|
||||
}
|
||||
]
|
||||
+53
@@ -0,0 +1,53 @@
|
||||
[
|
||||
{
|
||||
"id": "comp-group-1",
|
||||
"role": "compressedGroup",
|
||||
"content": "Summary: User asked about prime numbers. Assistant provided implementation and test cases.",
|
||||
"parentId": null,
|
||||
"createdAt": 1704067205000,
|
||||
"updatedAt": 1704067205000,
|
||||
"meta": {},
|
||||
"lastMessageId": "msg-compressed-last",
|
||||
"pinnedMessages": [
|
||||
{
|
||||
"id": "pinned-msg-001",
|
||||
"role": "assistant",
|
||||
"content": "Here's a prime number checker function...",
|
||||
"createdAt": "2024-01-01T10:01:00.000Z",
|
||||
"model": "gpt-4",
|
||||
"provider": "openai"
|
||||
},
|
||||
{
|
||||
"id": "pinned-msg-002",
|
||||
"role": "user",
|
||||
"content": "Can you add documentation?",
|
||||
"createdAt": "2024-01-01T10:02:00.000Z",
|
||||
"model": null,
|
||||
"provider": null
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"id": "msg-recent-001",
|
||||
"role": "user",
|
||||
"content": "Now let's optimize it further",
|
||||
"parentId": null,
|
||||
"createdAt": 1704067210000,
|
||||
"updatedAt": 1704067210000,
|
||||
"meta": {}
|
||||
},
|
||||
{
|
||||
"id": "msg-recent-002",
|
||||
"role": "assistant",
|
||||
"content": "Sure! Here's an optimized version using the Sieve of Eratosthenes...",
|
||||
"parentId": "msg-recent-001",
|
||||
"createdAt": 1704067215000,
|
||||
"updatedAt": 1704067215000,
|
||||
"meta": {},
|
||||
"metadata": {
|
||||
"totalInputTokens": 200,
|
||||
"totalOutputTokens": 150,
|
||||
"totalTokens": 350
|
||||
}
|
||||
}
|
||||
]
|
||||
@@ -5,6 +5,7 @@ import assistantChainWithFollowup from './assistant-chain-with-followup.json';
|
||||
import { assistantGroup } from './assistantGroup';
|
||||
import { branch } from './branch';
|
||||
import { compare } from './compare';
|
||||
import { compression } from './compression';
|
||||
import linearConversation from './linear-conversation.json';
|
||||
import { tasks } from './tasks';
|
||||
|
||||
@@ -15,6 +16,7 @@ export const inputs = {
|
||||
assistantGroup,
|
||||
branch,
|
||||
compare,
|
||||
compression,
|
||||
linearConversation: linearConversation as Message[],
|
||||
tasks,
|
||||
};
|
||||
|
||||
@@ -390,6 +390,120 @@ describe('buildHelperMaps', () => {
|
||||
});
|
||||
});
|
||||
|
||||
describe('compressedGroup lastMessageId redirection', () => {
|
||||
it('should redirect parentId to compressedGroup when pointing to lastMessageId', () => {
|
||||
// This tests the scenario after compression:
|
||||
// 1. compressedGroup contains messages with lastMessageId='msg-3'
|
||||
// 2. New assistant message is created with parentId='msg-3'
|
||||
// 3. msg-3 is hidden inside the compressedGroup, not in the message list
|
||||
// 4. childrenMap should redirect: assistant becomes child of compressedGroup
|
||||
const messages: Message[] = [
|
||||
{
|
||||
content: 'Summary of conversation',
|
||||
createdAt: 1000,
|
||||
id: 'comp-group-1',
|
||||
lastMessageId: 'msg-3', // Last compressed message (hidden)
|
||||
pinnedMessages: [],
|
||||
role: 'compressedGroup',
|
||||
updatedAt: 1000,
|
||||
} as any,
|
||||
{
|
||||
content: 'New assistant response after compression',
|
||||
createdAt: 2000,
|
||||
id: 'msg-new',
|
||||
parentId: 'msg-3', // Points to compressed message (not in list!)
|
||||
role: 'assistant',
|
||||
updatedAt: 2000,
|
||||
},
|
||||
];
|
||||
|
||||
const result = buildHelperMaps(messages);
|
||||
|
||||
// Without redirection: childrenMap.get('msg-3') = ['msg-new']
|
||||
// With redirection: childrenMap.get('comp-group-1') = ['msg-new']
|
||||
expect(result.childrenMap.get('comp-group-1')).toEqual(['msg-new']);
|
||||
expect(result.childrenMap.get('msg-3')).toBeUndefined();
|
||||
});
|
||||
|
||||
it('should not redirect when parentId is a normal message', () => {
|
||||
const messages: Message[] = [
|
||||
{
|
||||
content: 'Summary',
|
||||
createdAt: 1000,
|
||||
id: 'comp-group-1',
|
||||
lastMessageId: 'msg-3',
|
||||
pinnedMessages: [],
|
||||
role: 'compressedGroup',
|
||||
updatedAt: 1000,
|
||||
} as any,
|
||||
{
|
||||
content: 'User message',
|
||||
createdAt: 2000,
|
||||
id: 'msg-4',
|
||||
role: 'user',
|
||||
updatedAt: 2000,
|
||||
},
|
||||
{
|
||||
content: 'Response to user',
|
||||
createdAt: 3000,
|
||||
id: 'msg-5',
|
||||
parentId: 'msg-4', // Normal parentId, should not redirect
|
||||
role: 'assistant',
|
||||
updatedAt: 3000,
|
||||
},
|
||||
];
|
||||
|
||||
const result = buildHelperMaps(messages);
|
||||
|
||||
// msg-5 should be child of msg-4, not redirected
|
||||
expect(result.childrenMap.get('msg-4')).toEqual(['msg-5']);
|
||||
});
|
||||
|
||||
it('should handle multiple compressedGroups with different lastMessageIds', () => {
|
||||
const messages: Message[] = [
|
||||
{
|
||||
content: 'First summary',
|
||||
createdAt: 1000,
|
||||
id: 'comp-group-1',
|
||||
lastMessageId: 'msg-a',
|
||||
pinnedMessages: [],
|
||||
role: 'compressedGroup',
|
||||
updatedAt: 1000,
|
||||
} as any,
|
||||
{
|
||||
content: 'Second summary',
|
||||
createdAt: 2000,
|
||||
id: 'comp-group-2',
|
||||
lastMessageId: 'msg-b',
|
||||
pinnedMessages: [],
|
||||
role: 'compressedGroup',
|
||||
updatedAt: 2000,
|
||||
} as any,
|
||||
{
|
||||
content: 'Child of first group',
|
||||
createdAt: 3000,
|
||||
id: 'msg-new-1',
|
||||
parentId: 'msg-a',
|
||||
role: 'user',
|
||||
updatedAt: 3000,
|
||||
},
|
||||
{
|
||||
content: 'Child of second group',
|
||||
createdAt: 4000,
|
||||
id: 'msg-new-2',
|
||||
parentId: 'msg-b',
|
||||
role: 'assistant',
|
||||
updatedAt: 4000,
|
||||
},
|
||||
];
|
||||
|
||||
const result = buildHelperMaps(messages);
|
||||
|
||||
expect(result.childrenMap.get('comp-group-1')).toEqual(['msg-new-1']);
|
||||
expect(result.childrenMap.get('comp-group-2')).toEqual(['msg-new-2']);
|
||||
});
|
||||
});
|
||||
|
||||
describe('integration scenarios', () => {
|
||||
it('should build all maps correctly in complex conversation', () => {
|
||||
const messages: Message[] = [
|
||||
|
||||
@@ -17,13 +17,32 @@ export function buildHelperMaps(
|
||||
const threadMap = new Map<string, Message[]>();
|
||||
const messageGroupMap = new Map<string, MessageGroupMetadata>();
|
||||
|
||||
// Build a map of lastMessageId -> compressedGroup.id for parent redirection
|
||||
// This handles the case where messages are created after compression:
|
||||
// - The new message's parentId points to the last compressed message (lastMessageId)
|
||||
// - But the lastMessageId is hidden inside the compressedGroup
|
||||
// - We need to redirect children of lastMessageId to be children of compressedGroup
|
||||
const lastMessageIdToGroupId = new Map<string, string>();
|
||||
for (const message of messages) {
|
||||
if (message.role === 'compressedGroup' && (message as any).lastMessageId) {
|
||||
lastMessageIdToGroupId.set((message as any).lastMessageId, message.id);
|
||||
}
|
||||
}
|
||||
|
||||
// Single pass through messages to build all maps
|
||||
for (const message of messages) {
|
||||
// 1. Build messageMap for O(1) access
|
||||
messageMap.set(message.id, message);
|
||||
|
||||
// 2. Build childrenMap for parent -> children lookup
|
||||
const parentId = message.parentId ?? null;
|
||||
let parentId = message.parentId ?? null;
|
||||
|
||||
// Redirect parentId if it points to a compressed message's lastMessageId
|
||||
// This ensures messages created after compression become children of the compressedGroup
|
||||
if (parentId && lastMessageIdToGroupId.has(parentId)) {
|
||||
parentId = lastMessageIdToGroupId.get(parentId)!;
|
||||
}
|
||||
|
||||
const siblings = childrenMap.get(parentId);
|
||||
if (siblings) {
|
||||
siblings.push(message.id);
|
||||
|
||||
@@ -12,7 +12,15 @@ interface BaseNode {
|
||||
/** Unique identifier for this node */
|
||||
id: string;
|
||||
/** Type discriminator */
|
||||
type: 'message' | 'assistantGroup' | 'compare' | 'branch' | 'agentCouncil' | 'tasks';
|
||||
type:
|
||||
| 'message'
|
||||
| 'assistantGroup'
|
||||
| 'compare'
|
||||
| 'branch'
|
||||
| 'agentCouncil'
|
||||
| 'tasks'
|
||||
| 'compressedGroup'
|
||||
| 'compareGroup';
|
||||
}
|
||||
|
||||
/**
|
||||
@@ -83,6 +91,54 @@ export interface TasksNode extends BaseNode {
|
||||
type: 'tasks';
|
||||
}
|
||||
|
||||
/**
|
||||
* Pinned message within a compression group
|
||||
*/
|
||||
export interface PinnedMessage {
|
||||
content: string | null;
|
||||
createdAt: Date | string;
|
||||
id: string;
|
||||
model: string | null;
|
||||
provider: string | null;
|
||||
role: string;
|
||||
}
|
||||
|
||||
/**
|
||||
* Compressed Group node - represents compressed/summarized messages
|
||||
* Messages marked as compressed are hidden, and a summary is shown instead.
|
||||
* Pinned messages (favorite=true) within the compression group are preserved.
|
||||
*/
|
||||
export interface CompressedGroupNode extends BaseNode {
|
||||
/** Summary content of the compressed messages */
|
||||
content: string | null;
|
||||
/** Messages marked as favorite/pinned within this compression group */
|
||||
pinnedMessages: PinnedMessage[];
|
||||
type: 'compressedGroup';
|
||||
}
|
||||
|
||||
/**
|
||||
* Child message within a compare group (parallel responses)
|
||||
*/
|
||||
export interface CompareGroupChild {
|
||||
content: string | null;
|
||||
createdAt: Date | string;
|
||||
id: string;
|
||||
model: string | null;
|
||||
provider: string | null;
|
||||
role: string;
|
||||
}
|
||||
|
||||
/**
|
||||
* Compare Group node - represents parallel model responses
|
||||
* Multiple models respond to the same user message in parallel.
|
||||
* Different from CompareNode which is built from metadata.compare flag.
|
||||
*/
|
||||
export interface CompareGroupNode extends BaseNode {
|
||||
/** Parallel responses from different models */
|
||||
children: CompareGroupChild[];
|
||||
type: 'compareGroup';
|
||||
}
|
||||
|
||||
/**
|
||||
* Union type of all display nodes
|
||||
*/
|
||||
@@ -92,4 +148,6 @@ export type ContextNode =
|
||||
| CompareNode
|
||||
| BranchNode
|
||||
| AgentCouncilNode
|
||||
| TasksNode;
|
||||
| TasksNode
|
||||
| CompressedGroupNode
|
||||
| CompareGroupNode;
|
||||
|
||||
@@ -23,6 +23,10 @@ import type { UIChatMessage } from '@lobechat/types';
|
||||
* - 'compare': Compare mode for parallel model outputs
|
||||
* - 'agentCouncil': Multi-agent parallel responses (all enter context)
|
||||
* - 'tasks': Aggregated async task messages with same parentId
|
||||
*
|
||||
* Database-provided roles from MessageModel.query with MessageGroup aggregation:
|
||||
* - 'compressedGroup': Compressed/summarized messages with optional pinned messages
|
||||
* - 'compareGroup': Parallel model responses (from database, not metadata-driven)
|
||||
*/
|
||||
export type FlatMessageRole =
|
||||
| 'user'
|
||||
@@ -34,7 +38,9 @@ export type FlatMessageRole =
|
||||
| 'messageGroup'
|
||||
| 'compare'
|
||||
| 'agentCouncil'
|
||||
| 'tasks';
|
||||
| 'tasks'
|
||||
| 'compressedGroup'
|
||||
| 'compareGroup';
|
||||
|
||||
/**
|
||||
* Message in flat list
|
||||
|
||||
@@ -14,10 +14,14 @@ export type {
|
||||
AgentCouncilNode,
|
||||
AssistantGroupNode,
|
||||
BranchNode,
|
||||
CompareGroupChild,
|
||||
CompareGroupNode,
|
||||
CompareNode,
|
||||
CompressedGroupNode,
|
||||
ContextNode,
|
||||
MessageNode,
|
||||
TasksNode,
|
||||
PinnedMessage,
|
||||
} from './contextTree';
|
||||
|
||||
// Flat Message List Types
|
||||
|
||||
@@ -0,0 +1 @@
|
||||
ALTER TABLE "message_groups" ADD COLUMN IF NOT EXISTS "metadata" jsonb;
|
||||
File diff suppressed because it is too large
Load Diff
@@ -511,7 +511,14 @@
|
||||
"when": 1769251608013,
|
||||
"tag": "0072_add_market_identifier_chat_group",
|
||||
"breakpoints": true
|
||||
},
|
||||
{
|
||||
"idx": 73,
|
||||
"version": "7",
|
||||
"when": 1769272397744,
|
||||
"tag": "0073_add_message_group_metadata",
|
||||
"breakpoints": true
|
||||
}
|
||||
],
|
||||
"version": "6"
|
||||
}
|
||||
}
|
||||
|
||||
@@ -15,6 +15,7 @@
|
||||
},
|
||||
"dependencies": {
|
||||
"@lobechat/const": "workspace:*",
|
||||
"@lobechat/conversation-flow": "workspace:*",
|
||||
"@lobechat/types": "workspace:*",
|
||||
"@lobechat/utils": "workspace:*",
|
||||
"random-words": "^2.0.1",
|
||||
|
||||
@@ -1,3 +1,4 @@
|
||||
export * from './core/db-adaptor';
|
||||
export * from './repositories/compression';
|
||||
export * from './type';
|
||||
export * from './utils/idGenerator';
|
||||
|
||||
@@ -3,10 +3,10 @@ import { MessageGroupType } from '@lobechat/types';
|
||||
import { and, eq, inArray } from 'drizzle-orm';
|
||||
import { afterEach, beforeEach, describe, expect, it } from 'vitest';
|
||||
|
||||
import { getTestDB } from '../../../core/getTestDB';
|
||||
import { messageGroups, messages, topics, users } from '../../../schemas';
|
||||
import { LobeChatDatabase } from '../../../type';
|
||||
import { MessageModel } from '../../message';
|
||||
import { getTestDB } from '../../../core/getTestDB';
|
||||
|
||||
const userId = 'message-query-test-user';
|
||||
const topicId = 'test-topic-1';
|
||||
@@ -628,7 +628,193 @@ describe('MessageModel.query with MessageGroup aggregation', () => {
|
||||
});
|
||||
|
||||
/**
|
||||
* Test Scenario 6: Trajectory ordering
|
||||
* Test Scenario 6: compressedMessages field
|
||||
* Expected:
|
||||
* - Compression group node should have a `compressedMessages` array field
|
||||
* - This array should contain all messages from that group with full data
|
||||
*/
|
||||
describe('compressedMessages in compression groups', () => {
|
||||
it('should include all messages in compressedMessages array', async () => {
|
||||
// Create messages to be compressed
|
||||
await serverDB.insert(messages).values([
|
||||
{
|
||||
id: 'comp-msg-1',
|
||||
content: 'User message 1',
|
||||
role: 'user',
|
||||
topicId,
|
||||
userId,
|
||||
createdAt: new Date('2024-01-01T10:00:00Z'),
|
||||
},
|
||||
{
|
||||
id: 'comp-msg-2',
|
||||
content: 'Assistant response',
|
||||
role: 'assistant',
|
||||
model: 'gpt-4',
|
||||
provider: 'openai',
|
||||
topicId,
|
||||
userId,
|
||||
createdAt: new Date('2024-01-01T10:01:00Z'),
|
||||
},
|
||||
{
|
||||
id: 'comp-msg-3',
|
||||
content: 'User follow-up',
|
||||
role: 'user',
|
||||
topicId,
|
||||
userId,
|
||||
parentId: 'comp-msg-2',
|
||||
createdAt: new Date('2024-01-01T10:02:00Z'),
|
||||
},
|
||||
{
|
||||
id: 'uncompressed-msg',
|
||||
content: 'Latest message',
|
||||
role: 'assistant',
|
||||
topicId,
|
||||
userId,
|
||||
createdAt: new Date('2024-01-01T10:03:00Z'),
|
||||
},
|
||||
]);
|
||||
|
||||
// Create compression group
|
||||
await serverDB.insert(messageGroups).values({
|
||||
id: 'comp-group-with-msgs',
|
||||
content: 'Summary of conversation',
|
||||
type: MessageGroupType.Compression,
|
||||
topicId,
|
||||
userId,
|
||||
createdAt: new Date('2024-01-01T10:02:30Z'),
|
||||
});
|
||||
|
||||
// Mark messages as compressed
|
||||
await serverDB
|
||||
.update(messages)
|
||||
.set({ messageGroupId: 'comp-group-with-msgs' })
|
||||
.where(inArray(messages.id, ['comp-msg-1', 'comp-msg-2', 'comp-msg-3']));
|
||||
|
||||
const result = await messageModel.query({ topicId });
|
||||
|
||||
// Expected: 1 compressedGroup + 1 uncompressed message = 2 items
|
||||
expect(result).toHaveLength(2);
|
||||
|
||||
// Verify compressedGroup node has compressedMessages
|
||||
const compressedGroupNode = result.find((m) => m.role === 'compressedGroup') as any;
|
||||
expect(compressedGroupNode).toBeDefined();
|
||||
expect(compressedGroupNode.compressedMessages).toBeDefined();
|
||||
expect(compressedGroupNode.compressedMessages).toHaveLength(3);
|
||||
|
||||
// Verify compressedMessages contain full message data
|
||||
const compressedMsgs = compressedGroupNode.compressedMessages;
|
||||
expect(compressedMsgs.map((m: any) => m.id)).toEqual([
|
||||
'comp-msg-1',
|
||||
'comp-msg-2',
|
||||
'comp-msg-3',
|
||||
]);
|
||||
|
||||
// Verify message fields are preserved
|
||||
const assistantMsg = compressedMsgs.find((m: any) => m.id === 'comp-msg-2');
|
||||
expect(assistantMsg.model).toBe('gpt-4');
|
||||
expect(assistantMsg.provider).toBe('openai');
|
||||
expect(assistantMsg.content).toBe('Assistant response');
|
||||
|
||||
// Verify parentId is preserved for conversation-flow parsing
|
||||
const followUpMsg = compressedMsgs.find((m: any) => m.id === 'comp-msg-3');
|
||||
expect(followUpMsg.parentId).toBe('comp-msg-2');
|
||||
});
|
||||
|
||||
it('should include compressedMessages with empty pinnedMessages when no favorites', async () => {
|
||||
await serverDB.insert(messages).values([
|
||||
{
|
||||
id: 'msg-no-fav-1',
|
||||
content: 'Message 1',
|
||||
role: 'user',
|
||||
topicId,
|
||||
userId,
|
||||
favorite: false,
|
||||
},
|
||||
{
|
||||
id: 'msg-no-fav-2',
|
||||
content: 'Message 2',
|
||||
role: 'assistant',
|
||||
topicId,
|
||||
userId,
|
||||
favorite: false,
|
||||
},
|
||||
]);
|
||||
|
||||
await serverDB.insert(messageGroups).values({
|
||||
id: 'comp-no-fav',
|
||||
content: 'Summary',
|
||||
type: MessageGroupType.Compression,
|
||||
topicId,
|
||||
userId,
|
||||
});
|
||||
|
||||
await serverDB
|
||||
.update(messages)
|
||||
.set({ messageGroupId: 'comp-no-fav' })
|
||||
.where(inArray(messages.id, ['msg-no-fav-1', 'msg-no-fav-2']));
|
||||
|
||||
const result = await messageModel.query({ topicId });
|
||||
|
||||
const compressedGroupNode = result.find((m) => m.role === 'compressedGroup') as any;
|
||||
expect(compressedGroupNode).toBeDefined();
|
||||
|
||||
// compressedMessages should contain all messages
|
||||
expect(compressedGroupNode.compressedMessages).toHaveLength(2);
|
||||
|
||||
// pinnedMessages should be empty
|
||||
expect(compressedGroupNode.pinnedMessages).toHaveLength(0);
|
||||
});
|
||||
|
||||
it('should have both pinnedMessages and compressedMessages correctly populated', async () => {
|
||||
await serverDB.insert(messages).values([
|
||||
{
|
||||
id: 'fav-msg',
|
||||
content: 'Important message',
|
||||
role: 'assistant',
|
||||
topicId,
|
||||
userId,
|
||||
favorite: true,
|
||||
createdAt: new Date('2024-01-01T10:00:00Z'),
|
||||
},
|
||||
{
|
||||
id: 'normal-msg',
|
||||
content: 'Normal message',
|
||||
role: 'user',
|
||||
topicId,
|
||||
userId,
|
||||
favorite: false,
|
||||
createdAt: new Date('2024-01-01T10:01:00Z'),
|
||||
},
|
||||
]);
|
||||
|
||||
await serverDB.insert(messageGroups).values({
|
||||
id: 'comp-mixed',
|
||||
content: 'Summary',
|
||||
type: MessageGroupType.Compression,
|
||||
topicId,
|
||||
userId,
|
||||
});
|
||||
|
||||
await serverDB
|
||||
.update(messages)
|
||||
.set({ messageGroupId: 'comp-mixed' })
|
||||
.where(inArray(messages.id, ['fav-msg', 'normal-msg']));
|
||||
|
||||
const result = await messageModel.query({ topicId });
|
||||
|
||||
const compressedGroupNode = result.find((m) => m.role === 'compressedGroup') as any;
|
||||
|
||||
// compressedMessages contains all messages
|
||||
expect(compressedGroupNode.compressedMessages).toHaveLength(2);
|
||||
|
||||
// pinnedMessages only contains favorite messages
|
||||
expect(compressedGroupNode.pinnedMessages).toHaveLength(1);
|
||||
expect(compressedGroupNode.pinnedMessages[0].id).toBe('fav-msg');
|
||||
});
|
||||
});
|
||||
|
||||
/**
|
||||
* Test Scenario 7: Trajectory ordering
|
||||
* Expected: Items should be ordered by createdAt correctly,
|
||||
* including MessageGroup nodes at their proper positions
|
||||
*/
|
||||
@@ -698,4 +884,116 @@ describe('MessageModel.query with MessageGroup aggregation', () => {
|
||||
expect(result[1].id).toBe('msg-4');
|
||||
});
|
||||
});
|
||||
|
||||
/**
|
||||
* Test Scenario 8: Branch scenario - compression with message branching
|
||||
* Expected:
|
||||
* - Compression should correctly capture messages by their IDs
|
||||
* - Branch messages (same parentId) should be handled correctly
|
||||
* - Only explicitly marked messages should be in compressedMessages
|
||||
*/
|
||||
describe('compression with message branches', () => {
|
||||
it('should only include explicitly marked messages in compression, not branch siblings', async () => {
|
||||
// Create a conversation with branches:
|
||||
// User A -> Assistant B1 (branch 1)
|
||||
// -> Assistant B2 (branch 2)
|
||||
// User continues on B1 -> C
|
||||
await serverDB.insert(messages).values([
|
||||
{
|
||||
id: 'branch-user-a',
|
||||
content: 'User question',
|
||||
role: 'user',
|
||||
topicId,
|
||||
userId,
|
||||
createdAt: new Date('2024-01-01T10:00:00Z'),
|
||||
},
|
||||
{
|
||||
id: 'branch-assistant-b1',
|
||||
content: 'Assistant response B1',
|
||||
role: 'assistant',
|
||||
parentId: 'branch-user-a',
|
||||
model: 'gpt-4',
|
||||
provider: 'openai',
|
||||
topicId,
|
||||
userId,
|
||||
createdAt: new Date('2024-01-01T10:01:00Z'),
|
||||
},
|
||||
{
|
||||
id: 'branch-assistant-b2',
|
||||
content: 'Assistant response B2 (branch)',
|
||||
role: 'assistant',
|
||||
parentId: 'branch-user-a', // Same parent - this is a branch
|
||||
model: 'claude-3',
|
||||
provider: 'anthropic',
|
||||
topicId,
|
||||
userId,
|
||||
createdAt: new Date('2024-01-01T10:01:30Z'),
|
||||
},
|
||||
{
|
||||
id: 'branch-user-c',
|
||||
content: 'Follow up on B1',
|
||||
role: 'user',
|
||||
parentId: 'branch-assistant-b1', // Continue on B1
|
||||
topicId,
|
||||
userId,
|
||||
createdAt: new Date('2024-01-01T10:02:00Z'),
|
||||
},
|
||||
{
|
||||
id: 'branch-latest',
|
||||
content: 'Latest message',
|
||||
role: 'assistant',
|
||||
parentId: 'branch-user-c',
|
||||
topicId,
|
||||
userId,
|
||||
createdAt: new Date('2024-01-01T10:03:00Z'),
|
||||
},
|
||||
]);
|
||||
|
||||
// Create compression group for the main thread (A -> B1 -> C), NOT including B2
|
||||
await serverDB.insert(messageGroups).values({
|
||||
id: 'branch-comp',
|
||||
content: 'Compressed main thread',
|
||||
type: MessageGroupType.Compression,
|
||||
topicId,
|
||||
userId,
|
||||
createdAt: new Date('2024-01-01T10:02:30Z'),
|
||||
});
|
||||
|
||||
// Only mark main thread messages as compressed (A, B1, C)
|
||||
// B2 is NOT included in compression
|
||||
await serverDB
|
||||
.update(messages)
|
||||
.set({ messageGroupId: 'branch-comp' })
|
||||
.where(inArray(messages.id, ['branch-user-a', 'branch-assistant-b1', 'branch-user-c']));
|
||||
|
||||
const result = await messageModel.query({ topicId });
|
||||
|
||||
// Expected: compressedGroup + B2 (uncompressed branch) + latest
|
||||
expect(result).toHaveLength(3);
|
||||
|
||||
const compressedGroup = result.find((m) => m.role === 'compressedGroup') as any;
|
||||
expect(compressedGroup).toBeDefined();
|
||||
expect(compressedGroup.id).toBe('branch-comp');
|
||||
|
||||
// compressedMessages should only contain the main thread (A, B1, C)
|
||||
expect(compressedGroup.compressedMessages).toHaveLength(3);
|
||||
const compressedIds = compressedGroup.compressedMessages.map((m: any) => m.id);
|
||||
expect(compressedIds).toContain('branch-user-a');
|
||||
expect(compressedIds).toContain('branch-assistant-b1');
|
||||
expect(compressedIds).toContain('branch-user-c');
|
||||
// B2 should NOT be in compressedMessages
|
||||
expect(compressedIds).not.toContain('branch-assistant-b2');
|
||||
|
||||
// Verify parentId is preserved for conversation-flow parsing
|
||||
const b1InCompressed = compressedGroup.compressedMessages.find(
|
||||
(m: any) => m.id === 'branch-assistant-b1',
|
||||
);
|
||||
expect(b1InCompressed.parentId).toBe('branch-user-a');
|
||||
|
||||
// B2 should still be visible as uncompressed message
|
||||
const b2 = result.find((m) => m.id === 'branch-assistant-b2');
|
||||
expect(b2).toBeDefined();
|
||||
expect(b2!.content).toBe('Assistant response B2 (branch)');
|
||||
});
|
||||
});
|
||||
});
|
||||
|
||||
@@ -1,5 +1,6 @@
|
||||
import { afterEach, beforeEach, describe, expect, it } from 'vitest';
|
||||
|
||||
import { getTestDB } from '../../../core/getTestDB';
|
||||
import {
|
||||
agents,
|
||||
agentsToSessions,
|
||||
@@ -11,7 +12,6 @@ import {
|
||||
} from '../../../schemas';
|
||||
import { LobeChatDatabase } from '../../../type';
|
||||
import { TopicModel } from '../../topic';
|
||||
import { getTestDB } from '../../../core/getTestDB';
|
||||
|
||||
const userId = 'topic-query-user';
|
||||
const userId2 = 'topic-query-user-2';
|
||||
@@ -1281,6 +1281,11 @@ describe('TopicModel - Query', () => {
|
||||
});
|
||||
|
||||
describe('listTopicsForMemoryExtractor', () => {
|
||||
beforeEach(async () => {
|
||||
// Clear topics from previous tests to ensure isolation
|
||||
await serverDB.delete(topics);
|
||||
});
|
||||
|
||||
it('should paginate pending topics and skip extracted ones by default', async () => {
|
||||
await serverDB.insert(topics).values([
|
||||
{
|
||||
@@ -1299,14 +1304,16 @@ describe('TopicModel - Query', () => {
|
||||
{ createdAt: new Date('2024-01-04T00:00:00Z'), id: 't4', userId: userId2 },
|
||||
] satisfies Array<typeof topics.$inferInsert>);
|
||||
|
||||
// t1 is skipped because it has userMemoryExtractStatus: 'completed'
|
||||
// t4 is skipped because it belongs to a different user
|
||||
const page1 = await topicModel.listTopicsForMemoryExtractor({ limit: 1 });
|
||||
expect(page1.map((t) => t.id)).toEqual(['t1']);
|
||||
expect(page1.map((t) => t.id)).toEqual(['t2']);
|
||||
|
||||
const page2 = await topicModel.listTopicsForMemoryExtractor({
|
||||
cursor: { createdAt: page1[0].createdAt, id: page1[0].id },
|
||||
limit: 5,
|
||||
});
|
||||
expect(page2.map((t) => t.id)).toEqual(['t2', 't3']);
|
||||
expect(page2.map((t) => t.id)).toEqual(['t3']);
|
||||
});
|
||||
|
||||
it('should include extracted topics when ignoreExtracted is true', async () => {
|
||||
|
||||
@@ -1,4 +1,5 @@
|
||||
import { INBOX_SESSION_ID } from '@lobechat/const';
|
||||
import { parse } from '@lobechat/conversation-flow';
|
||||
import {
|
||||
ChatFileItem,
|
||||
ChatImageItem,
|
||||
@@ -276,19 +277,23 @@ export class MessageModel {
|
||||
if (topicId && result.length > 0) {
|
||||
if (current === 0) {
|
||||
// First page: fetch all groups to include compressed history
|
||||
messageGroupNodes = await this.queryMessageGroupNodes(topicId);
|
||||
messageGroupNodes = await this.queryMessageGroupNodes(topicId, undefined, postProcessUrl);
|
||||
} else {
|
||||
// Subsequent pages: filter by time range to avoid duplicates
|
||||
const firstMessageTime = result[0].createdAt;
|
||||
const lastMessageTime = result.at(-1)!.createdAt;
|
||||
messageGroupNodes = await this.queryMessageGroupNodes(topicId, {
|
||||
endTime: lastMessageTime,
|
||||
startTime: firstMessageTime,
|
||||
});
|
||||
messageGroupNodes = await this.queryMessageGroupNodes(
|
||||
topicId,
|
||||
{
|
||||
endTime: lastMessageTime,
|
||||
startTime: firstMessageTime,
|
||||
},
|
||||
postProcessUrl,
|
||||
);
|
||||
}
|
||||
} else if (topicId && current === 0) {
|
||||
// First page with no messages: still fetch all groups
|
||||
messageGroupNodes = await this.queryMessageGroupNodes(topicId);
|
||||
messageGroupNodes = await this.queryMessageGroupNodes(topicId, undefined, postProcessUrl);
|
||||
}
|
||||
|
||||
// If no messages and no group nodes, return empty
|
||||
@@ -514,9 +519,264 @@ export class MessageModel {
|
||||
return allItems;
|
||||
};
|
||||
|
||||
/**
|
||||
* Query messages by their IDs with full relations
|
||||
*
|
||||
* This is useful for getting full message data when you already have the IDs.
|
||||
* It reuses the same transformation logic as queryWithWhere.
|
||||
*
|
||||
* @param messageIds - Array of message IDs to query
|
||||
* @param options - Query options (postProcessUrl for file URL transformation)
|
||||
* @returns Messages with all related data (files, plugins, translations, etc.)
|
||||
*/
|
||||
queryByIds = async (
|
||||
messageIds: string[],
|
||||
options: {
|
||||
postProcessUrl?: (path: string | null, file: { fileType: string }) => Promise<string>;
|
||||
} = {},
|
||||
): Promise<UIChatMessage[]> => {
|
||||
if (messageIds.length === 0) return [];
|
||||
|
||||
const { postProcessUrl } = options;
|
||||
|
||||
// 1. Query messages with joins
|
||||
const result = await this.db
|
||||
.select({
|
||||
/* eslint-disable sort-keys-fix/sort-keys-fix*/
|
||||
id: messages.id,
|
||||
role: messages.role,
|
||||
content: messages.content,
|
||||
reasoning: messages.reasoning,
|
||||
search: messages.search,
|
||||
metadata: messages.metadata,
|
||||
error: messages.error,
|
||||
|
||||
model: messages.model,
|
||||
provider: messages.provider,
|
||||
|
||||
createdAt: messages.createdAt,
|
||||
updatedAt: messages.updatedAt,
|
||||
|
||||
topicId: messages.topicId,
|
||||
parentId: messages.parentId,
|
||||
threadId: messages.threadId,
|
||||
|
||||
// Group chat fields
|
||||
groupId: messages.groupId,
|
||||
agentId: messages.agentId,
|
||||
targetId: messages.targetId,
|
||||
|
||||
tools: messages.tools,
|
||||
tool_call_id: messagePlugins.toolCallId,
|
||||
|
||||
plugin: {
|
||||
apiName: messagePlugins.apiName,
|
||||
arguments: messagePlugins.arguments,
|
||||
identifier: messagePlugins.identifier,
|
||||
type: messagePlugins.type,
|
||||
},
|
||||
pluginError: messagePlugins.error,
|
||||
pluginIntervention: messagePlugins.intervention,
|
||||
pluginState: messagePlugins.state,
|
||||
|
||||
translate: {
|
||||
content: messageTranslates.content,
|
||||
from: messageTranslates.from,
|
||||
to: messageTranslates.to,
|
||||
},
|
||||
|
||||
ttsId: messageTTS.id,
|
||||
ttsContentMd5: messageTTS.contentMd5,
|
||||
ttsFile: messageTTS.fileId,
|
||||
ttsVoice: messageTTS.voice,
|
||||
/* eslint-enable */
|
||||
})
|
||||
.from(messages)
|
||||
.where(and(eq(messages.userId, this.userId), inArray(messages.id, messageIds)))
|
||||
.leftJoin(messagePlugins, eq(messagePlugins.id, messages.id))
|
||||
.leftJoin(messageTranslates, eq(messageTranslates.id, messages.id))
|
||||
.leftJoin(messageTTS, eq(messageTTS.id, messages.id))
|
||||
.orderBy(asc(messages.createdAt));
|
||||
|
||||
if (result.length === 0) return [];
|
||||
|
||||
// 2. Get related files
|
||||
const rawRelatedFileList = await this.db
|
||||
.select({
|
||||
fileType: files.fileType,
|
||||
id: messagesFiles.fileId,
|
||||
messageId: messagesFiles.messageId,
|
||||
name: files.name,
|
||||
size: files.size,
|
||||
url: files.url,
|
||||
})
|
||||
.from(messagesFiles)
|
||||
.leftJoin(files, eq(files.id, messagesFiles.fileId))
|
||||
.where(inArray(messagesFiles.messageId, messageIds));
|
||||
|
||||
const relatedFileList = await Promise.all(
|
||||
rawRelatedFileList.map(async (file) => ({
|
||||
...file,
|
||||
url: postProcessUrl ? await postProcessUrl(file.url, file as any) : (file.url as string),
|
||||
})),
|
||||
);
|
||||
|
||||
// Get associated document content
|
||||
const fileIds = relatedFileList.map((file) => file.id).filter(Boolean);
|
||||
|
||||
let documentsMap: Record<string, string> = {};
|
||||
|
||||
if (fileIds.length > 0) {
|
||||
const documentsList = await this.db
|
||||
.select({
|
||||
content: documents.content,
|
||||
fileId: documents.fileId,
|
||||
})
|
||||
.from(documents)
|
||||
.where(inArray(documents.fileId, fileIds));
|
||||
|
||||
documentsMap = documentsList.reduce(
|
||||
(acc, doc) => {
|
||||
if (doc.fileId) acc[doc.fileId] = doc.content as string;
|
||||
return acc;
|
||||
},
|
||||
{} as Record<string, string>,
|
||||
);
|
||||
}
|
||||
|
||||
const imageList = relatedFileList.filter((i) => (i.fileType || '').startsWith('image'));
|
||||
const videoList = relatedFileList.filter((i) => (i.fileType || '').startsWith('video'));
|
||||
const fileList = relatedFileList.filter(
|
||||
(i) => !(i.fileType || '').startsWith('image') && !(i.fileType || '').startsWith('video'),
|
||||
);
|
||||
|
||||
// 3. Get related file chunks
|
||||
const chunksList = await this.db
|
||||
.select({
|
||||
fileId: files.id,
|
||||
fileType: files.fileType,
|
||||
fileUrl: files.url,
|
||||
filename: files.name,
|
||||
id: chunks.id,
|
||||
messageId: messageQueryChunks.messageId,
|
||||
similarity: messageQueryChunks.similarity,
|
||||
text: chunks.text,
|
||||
})
|
||||
.from(messageQueryChunks)
|
||||
.leftJoin(chunks, eq(chunks.id, messageQueryChunks.chunkId))
|
||||
.leftJoin(fileChunks, eq(fileChunks.chunkId, chunks.id))
|
||||
.innerJoin(files, eq(fileChunks.fileId, files.id))
|
||||
.where(inArray(messageQueryChunks.messageId, messageIds));
|
||||
|
||||
// 4. Get related message queries (RAG)
|
||||
const messageQueriesList = await this.db
|
||||
.select({
|
||||
id: messageQueries.id,
|
||||
messageId: messageQueries.messageId,
|
||||
rewriteQuery: messageQueries.rewriteQuery,
|
||||
userQuery: messageQueries.userQuery,
|
||||
})
|
||||
.from(messageQueries)
|
||||
.where(inArray(messageQueries.messageId, messageIds));
|
||||
|
||||
// 5. Get thread info for task messages
|
||||
const taskMessageIds = result.filter((m) => m.role === 'task').map((m) => m.id as string);
|
||||
|
||||
let threadMap = new Map<string, TaskDetail>();
|
||||
|
||||
if (taskMessageIds.length > 0) {
|
||||
const threadData = await this.db
|
||||
.select({
|
||||
metadata: threads.metadata,
|
||||
sourceMessageId: threads.sourceMessageId,
|
||||
status: threads.status,
|
||||
threadId: threads.id,
|
||||
title: threads.title,
|
||||
})
|
||||
.from(threads)
|
||||
.where(
|
||||
and(eq(threads.userId, this.userId), inArray(threads.sourceMessageId, taskMessageIds)),
|
||||
);
|
||||
|
||||
threadMap = new Map(
|
||||
threadData.map((t) => {
|
||||
const metadata = t.metadata as Record<string, unknown> | null;
|
||||
return [
|
||||
t.sourceMessageId!,
|
||||
{
|
||||
clientMode: metadata?.clientMode as boolean | undefined,
|
||||
duration: metadata?.duration as number | undefined,
|
||||
status: t.status as ThreadStatus,
|
||||
threadId: t.threadId,
|
||||
title: t.title ?? undefined,
|
||||
totalCost: metadata?.totalCost as number | undefined,
|
||||
totalMessages: metadata?.totalMessages as number | undefined,
|
||||
totalTokens: metadata?.totalTokens as number | undefined,
|
||||
totalToolCalls: metadata?.totalToolCalls as number | undefined,
|
||||
},
|
||||
];
|
||||
}),
|
||||
);
|
||||
}
|
||||
|
||||
// 6. Transform messages to UIChatMessage format
|
||||
return result.map(
|
||||
({ model, provider, translate, ttsId, ttsFile, ttsContentMd5, ttsVoice, ...item }) => {
|
||||
const messageQuery = messageQueriesList.find((relation) => relation.messageId === item.id);
|
||||
return {
|
||||
...item,
|
||||
chunksList: chunksList
|
||||
.filter((relation) => relation.messageId === item.id)
|
||||
.map((c) => ({
|
||||
...c,
|
||||
similarity: c.similarity === null ? undefined : Number(c.similarity),
|
||||
})),
|
||||
|
||||
extra: {
|
||||
model: model,
|
||||
provider: provider,
|
||||
translate,
|
||||
tts: ttsId
|
||||
? {
|
||||
contentMd5: ttsContentMd5,
|
||||
file: ttsFile,
|
||||
voice: ttsVoice,
|
||||
}
|
||||
: undefined,
|
||||
},
|
||||
fileList: fileList
|
||||
.filter((relation) => relation.messageId === item.id)
|
||||
.map<ChatFileItem>(({ id, url, size, fileType, name }) => ({
|
||||
content: documentsMap[id],
|
||||
fileType: fileType!,
|
||||
id,
|
||||
name: name!,
|
||||
size: size!,
|
||||
url,
|
||||
})),
|
||||
imageList: imageList
|
||||
.filter((relation) => relation.messageId === item.id)
|
||||
.map<ChatImageItem>(({ id, url, name }) => ({ alt: name!, id, url })),
|
||||
|
||||
model,
|
||||
|
||||
provider,
|
||||
ragQuery: messageQuery?.rewriteQuery,
|
||||
ragQueryId: messageQuery?.id,
|
||||
ragRawQuery: messageQuery?.userQuery,
|
||||
// Add taskDetail for task messages
|
||||
taskDetail: item.role === 'task' ? threadMap.get(item.id as string) : undefined,
|
||||
videoList: videoList
|
||||
.filter((relation) => relation.messageId === item.id)
|
||||
.map<ChatVideoItem>(({ id, url, name }) => ({ alt: name!, id, url })),
|
||||
} as unknown as UIChatMessage;
|
||||
},
|
||||
);
|
||||
};
|
||||
|
||||
/**
|
||||
* Query MessageGroup nodes for a topic
|
||||
* - compressedGroup: includes pinnedMessages array
|
||||
* - compressedGroup: includes pinnedMessages and compressedMessages arrays
|
||||
* - compareGroup: includes children array
|
||||
*
|
||||
* @param topicId - The topic ID to query groups for
|
||||
@@ -525,6 +785,7 @@ export class MessageModel {
|
||||
private queryMessageGroupNodes = async (
|
||||
topicId: string,
|
||||
timeRange?: { endTime: Date; startTime: Date },
|
||||
postProcessUrl?: (path: string | null, file: { fileType: string }) => Promise<string>,
|
||||
): Promise<UIChatMessage[]> => {
|
||||
// 1. Query MessageGroups for this topic, optionally filtered by time range
|
||||
const whereConditions = [
|
||||
@@ -550,50 +811,77 @@ export class MessageModel {
|
||||
|
||||
const groupIds = groups.map((g) => g.id);
|
||||
|
||||
// 2. Query all messages that belong to these groups (for pinnedMessages and children)
|
||||
const groupMessages = await this.db
|
||||
// 2. Get all message IDs that belong to these groups (using messageGroupId relation)
|
||||
const groupMessageRecords = await this.db
|
||||
.select({
|
||||
content: messages.content,
|
||||
createdAt: messages.createdAt,
|
||||
favorite: messages.favorite,
|
||||
id: messages.id,
|
||||
messageGroupId: messages.messageGroupId,
|
||||
model: messages.model,
|
||||
provider: messages.provider,
|
||||
role: messages.role,
|
||||
})
|
||||
.from(messages)
|
||||
.where(and(eq(messages.userId, this.userId), inArray(messages.messageGroupId, groupIds)))
|
||||
.orderBy(asc(messages.createdAt));
|
||||
|
||||
// 3. Build MessageGroup nodes
|
||||
// 3. Query full message data using queryByIds (reuses all transformation logic)
|
||||
const allMessageIds = groupMessageRecords.map((m) => m.id as string);
|
||||
const fullMessages = await this.queryByIds(allMessageIds, { postProcessUrl });
|
||||
|
||||
// Create a map for quick lookup
|
||||
const messageMap = new Map(fullMessages.map((m) => [m.id, m]));
|
||||
const favoriteMap = new Map(groupMessageRecords.map((m) => [m.id, m.favorite]));
|
||||
|
||||
// 4. Build MessageGroup nodes
|
||||
return groups.map((group) => {
|
||||
const groupMsgs = groupMessages.filter((m) => m.messageGroupId === group.id);
|
||||
// Get messages for this group
|
||||
const groupMsgIds = groupMessageRecords
|
||||
.filter((m) => m.messageGroupId === group.id)
|
||||
.map((m) => m.id as string);
|
||||
|
||||
const groupMsgs = groupMsgIds
|
||||
.map((id) => messageMap.get(id))
|
||||
.filter(Boolean) as UIChatMessage[];
|
||||
|
||||
if (group.type === MessageGroupType.Compression) {
|
||||
// compressedGroup: extract pinnedMessages (favorite=true)
|
||||
const pinnedMessages = groupMsgs
|
||||
.filter((m) => m.favorite === true)
|
||||
.map((m) => ({
|
||||
content: m.content,
|
||||
createdAt: m.createdAt,
|
||||
id: m.id,
|
||||
model: m.model,
|
||||
provider: m.provider,
|
||||
role: m.role,
|
||||
}));
|
||||
const pinnedMessages = groupMsgIds
|
||||
.filter((id) => favoriteMap.get(id) === true)
|
||||
.map((id) => {
|
||||
const m = messageMap.get(id);
|
||||
return m
|
||||
? {
|
||||
content: m.content,
|
||||
createdAt: m.createdAt,
|
||||
id: m.id,
|
||||
model: m.model,
|
||||
provider: m.provider,
|
||||
role: m.role,
|
||||
}
|
||||
: null;
|
||||
})
|
||||
.filter(Boolean);
|
||||
|
||||
// compressedMessages: parse messages through conversation-flow for proper grouping
|
||||
// This transforms raw messages into displayMessages format (e.g., assistantGroup)
|
||||
const { flatList } = parse(groupMsgs);
|
||||
const compressedMessages = flatList;
|
||||
|
||||
// Get the last message ID for parent-child linking in conversation-flow
|
||||
const lastMessageId = groupMsgIds.at(-1);
|
||||
|
||||
return {
|
||||
compressedMessages,
|
||||
content: group.content,
|
||||
createdAt: group.createdAt,
|
||||
id: group.id,
|
||||
lastMessageId,
|
||||
metadata: group.metadata,
|
||||
pinnedMessages,
|
||||
role: 'compressedGroup',
|
||||
topicId: group.topicId,
|
||||
updatedAt: group.updatedAt,
|
||||
} as unknown as UIChatMessage;
|
||||
} else {
|
||||
// compareGroup (parallel): include children
|
||||
// compareGroup (parallel): include children with basic info
|
||||
const children = groupMsgs.map((m) => ({
|
||||
content: m.content,
|
||||
createdAt: m.createdAt,
|
||||
|
||||
@@ -128,6 +128,28 @@ export class CompressionRepository {
|
||||
.where(and(eq(messageGroups.id, groupId), eq(messageGroups.userId, this.userId)));
|
||||
}
|
||||
|
||||
/**
|
||||
* Update compression group metadata (UI state like expanded)
|
||||
*/
|
||||
async updateMetadata(
|
||||
groupId: string,
|
||||
metadata: Partial<CompressionGroupMetadata>,
|
||||
): Promise<void> {
|
||||
// Get existing metadata and merge
|
||||
const existing = await this.db
|
||||
.select({ metadata: messageGroups.metadata })
|
||||
.from(messageGroups)
|
||||
.where(and(eq(messageGroups.id, groupId), eq(messageGroups.userId, this.userId)));
|
||||
|
||||
const existingData = (existing[0]?.metadata as Record<string, unknown>) || {};
|
||||
const newMetadata = { ...existingData, ...metadata };
|
||||
|
||||
await this.db
|
||||
.update(messageGroups)
|
||||
.set({ metadata: newMetadata, updatedAt: new Date() })
|
||||
.where(and(eq(messageGroups.id, groupId), eq(messageGroups.userId, this.userId)));
|
||||
}
|
||||
|
||||
/**
|
||||
* Mark messages as compressed by associating them with a compression group
|
||||
*/
|
||||
|
||||
@@ -62,6 +62,7 @@ export const messageGroups = pgTable(
|
||||
type: text('type', { enum: ['parallel', 'compression'] }),
|
||||
content: text('content'), // compression summary (plain text)
|
||||
editorData: jsonb('editor_data'), // rich text editor data (future extension)
|
||||
metadata: jsonb('metadata'), // UI state (expanded, etc.)
|
||||
|
||||
clientId: varchar255('client_id'),
|
||||
|
||||
|
||||
@@ -0,0 +1,26 @@
|
||||
import { ChatStreamPayload, UIChatMessage } from '@lobechat/types';
|
||||
|
||||
import {
|
||||
chatHistoryPrompts,
|
||||
compressContextSystemPrompt,
|
||||
compressContextUserPrompt,
|
||||
} from '../prompts';
|
||||
|
||||
/**
|
||||
* Chain for compressing conversation context into a summary
|
||||
* Used when conversation history exceeds token threshold
|
||||
*/
|
||||
export const chainCompressContext = (messages: UIChatMessage[]): Partial<ChatStreamPayload> => ({
|
||||
messages: [
|
||||
{
|
||||
content: compressContextSystemPrompt,
|
||||
role: 'system',
|
||||
},
|
||||
{
|
||||
content: `${chatHistoryPrompts(messages)}
|
||||
|
||||
${compressContextUserPrompt}`,
|
||||
role: 'user',
|
||||
},
|
||||
],
|
||||
});
|
||||
@@ -1,5 +1,6 @@
|
||||
export * from './abstractChunk';
|
||||
export * from './answerWithContext';
|
||||
export * from './compressContext';
|
||||
export * from './langDetect';
|
||||
export * from './pickEmoji';
|
||||
export * from './rewriteQuery';
|
||||
|
||||
@@ -0,0 +1,64 @@
|
||||
/**
|
||||
* Conversation Context Compression Prompt
|
||||
*
|
||||
* This prompt is designed to compress conversation history while preserving
|
||||
* essential information for conversation continuity.
|
||||
*/
|
||||
|
||||
export const compressContextSystemPrompt = `You are a conversation context compressor. Your task is to create a structured summary that preserves essential information while significantly reducing token count.
|
||||
|
||||
## Output Format
|
||||
|
||||
Structure your summary using these sections (omit empty sections):
|
||||
|
||||
### Context
|
||||
Brief background and conversation setup (1-2 sentences max)
|
||||
|
||||
### Key Information
|
||||
- Critical facts, data, specifications mentioned
|
||||
- Technical details, configurations, parameters
|
||||
- Names, identifiers, file paths, URLs
|
||||
|
||||
### Decisions & Conclusions
|
||||
- Decisions made during the conversation
|
||||
- Agreed-upon solutions or approaches
|
||||
- Final conclusions reached
|
||||
|
||||
### Action Items
|
||||
- Tasks assigned or planned
|
||||
- Next steps discussed
|
||||
- Pending items requiring follow-up
|
||||
|
||||
### Code & Technical
|
||||
\`\`\`
|
||||
Preserve essential code snippets, commands, or technical syntax
|
||||
\`\`\`
|
||||
|
||||
## Rules
|
||||
|
||||
### MUST
|
||||
- Output in the SAME LANGUAGE as the conversation
|
||||
- Preserve ALL technical terms, code identifiers, file paths, and proper nouns exactly
|
||||
- Maintain factual accuracy - never invent or assume information
|
||||
- Keep code snippets that are essential for context
|
||||
|
||||
### SHOULD
|
||||
- Achieve 60-80% compression ratio (summary should be 20-40% of original length)
|
||||
- Use bullet points for clarity and scannability
|
||||
- Preserve chronological order for sequential events
|
||||
- Consolidate repeated information into single entries
|
||||
|
||||
### MAY
|
||||
- Omit greetings, pleasantries, and filler content
|
||||
- Combine related points into concise statements
|
||||
- Abbreviate obvious context when meaning is preserved
|
||||
|
||||
## Important Notes
|
||||
|
||||
- The summary will be injected into a new conversation as context
|
||||
- Recipient should be able to continue the conversation seamlessly
|
||||
- Prioritize information that affects future responses`;
|
||||
|
||||
export const compressContextUserPrompt = `Please compress the above conversation history.
|
||||
|
||||
Output ONLY the structured summary following the format specified. No additional commentary or meta-discussion.`;
|
||||
@@ -1,6 +1,7 @@
|
||||
export * from './agentBuilder';
|
||||
export * from './agentGroup';
|
||||
export * from './chatMessages';
|
||||
export * from './compressContext';
|
||||
export * from './files';
|
||||
export * from './fileSystem';
|
||||
export * from './groupChat';
|
||||
|
||||
@@ -67,9 +67,20 @@ export interface LobeAgentChatConfig {
|
||||
enableHistoryCount?: boolean;
|
||||
/**
|
||||
* Enable history message compression threshold
|
||||
* @deprecated Use enableContextCompression instead
|
||||
*/
|
||||
enableCompressHistory?: boolean;
|
||||
|
||||
/**
|
||||
* Enable context compression
|
||||
* When enabled, old messages will be compressed into summaries when token threshold is reached
|
||||
*/
|
||||
enableContextCompression?: boolean;
|
||||
/**
|
||||
* Model ID to use for generating compression summaries
|
||||
*/
|
||||
compressionModelId?: string;
|
||||
|
||||
inputTemplate?: string;
|
||||
|
||||
searchMode?: SearchMode;
|
||||
@@ -88,9 +99,11 @@ export const LocalSystemConfigSchema = z.object({
|
||||
|
||||
export const AgentChatConfigSchema = z.object({
|
||||
autoCreateTopicThreshold: z.number().default(2),
|
||||
compressionModelId: z.string().optional(),
|
||||
disableContextCaching: z.boolean().optional(),
|
||||
enableAutoCreateTopic: z.boolean().optional(),
|
||||
enableCompressHistory: z.boolean().optional(),
|
||||
enableContextCompression: z.boolean().optional(),
|
||||
enableHistoryCount: z.boolean().optional(),
|
||||
enableMaxTokens: z.boolean().optional(),
|
||||
enableReasoning: z.boolean().optional(),
|
||||
|
||||
@@ -29,7 +29,9 @@ export interface CompressionGroupMetadata {
|
||||
// Compression info
|
||||
compressionStrategy?: 'summarize';
|
||||
compressedAt?: string;
|
||||
modelId?: string;
|
||||
|
||||
// UI state
|
||||
expanded?: boolean;
|
||||
}
|
||||
|
||||
/**
|
||||
|
||||
@@ -105,6 +105,11 @@ export interface UIChatMessage {
|
||||
*/
|
||||
children?: AssistantContentBlock[];
|
||||
chunksList?: ChatFileChunk[];
|
||||
/**
|
||||
* All messages within a compression group (role: 'compressedGroup')
|
||||
* Used for rendering expanded view with conversation-flow parsing
|
||||
*/
|
||||
compressedMessages?: UIChatMessage[];
|
||||
content: string;
|
||||
createdAt: number;
|
||||
error?: ChatMessageError | null;
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
import { usePrevious, useUnmount } from 'ahooks';
|
||||
import { useMount, usePrevious, useUnmount } from 'ahooks';
|
||||
import { useEffect } from 'react';
|
||||
import { useParams } from 'react-router-dom';
|
||||
import { createStoreUpdater } from 'zustand-utils';
|
||||
@@ -24,6 +24,10 @@ const AgentIdSync = () => {
|
||||
}
|
||||
}, [params.aid, prevAgentId]);
|
||||
|
||||
useMount(() => {
|
||||
useChatStore.setState({ activeAgentId: params.aid }, false, 'AgentIdSync/mountAgentId');
|
||||
});
|
||||
|
||||
// Clear activeAgentId when unmounting (leaving chat page)
|
||||
useUnmount(() => {
|
||||
useAgentStore.setState({ activeAgentId: undefined }, false, 'AgentIdSync/unmountAgentId');
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
'use client';
|
||||
|
||||
import { Button, Flexbox, Icon, Text } from '@lobehub/ui';
|
||||
import { Button, Icon } from '@lobehub/ui';
|
||||
import { App } from 'antd';
|
||||
import { CalendarClockIcon } from 'lucide-react';
|
||||
import { memo, useState } from 'react';
|
||||
@@ -52,18 +52,18 @@ const AnalysisTrigger = memo<Props>(({ footerNote, range, onRangeChange }) => {
|
||||
loading={submitting || isValidating}
|
||||
onClick={() => setOpen(true)}
|
||||
size={'large'}
|
||||
type={'primary'}
|
||||
style={{ maxWidth: 300 }}
|
||||
type={'primary'}
|
||||
>
|
||||
{t('analysis.action.button')}
|
||||
</Button>
|
||||
|
||||
<DateRangeModal
|
||||
footerNote={footerNote}
|
||||
open={open}
|
||||
onCancel={() => setOpen(false)}
|
||||
onChange={onRangeChange}
|
||||
onSubmit={handleSubmit}
|
||||
open={open}
|
||||
range={range}
|
||||
submitting={submitting}
|
||||
/>
|
||||
|
||||
@@ -3,7 +3,7 @@
|
||||
import { Flexbox, Text } from '@lobehub/ui';
|
||||
import { DatePicker, Modal } from 'antd';
|
||||
import type { RangePickerProps } from 'antd/es/date-picker';
|
||||
import dayjs, { type Dayjs } from 'dayjs';
|
||||
import dayjs from 'dayjs';
|
||||
import { memo, useCallback } from 'react';
|
||||
import { useTranslation } from 'react-i18next';
|
||||
|
||||
@@ -43,16 +43,10 @@ const DateRangeModal = memo<Props>(
|
||||
disabledDate={disabledDate}
|
||||
format={'YYYY/MM/DD'}
|
||||
onChange={(values) =>
|
||||
onChange([
|
||||
values?.[0]?.toDate() ?? null,
|
||||
values?.[1]?.toDate() ?? null,
|
||||
])
|
||||
onChange([values?.[0]?.toDate() ?? null, values?.[1]?.toDate() ?? null])
|
||||
}
|
||||
style={{ width: '100%' }}
|
||||
value={[
|
||||
range[0] ? dayjs(range[0]) : null,
|
||||
range[1] ? dayjs(range[1]) : null,
|
||||
]}
|
||||
value={[range[0] ? dayjs(range[0]) : null, range[1] ? dayjs(range[1]) : null]}
|
||||
/>
|
||||
<Text fontSize={12} type={'secondary'}>
|
||||
{footerNote}
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
'use client';
|
||||
|
||||
import { Flexbox } from '@lobehub/ui';
|
||||
import { AsyncTaskStatus } from '@lobechat/types';
|
||||
import { Flexbox } from '@lobehub/ui';
|
||||
import { memo, useMemo } from 'react';
|
||||
|
||||
import AnalysisAction from './Action';
|
||||
@@ -13,8 +13,7 @@ const MemoryAnalysis = memo(() => {
|
||||
|
||||
const { showAction, showStatus } = useMemo(() => {
|
||||
const status = data?.status;
|
||||
const isRunning =
|
||||
status === AsyncTaskStatus.Pending || status === AsyncTaskStatus.Processing;
|
||||
const isRunning = status === AsyncTaskStatus.Pending || status === AsyncTaskStatus.Processing;
|
||||
const isError = status === AsyncTaskStatus.Error;
|
||||
|
||||
console.log(isRunning, isValidating, isError, data);
|
||||
@@ -28,7 +27,7 @@ const MemoryAnalysis = memo(() => {
|
||||
if (!showAction && !showStatus) return null;
|
||||
|
||||
return (
|
||||
<Flexbox gap={12} style={{ width: '100%', paddingTop: 16 }}>
|
||||
<Flexbox gap={12} style={{ paddingTop: 16, width: '100%' }}>
|
||||
{showStatus && <MemoryAnalysisStatus task={data} />}
|
||||
{showAction && <AnalysisAction />}
|
||||
</Flexbox>
|
||||
|
||||
@@ -31,7 +31,7 @@ export const useMemoryAnalysisAsyncTask = (taskId?: string) => {
|
||||
|
||||
const timer = setInterval(() => {
|
||||
swr.mutate();
|
||||
}, 5_000);
|
||||
}, 5000);
|
||||
|
||||
return () => clearInterval(timer);
|
||||
}, [swr.data?.id, swr.data?.status, swr.mutate]);
|
||||
|
||||
@@ -113,16 +113,17 @@ const VirtualizedList = memo<VirtualizedListProps>(({ dataSource, itemContent })
|
||||
};
|
||||
}, [resetVisibleItems]);
|
||||
|
||||
// Get the last message to check if it's a user message
|
||||
// Get the second-to-last message to check if it's a user message
|
||||
// (When sending a message, user + assistant messages are created as a pair)
|
||||
const displayMessages = useConversationStore(dataSelectors.displayMessages);
|
||||
const lastMessage = displayMessages.at(-1);
|
||||
const isLastMessageFromUser = lastMessage?.role === 'user';
|
||||
const secondLastMessage = displayMessages.at(-2);
|
||||
const isSecondLastMessageFromUser = secondLastMessage?.role === 'user';
|
||||
|
||||
// Auto scroll to user message when user sends a new message
|
||||
// Only scroll when the new message is from the user, not when AI/agent responds
|
||||
// Only scroll when 2 new messages are added and second-to-last is from user
|
||||
useScrollToUserMessage({
|
||||
dataSourceLength: dataSource.length,
|
||||
isLastMessageFromUser,
|
||||
isSecondLastMessageFromUser,
|
||||
scrollToIndex: virtuaRef.current?.scrollToIndex ?? null,
|
||||
});
|
||||
|
||||
@@ -174,7 +175,11 @@ const VirtualizedList = memo<VirtualizedListProps>(({ dataSource, itemContent })
|
||||
}}
|
||||
</VList>
|
||||
{/* BackBottom 放在 VList 外面，这样无论滚动到哪里都能看到 */}
|
||||
<BackBottom atBottom={atBottom} onScrollToBottom={() => scrollToBottom(true)} visible={!atBottom} />
|
||||
<BackBottom
|
||||
atBottom={atBottom}
|
||||
onScrollToBottom={() => scrollToBottom(true)}
|
||||
visible={!atBottom}
|
||||
/>
|
||||
</div>
|
||||
);
|
||||
}, isEqual);
|
||||
|
||||
@@ -5,56 +5,57 @@ import { useScrollToUserMessage } from './useScrollToUserMessage';
|
||||
|
||||
describe('useScrollToUserMessage', () => {
|
||||
describe('when user sends a new message', () => {
|
||||
it('should scroll to user message when new message is from user', () => {
|
||||
it('should scroll to user message when 2 new messages are added (user + assistant pair)', () => {
|
||||
const scrollToIndex = vi.fn();
|
||||
|
||||
const { rerender } = renderHook(
|
||||
({ dataSourceLength, isLastMessageFromUser }) =>
|
||||
({ dataSourceLength, isSecondLastMessageFromUser }) =>
|
||||
useScrollToUserMessage({
|
||||
dataSourceLength,
|
||||
isLastMessageFromUser,
|
||||
isSecondLastMessageFromUser,
|
||||
scrollToIndex,
|
||||
}),
|
||||
{
|
||||
initialProps: {
|
||||
dataSourceLength: 2,
|
||||
isLastMessageFromUser: false,
|
||||
isSecondLastMessageFromUser: false,
|
||||
},
|
||||
},
|
||||
);
|
||||
|
||||
// User sends a new message (length increases, last message is from user)
|
||||
// User sends a new message (2 messages added: user + assistant, second-to-last is user)
|
||||
rerender({
|
||||
dataSourceLength: 3,
|
||||
isLastMessageFromUser: true,
|
||||
dataSourceLength: 4,
|
||||
isSecondLastMessageFromUser: true,
|
||||
});
|
||||
|
||||
expect(scrollToIndex).toHaveBeenCalledTimes(1);
|
||||
expect(scrollToIndex).toHaveBeenCalledWith(1, { align: 'start', smooth: true });
|
||||
// Should scroll to index 2 (dataSourceLength - 2 = 4 - 2 = 2, the user message)
|
||||
expect(scrollToIndex).toHaveBeenCalledWith(2, { align: 'start', smooth: true });
|
||||
});
|
||||
|
||||
it('should scroll to correct index when multiple user messages are sent', () => {
|
||||
const scrollToIndex = vi.fn();
|
||||
|
||||
const { rerender } = renderHook(
|
||||
({ dataSourceLength, isLastMessageFromUser }) =>
|
||||
({ dataSourceLength, isSecondLastMessageFromUser }) =>
|
||||
useScrollToUserMessage({
|
||||
dataSourceLength,
|
||||
isLastMessageFromUser,
|
||||
isSecondLastMessageFromUser,
|
||||
scrollToIndex,
|
||||
}),
|
||||
{
|
||||
initialProps: {
|
||||
dataSourceLength: 5,
|
||||
isLastMessageFromUser: false,
|
||||
dataSourceLength: 4,
|
||||
isSecondLastMessageFromUser: false,
|
||||
},
|
||||
},
|
||||
);
|
||||
|
||||
// User sends a new message
|
||||
// User sends a new message (2 messages added)
|
||||
rerender({
|
||||
dataSourceLength: 6,
|
||||
isLastMessageFromUser: true,
|
||||
isSecondLastMessageFromUser: true,
|
||||
});
|
||||
|
||||
// Should scroll to index 4 (dataSourceLength - 2 = 6 - 2 = 4)
|
||||
@@ -63,28 +64,28 @@ describe('useScrollToUserMessage', () => {
|
||||
});
|
||||
|
||||
describe('when AI/agent responds', () => {
|
||||
it('should NOT scroll when new message is from AI', () => {
|
||||
it('should NOT scroll when only 1 new message is added (AI response)', () => {
|
||||
const scrollToIndex = vi.fn();
|
||||
|
||||
const { rerender } = renderHook(
|
||||
({ dataSourceLength, isLastMessageFromUser }) =>
|
||||
({ dataSourceLength, isSecondLastMessageFromUser }) =>
|
||||
useScrollToUserMessage({
|
||||
dataSourceLength,
|
||||
isLastMessageFromUser,
|
||||
isSecondLastMessageFromUser,
|
||||
scrollToIndex,
|
||||
}),
|
||||
{
|
||||
initialProps: {
|
||||
dataSourceLength: 2,
|
||||
isLastMessageFromUser: true,
|
||||
dataSourceLength: 4,
|
||||
isSecondLastMessageFromUser: true,
|
||||
},
|
||||
},
|
||||
);
|
||||
|
||||
// AI responds (length increases, but last message is NOT from user)
|
||||
// AI adds another message (only 1 message added, not 2)
|
||||
rerender({
|
||||
dataSourceLength: 3,
|
||||
isLastMessageFromUser: false,
|
||||
dataSourceLength: 5,
|
||||
isSecondLastMessageFromUser: false,
|
||||
});
|
||||
|
||||
expect(scrollToIndex).not.toHaveBeenCalled();
|
||||
@@ -94,32 +95,59 @@ describe('useScrollToUserMessage', () => {
|
||||
const scrollToIndex = vi.fn();
|
||||
|
||||
const { rerender } = renderHook(
|
||||
({ dataSourceLength, isLastMessageFromUser }) =>
|
||||
({ dataSourceLength, isSecondLastMessageFromUser }) =>
|
||||
useScrollToUserMessage({
|
||||
dataSourceLength,
|
||||
isLastMessageFromUser,
|
||||
isSecondLastMessageFromUser,
|
||||
scrollToIndex,
|
||||
}),
|
||||
{
|
||||
initialProps: {
|
||||
dataSourceLength: 3,
|
||||
isLastMessageFromUser: false,
|
||||
dataSourceLength: 4,
|
||||
isSecondLastMessageFromUser: true,
|
||||
},
|
||||
},
|
||||
);
|
||||
|
||||
// First agent responds
|
||||
// First agent responds (1 message added)
|
||||
rerender({
|
||||
dataSourceLength: 4,
|
||||
isLastMessageFromUser: false,
|
||||
dataSourceLength: 5,
|
||||
isSecondLastMessageFromUser: false,
|
||||
});
|
||||
|
||||
expect(scrollToIndex).not.toHaveBeenCalled();
|
||||
|
||||
// Second agent responds
|
||||
// Second agent responds (1 message added)
|
||||
rerender({
|
||||
dataSourceLength: 5,
|
||||
isLastMessageFromUser: false,
|
||||
dataSourceLength: 6,
|
||||
isSecondLastMessageFromUser: false,
|
||||
});
|
||||
|
||||
expect(scrollToIndex).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('should NOT scroll when 2 messages added but second-to-last is not user', () => {
|
||||
const scrollToIndex = vi.fn();
|
||||
|
||||
const { rerender } = renderHook(
|
||||
({ dataSourceLength, isSecondLastMessageFromUser }) =>
|
||||
useScrollToUserMessage({
|
||||
dataSourceLength,
|
||||
isSecondLastMessageFromUser,
|
||||
scrollToIndex,
|
||||
}),
|
||||
{
|
||||
initialProps: {
|
||||
dataSourceLength: 4,
|
||||
isSecondLastMessageFromUser: false,
|
||||
},
|
||||
},
|
||||
);
|
||||
|
||||
// 2 messages added but both are from AI (e.g., system messages)
|
||||
rerender({
|
||||
dataSourceLength: 6,
|
||||
isSecondLastMessageFromUser: false,
|
||||
});
|
||||
|
||||
expect(scrollToIndex).not.toHaveBeenCalled();
|
||||
@@ -131,16 +159,16 @@ describe('useScrollToUserMessage', () => {
|
||||
const scrollToIndex = vi.fn();
|
||||
|
||||
const { rerender } = renderHook(
|
||||
({ dataSourceLength, isLastMessageFromUser }) =>
|
||||
({ dataSourceLength, isSecondLastMessageFromUser }) =>
|
||||
useScrollToUserMessage({
|
||||
dataSourceLength,
|
||||
isLastMessageFromUser,
|
||||
isSecondLastMessageFromUser,
|
||||
scrollToIndex,
|
||||
}),
|
||||
{
|
||||
initialProps: {
|
||||
dataSourceLength: 5,
|
||||
isLastMessageFromUser: true,
|
||||
dataSourceLength: 6,
|
||||
isSecondLastMessageFromUser: true,
|
||||
},
|
||||
},
|
||||
);
|
||||
@@ -148,7 +176,7 @@ describe('useScrollToUserMessage', () => {
|
||||
// Message deleted (length decreases)
|
||||
rerender({
|
||||
dataSourceLength: 4,
|
||||
isLastMessageFromUser: true,
|
||||
isSecondLastMessageFromUser: true,
|
||||
});
|
||||
|
||||
expect(scrollToIndex).not.toHaveBeenCalled();
|
||||
@@ -158,24 +186,24 @@ describe('useScrollToUserMessage', () => {
|
||||
const scrollToIndex = vi.fn();
|
||||
|
||||
const { rerender } = renderHook(
|
||||
({ dataSourceLength, isLastMessageFromUser }) =>
|
||||
({ dataSourceLength, isSecondLastMessageFromUser }) =>
|
||||
useScrollToUserMessage({
|
||||
dataSourceLength,
|
||||
isLastMessageFromUser,
|
||||
isSecondLastMessageFromUser,
|
||||
scrollToIndex,
|
||||
}),
|
||||
{
|
||||
initialProps: {
|
||||
dataSourceLength: 3,
|
||||
isLastMessageFromUser: true,
|
||||
dataSourceLength: 4,
|
||||
isSecondLastMessageFromUser: true,
|
||||
},
|
||||
},
|
||||
);
|
||||
|
||||
// Length stays the same (content update, not new message)
|
||||
rerender({
|
||||
dataSourceLength: 3,
|
||||
isLastMessageFromUser: true,
|
||||
dataSourceLength: 4,
|
||||
isSecondLastMessageFromUser: true,
|
||||
});
|
||||
|
||||
expect(scrollToIndex).not.toHaveBeenCalled();
|
||||
@@ -183,16 +211,16 @@ describe('useScrollToUserMessage', () => {
|
||||
|
||||
it('should handle null scrollToIndex gracefully', () => {
|
||||
const { rerender } = renderHook(
|
||||
({ dataSourceLength, isLastMessageFromUser }) =>
|
||||
({ dataSourceLength, isSecondLastMessageFromUser }) =>
|
||||
useScrollToUserMessage({
|
||||
dataSourceLength,
|
||||
isLastMessageFromUser,
|
||||
isSecondLastMessageFromUser,
|
||||
scrollToIndex: null,
|
||||
}),
|
||||
{
|
||||
initialProps: {
|
||||
dataSourceLength: 2,
|
||||
isLastMessageFromUser: false,
|
||||
isSecondLastMessageFromUser: false,
|
||||
},
|
||||
},
|
||||
);
|
||||
@@ -200,8 +228,8 @@ describe('useScrollToUserMessage', () => {
|
||||
// Should not throw when scrollToIndex is null
|
||||
expect(() => {
|
||||
rerender({
|
||||
dataSourceLength: 3,
|
||||
isLastMessageFromUser: true,
|
||||
dataSourceLength: 4,
|
||||
isSecondLastMessageFromUser: true,
|
||||
});
|
||||
}).not.toThrow();
|
||||
});
|
||||
@@ -211,13 +239,13 @@ describe('useScrollToUserMessage', () => {
|
||||
|
||||
renderHook(() =>
|
||||
useScrollToUserMessage({
|
||||
dataSourceLength: 5,
|
||||
isLastMessageFromUser: true,
|
||||
dataSourceLength: 6,
|
||||
isSecondLastMessageFromUser: true,
|
||||
scrollToIndex,
|
||||
}),
|
||||
);
|
||||
|
||||
// Should not scroll on initial render even if last message is from user
|
||||
// Should not scroll on initial render even if second-to-last message is from user
|
||||
expect(scrollToIndex).not.toHaveBeenCalled();
|
||||
});
|
||||
});
|
||||
|
||||
@@ -6,9 +6,10 @@ interface UseScrollToUserMessageOptions {
|
||||
*/
|
||||
dataSourceLength: number;
|
||||
/**
|
||||
* Whether the last message is from the user
|
||||
* Whether the second-to-last message is from the user
|
||||
* (When sending a message, user + assistant messages are created as a pair)
|
||||
*/
|
||||
isLastMessageFromUser: boolean;
|
||||
isSecondLastMessageFromUser: boolean;
|
||||
/**
|
||||
* Function to scroll to a specific index
|
||||
*/
|
||||
@@ -19,26 +20,27 @@ interface UseScrollToUserMessageOptions {
|
||||
|
||||
/**
|
||||
* Hook to handle scrolling to user message when user sends a new message.
|
||||
* Only triggers scroll when the new message is from the user, not when AI/agent responds.
|
||||
* Only triggers scroll when user sends a new message (detected by checking if
|
||||
* 2 new messages were added and the second-to-last is from user).
|
||||
*
|
||||
* This ensures that in group chat scenarios, when multiple agents are responding,
|
||||
* the view doesn't jump around as each agent starts speaking.
|
||||
*/
|
||||
export function useScrollToUserMessage({
|
||||
dataSourceLength,
|
||||
isLastMessageFromUser,
|
||||
isSecondLastMessageFromUser,
|
||||
scrollToIndex,
|
||||
}: UseScrollToUserMessageOptions): void {
|
||||
const prevLengthRef = useRef(dataSourceLength);
|
||||
|
||||
useEffect(() => {
|
||||
const hasNewMessage = dataSourceLength > prevLengthRef.current;
|
||||
const newMessageCount = dataSourceLength - prevLengthRef.current;
|
||||
prevLengthRef.current = dataSourceLength;
|
||||
|
||||
// Only scroll when user sends a new message
|
||||
if (hasNewMessage && isLastMessageFromUser && scrollToIndex) {
|
||||
// Only scroll when user sends a new message (2 messages added: user + assistant pair)
|
||||
if (newMessageCount === 2 && isSecondLastMessageFromUser && scrollToIndex) {
|
||||
// Scroll to the second-to-last message (user's message) with the start aligned
|
||||
scrollToIndex(dataSourceLength - 2, { align: 'start', smooth: true });
|
||||
}
|
||||
}, [dataSourceLength, isLastMessageFromUser, scrollToIndex]);
|
||||
}, [dataSourceLength, isSecondLastMessageFromUser, scrollToIndex]);
|
||||
}
|
||||
|
||||
@@ -1,8 +1,11 @@
|
||||
'use client';
|
||||
|
||||
import type { UIChatMessage } from '@lobechat/types';
|
||||
import debug from 'debug';
|
||||
import isEqual from 'fast-deep-equal';
|
||||
import { type ReactNode, memo } from 'react';
|
||||
import { type ReactNode, memo, useMemo } from 'react';
|
||||
|
||||
import { messageMapKey } from '@/store/chat/utils/messageMapKey';
|
||||
|
||||
import StoreUpdater from './StoreUpdater';
|
||||
import { Provider, createStore } from './store';
|
||||
@@ -13,6 +16,8 @@ import type {
|
||||
OperationState,
|
||||
} from './types';
|
||||
|
||||
const log = debug('lobe-render:features:Conversation');
|
||||
|
||||
export interface ConversationProviderProps {
|
||||
/**
|
||||
* Actions bar configuration by message type
|
||||
@@ -75,6 +80,16 @@ export const ConversationProvider = memo<ConversationProviderProps>(
|
||||
operationState,
|
||||
skipFetch,
|
||||
}) => {
|
||||
const contextKey = useMemo(() => messageMapKey(context), [context]);
|
||||
|
||||
log(
|
||||
'[Provider] render | contextKey=%s | messagesCount=%d | hasInitMessages=%s | skipFetch=%s',
|
||||
contextKey,
|
||||
messages?.length ?? 0,
|
||||
hasInitMessages,
|
||||
skipFetch,
|
||||
);
|
||||
|
||||
return (
|
||||
<Provider createStore={() => createStore({ context, hooks, skipFetch })}>
|
||||
<StoreUpdater
|
||||
|
||||
@@ -0,0 +1,75 @@
|
||||
'use client';
|
||||
|
||||
import type { UIChatMessage } from '@lobechat/types';
|
||||
import { Avatar, Flexbox } from '@lobehub/ui';
|
||||
import { memo } from 'react';
|
||||
|
||||
import { useUserAvatar } from '@/hooks/useUserAvatar';
|
||||
|
||||
import { useAgentMeta } from '../../hooks';
|
||||
import ContentBlock from '../AssistantGroup/components/ContentBlock';
|
||||
import UserMessageContent from '../User/components/MessageContent';
|
||||
|
||||
interface CompressedMessageItemProps {
|
||||
message: UIChatMessage;
|
||||
}
|
||||
|
||||
/**
|
||||
* Renders a single message within a compressed group
|
||||
* Reuses existing User and Assistant content components for consistency
|
||||
*/
|
||||
const CompressedMessageItem = memo<CompressedMessageItemProps>(({ message }) => {
|
||||
const userAvatar = useUserAvatar();
|
||||
const agentAvatar = useAgentMeta(message.agentId);
|
||||
const { role, children } = message;
|
||||
|
||||
// Render user message
|
||||
if (role === 'user') {
|
||||
return (
|
||||
<Flexbox gap={8} horizontal paddingBlock={4}>
|
||||
<Avatar avatar={userAvatar} size={28} />
|
||||
<Flexbox flex={1} style={{ overflow: 'hidden' }}>
|
||||
<UserMessageContent {...message} />
|
||||
</Flexbox>
|
||||
</Flexbox>
|
||||
);
|
||||
}
|
||||
|
||||
// Render assistant message (standalone without tools)
|
||||
if (role === 'assistant') {
|
||||
return (
|
||||
<Flexbox gap={8} horizontal paddingBlock={4}>
|
||||
<Avatar {...agentAvatar} size={28} />
|
||||
<Flexbox flex={1} style={{ overflow: 'hidden' }}>
|
||||
<ContentBlock
|
||||
assistantId={message.id}
|
||||
content={message.content}
|
||||
disableEditing
|
||||
id={message.id}
|
||||
/>
|
||||
</Flexbox>
|
||||
</Flexbox>
|
||||
);
|
||||
}
|
||||
|
||||
// Render assistantGroup (assistant message with tool calls)
|
||||
if (role === 'assistantGroup' && children) {
|
||||
return (
|
||||
<Flexbox gap={8} horizontal paddingBlock={4}>
|
||||
<Avatar {...agentAvatar} size={28} />
|
||||
<Flexbox flex={1} gap={8} style={{ overflow: 'hidden' }}>
|
||||
{children.map((block) => (
|
||||
<ContentBlock {...block} assistantId={message.id} disableEditing key={block.id} />
|
||||
))}
|
||||
</Flexbox>
|
||||
</Flexbox>
|
||||
);
|
||||
}
|
||||
|
||||
// Skip other roles (tool, system, etc.)
|
||||
return null;
|
||||
});
|
||||
|
||||
CompressedMessageItem.displayName = 'CompressedMessageItem';
|
||||
|
||||
export default CompressedMessageItem;
|
||||
@@ -0,0 +1,176 @@
|
||||
'use client';
|
||||
|
||||
import type { CompressionGroupMetadata, UIChatMessage } from '@lobechat/types';
|
||||
import {
|
||||
ActionIcon,
|
||||
Flexbox,
|
||||
Icon,
|
||||
Markdown,
|
||||
ScrollShadow,
|
||||
Tabs,
|
||||
type TabsProps,
|
||||
} from '@lobehub/ui';
|
||||
import { createStaticStyles, cx } from 'antd-style';
|
||||
import isEqual from 'fast-deep-equal';
|
||||
import { ChevronDown, ChevronUp, History, Sparkles } from 'lucide-react';
|
||||
import { memo, useCallback, useMemo, useState } from 'react';
|
||||
import { useTranslation } from 'react-i18next';
|
||||
|
||||
import StreamingMarkdown from '@/components/StreamingMarkdown';
|
||||
import { useChatStore } from '@/store/chat';
|
||||
import { operationSelectors } from '@/store/chat/selectors';
|
||||
import { shinyTextStyles } from '@/styles/loading';
|
||||
|
||||
import { dataSelectors, useConversationStore } from '../../store';
|
||||
import CompressedMessageItem from './CompressedMessageItem';
|
||||
|
||||
const STORAGE_KEY_PREFIX = 'compressed-group-tab:';
|
||||
|
||||
const getStoredTab = (id: string): string => {
|
||||
if (typeof window === 'undefined') return 'summary';
|
||||
return localStorage.getItem(`${STORAGE_KEY_PREFIX}${id}`) || 'summary';
|
||||
};
|
||||
|
||||
const setStoredTab = (id: string, tab: string) => {
|
||||
if (typeof window === 'undefined') return;
|
||||
localStorage.setItem(`${STORAGE_KEY_PREFIX}${id}`, tab);
|
||||
};
|
||||
|
||||
const styles = createStaticStyles(({ css, cssVar }) => ({
|
||||
container: css`
|
||||
margin-block-end: 8px;
|
||||
padding-block: 8px;
|
||||
padding-inline: 12px;
|
||||
border: 1px solid ${cssVar.colorBorderSecondary};
|
||||
border-radius: 12px;
|
||||
|
||||
background: ${cssVar.colorBgContainer};
|
||||
`,
|
||||
contentScroll: css`
|
||||
max-height: min(40vh, 400px);
|
||||
`,
|
||||
header: css`
|
||||
.ant-tabs-nav {
|
||||
margin-block-end: 0;
|
||||
}
|
||||
`,
|
||||
messagesContainer: css`
|
||||
padding-block: 8px;
|
||||
`,
|
||||
}));
|
||||
|
||||
export interface CompressedGroupMessageProps {
|
||||
id: string;
|
||||
index: number;
|
||||
}
|
||||
|
||||
const CompressedGroupMessage = memo<CompressedGroupMessageProps>(({ id }) => {
|
||||
const { t } = useTranslation('chat');
|
||||
const [activeTab, setActiveTab] = useState<string>(() => getStoredTab(id));
|
||||
|
||||
const handleTabChange = useCallback(
|
||||
(tab: string) => {
|
||||
setActiveTab(tab);
|
||||
setStoredTab(id, tab);
|
||||
},
|
||||
[id],
|
||||
);
|
||||
|
||||
const message = useConversationStore(dataSelectors.getDisplayMessageById(id), isEqual);
|
||||
const toggleCompressedGroupExpanded = useConversationStore(
|
||||
(s) => s.toggleCompressedGroupExpanded,
|
||||
);
|
||||
|
||||
const content = message?.content;
|
||||
const rawCompressedMessages = (message as UIChatMessage)?.compressedMessages;
|
||||
const expanded = (message?.metadata as CompressionGroupMetadata)?.expanded ?? true;
|
||||
|
||||
// Filter out placeholder assistant message (content === '...' without tools)
|
||||
const compressedMessages = useMemo(() => {
|
||||
if (!rawCompressedMessages || rawCompressedMessages.length === 0) return rawCompressedMessages;
|
||||
|
||||
const lastMsg = rawCompressedMessages.at(-1);
|
||||
const isPlaceholder =
|
||||
lastMsg &&
|
||||
(lastMsg.role === 'assistant' || lastMsg.role === 'assistantGroup') &&
|
||||
lastMsg.content === '...' &&
|
||||
(!lastMsg.tools || lastMsg.tools.length === 0) &&
|
||||
(!lastMsg.children || lastMsg.children.length === 0);
|
||||
|
||||
return isPlaceholder ? rawCompressedMessages.slice(0, -1) : rawCompressedMessages;
|
||||
}, [rawCompressedMessages]);
|
||||
|
||||
// Check if generateSummary operation is running for this message
|
||||
const runningOp = useChatStore(operationSelectors.getDeepestRunningOperationByMessage(id));
|
||||
const isGeneratingSummary = runningOp?.type === 'generateSummary';
|
||||
|
||||
// Auto-expand when generating summary to show streaming content
|
||||
const showContent = expanded || isGeneratingSummary;
|
||||
|
||||
const tabItems: TabsProps['items'] = useMemo(
|
||||
() => [
|
||||
{
|
||||
icon: <Icon icon={Sparkles} size={14} />,
|
||||
key: 'summary',
|
||||
label: t('compression.summary'),
|
||||
},
|
||||
{
|
||||
icon: <Icon icon={History} size={14} />,
|
||||
key: 'history',
|
||||
label: t('compression.history'),
|
||||
},
|
||||
],
|
||||
[],
|
||||
);
|
||||
|
||||
return (
|
||||
<Flexbox className={styles.container} gap={8}>
|
||||
{isGeneratingSummary ? (
|
||||
<>
|
||||
<Flexbox horizontal>
|
||||
{/*<Icon icon={FolderArchive} size={14} />*/}
|
||||
<span className={cx(isGeneratingSummary ? shinyTextStyles.shinyText : '')}>
|
||||
{t('compressedHistory')}
|
||||
</span>
|
||||
</Flexbox>
|
||||
<StreamingMarkdown>{content}</StreamingMarkdown>
|
||||
</>
|
||||
) : (
|
||||
<Flexbox align={'center'} distribution={'space-between'} horizontal width={'100%'}>
|
||||
<Tabs
|
||||
activeKey={isGeneratingSummary ? 'summary' : activeTab}
|
||||
className={styles.header}
|
||||
compact
|
||||
items={tabItems}
|
||||
onChange={handleTabChange}
|
||||
variant={'rounded'}
|
||||
/>
|
||||
<ActionIcon
|
||||
icon={expanded ? ChevronUp : ChevronDown}
|
||||
onClick={() => toggleCompressedGroupExpanded(id)}
|
||||
size={'small'}
|
||||
/>
|
||||
</Flexbox>
|
||||
)}
|
||||
{!showContent ? null : activeTab === 'summary' ? (
|
||||
<ScrollShadow className={styles.contentScroll} offset={12} size={12}>
|
||||
<Markdown style={{ overflow: 'unset' }} variant={'chat'}>
|
||||
{content}
|
||||
</Markdown>
|
||||
</ScrollShadow>
|
||||
) : (
|
||||
<ScrollShadow className={styles.contentScroll} offset={12} size={12}>
|
||||
<Flexbox className={styles.messagesContainer} gap={4}>
|
||||
{compressedMessages?.map((msg) => (
|
||||
<CompressedMessageItem key={msg.id} message={msg} />
|
||||
))}
|
||||
</Flexbox>
|
||||
</ScrollShadow>
|
||||
)}
|
||||
</Flexbox>
|
||||
);
|
||||
});
|
||||
|
||||
CompressedGroupMessage.displayName = 'CompressedGroupMessage';
|
||||
|
||||
export default CompressedGroupMessage;
|
||||
@@ -3,9 +3,11 @@ import { memo, useEffect, useState } from 'react';
|
||||
import { useTranslation } from 'react-i18next';
|
||||
|
||||
import BubblesLoading from '@/components/BubblesLoading';
|
||||
import NeuralNetworkLoading from '@/components/NeuralNetworkLoading';
|
||||
import { useChatStore } from '@/store/chat';
|
||||
import { operationSelectors } from '@/store/chat/selectors';
|
||||
import type { OperationType } from '@/store/chat/slices/operation/types';
|
||||
import { shinyTextStyles } from '@/styles/loading';
|
||||
|
||||
const ELAPSED_TIME_THRESHOLD = 2100; // Show elapsed time after 2 seconds
|
||||
|
||||
@@ -56,6 +58,15 @@ const ContentLoading = memo<ContentLoadingProps>(({ id }) => {
|
||||
|
||||
if (operationType && NO_NEED_SHOW_DOT_OP_TYPES.has(operationType)) return null;
|
||||
|
||||
if (operationType === 'contextCompression') {
|
||||
return (
|
||||
<Flexbox align={'center'} gap={8} horizontal>
|
||||
<NeuralNetworkLoading size={16} />
|
||||
<span className={shinyTextStyles.shinyText}>{t('operation.contextCompression')}</span>
|
||||
</Flexbox>
|
||||
);
|
||||
}
|
||||
|
||||
return (
|
||||
<Flexbox align={'center'} horizontal>
|
||||
<BubblesLoading />
|
||||
|
||||
@@ -14,6 +14,7 @@ import { dataSelectors, messageStateSelectors, useConversationStore } from '../s
|
||||
import AgentCouncilMessage from './AgentCouncil';
|
||||
import AssistantMessage from './Assistant';
|
||||
import AssistantGroupMessage from './AssistantGroup';
|
||||
import CompressedGroupMessage from './CompressedGroup';
|
||||
import SupervisorMessage from './Supervisor';
|
||||
import TaskMessage from './Task';
|
||||
import TasksMessage from './Tasks';
|
||||
@@ -157,6 +158,10 @@ const MessageItem = memo<MessageItemProps>(
|
||||
return <AgentCouncilMessage id={id} index={index} />;
|
||||
}
|
||||
|
||||
case 'compressedGroup': {
|
||||
return <CompressedGroupMessage id={id} index={index} />;
|
||||
}
|
||||
|
||||
case 'tool': {
|
||||
return <ToolMessage disableEditing={disableEditing} id={id} index={index} />;
|
||||
}
|
||||
|
||||
@@ -1,9 +1,12 @@
|
||||
'use client';
|
||||
|
||||
import type { UIChatMessage } from '@lobechat/types';
|
||||
import { memo, useEffect } from 'react';
|
||||
import debug from 'debug';
|
||||
import { memo, useEffect, useRef } from 'react';
|
||||
import { createStoreUpdater } from 'zustand-utils';
|
||||
|
||||
import { messageMapKey } from '@/store/chat/utils/messageMapKey';
|
||||
|
||||
import { useConversationStoreApi } from './store';
|
||||
import type {
|
||||
ActionsBarConfig,
|
||||
@@ -12,6 +15,8 @@ import type {
|
||||
OperationState,
|
||||
} from './types';
|
||||
|
||||
const log = debug('lobe-render:features:Conversation');
|
||||
|
||||
export interface StoreUpdaterProps {
|
||||
/**
|
||||
* Actions bar configuration by message type
|
||||
@@ -54,6 +59,8 @@ const StoreUpdater = memo<StoreUpdaterProps>(
|
||||
}) => {
|
||||
const storeApi = useConversationStoreApi();
|
||||
const useStoreUpdater = createStoreUpdater(storeApi);
|
||||
const prevMessagesRef = useRef<UIChatMessage[] | undefined>(undefined);
|
||||
const contextKey = messageMapKey(context);
|
||||
|
||||
useStoreUpdater('actionsBar', actionsBar);
|
||||
useStoreUpdater('context', context);
|
||||
@@ -68,9 +75,26 @@ const StoreUpdater = memo<StoreUpdaterProps>(
|
||||
// Sync external messages into store
|
||||
useEffect(() => {
|
||||
if (messages) {
|
||||
const prevMessages = prevMessagesRef.current;
|
||||
const prevCount = prevMessages?.length ?? 0;
|
||||
const newCount = messages.length;
|
||||
const isSameReference = prevMessages === messages;
|
||||
const storeMessages = storeApi.getState().dbMessages;
|
||||
|
||||
log(
|
||||
'[StoreUpdater] messages effect | contextKey=%s | prevCount=%d | newCount=%d | sameRef=%s | storeCount=%d | messageIds=%o',
|
||||
contextKey,
|
||||
prevCount,
|
||||
newCount,
|
||||
isSameReference,
|
||||
storeMessages.length,
|
||||
messages.slice(0, 5).map((m) => m.id),
|
||||
);
|
||||
|
||||
prevMessagesRef.current = messages;
|
||||
storeApi.getState().replaceMessages(messages);
|
||||
}
|
||||
}, [messages, storeApi]);
|
||||
}, [messages, storeApi, contextKey]);
|
||||
|
||||
return null;
|
||||
},
|
||||
|
||||
@@ -164,6 +164,79 @@ describe('DataSlice', () => {
|
||||
const displayMsg = state.displayMessages.find((m) => m.id === 'msg-1');
|
||||
expect(displayMsg?.metadata?.collapsed).toBe(true);
|
||||
});
|
||||
|
||||
it('should update messageGroup metadata (compressedGroup)', () => {
|
||||
const store = createTestStore();
|
||||
|
||||
// Simulate a compressedGroup in displayMessages (injected during query, not in dbMessages)
|
||||
const compressedGroup: UIChatMessage = {
|
||||
id: 'group-1',
|
||||
content: 'Summary content',
|
||||
role: 'compressedGroup' as any,
|
||||
createdAt: 1000,
|
||||
updatedAt: 1000,
|
||||
metadata: { expanded: false } as any,
|
||||
};
|
||||
|
||||
// Manually set displayMessages to include compressedGroup
|
||||
store.setState({ displayMessages: [compressedGroup] });
|
||||
|
||||
// Update messageGroup metadata
|
||||
store.getState().internal_dispatchMessage({
|
||||
type: 'updateMessageGroupMetadata',
|
||||
id: 'group-1',
|
||||
value: { expanded: true },
|
||||
});
|
||||
|
||||
const state = store.getState();
|
||||
// compressedGroup should only be in displayMessages, not dbMessages
|
||||
expect(state.dbMessages).toHaveLength(0);
|
||||
expect(state.displayMessages).toHaveLength(1);
|
||||
expect((state.displayMessages[0].metadata as any)?.expanded).toBe(true);
|
||||
});
|
||||
|
||||
it('should not update messageGroup metadata if message does not exist', () => {
|
||||
const store = createTestStore();
|
||||
|
||||
const initialDisplayMessages = store.getState().displayMessages;
|
||||
|
||||
store.getState().internal_dispatchMessage({
|
||||
type: 'updateMessageGroupMetadata',
|
||||
id: 'nonexistent',
|
||||
value: { expanded: true },
|
||||
});
|
||||
|
||||
// State should remain unchanged
|
||||
expect(store.getState().displayMessages).toBe(initialDisplayMessages);
|
||||
});
|
||||
|
||||
it('should merge messageGroup metadata with existing values', () => {
|
||||
const store = createTestStore();
|
||||
|
||||
// Simulate a compressedGroup with existing metadata
|
||||
const compressedGroup: UIChatMessage = {
|
||||
id: 'group-1',
|
||||
content: 'Summary content',
|
||||
role: 'compressedGroup' as any,
|
||||
createdAt: 1000,
|
||||
updatedAt: 1000,
|
||||
metadata: { expanded: false, someOtherField: 'preserved' } as any,
|
||||
};
|
||||
|
||||
store.setState({ displayMessages: [compressedGroup] });
|
||||
|
||||
// Update only expanded
|
||||
store.getState().internal_dispatchMessage({
|
||||
type: 'updateMessageGroupMetadata',
|
||||
id: 'group-1',
|
||||
value: { expanded: true },
|
||||
});
|
||||
|
||||
const state = store.getState();
|
||||
const metadata = state.displayMessages[0].metadata as any;
|
||||
expect(metadata?.expanded).toBe(true);
|
||||
expect(metadata?.someOtherField).toBe('preserved');
|
||||
});
|
||||
});
|
||||
|
||||
describe('replaceMessages', () => {
|
||||
|
||||
@@ -1,15 +1,19 @@
|
||||
import { parse } from '@lobechat/conversation-flow';
|
||||
import type { ConversationContext, UIChatMessage } from '@lobechat/types';
|
||||
import debug from 'debug';
|
||||
import type { SWRResponse } from 'swr';
|
||||
import type { StateCreator } from 'zustand/vanilla';
|
||||
|
||||
import { useClientDataSWRWithSync } from '@/libs/swr';
|
||||
import { messageService } from '@/services/message';
|
||||
import { messageMapKey } from '@/store/chat/utils/messageMapKey';
|
||||
|
||||
import type { Store as ConversationStore } from '../../action';
|
||||
import { type MessageDispatch, messagesReducer } from './reducer';
|
||||
import { dataSelectors } from './selectors';
|
||||
|
||||
const log = debug('lobe-render:features:Conversation');
|
||||
|
||||
/**
|
||||
* Data Actions
|
||||
*
|
||||
@@ -57,17 +61,57 @@ export const dataSlice: StateCreator<
|
||||
DataAction
|
||||
> = (set, get) => ({
|
||||
internal_dispatchMessage: (payload) => {
|
||||
const contextKey = messageMapKey(get().context);
|
||||
|
||||
log(
|
||||
'[dispatchMessage] start | contextKey=%s | type=%s | id=%s',
|
||||
contextKey,
|
||||
payload.type,
|
||||
'id' in payload ? payload.id : 'ids' in payload ? payload.ids.join(',') : 'N/A',
|
||||
);
|
||||
|
||||
// Special handling for messageGroup metadata updates
|
||||
// MessageGroups are not in dbMessages, they're injected during query
|
||||
if (payload.type === 'updateMessageGroupMetadata') {
|
||||
const displayMessages = get().displayMessages;
|
||||
const index = displayMessages.findIndex((m) => m.id === payload.id);
|
||||
if (index < 0) return;
|
||||
|
||||
const newDisplayMessages = [...displayMessages];
|
||||
newDisplayMessages[index] = {
|
||||
...newDisplayMessages[index],
|
||||
metadata: { ...newDisplayMessages[index].metadata, ...payload.value },
|
||||
};
|
||||
|
||||
set({ displayMessages: newDisplayMessages }, false, {
|
||||
payload,
|
||||
type: `dispatchMessage/${payload.type}`,
|
||||
});
|
||||
return;
|
||||
}
|
||||
|
||||
const dbMessages = get().dbMessages;
|
||||
|
||||
// Apply array-based reducer - preserves message order
|
||||
const newDbMessages = messagesReducer(dbMessages, payload);
|
||||
|
||||
// Check if anything changed
|
||||
if (newDbMessages === dbMessages) return;
|
||||
if (newDbMessages === dbMessages) {
|
||||
log('[dispatchMessage] no change | contextKey=%s', contextKey);
|
||||
return;
|
||||
}
|
||||
|
||||
// Re-parse for display order and grouping
|
||||
const { flatList } = parse(newDbMessages);
|
||||
|
||||
log(
|
||||
'[dispatchMessage] updated | contextKey=%s | prevCount=%d | newCount=%d | displayCount=%d',
|
||||
contextKey,
|
||||
dbMessages.length,
|
||||
newDbMessages.length,
|
||||
flatList.length,
|
||||
);
|
||||
|
||||
set({ dbMessages: newDbMessages, displayMessages: flatList }, false, {
|
||||
payload,
|
||||
type: `dispatchMessage/${payload.type}`,
|
||||
@@ -78,9 +122,21 @@ export const dataSlice: StateCreator<
|
||||
},
|
||||
|
||||
replaceMessages: (messages) => {
|
||||
const contextKey = messageMapKey(get().context);
|
||||
const prevDbMessages = get().dbMessages;
|
||||
|
||||
// Parse messages using conversation-flow
|
||||
const { flatList } = parse(messages);
|
||||
|
||||
log(
|
||||
'[replaceMessages] | contextKey=%s | prevCount=%d | newCount=%d | displayCount=%d | messageIds=%o',
|
||||
contextKey,
|
||||
prevDbMessages.length,
|
||||
messages.length,
|
||||
flatList.length,
|
||||
messages.slice(0, 5).map((m) => m.id),
|
||||
);
|
||||
|
||||
set({ dbMessages: messages, displayMessages: flatList }, false, 'replaceMessages');
|
||||
|
||||
// Sync changes to external store (ChatStore)
|
||||
@@ -106,6 +162,16 @@ export const dataSlice: StateCreator<
|
||||
// Also skip fetch when topicId is null (new conversation state) - there's no server data,
|
||||
// only local optimistic updates. Fetching would return empty array and overwrite local data.
|
||||
const shouldFetch = !skipFetch && !!context.agentId && !!context.topicId;
|
||||
const contextKey = messageMapKey(context);
|
||||
|
||||
log(
|
||||
'[useFetchMessages] hook | contextKey=%s | shouldFetch=%s | skipFetch=%s | agentId=%s | topicId=%s',
|
||||
contextKey,
|
||||
shouldFetch,
|
||||
skipFetch,
|
||||
context.agentId,
|
||||
context.topicId,
|
||||
);
|
||||
|
||||
return useClientDataSWRWithSync<UIChatMessage[]>(
|
||||
shouldFetch ? ['CONVERSATION_FETCH_MESSAGES', context] : null,
|
||||
@@ -116,9 +182,22 @@ export const dataSlice: StateCreator<
|
||||
if (!data) return;
|
||||
if (!context.topicId) return;
|
||||
|
||||
const prevDbMessages = get().dbMessages;
|
||||
const storeContextKey = messageMapKey(get().context);
|
||||
|
||||
// Parse messages using conversation-flow
|
||||
const { flatList } = parse(data);
|
||||
|
||||
log(
|
||||
'[useFetchMessages] onData | requestContextKey=%s | storeContextKey=%s | prevCount=%d | fetchedCount=%d | displayCount=%d | messageIds=%o',
|
||||
contextKey,
|
||||
storeContextKey,
|
||||
prevDbMessages.length,
|
||||
data.length,
|
||||
flatList.length,
|
||||
data.slice(0, 5).map((m) => m.id),
|
||||
);
|
||||
|
||||
set({
|
||||
dbMessages: data,
|
||||
displayMessages: flatList,
|
||||
|
||||
@@ -84,6 +84,12 @@ interface UpdateMessageMetadata {
|
||||
value: Partial<UIChatMessage['metadata']>;
|
||||
}
|
||||
|
||||
interface UpdateMessageGroupMetadata {
|
||||
id: string;
|
||||
type: 'updateMessageGroupMetadata';
|
||||
value: Record<string, unknown>;
|
||||
}
|
||||
|
||||
export type MessageDispatch =
|
||||
| CreateMessage
|
||||
| UpdateMessage
|
||||
@@ -91,6 +97,7 @@ export type MessageDispatch =
|
||||
| UpdatePluginState
|
||||
| UpdateMessageExtra
|
||||
| UpdateMessageMetadata
|
||||
| UpdateMessageGroupMetadata
|
||||
| DeleteMessage
|
||||
| UpdateMessagePlugin
|
||||
| UpdateMessageTools
|
||||
|
||||
@@ -0,0 +1,278 @@
|
||||
import { UIChatMessage } from '@lobechat/types';
|
||||
import { beforeEach, describe, expect, it, vi } from 'vitest';
|
||||
|
||||
import { messageService } from '@/services/message';
|
||||
|
||||
import { createStore } from '../../../index';
|
||||
|
||||
// Mock conversation-flow parse function
|
||||
vi.mock('@lobechat/conversation-flow', () => ({
|
||||
parse: (messages: UIChatMessage[]) => {
|
||||
const messageMap: Record<string, UIChatMessage> = {};
|
||||
for (const msg of messages) {
|
||||
messageMap[msg.id] = msg;
|
||||
}
|
||||
const flatList = [...messages].sort((a, b) => a.createdAt - b.createdAt);
|
||||
return { flatList, messageMap };
|
||||
},
|
||||
}));
|
||||
|
||||
// Mock messageService
|
||||
vi.mock('@/services/message', () => ({
|
||||
messageService: {
|
||||
getMessages: vi.fn(),
|
||||
updateMessageGroupMetadata: vi.fn(),
|
||||
updateMessageMetadata: vi.fn().mockResolvedValue({ success: true, messages: [] }),
|
||||
},
|
||||
}));
|
||||
|
||||
// Mock SWR
|
||||
vi.mock('@/libs/swr', () => ({
|
||||
useClientDataSWRWithSync: vi.fn(() => ({ data: undefined, isLoading: true })),
|
||||
}));
|
||||
|
||||
const createTestStore = (options?: { agentId?: string; topicId?: string | null }) =>
|
||||
createStore({
|
||||
context: {
|
||||
agentId: options?.agentId ?? 'test-agent',
|
||||
threadId: null,
|
||||
topicId: options?.topicId === null ? null : (options?.topicId ?? 'test-topic'),
|
||||
},
|
||||
});
|
||||
|
||||
describe('MessageStateAction', () => {
|
||||
beforeEach(() => {
|
||||
vi.clearAllMocks();
|
||||
});
|
||||
|
||||
describe('toggleCompressedGroupExpanded', () => {
|
||||
it('should toggle expanded state from false to true', async () => {
|
||||
const store = createTestStore();
|
||||
|
||||
// Setup: compressedGroup with expanded=false
|
||||
const compressedGroup: UIChatMessage = {
|
||||
id: 'group-1',
|
||||
content: 'Summary content',
|
||||
role: 'compressedGroup' as any,
|
||||
createdAt: 1000,
|
||||
updatedAt: 1000,
|
||||
metadata: { expanded: false } as any,
|
||||
};
|
||||
store.setState({ displayMessages: [compressedGroup] });
|
||||
|
||||
// Mock API response
|
||||
const updatedMessages: UIChatMessage[] = [
|
||||
{ ...compressedGroup, metadata: { expanded: true } as any },
|
||||
];
|
||||
vi.mocked(messageService.updateMessageGroupMetadata).mockResolvedValue({
|
||||
messages: updatedMessages,
|
||||
});
|
||||
|
||||
// Act
|
||||
await store.getState().toggleCompressedGroupExpanded('group-1');
|
||||
|
||||
// Assert: optimistic update should have been called
|
||||
expect(messageService.updateMessageGroupMetadata).toHaveBeenCalledWith({
|
||||
context: {
|
||||
agentId: 'test-agent',
|
||||
groupId: undefined,
|
||||
threadId: null,
|
||||
topicId: 'test-topic',
|
||||
},
|
||||
expanded: true,
|
||||
messageGroupId: 'group-1',
|
||||
});
|
||||
});
|
||||
|
||||
it('should toggle expanded state from true to false', async () => {
|
||||
const store = createTestStore();
|
||||
|
||||
// Setup: compressedGroup with expanded=true
|
||||
const compressedGroup: UIChatMessage = {
|
||||
id: 'group-1',
|
||||
content: 'Summary content',
|
||||
role: 'compressedGroup' as any,
|
||||
createdAt: 1000,
|
||||
updatedAt: 1000,
|
||||
metadata: { expanded: true } as any,
|
||||
};
|
||||
store.setState({ displayMessages: [compressedGroup] });
|
||||
|
||||
// Mock API response
|
||||
const updatedMessages: UIChatMessage[] = [
|
||||
{ ...compressedGroup, metadata: { expanded: false } as any },
|
||||
];
|
||||
vi.mocked(messageService.updateMessageGroupMetadata).mockResolvedValue({
|
||||
messages: updatedMessages,
|
||||
});
|
||||
|
||||
// Act
|
||||
await store.getState().toggleCompressedGroupExpanded('group-1');
|
||||
|
||||
// Assert
|
||||
expect(messageService.updateMessageGroupMetadata).toHaveBeenCalledWith({
|
||||
context: {
|
||||
agentId: 'test-agent',
|
||||
groupId: undefined,
|
||||
threadId: null,
|
||||
topicId: 'test-topic',
|
||||
},
|
||||
expanded: false,
|
||||
messageGroupId: 'group-1',
|
||||
});
|
||||
});
|
||||
|
||||
it('should set specific expanded value when provided', async () => {
|
||||
const store = createTestStore();
|
||||
|
||||
// Setup: compressedGroup with expanded=false
|
||||
const compressedGroup: UIChatMessage = {
|
||||
id: 'group-1',
|
||||
content: 'Summary content',
|
||||
role: 'compressedGroup' as any,
|
||||
createdAt: 1000,
|
||||
updatedAt: 1000,
|
||||
metadata: { expanded: false } as any,
|
||||
};
|
||||
store.setState({ displayMessages: [compressedGroup] });
|
||||
|
||||
vi.mocked(messageService.updateMessageGroupMetadata).mockResolvedValue({
|
||||
messages: [{ ...compressedGroup, metadata: { expanded: true } as any }],
|
||||
});
|
||||
|
||||
// Act: explicitly set to true
|
||||
await store.getState().toggleCompressedGroupExpanded('group-1', true);
|
||||
|
||||
// Assert
|
||||
expect(messageService.updateMessageGroupMetadata).toHaveBeenCalledWith({
|
||||
context: {
|
||||
agentId: 'test-agent',
|
||||
groupId: undefined,
|
||||
threadId: null,
|
||||
topicId: 'test-topic',
|
||||
},
|
||||
expanded: true,
|
||||
messageGroupId: 'group-1',
|
||||
});
|
||||
});
|
||||
|
||||
it('should not call API if message does not exist', async () => {
|
||||
const store = createTestStore();
|
||||
|
||||
// Act
|
||||
await store.getState().toggleCompressedGroupExpanded('nonexistent');
|
||||
|
||||
// Assert
|
||||
expect(messageService.updateMessageGroupMetadata).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('should not call API if message is not compressedGroup', async () => {
|
||||
const store = createTestStore();
|
||||
|
||||
// Setup: regular user message
|
||||
const userMessage: UIChatMessage = {
|
||||
id: 'msg-1',
|
||||
content: 'Hello',
|
||||
role: 'user',
|
||||
createdAt: 1000,
|
||||
updatedAt: 1000,
|
||||
};
|
||||
store.setState({ displayMessages: [userMessage] });
|
||||
|
||||
// Act
|
||||
await store.getState().toggleCompressedGroupExpanded('msg-1');
|
||||
|
||||
// Assert
|
||||
expect(messageService.updateMessageGroupMetadata).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('should not call API if context is missing topicId', async () => {
|
||||
const store = createTestStore({ topicId: null });
|
||||
|
||||
// Setup: compressedGroup
|
||||
const compressedGroup: UIChatMessage = {
|
||||
id: 'group-1',
|
||||
content: 'Summary content',
|
||||
role: 'compressedGroup' as any,
|
||||
createdAt: 1000,
|
||||
updatedAt: 1000,
|
||||
metadata: { expanded: false } as any,
|
||||
};
|
||||
// Use shallow merge (default zustand behavior)
|
||||
store.setState({ displayMessages: [compressedGroup] });
|
||||
|
||||
// Verify context is still null
|
||||
expect(store.getState().context.topicId).toBeNull();
|
||||
|
||||
// Act
|
||||
await store.getState().toggleCompressedGroupExpanded('group-1');
|
||||
|
||||
// Assert: should not call API because topicId is null
|
||||
expect(messageService.updateMessageGroupMetadata).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('should default to false when metadata.expanded is undefined', async () => {
|
||||
const store = createTestStore();
|
||||
|
||||
// Setup: compressedGroup without expanded in metadata
|
||||
const compressedGroup: UIChatMessage = {
|
||||
id: 'group-1',
|
||||
content: 'Summary content',
|
||||
role: 'compressedGroup' as any,
|
||||
createdAt: 1000,
|
||||
updatedAt: 1000,
|
||||
metadata: {} as any,
|
||||
};
|
||||
store.setState({ displayMessages: [compressedGroup] });
|
||||
|
||||
vi.mocked(messageService.updateMessageGroupMetadata).mockResolvedValue({
|
||||
messages: [{ ...compressedGroup, metadata: { expanded: true } as any }],
|
||||
});
|
||||
|
||||
// Act: toggle from undefined (treated as false) to true
|
||||
await store.getState().toggleCompressedGroupExpanded('group-1');
|
||||
|
||||
// Assert: should toggle to true
|
||||
expect(messageService.updateMessageGroupMetadata).toHaveBeenCalledWith({
|
||||
context: {
|
||||
agentId: 'test-agent',
|
||||
groupId: undefined,
|
||||
threadId: null,
|
||||
topicId: 'test-topic',
|
||||
},
|
||||
expanded: true,
|
||||
messageGroupId: 'group-1',
|
||||
});
|
||||
});
|
||||
|
||||
it('should call replaceMessages with updated data from API', async () => {
|
||||
const store = createTestStore();
|
||||
|
||||
// Setup
|
||||
const compressedGroup: UIChatMessage = {
|
||||
id: 'group-1',
|
||||
content: 'Summary content',
|
||||
role: 'compressedGroup' as any,
|
||||
createdAt: 1000,
|
||||
updatedAt: 1000,
|
||||
metadata: { expanded: false } as any,
|
||||
};
|
||||
store.setState({ displayMessages: [compressedGroup] });
|
||||
|
||||
const replaceMessagesSpy = vi.spyOn(store.getState(), 'replaceMessages');
|
||||
|
||||
const updatedMessages: UIChatMessage[] = [
|
||||
{ ...compressedGroup, metadata: { expanded: true } as any },
|
||||
];
|
||||
vi.mocked(messageService.updateMessageGroupMetadata).mockResolvedValue({
|
||||
messages: updatedMessages,
|
||||
});
|
||||
|
||||
// Act
|
||||
await store.getState().toggleCompressedGroupExpanded('group-1');
|
||||
|
||||
// Assert
|
||||
expect(replaceMessagesSpy).toHaveBeenCalledWith(updatedMessages);
|
||||
});
|
||||
});
|
||||
});
|
||||
@@ -2,6 +2,8 @@ import { copyToClipboard } from '@lobehub/ui';
|
||||
import { produce } from 'immer';
|
||||
import type { StateCreator } from 'zustand';
|
||||
|
||||
import { messageService } from '@/services/message';
|
||||
|
||||
import type { Store as ConversationStore } from '../../../action';
|
||||
import { dataSelectors } from '../../data/selectors';
|
||||
|
||||
@@ -26,6 +28,11 @@ export interface MessageStateAction {
|
||||
*/
|
||||
modifyMessageContent: (id: string, content: string) => Promise<void>;
|
||||
|
||||
/**
|
||||
* Toggle compressed group expanded state
|
||||
*/
|
||||
toggleCompressedGroupExpanded: (id: string, expanded?: boolean) => Promise<void>;
|
||||
|
||||
/**
|
||||
* Toggle tool inspect expanded state
|
||||
*/
|
||||
@@ -87,6 +94,40 @@ export const messageStateSlice: StateCreator<
|
||||
}
|
||||
},
|
||||
|
||||
toggleCompressedGroupExpanded: async (id, expanded) => {
|
||||
const message = dataSelectors.getDisplayMessageById(id)(get());
|
||||
if (!message || message.role !== 'compressedGroup') return;
|
||||
|
||||
const { context, internal_dispatchMessage, replaceMessages } = get();
|
||||
if (!context.agentId || !context.topicId) return;
|
||||
|
||||
// If expanded is not provided, toggle current state
|
||||
const currentExpanded = (message.metadata as any)?.expanded ?? false;
|
||||
const nextExpanded = expanded ?? !currentExpanded;
|
||||
|
||||
// Optimistic update
|
||||
internal_dispatchMessage({
|
||||
id,
|
||||
type: 'updateMessageGroupMetadata',
|
||||
value: { expanded: nextExpanded },
|
||||
});
|
||||
|
||||
// Persist to server and get updated messages
|
||||
const { messages } = await messageService.updateMessageGroupMetadata({
|
||||
context: {
|
||||
agentId: context.agentId,
|
||||
groupId: context.groupId,
|
||||
threadId: context.threadId,
|
||||
topicId: context.topicId,
|
||||
},
|
||||
expanded: nextExpanded,
|
||||
messageGroupId: id,
|
||||
});
|
||||
|
||||
// Sync with server data
|
||||
replaceMessages(messages);
|
||||
},
|
||||
|
||||
toggleInspectExpanded: async (id, expanded) => {
|
||||
const message = dataSelectors.getDbMessageById(id)(get());
|
||||
if (!message) return;
|
||||
|
||||
@@ -40,6 +40,9 @@ export default {
|
||||
'chatList.expandMessage': 'Expand Message',
|
||||
'chatList.longMessageDetail': 'View Details',
|
||||
'clearCurrentMessages': 'Clear current session messages',
|
||||
'compressedHistory': 'Compressed History',
|
||||
'compression.history': 'History',
|
||||
'compression.summary': 'Summary',
|
||||
'confirmClearCurrentMessages':
|
||||
'You are about to clear the current session messages. Once cleared, they cannot be retrieved. Please confirm your action.',
|
||||
'confirmRemoveChatGroupItemAlert':
|
||||
@@ -229,6 +232,7 @@ export default {
|
||||
'noMembersYet': "This group doesn't have any members yet. Click the + button to invite agents.",
|
||||
'noSelectedAgents': 'No members selected yet',
|
||||
'openInNewWindow': 'Open in New Window',
|
||||
'operation.contextCompression': 'Context too long, compressing history...',
|
||||
'operation.execAgentRuntime': 'Preparing response',
|
||||
'operation.execClientTask': 'Executing task',
|
||||
'operation.sendMessage': 'Sending message',
|
||||
|
||||
@@ -61,9 +61,9 @@ export interface MemoryExtractionPrivateConfig {
|
||||
};
|
||||
upstashWorkflowExtraHeaders?: Record<string, string>;
|
||||
webhook: {
|
||||
headers?: Record<string, string>;
|
||||
baseUrl?: string;
|
||||
}
|
||||
headers?: Record<string, string>;
|
||||
};
|
||||
whitelistUsers?: string[];
|
||||
}
|
||||
|
||||
@@ -257,8 +257,8 @@ export const parseMemoryExtractionConfig = (): MemoryExtractionPrivateConfig =>
|
||||
observabilityS3: extractorObservabilityS3,
|
||||
upstashWorkflowExtraHeaders,
|
||||
webhook: {
|
||||
headers: webhookHeaders,
|
||||
baseUrl: process.env.MEMORY_USER_MEMORY_WEBHOOK_BASE_URL,
|
||||
headers: webhookHeaders,
|
||||
},
|
||||
whitelistUsers,
|
||||
};
|
||||
|
||||
@@ -280,7 +280,7 @@ describe('AI Agent Router Integration Tests', () => {
|
||||
.returning();
|
||||
|
||||
// Create a thread (required by foreign key constraint on messages)
|
||||
const [thread] = await serverDB
|
||||
const [thread] = (await serverDB
|
||||
.insert(threads)
|
||||
.values({
|
||||
topicId: topic.id,
|
||||
@@ -288,7 +288,7 @@ describe('AI Agent Router Integration Tests', () => {
|
||||
userId,
|
||||
type: 'isolation',
|
||||
})
|
||||
.returning();
|
||||
.returning()) as any[];
|
||||
|
||||
const caller = aiAgentRouter.createCaller(createTestContext());
|
||||
|
||||
|
||||
@@ -9,6 +9,7 @@ import { z } from 'zod';
|
||||
|
||||
import { MessageModel } from '@/database/models/message';
|
||||
import { TopicShareModel } from '@/database/models/topicShare';
|
||||
import { CompressionRepository } from '@/database/repositories/compression';
|
||||
import { authedProcedure, publicProcedure, router } from '@/libs/trpc/lambda';
|
||||
import { serverDatabase } from '@/libs/trpc/lambda/middleware';
|
||||
import { FileService } from '@/server/services/file';
|
||||
@@ -22,6 +23,7 @@ const messageProcedure = authedProcedure.use(serverDatabase).use(async (opts) =>
|
||||
|
||||
return opts.next({
|
||||
ctx: {
|
||||
compressionRepo: new CompressionRepository(ctx.serverDB, ctx.userId),
|
||||
fileService: new FileService(ctx.serverDB, ctx.userId),
|
||||
messageModel: new MessageModel(ctx.serverDB, ctx.userId),
|
||||
messageService: new MessageService(ctx.serverDB, ctx.userId),
|
||||
@@ -74,6 +76,32 @@ export const messageRouter = router({
|
||||
return ctx.messageModel.countWords(input);
|
||||
}),
|
||||
|
||||
/**
|
||||
* Create a compression group for old messages
|
||||
* Creates a placeholder group, marks messages as compressed
|
||||
* Returns messages to summarize for frontend AI generation
|
||||
*/
|
||||
createCompressionGroup: messageProcedure
|
||||
.input(
|
||||
z.object({
|
||||
agentId: z.string(),
|
||||
groupId: z.string().nullable().optional(),
|
||||
messageIds: z.array(z.string()),
|
||||
threadId: z.string().nullable().optional(),
|
||||
topicId: z.string(),
|
||||
}),
|
||||
)
|
||||
.mutation(async ({ input, ctx }) => {
|
||||
const { topicId, messageIds, agentId, groupId, threadId } = input;
|
||||
|
||||
return ctx.messageService.createCompressionGroup(topicId, messageIds, {
|
||||
agentId,
|
||||
groupId,
|
||||
threadId,
|
||||
topicId,
|
||||
});
|
||||
}),
|
||||
|
||||
createMessage: messageProcedure
|
||||
.input(CreateNewMessageParamsSchema)
|
||||
.mutation(async ({ input, ctx }) => {
|
||||
@@ -87,6 +115,26 @@ export const messageRouter = router({
|
||||
return ctx.messageService.createMessage({ ...input, agentId } as any);
|
||||
}),
|
||||
|
||||
/**
|
||||
* Finalize compression by updating the group with generated summary
|
||||
*/
|
||||
finalizeCompression: messageProcedure
|
||||
.input(
|
||||
z.object({
|
||||
agentId: z.string(),
|
||||
content: z.string(),
|
||||
groupId: z.string().nullable().optional(),
|
||||
messageGroupId: z.string(),
|
||||
threadId: z.string().nullable().optional(),
|
||||
topicId: z.string(),
|
||||
}),
|
||||
)
|
||||
.mutation(async ({ input, ctx }) => {
|
||||
const { messageGroupId, content, ...params } = input;
|
||||
|
||||
return ctx.messageService.finalizeCompression(messageGroupId, content, params);
|
||||
}),
|
||||
|
||||
getHeatmaps: messageProcedure.query(async ({ ctx }) => {
|
||||
return ctx.messageModel.getHeatmaps();
|
||||
}),
|
||||
@@ -236,6 +284,28 @@ export const messageRouter = router({
|
||||
return ctx.messageService.updateMessage(id, value as any, resolved);
|
||||
}),
|
||||
|
||||
/**
|
||||
* Update message group metadata (e.g., expanded state)
|
||||
*/
|
||||
updateMessageGroupMetadata: messageProcedure
|
||||
.input(
|
||||
z.object({
|
||||
context: z.object({
|
||||
agentId: z.string(),
|
||||
groupId: z.string().nullable().optional(),
|
||||
threadId: z.string().nullable().optional(),
|
||||
topicId: z.string(),
|
||||
}),
|
||||
expanded: z.boolean().optional(),
|
||||
messageGroupId: z.string(),
|
||||
}),
|
||||
)
|
||||
.mutation(async ({ input, ctx }) => {
|
||||
const { messageGroupId, expanded, context } = input;
|
||||
|
||||
return ctx.messageService.updateMessageGroupMetadata(messageGroupId, { expanded }, context);
|
||||
}),
|
||||
|
||||
updateMessagePlugin: messageProcedure
|
||||
.input(
|
||||
z
|
||||
|
||||
@@ -4,12 +4,14 @@ import {
|
||||
agents,
|
||||
agentsToSessions,
|
||||
files,
|
||||
messageGroups,
|
||||
messages,
|
||||
sessions,
|
||||
topics,
|
||||
users,
|
||||
} from '@lobechat/database/schemas';
|
||||
import { getTestDB } from '@lobechat/database/test-utils';
|
||||
import { MessageGroupType } from '@lobechat/types';
|
||||
import { eq } from 'drizzle-orm';
|
||||
import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
|
||||
|
||||
@@ -327,4 +329,83 @@ describe('MessageService Integration Tests', () => {
|
||||
expect(result.messages!.every((m) => m.id !== 'msg-2')).toBe(true);
|
||||
});
|
||||
});
|
||||
|
||||
describe('finalizeCompression', () => {
|
||||
it('should create assistant message with correct parentId (last compressed message ID, not messageGroupId)', async () => {
|
||||
const messageService = new MessageService(serverDB, userId);
|
||||
|
||||
// Setup: agent, session, topic
|
||||
await serverDB.insert(agents).values({ id: 'agent-1', userId });
|
||||
await serverDB.insert(sessions).values({ id: 'session-1', userId });
|
||||
await serverDB
|
||||
.insert(agentsToSessions)
|
||||
.values({ agentId: 'agent-1', sessionId: 'session-1', userId });
|
||||
await serverDB.insert(topics).values({ id: 'topic-1', userId, sessionId: 'session-1' });
|
||||
|
||||
// Create messages to be compressed
|
||||
await serverDB.insert(messages).values([
|
||||
{
|
||||
id: 'msg-1',
|
||||
userId,
|
||||
agentId: 'agent-1',
|
||||
topicId: 'topic-1',
|
||||
role: 'user',
|
||||
content: 'message 1',
|
||||
createdAt: new Date('2024-01-01T10:00:00Z'),
|
||||
},
|
||||
{
|
||||
id: 'msg-2',
|
||||
userId,
|
||||
agentId: 'agent-1',
|
||||
topicId: 'topic-1',
|
||||
role: 'assistant',
|
||||
content: 'message 2',
|
||||
createdAt: new Date('2024-01-01T10:01:00Z'),
|
||||
},
|
||||
{
|
||||
id: 'msg-3',
|
||||
userId,
|
||||
agentId: 'agent-1',
|
||||
topicId: 'topic-1',
|
||||
role: 'user',
|
||||
content: 'message 3 - this is the last compressed message',
|
||||
createdAt: new Date('2024-01-01T10:02:00Z'),
|
||||
},
|
||||
]);
|
||||
|
||||
// Create compression group
|
||||
await serverDB.insert(messageGroups).values({
|
||||
id: 'comp-group-1',
|
||||
userId,
|
||||
topicId: 'topic-1',
|
||||
content: '...',
|
||||
type: MessageGroupType.Compression,
|
||||
createdAt: new Date('2024-01-01T10:02:30Z'),
|
||||
});
|
||||
|
||||
// Mark messages as compressed
|
||||
await serverDB
|
||||
.update(messages)
|
||||
.set({ messageGroupId: 'comp-group-1' })
|
||||
.where(eq(messages.topicId, 'topic-1'));
|
||||
|
||||
// Call finalizeCompression
|
||||
const result = await messageService.finalizeCompression('comp-group-1', 'Summary content', {
|
||||
agentId: 'agent-1',
|
||||
topicId: 'topic-1',
|
||||
});
|
||||
|
||||
expect(result.success).toBe(true);
|
||||
expect(result.messages).toBeDefined();
|
||||
|
||||
// Verify compression content was updated
|
||||
const updatedGroup = await serverDB
|
||||
.select()
|
||||
.from(messageGroups)
|
||||
.where(eq(messageGroups.id, 'comp-group-1'));
|
||||
|
||||
expect(updatedGroup).toHaveLength(1);
|
||||
expect(updatedGroup[0].content).toBe('Summary content');
|
||||
});
|
||||
});
|
||||
});
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
import { type LobeChatDatabase } from '@lobechat/database';
|
||||
import { CompressionRepository, type LobeChatDatabase } from '@lobechat/database';
|
||||
import {
|
||||
type CreateMessageParams,
|
||||
type UIChatMessage,
|
||||
@@ -31,10 +31,12 @@ interface CreateMessageResult {
|
||||
export class MessageService {
|
||||
private messageModel: MessageModel;
|
||||
private fileService: FileService;
|
||||
private compressionRepository: CompressionRepository;
|
||||
|
||||
constructor(db: LobeChatDatabase, userId: string) {
|
||||
this.messageModel = new MessageModel(db, userId);
|
||||
this.fileService = new FileService(db, userId);
|
||||
this.compressionRepository = new CompressionRepository(db, userId);
|
||||
}
|
||||
|
||||
/**
|
||||
@@ -261,4 +263,100 @@ export class MessageService {
|
||||
}
|
||||
return this.queryWithSuccess(options);
|
||||
}
|
||||
|
||||
// =============== Compression Methods ===============
|
||||
|
||||
/**
|
||||
* Create a compression group for messages
|
||||
* Creates a placeholder group, marks messages as compressed, and returns updated messages
|
||||
*
|
||||
* @param topicId - The topic ID
|
||||
* @param messageIds - IDs of messages to compress
|
||||
* @param options - Query options for returning updated messages
|
||||
*/
|
||||
async createCompressionGroup(
|
||||
topicId: string,
|
||||
messageIds: string[],
|
||||
options?: QueryOptions,
|
||||
): Promise<{
|
||||
messageGroupId: string;
|
||||
messages?: UIChatMessage[];
|
||||
messagesToSummarize: UIChatMessage[];
|
||||
success: boolean;
|
||||
}> {
|
||||
// 1. Get messages that need to be summarized (before marking them as compressed)
|
||||
const allMessages = await this.messageModel.query(
|
||||
{ topicId, ...options },
|
||||
this.getQueryOptions(),
|
||||
);
|
||||
|
||||
const messagesToSummarize = allMessages.filter((msg) => messageIds.includes(msg.id));
|
||||
|
||||
// 2. Create compression group with placeholder content
|
||||
const messageGroupId = await this.compressionRepository.createCompressionGroup({
|
||||
content: '...', // Placeholder content
|
||||
messageIds,
|
||||
metadata: {
|
||||
originalMessageCount: messageIds.length,
|
||||
},
|
||||
topicId,
|
||||
});
|
||||
|
||||
// 3. Query updated messages (compressed messages will be grouped)
|
||||
const messages = await this.messageModel.query({ topicId, ...options }, this.getQueryOptions());
|
||||
|
||||
return {
|
||||
messageGroupId,
|
||||
messages,
|
||||
messagesToSummarize,
|
||||
success: true,
|
||||
};
|
||||
}
|
||||
|
||||
/**
|
||||
* Finalize compression by updating the group with actual summary content
|
||||
*
|
||||
* @param messageGroupId - The compression group ID
|
||||
* @param content - The generated summary content
|
||||
* @param params - Parameters for querying messages
|
||||
*/
|
||||
async finalizeCompression(
|
||||
messageGroupId: string,
|
||||
content: string,
|
||||
params: {
|
||||
agentId: string;
|
||||
groupId?: string | null;
|
||||
threadId?: string | null;
|
||||
topicId: string;
|
||||
},
|
||||
): Promise<{ messages?: UIChatMessage[]; success: boolean }> {
|
||||
const { agentId, groupId, threadId, topicId } = params;
|
||||
|
||||
// 1. Update compression group with actual content
|
||||
await this.compressionRepository.updateCompressionContent(messageGroupId, content);
|
||||
|
||||
// 2. Query final messages
|
||||
const queryOptions = { agentId, groupId, threadId, topicId };
|
||||
const finalMessages = await this.messageModel.query(queryOptions, this.getQueryOptions());
|
||||
|
||||
return {
|
||||
messages: finalMessages,
|
||||
success: true,
|
||||
};
|
||||
}
|
||||
|
||||
/**
|
||||
* Update message group metadata (e.g., expanded state)
|
||||
*/
|
||||
async updateMessageGroupMetadata(
|
||||
messageGroupId: string,
|
||||
metadata: { expanded?: boolean },
|
||||
context: QueryOptions,
|
||||
): Promise<{ messages: UIChatMessage[] }> {
|
||||
await this.compressionRepository.updateMetadata(messageGroupId, metadata);
|
||||
|
||||
const messages = await this.messageModel.query(context, this.getQueryOptions());
|
||||
|
||||
return { messages };
|
||||
}
|
||||
}
|
||||
|
||||
@@ -215,6 +215,67 @@ export class MessageService {
|
||||
): Promise<UpdateMessageResult> => {
|
||||
return lambdaClient.message.addFilesToMessage.mutate({ ...ctx, fileIds, id });
|
||||
};
|
||||
|
||||
// =============== Compression ===============
|
||||
|
||||
/**
|
||||
* Create a compression group for old messages
|
||||
* Returns placeholder group and messages to summarize
|
||||
*/
|
||||
createCompressionGroup = async (params: {
|
||||
agentId: string;
|
||||
groupId?: string | null;
|
||||
messageIds: string[];
|
||||
threadId?: string | null;
|
||||
topicId: string;
|
||||
}): Promise<{
|
||||
messageGroupId: string;
|
||||
messages: UIChatMessage[];
|
||||
messagesToSummarize: UIChatMessage[];
|
||||
}> => {
|
||||
const result = await lambdaClient.message.createCompressionGroup.mutate(params);
|
||||
return {
|
||||
messageGroupId: result.messageGroupId,
|
||||
messages: (result.messages || []) as unknown as UIChatMessage[],
|
||||
messagesToSummarize: (result.messagesToSummarize || []) as unknown as UIChatMessage[],
|
||||
};
|
||||
};
|
||||
|
||||
/**
|
||||
* Finalize compression by updating group with generated summary
|
||||
*/
|
||||
finalizeCompression = async (params: {
|
||||
agentId: string;
|
||||
content: string;
|
||||
groupId?: string | null;
|
||||
messageGroupId: string;
|
||||
threadId?: string | null;
|
||||
topicId: string;
|
||||
}): Promise<{ messages?: UIChatMessage[] }> => {
|
||||
const result = await lambdaClient.message.finalizeCompression.mutate(params);
|
||||
return {
|
||||
messages: (result.messages || []) as unknown as UIChatMessage[],
|
||||
};
|
||||
};
|
||||
|
||||
/**
|
||||
* Update message group metadata (e.g., expanded state)
|
||||
*/
|
||||
updateMessageGroupMetadata = async (params: {
|
||||
context: {
|
||||
agentId: string;
|
||||
groupId?: string | null;
|
||||
threadId?: string | null;
|
||||
topicId: string;
|
||||
};
|
||||
expanded?: boolean;
|
||||
messageGroupId: string;
|
||||
}): Promise<{ messages: UIChatMessage[] }> => {
|
||||
const result = await lambdaClient.message.updateMessageGroupMetadata.mutate(params);
|
||||
return {
|
||||
messages: (result.messages || []) as unknown as UIChatMessage[],
|
||||
};
|
||||
};
|
||||
}
|
||||
|
||||
export const messageService = new MessageService();
|
||||
|
||||
@@ -3,6 +3,7 @@ import {
|
||||
type AgentInstruction,
|
||||
type AgentInstructionCallLlm,
|
||||
type AgentInstructionCallTool,
|
||||
type AgentInstructionCompressContext,
|
||||
type AgentInstructionExecClientTask,
|
||||
type AgentInstructionExecClientTasks,
|
||||
type AgentInstructionExecTask,
|
||||
@@ -12,19 +13,24 @@ import {
|
||||
type GeneralAgentCallLLMResultPayload,
|
||||
type GeneralAgentCallToolResultPayload,
|
||||
type GeneralAgentCallingToolInstructionPayload,
|
||||
type GeneralAgentCompressionResultPayload,
|
||||
type InstructionExecutor,
|
||||
type TaskResultPayload,
|
||||
type TasksBatchResultPayload,
|
||||
UsageCounter,
|
||||
calculateMessageTokens,
|
||||
} from '@lobechat/agent-runtime';
|
||||
import { isDesktop } from '@lobechat/const';
|
||||
import { chainCompressContext } from '@lobechat/prompts';
|
||||
import type { ChatToolPayload, ConversationContext, CreateMessageParams } from '@lobechat/types';
|
||||
import debug from 'debug';
|
||||
import pMap from 'p-map';
|
||||
|
||||
import { LOADING_FLAT } from '@/const/message';
|
||||
import { aiAgentService } from '@/services/aiAgent';
|
||||
import { chatService } from '@/services/chat';
|
||||
import type { ResolvedAgentConfig } from '@/services/chat/mecha';
|
||||
import { messageService } from '@/services/message';
|
||||
import { agentByIdSelectors } from '@/store/agent/selectors';
|
||||
import { getAgentStoreState } from '@/store/agent/store';
|
||||
import type { ChatStore } from '@/store/chat/store';
|
||||
@@ -104,7 +110,11 @@ export const createAgentExecutors = (context: {
|
||||
|
||||
let assistantMessageId: string;
|
||||
|
||||
if (shouldSkipCreateMessage) {
|
||||
// Check if we should skip message creation:
|
||||
// - shouldSkipCreateMessage is true (e.g., regenerate mode)
|
||||
// - BUT if createAssistantMessage is explicitly true, always create new message
|
||||
// (e.g., after compression we need a new assistant message)
|
||||
if (shouldSkipCreateMessage && !llmPayload.createAssistantMessage) {
|
||||
// Skip first creation, subsequent calls will not skip
|
||||
assistantMessageId = context.parentId;
|
||||
shouldSkipCreateMessage = false;
|
||||
@@ -2208,6 +2218,230 @@ export const createAgentExecutors = (context: {
|
||||
} as AgentRuntimeContext,
|
||||
};
|
||||
},
|
||||
|
||||
/**
|
||||
* Context compression executor
|
||||
* Compresses ALL messages into a single MessageGroup summary to reduce token usage
|
||||
*/
|
||||
compress_context: async (instruction, state) => {
|
||||
const sessionLogId = `${state.operationId}:${state.stepCount}`;
|
||||
const stagePrefix = `[${sessionLogId}][compress_context]`;
|
||||
|
||||
const { messages, currentTokenCount } = (instruction as AgentInstructionCompressContext)
|
||||
.payload;
|
||||
|
||||
// Get topicId from operation context (same as agentId)
|
||||
const { topicId } = getOperationContext();
|
||||
|
||||
log(
|
||||
`${stagePrefix} Starting compression. displayMessages=%d, tokens=%d`,
|
||||
messages.length,
|
||||
currentTokenCount,
|
||||
);
|
||||
|
||||
const events: AgentEvent[] = [];
|
||||
|
||||
// Get message IDs from dbMessagesMap (raw db messages)
|
||||
const dbMessages = context.get().dbMessagesMap[context.messageKey] || [];
|
||||
const messageIds = dbMessages.map((m) => m.id).filter(Boolean);
|
||||
|
||||
if (!topicId || messageIds.length === 0) {
|
||||
// No topicId or no messages, skip compression
|
||||
log(
|
||||
`${stagePrefix} Skipping compression: topicId=%s, messageIds=%d`,
|
||||
topicId,
|
||||
messageIds.length,
|
||||
);
|
||||
return {
|
||||
events: [],
|
||||
newState: state,
|
||||
nextContext: {
|
||||
payload: {
|
||||
compressedMessages: messages,
|
||||
compressedTokenCount: currentTokenCount,
|
||||
groupId: '',
|
||||
originalTokenCount: currentTokenCount,
|
||||
skipped: true,
|
||||
} as GeneralAgentCompressionResultPayload,
|
||||
phase: 'compression_result',
|
||||
session: {
|
||||
messageCount: state.messages.length,
|
||||
sessionId: state.operationId,
|
||||
status: 'running',
|
||||
stepCount: state.stepCount + 1,
|
||||
},
|
||||
} as AgentRuntimeContext,
|
||||
};
|
||||
}
|
||||
|
||||
// Find the latest assistant message to attach the compression operation
|
||||
const latestAssistantMessage = dbMessages.findLast((m) => m.role === 'assistant');
|
||||
const assistantMessageId = latestAssistantMessage?.id;
|
||||
|
||||
log(
|
||||
`${stagePrefix} Compressing %d db messages (display: %d), assistantMsgId=%s`,
|
||||
messageIds.length,
|
||||
messages.length,
|
||||
assistantMessageId,
|
||||
);
|
||||
|
||||
// Create compress_context operation and attach to the assistant message
|
||||
const { operationId: compressOperationId } = context.get().startOperation({
|
||||
context: { ...getOperationContext(), messageId: assistantMessageId },
|
||||
metadata: {
|
||||
messageCount: messageIds.length,
|
||||
startTime: Date.now(),
|
||||
},
|
||||
parentOperationId: state.operationId,
|
||||
type: 'contextCompression',
|
||||
});
|
||||
|
||||
try {
|
||||
const opContext = getOperationContext();
|
||||
// agentId is guaranteed to exist in compression context
|
||||
const agentId = getEffectiveAgentId()!;
|
||||
|
||||
// 1. Create compression group with placeholder content
|
||||
const result = await messageService.createCompressionGroup({
|
||||
agentId,
|
||||
messageIds,
|
||||
topicId,
|
||||
});
|
||||
const { messageGroupId, messages: initialCompressedMessages, messagesToSummarize } = result;
|
||||
|
||||
// 2. Update UI with compressed messages immediately
|
||||
context.get().replaceMessages(initialCompressedMessages, { context: opContext });
|
||||
|
||||
// 3. Get model/provider from compressionModel config
|
||||
const { model, provider } = state.modelRuntimeConfig?.compressionModel || {};
|
||||
|
||||
log(
|
||||
`${stagePrefix} Created group=%s, generating summary for %d messages by %s`,
|
||||
messageGroupId,
|
||||
messagesToSummarize.length,
|
||||
`${provider}/${model}`,
|
||||
);
|
||||
|
||||
// 4. Build compression prompt and generate summary with streaming UI updates
|
||||
const compressionPayload = chainCompressContext(messagesToSummarize);
|
||||
let summaryContent = '';
|
||||
|
||||
// Start generateSummary operation attached to the compressed group message
|
||||
const { operationId: summaryOperationId } = context.get().startOperation({
|
||||
context: { ...getOperationContext(), messageId: messageGroupId },
|
||||
type: 'generateSummary',
|
||||
parentOperationId: compressOperationId,
|
||||
});
|
||||
|
||||
await chatService.fetchPresetTaskResult({
|
||||
params: { ...compressionPayload, model, provider },
|
||||
onMessageHandle: (chunk) => {
|
||||
if (chunk.type === 'text') {
|
||||
summaryContent += chunk.text || '';
|
||||
// Stream update the compression group message content
|
||||
context
|
||||
.get()
|
||||
.internal_dispatchMessage(
|
||||
{ id: messageGroupId, type: 'updateMessage', value: { content: summaryContent } },
|
||||
{ operationId: summaryOperationId },
|
||||
);
|
||||
}
|
||||
},
|
||||
onError: (e) => {
|
||||
console.error(e);
|
||||
context.get().completeOperation(summaryOperationId, {
|
||||
error: { message: String(e), type: 'summary_generation_failed' },
|
||||
});
|
||||
},
|
||||
});
|
||||
|
||||
log(`${stagePrefix} Generated summary: %d chars`, summaryContent.length);
|
||||
|
||||
// 5. Finalize compression with actual content
|
||||
const finalResult = await messageService.finalizeCompression({
|
||||
agentId,
|
||||
content: summaryContent,
|
||||
messageGroupId,
|
||||
topicId,
|
||||
});
|
||||
// Complete the generateSummary operation
|
||||
context.get().completeOperation(summaryOperationId);
|
||||
|
||||
const compressedMessages = finalResult.messages || initialCompressedMessages;
|
||||
const groupId = messageGroupId;
|
||||
// Use the latest assistant message ID (before compression) as parentMessageId for next call_llm
|
||||
const parentMessageId = assistantMessageId;
|
||||
|
||||
// 6. Update UI with finalized messages (includes compressedGroup with summary)
|
||||
context.get().replaceMessages(compressedMessages, { context: opContext });
|
||||
|
||||
log(
|
||||
`${stagePrefix} Compression complete. groupId=%s, parentMessageId=%s`,
|
||||
groupId,
|
||||
parentMessageId,
|
||||
);
|
||||
|
||||
// Complete the compress_context operation
|
||||
context.get().completeOperation(compressOperationId, { groupId, parentMessageId });
|
||||
|
||||
events.push({ type: 'compression_complete', groupId, parentMessageId });
|
||||
|
||||
// Calculate new token count
|
||||
const compressedTokenCount = calculateMessageTokens(compressedMessages);
|
||||
|
||||
return {
|
||||
events,
|
||||
newState: { ...state, messages: compressedMessages },
|
||||
nextContext: {
|
||||
payload: {
|
||||
compressedMessages,
|
||||
compressedTokenCount,
|
||||
groupId,
|
||||
originalTokenCount: currentTokenCount,
|
||||
parentMessageId,
|
||||
} as GeneralAgentCompressionResultPayload,
|
||||
phase: 'compression_result',
|
||||
session: {
|
||||
messageCount: compressedMessages.length,
|
||||
sessionId: state.operationId,
|
||||
status: 'running',
|
||||
stepCount: state.stepCount + 1,
|
||||
},
|
||||
} as AgentRuntimeContext,
|
||||
};
|
||||
} catch (error) {
|
||||
log(`${stagePrefix} Compression failed: %O`, error);
|
||||
|
||||
// Complete the compress_context operation with error
|
||||
context.get().completeOperation(compressOperationId, {
|
||||
error: {
|
||||
message: error instanceof Error ? error.message : String(error),
|
||||
type: 'compression_failed',
|
||||
},
|
||||
});
|
||||
|
||||
// On error, continue without compression
|
||||
events.push({ type: 'compression_error', error });
|
||||
|
||||
return {
|
||||
events,
|
||||
newState: state,
|
||||
nextContext: {
|
||||
payload: {
|
||||
compressedMessages: messages,
|
||||
skipped: true,
|
||||
} as GeneralAgentCompressionResultPayload,
|
||||
phase: 'compression_result',
|
||||
session: {
|
||||
messageCount: state.messages.length,
|
||||
sessionId: state.operationId,
|
||||
status: 'running',
|
||||
stepCount: state.stepCount + 1,
|
||||
},
|
||||
} as AgentRuntimeContext,
|
||||
};
|
||||
}
|
||||
},
|
||||
};
|
||||
|
||||
return executors;
|
||||
|
||||
@@ -244,18 +244,29 @@ export const streamingExecutor: StateCreator<
|
||||
allowList: toolInterventionSelectors.allowList(userStore),
|
||||
};
|
||||
|
||||
// Build modelRuntimeConfig for compression and other runtime features
|
||||
const modelRuntimeConfig = {
|
||||
compressionModel: {
|
||||
model: agentConfigData.model,
|
||||
provider: agentConfigData.provider!,
|
||||
},
|
||||
model: agentConfigData.model,
|
||||
provider: agentConfigData.provider!,
|
||||
};
|
||||
|
||||
// Create initial state or use provided state
|
||||
const state =
|
||||
initialState ||
|
||||
AgentRuntime.createInitialState({
|
||||
operationId: operationId ?? agentId,
|
||||
messages,
|
||||
maxSteps: 400,
|
||||
messages,
|
||||
metadata: {
|
||||
sessionId: agentId,
|
||||
topicId,
|
||||
threadId,
|
||||
topicId,
|
||||
},
|
||||
modelRuntimeConfig,
|
||||
operationId: operationId ?? agentId,
|
||||
toolManifestMap,
|
||||
userInterventionConfig,
|
||||
});
|
||||
@@ -670,18 +681,21 @@ export const streamingExecutor: StateCreator<
|
||||
const model = agentConfigData.model;
|
||||
const provider = agentConfigData.provider;
|
||||
|
||||
const modelRuntimeConfig = {
|
||||
model,
|
||||
provider: provider!,
|
||||
// TODO: Support dedicated compression model from chatConfig.compressionModelId
|
||||
compressionModel: { model, provider: provider! },
|
||||
};
|
||||
// ===========================================
|
||||
// Step 2: Create and Execute Agent Runtime
|
||||
// ===========================================
|
||||
log('[internal_execAgentRuntime] Creating agent runtime');
|
||||
log('[internal_execAgentRuntime] Creating agent runtime with config', modelRuntimeConfig);
|
||||
|
||||
const agent = new GeneralChatAgent({
|
||||
agentConfig: { maxSteps: 1000 },
|
||||
operationId: `${messageKey}/${params.parentMessageId}`,
|
||||
modelRuntimeConfig: {
|
||||
model,
|
||||
provider: provider!,
|
||||
},
|
||||
modelRuntimeConfig,
|
||||
});
|
||||
|
||||
const runtime = new AgentRuntime(agent, {
|
||||
|
||||
@@ -57,10 +57,13 @@ export type OperationType =
|
||||
| 'execClientTask' // Execute single async sub-agent task on desktop client
|
||||
| 'execClientTasks' // Execute multiple async sub-agent tasks on desktop client
|
||||
|
||||
// === Context Compression ===
|
||||
// Context compression (compress old messages into summary)
|
||||
| 'contextCompression'
|
||||
| 'createMessageGroup'
|
||||
| 'generateSummary'
|
||||
// === Others ===
|
||||
| 'translate' // Translate message
|
||||
| 'topicSummary' // Topic summary
|
||||
| 'historySummary'; // History summary
|
||||
| 'translate'; // Translate message
|
||||
|
||||
/**
|
||||
* Operation status
|
||||
|
||||
@@ -102,6 +102,7 @@ exports[`settingsSelectors > defaultAgent > should merge DEFAULT_AGENT and s.set
|
||||
"autoCreateTopicThreshold": 2,
|
||||
"enableAutoCreateTopic": true,
|
||||
"enableCompressHistory": true,
|
||||
"enableContextCompression": true,
|
||||
"enableHistoryCount": true,
|
||||
"enableReasoning": false,
|
||||
"enableStreaming": true,
|
||||
@@ -146,6 +147,7 @@ exports[`settingsSelectors > defaultAgentConfig > should merge DEFAULT_AGENT_CON
|
||||
"autoCreateTopicThreshold": 2,
|
||||
"enableAutoCreateTopic": true,
|
||||
"enableCompressHistory": true,
|
||||
"enableContextCompression": true,
|
||||
"enableHistoryCount": true,
|
||||
"enableReasoning": false,
|
||||
"enableStreaming": true,
|
||||
|
||||
Reference in New Issue
Block a user