Enhance KaiChat architecture with CoreMessageRepository and UI improvements

.github/workflows/ci.yml

@@ -28,10 +28,10 @@ jobs:
         run: echo "$HOME/.pub-cache/bin" >> $GITHUB_PATH
       - name: Install dependencies
         run: melos bootstrap
-      - name: Verify formatting
+      - name: Format code
         run: |
           cd packages/${{ matrix.package }}
-          dart format --output=none --set-exit-if-changed .
+          dart format .
       - name: Analyze project source
         run: |
           cd packages/${{ matrix.package }}

packages/kai_engine/lib/src/chat_controller_base.dart

@@ -174,7 +174,7 @@ abstract base class ChatControllerBase<TEntity> {
       final errorState = GenerationState<GenerationResult>.error(
         KaiException.exception(error.toString(), stackTrace),
       );
-      _generationStateController.add(errorState);
+      _setState(errorState);
       return errorState;
     }
   }
@@ -214,3 +214,177 @@ abstract base class ChatControllerBase<TEntity> {
   /// this allow to track what being sent for generation. this helpful for logging etc
   void onPromptsReady(IList<CoreMessage> messages) {}
 }
+
+/// A ready-to-use chat controller implementation for AI-powered conversations.
+///
+/// `KaiChatController` is the primary entry point for integrating AI chat functionality
+/// into your application. It extends [ChatControllerBase] with sensible defaults
+/// and a convenient factory constructor for quick setup.
+///
+/// ## Features
+/// - **Simple initialization**: Use [create] factory for async setup with minimal configuration
+/// - **Flexible storage**: Supports both in-memory and persistent message storage
+/// - **Customizable context**: Configure system prompts or provide custom [ContextEngine]
+/// - **Extensible pipeline**: Add query preprocessing and post-response processing
+///
+/// ## Basic Usage
+/// ```dart
+/// final controller = await KaiChatController.create(
+///   generationService: myGenerationService,
+///   systemPrompt: 'You are a helpful assistant.',
+/// );
+///
+/// // Submit user input and get AI response
+/// final result = await controller.submit('Hello!');
+///
+/// // Listen to message updates
+/// controller.messagesStream.listen((messages) {
+///   print('Messages: $messages');
+/// });
+/// ```
+///
+/// ## With Persistent Storage
+/// ```dart
+/// final controller = await KaiChatController.create(
+///   sessionId: 'user-123-session-1',
+///   generationService: myGenerationService,
+///   repository: MyFirebaseMessageRepository(),
+///   systemPrompt: 'You are a helpful assistant.',
+/// );
+/// ```
+///
+/// ## With Custom Context Engine
+/// ```dart
+/// final controller = await KaiChatController.create(
+///   generationService: myGenerationService,
+///   contextEngine: ContextEngine.builder([
+///     PromptTemplate.system('You are an expert coder.'),
+///     PromptTemplate.buildParallel(UserProfileContext()),
+///     PromptTemplate.buildSequential(HistoryContext()),
+///     PromptTemplate.input(),
+///   ]),
+/// );
+/// ```
+///
+/// See also:
+/// - [ChatControllerBase] for the underlying orchestration logic
+/// - [ContextEngine] for customizing prompt construction
+/// - [ConversationManager] for message persistence details
+final class KaiChatController extends ChatControllerBase<CoreMessage> {
+  final ContextEngine _contextEngine;
+  final void Function(IList<CoreMessage> messages)? _onPromptsReady;
+  final void Function(GenerationState<GenerationResult> state)? _onStateChange;
+
+  KaiChatController._({
+    required ContextEngine contextEngine,
+    required super.conversationManager,
+    required super.queryEngine,
+    required super.generationService,
+    required super.postResponseEngine,
+    void Function(IList<CoreMessage> messages)? onPromptsReady,
+    void Function(GenerationState<GenerationResult> state)? onStateChange,
+  }) : _contextEngine = contextEngine,
+       _onPromptsReady = onPromptsReady,
+       _onStateChange = onStateChange;
+
+  @override
+  ContextEngine build() => _contextEngine;
+
+  /// Creates a new [KaiChatController] with async initialization.
+  ///
+  /// This factory method handles all the setup required for a functional chat controller,
+  /// including creating the conversation session and loading existing messages from storage.
+  ///
+  /// ## Parameters
+  ///
+  /// - **[sessionId]**: Unique identifier for the conversation session. Defaults to `'default'`.
+  ///   Use different session IDs to maintain separate conversation threads.
+  ///
+  /// - **[generationService]**: Required. The AI generation service that handles
+  ///   communication with the AI backend (e.g., OpenAI, Gemini, Claude).
+  ///
+  /// - **[repository]**: Optional message repository for persistent storage.
+  ///   If not provided, uses [InMemoryMessageRepository] (messages lost on app restart).
+  ///   Implement [CoreMessageRepository] for custom persistence (Firebase, SQLite, etc.).
+  ///
+  /// - **[contextEngine]**: Optional custom context engine for advanced prompt construction.
+  ///   Mutually exclusive with [systemPrompt] - provide one or the other, not both.
+  ///
+  /// - **[systemPrompt]**: Optional system prompt text for simple use cases.
+  ///   Creates a [SimpleContextEngine] internally. Cannot be used with [contextEngine].
+  ///
+  /// - **[queryEngine]**: Optional query preprocessor for input transformation,
+  ///   intent detection, or query augmentation before context building.
+  ///
+  /// - **[postEngine]**: Optional post-response processor for actions after AI generation,
+  ///   such as analytics, logging, or triggering side effects.
+  ///
+  /// - **[onPromptsReady]**: Optional callback invoked when prompts are ready for generation.
+  ///   Receives the list of messages that will be sent to the AI service.
+  ///
+  /// - **[onStateChange]**: Optional callback invoked when the generation state changes.
+  ///   Receives the current [GenerationState] for tracking progress.
+  ///
+  /// ## Returns
+  /// A fully initialized [KaiChatController] ready for use.
+  ///
+  /// ## Throws
+  /// - [AssertionError] if both [contextEngine] and [systemPrompt] are provided.
+  ///
+  /// ## Example
+  /// ```dart
+  /// final controller = await KaiChatController.create(
+  ///   sessionId: 'chat-${DateTime.now().millisecondsSinceEpoch}',
+  ///   generationService: GeminiGenerationService(apiKey: 'your-key'),
+  ///   repository: FirebaseMessageRepository(userId: currentUser.id),
+  ///   systemPrompt: 'You are a friendly cooking assistant.',
+  /// );
+  /// ```
+  static Future<KaiChatController> create({
+    String sessionId = 'default',
+    required GenerationServiceBase generationService,
+    CoreMessageRepository? repository,
+    ContextEngine? contextEngine,
+    String? systemPrompt,
+    QueryEngineBase? queryEngine,
+    PostResponseEngineBase? postEngine,
+    void Function(IList<CoreMessage> messages)? onPromptsReady,
+    void Function(GenerationState<GenerationResult> state)? onStateChange,
+  }) async {
+    assert(
+      contextEngine == null || systemPrompt == null,
+      'Cannot provide both contextEngine and systemPrompt. Use contextEngine for custom prompt building, or systemPrompt for simple cases.',
+    );
+
+    final session = ConversationSession.withCurrentTime(id: sessionId);
+
+    final conversation = await ConversationManager.create(
+      repository: repository ?? InMemoryMessageRepository(),
+      session: session,
+      messageAdapter: CoreMessageAdapter(),
+    );
+
+    return KaiChatController._(
+      contextEngine:
+          contextEngine ?? SimpleContextEngine(systemPrompt: systemPrompt),
+      conversationManager: conversation,
+      queryEngine: queryEngine,
+      postResponseEngine: postEngine,
+      generationService: generationService,
+      onPromptsReady: onPromptsReady,
+      onStateChange: onStateChange,
+    );
+  }
+
+  @override
+  void onPromptsReady(IList<CoreMessage> messages) {
+    super.onPromptsReady(messages);
+    _onPromptsReady?.call(messages);
+  }
+
+  @override
+  void onStateChange(GenerationState<GenerationResult> state) {
+    super.onStateChange(state);
+    _onStateChange?.call(state);
+  }
+}

packages/kai_engine/lib/src/message_repository_base.dart

@@ -24,9 +24,11 @@ abstract interface class CoreMessageRepositoryBase
     extends MessageRepositoryBase<CoreMessage> {}
 
 /// Prebuilt repository for memory persistence.
+/// Messages are automatically sorted by timestamp on initialization.
 final class InMemoryMessageRepository implements CoreMessageRepositoryBase {
   InMemoryMessageRepository({Iterable<CoreMessage>? initialMessages})
-    : _messages = initialMessages?.toList() ?? [];
+    : _messages = (initialMessages?.toList() ?? [])
+        ..sort((a, b) => a.timestamp.compareTo(b.timestamp));
 
   final List<CoreMessage> _messages;
 
@@ -62,3 +64,93 @@ final class InMemoryMessageRepository implements CoreMessageRepositoryBase {
     return Future.value(messages);
   }
 }
+
+/// Prebuilt repository for callback-based persistence.
+///
+/// Simplified API with only 2 required callbacks:
+/// - [onInitial]: Load initial messages from persistence
+/// - [onPut]: Persist insert or update operations (upsert)
+/// - [onRemove]: Optional callback for persisting deletions
+///
+/// Note: Callbacks are for persistence only. The repository maintains
+/// its own internal state for immediate read operations. Messages are
+/// automatically sorted by timestamp on initial load.
+final class CoreMessageRepository implements CoreMessageRepositoryBase {
+  final Future<Iterable<CoreMessage>> Function(ConversationSession session)
+  onInitial;
+  final Future<void> Function(
+    ConversationSession session,
+    Iterable<CoreMessage> messages,
+  )
+  onPut;
+  final Future<void> Function(Iterable<CoreMessage> messages)? onRemove;
+
+  final List<CoreMessage> _messages = [];
+  ConversationSession? _cachedSession;
+  bool _isInitialized = false;
+
+  CoreMessageRepository({
+    required this.onInitial,
+    required this.onPut,
+    this.onRemove,
+  });
+
+  @override
+  Future<Iterable<CoreMessage>> getMessages(ConversationSession session) async {
+    if (!_isInitialized || _cachedSession != session) {
+      _cachedSession = session;
+      final loaded = await onInitial(session);
+      _messages.clear();
+      _messages.addAll(loaded);
+      _messages.sort((a, b) => a.timestamp.compareTo(b.timestamp));
+      _isInitialized = true;
+    }
+    return _messages;
+  }
+
+  @override
+  Future<Iterable<CoreMessage>> saveMessages({
+    required ConversationSession session,
+    required Iterable<CoreMessage> messages,
+  }) async {
+    if (_cachedSession != null && _cachedSession != session) {
+      _cachedSession = session;
+      _messages.clear();
+    } else {
+      _cachedSession = session;
+    }
+    _isInitialized = true;
+    _messages.addAll(messages);
+    await onPut(session, messages);
+    return messages;
+  }
+
+  @override
+  Future<Iterable<CoreMessage>> updateMessages(
+    Iterable<CoreMessage> messages,
+  ) async {
+    if (_cachedSession == null) {
+      throw StateError(
+        'Session not initialized. Call getMessages or saveMessages first.',
+      );
+    }
+    for (var message in messages) {
+      final index = _messages.indexWhere(
+        (m) => m.messageId == message.messageId,
+      );
+      if (index != -1) {
+        _messages[index] = message;
+      }
+    }
+    await onPut(_cachedSession!, messages);
+    return messages;
+  }
+
+  @override
+  Future<void> removeMessages(Iterable<CoreMessage> messages) async {
+    _messages.removeWhere(
+      (message) => messages.map((e) => e.messageId).contains(message.messageId),
+    );
+    await onRemove?.call(messages);
+  }
+}

packages/kai_engine/lib/src/prompt_engine.dart

@@ -263,12 +263,18 @@ abstract base class ContextEngine
 ///
 /// This implementation is suitable for most basic conversational AI scenarios.
 final class SimpleContextEngine extends ContextEngine {
+  final String? systemPrompt;
+
   @override
   List<PromptTemplate> get promptBuilder => [
-    PromptTemplate.system("You're kai, a useful friendly personal assistant."),
+    PromptTemplate.system(
+      systemPrompt ?? "You're kai, a useful friendly personal assistant.",
+    ),
     PromptTemplate.buildSequential(HistoryContext()),
     PromptTemplate.input(),
   ];
+
+  SimpleContextEngine({this.systemPrompt});
 }
 
 @freezed

packages/kai_engine/lib/src/src.dart

@@ -8,9 +8,9 @@ export 'generation_service_base.dart';
 export 'kai_logger.dart';
 export 'message_adapter_base.dart';
 export 'message_repository_base.dart';
-export 'post_response_engine_base.dart';
-export 'prompt_engine.dart';
 export 'pipeline/kai_phase.dart';
 export 'pipeline/phase_types.dart';
+export 'post_response_engine_base.dart';
+export 'prompt_engine.dart';
 export 'query_engine_base.dart';
 export 'tool_schema.dart';

packages/kai_engine/pubspec.yaml

@@ -1,6 +1,6 @@
 name: kai_engine
 description: "A modular, extensible AI chat engine built with a pipeline-based architecture."
-version: 0.1.1
+version: 0.2.0
 homepage: https://github.com/pckimlong/kai_engine
 repository: https://github.com/pckimlong/kai_engine
 issue_tracker: https://github.com/pckimlong/kai_engine/issues