import 'package:dio/dio.dart'; import 'package:riverpod_annotation/riverpod_annotation.dart'; import 'package:trainhub_flutter/core/constants/ai_constants.dart'; import 'package:trainhub_flutter/domain/repositories/chat_repository.dart'; import 'package:trainhub_flutter/domain/repositories/exercise_repository.dart'; import 'package:trainhub_flutter/domain/repositories/note_repository.dart'; import 'package:trainhub_flutter/domain/repositories/training_plan_repository.dart'; import 'package:trainhub_flutter/data/services/ai_process_manager.dart'; import 'package:trainhub_flutter/data/services/ai_settings_service.dart'; import 'package:trainhub_flutter/data/services/llm_client.dart'; import 'package:trainhub_flutter/injection.dart'; import 'package:trainhub_flutter/presentation/chat/chat_state.dart'; import 'package:uuid/uuid.dart'; part 'chat_controller.g.dart'; @riverpod AiProcessManager aiProcessManager(AiProcessManagerRef ref) { final manager = getIt(); manager.addListener(() => ref.notifyListeners()); return manager; } @riverpod AiSettingsService aiSettingsService(AiSettingsServiceRef ref) { final service = getIt(); service.addListener(() => ref.notifyListeners()); return service; } @riverpod class ChatController extends _$ChatController { late ChatRepository _repo; late NoteRepository _noteRepo; late LlmClient _llm; CancelToken? _cancelToken; @override Future build() async { _repo = getIt(); _noteRepo = getIt(); _llm = getIt(); // Abort any in-flight generation when the user leaves the chat page. ref.onDispose(() => _cancelToken?.cancel()); final aiManager = ref.read(aiProcessManagerProvider); if (aiManager.status == AiServerStatus.offline) { aiManager.startServers(); } final sessions = await _repo.getAllSessions(); return ChatState(sessions: sessions); } Future createSession() async { final session = await _repo.createSession(); final sessions = await _repo.getAllSessions(); state = AsyncValue.data( ChatState(sessions: sessions, activeSession: session), ); } Future loadSession(String id) async { final session = await _repo.getSession(id); if (session == null) return; final messages = await _repo.getMessages(id); final current = state.valueOrNull ?? const ChatState(); state = AsyncValue.data( current.copyWith(activeSession: session, messages: messages), ); } Future deleteSession(String id) async { await _repo.deleteSession(id); final sessions = await _repo.getAllSessions(); final current = state.valueOrNull ?? const ChatState(); state = AsyncValue.data( current.copyWith( sessions: sessions, activeSession: current.activeSession?.id == id ? null : current.activeSession, messages: current.activeSession?.id == id ? [] : current.messages, ), ); } Future sendMessage(String content) async { final current = state.valueOrNull; if (current == null) return; final sessionId = await _resolveSession(current, content); await _persistUserMessage(sessionId, content); final contextChunks = await _searchKnowledgeBase(content); final trainingContext = await _buildTrainingContext(); final systemPrompt = _buildSystemPrompt(contextChunks, trainingContext); final history = _buildHistory(); final fullAiResponse = await _streamResponse(systemPrompt, history); await _persistAssistantResponse(sessionId, content, fullAiResponse); } /// Summarizes the user's exercise library and training plans so the model /// can reference and plan around real data. Only attached when a cloud /// provider is active — the local 4B model's context is too small for it. Future _buildTrainingContext() async { if (getIt().activeProvider == AiProvider.local) return ''; try { final exercises = await getIt().getAll(); final plans = await getIt().getAll(); final buffer = StringBuffer(); if (exercises.isNotEmpty) { buffer.writeln("### The trainer's exercise library:"); for (final e in exercises.take(150)) { buffer.write('- ${e.name}'); final tags = e.tags; if (tags != null && tags.isNotEmpty) buffer.write(' [$tags]'); buffer.writeln(); } } if (plans.isNotEmpty) { buffer.writeln("\n### The trainer's training plans:"); for (final plan in plans.take(30)) { buffer.writeln('- ${plan.name}'); } } return buffer.toString(); } catch (_) { return ''; } } Future _resolveSession(ChatState current, String content) async { if (current.activeSession != null) return current.activeSession!.id; final session = await _repo.createSession(); final sessions = await _repo.getAllSessions(); state = AsyncValue.data( current.copyWith(sessions: sessions, activeSession: session), ); return session.id; } Future _persistUserMessage(String sessionId, String content) async { await _repo.addMessage( sessionId: sessionId, role: 'user', content: content, ); final messagesAfterUser = await _repo.getMessages(sessionId); state = AsyncValue.data( state.valueOrNull!.copyWith( messages: messagesAfterUser, isTyping: true, thinkingSteps: [], streamingContent: '', ), ); } Future> _searchKnowledgeBase(String query) async { final searchStep = _createStep('Searching knowledge base...'); List contextChunks = []; try { contextChunks = await _noteRepo.searchSimilar(query, topK: 3); if (contextChunks.isNotEmpty) { _updateStep( searchStep.id, status: ThinkingStepStatus.completed, title: 'Found ${contextChunks.length} documents', details: 'Context added for assistant.', ); } else { _updateStep( searchStep.id, status: ThinkingStepStatus.completed, title: 'No matching documents in knowledge base', details: 'Responding based on general knowledge.', ); } } catch (e) { _updateStep( searchStep.id, status: ThinkingStepStatus.error, title: 'Knowledge base search error', details: e.toString(), ); } return contextChunks; } /// Most recent messages only — an unbounded history would eventually /// overflow the model context and slow every request down. List> _buildHistory() { final messages = state.valueOrNull?.messages ?? []; final recent = messages.length > AiConstants.chatHistoryLimit ? messages.sublist(messages.length - AiConstants.chatHistoryLimit) : messages; return recent .map( (m) => { 'role': m.isUser ? 'user' : 'assistant', 'content': m.content, }, ) .toList(); } /// Stops an in-flight generation. The partial response streamed so far is /// kept and persisted like a normal reply. void stopGeneration() => _cancelToken?.cancel(); Future _streamResponse( String systemPrompt, List> history, ) async { final generateStep = _createStep('Generating response...'); String fullAiResponse = ''; _cancelToken = CancelToken(); try { final stream = _llm.streamChat([ {'role': 'system', 'content': systemPrompt}, ...history, ], cancelToken: _cancelToken); _updateStep( generateStep.id, status: ThinkingStepStatus.running, title: 'Writing...', ); await for (final delta in stream) { fullAiResponse += delta; final updatedState = state.valueOrNull; if (updatedState != null) { state = AsyncValue.data( updatedState.copyWith(streamingContent: fullAiResponse), ); } } _updateStep( generateStep.id, status: ThinkingStepStatus.completed, title: 'Response generated', ); } on DioException catch (e) { if (CancelToken.isCancel(e)) { _updateStep( generateStep.id, status: ThinkingStepStatus.completed, title: 'Stopped by user', ); } else { _updateStep( generateStep.id, status: ThinkingStepStatus.error, title: 'Generation failed', details: e.message ?? e.toString(), ); } } catch (e) { _updateStep( generateStep.id, status: ThinkingStepStatus.error, title: 'Generation failed', details: e.toString(), ); } finally { _cancelToken = null; } return fullAiResponse; } Future _persistAssistantResponse( String sessionId, String userContent, String aiResponse, ) async { // A failed generation yields an empty response — leave the error visible // in the thinking steps instead of saving an empty assistant message. if (aiResponse.trim().isEmpty) { final current = state.valueOrNull; if (current != null) { state = AsyncValue.data( current.copyWith(isTyping: false, streamingContent: null), ); } return; } await _repo.addMessage( sessionId: sessionId, role: 'assistant', content: aiResponse, ); final messagesAfterAi = await _repo.getMessages(sessionId); if (messagesAfterAi.length <= 2) { final title = userContent.length > 30 ? '${userContent.substring(0, 30)}…' : userContent; await _repo.updateSessionTitle(sessionId, title); } final sessions = await _repo.getAllSessions(); state = AsyncValue.data( state.valueOrNull!.copyWith( messages: messagesAfterAi, isTyping: false, streamingContent: null, thinkingSteps: [], sessions: sessions, ), ); } ThinkingStep _createStep(String title) { final step = ThinkingStep( id: const Uuid().v4(), title: title, status: ThinkingStepStatus.pending, ); final current = state.valueOrNull; if (current != null) { state = AsyncValue.data( current.copyWith(thinkingSteps: [...current.thinkingSteps, step]), ); } return step; } void _updateStep( String id, { ThinkingStepStatus? status, String? title, String? details, }) { final current = state.valueOrNull; if (current == null) return; final updatedSteps = current.thinkingSteps.map((s) { if (s.id != id) return s; return s.copyWith( status: status ?? s.status, title: title ?? s.title, details: details ?? s.details, ); }).toList(); state = AsyncValue.data(current.copyWith(thinkingSteps: updatedSteps)); } static String _buildSystemPrompt( List chunks, String trainingContext, ) { final buffer = StringBuffer(AiConstants.baseSystemPrompt); if (trainingContext.isNotEmpty) { buffer.write('\n\n$trainingContext'); buffer.write( '\nWhen designing or discussing training plans, prefer exercises ' 'from the library above and reference existing plans by name.', ); } if (chunks.isNotEmpty) { final contextBlock = chunks .asMap() .entries .map((e) => '[${e.key + 1}] ${e.value}') .join('\n\n'); buffer.write( '\n\n### Relevant notes from the trainer\'s knowledge base:\n' '$contextBlock\n\n' 'Use the above context to inform your response when relevant. ' 'If the context is not directly applicable, rely on your general ' 'fitness knowledge.', ); } return buffer.toString(); } }