From b9ec730fad14b43227f3f548b2177c93d6857289 Mon Sep 17 00:00:00 2001
From: cogwheel0 <172976095+cogwheel0@users.noreply.github.com>
Date: Wed, 5 Nov 2025 13:10:00 +0530
Subject: [PATCH] feat(conversation): Enhance source parsing and normalization
 for OpenWebUI messages

---
 lib/core/services/conversation_parsing.dart | 45 +++++++++++++++++----
 lib/core/services/streaming_helper.dart     | 40 ++++++++++++++++++
 2 files changed, 77 insertions(+), 8 deletions(-)

diff --git a/lib/core/services/conversation_parsing.dart b/lib/core/services/conversation_parsing.dart
index 6cc8ee6..fe3fc1f 100644
--- a/lib/core/services/conversation_parsing.dart
+++ b/lib/core/services/conversation_parsing.dart
@@ -2,6 +2,8 @@ import 'dart:convert';
 
 import 'package:uuid/uuid.dart';
 
+import '../utils/openwebui_source_parser.dart';
+
 /// Utilities for converting OpenWebUI conversation payloads into JSON maps
 /// that match the app's `Conversation` / `ChatMessage` schemas. All helpers
 /// here are isolate-safe (they only work with primitive JSON types) so they
@@ -293,9 +295,11 @@ Map<String, dynamic> _parseOpenWebUIMessageToJson(
   final codeExecRaw = historyMsg != null
       ? historyMsg['code_executions'] ?? historyMsg['codeExecutions']
       : msgData['code_executions'] ?? msgData['codeExecutions'];
-  final sourcesRaw = historyMsg != null && historyMsg.containsKey('sources')
-      ? historyMsg['sources']
-      : msgData['sources'];
+  final sourcesRaw = historyMsg != null
+      ? historyMsg.containsKey('sources')
+            ? historyMsg['sources']
+            : historyMsg['citations']
+      : msgData['sources'] ?? msgData['citations'];
 
   return <String, dynamic>{
     'id': (msgData['id'] ?? _uuid.v4()).toString(),
@@ -409,21 +413,46 @@ List<Map<String, dynamic>> _parseCodeExecutionsField(dynamic raw) {
 }
 
 List<Map<String, dynamic>> _parseSourcesField(dynamic raw) {
+  final normalized = _coerceSourcesList(raw);
+  if (normalized == null || normalized.isEmpty) {
+    return const <Map<String, dynamic>>[];
+  }
+
+  final parsed = parseOpenWebUISourceList(normalized);
+  if (parsed.isNotEmpty) {
+    return parsed
+        .map((reference) => reference.toJson())
+        .toList(growable: false);
+  }
+
+  return normalized
+      .whereType<Map>()
+      .map(_coerceJsonMap)
+      .toList(growable: false);
+}
+
+List<dynamic>? _coerceSourcesList(dynamic raw) {
   if (raw is List) {
-    return raw.whereType<Map>().map(_coerceJsonMap).toList(growable: false);
+    return raw;
+  }
+  if (raw is Iterable) {
+    return raw.toList(growable: false);
   }
   if (raw is Map) {
-    return [_coerceJsonMap(raw)];
+    return [raw];
   }
-  if (raw is String) {
+  if (raw is String && raw.isNotEmpty) {
     try {
       final decoded = jsonDecode(raw);
       if (decoded is List) {
-        return decoded.whereType<Map>().map(_coerceJsonMap).toList();
+        return decoded;
+      }
+      if (decoded is Map) {
+        return [decoded];
       }
     } catch (_) {}
   }
-  return const <Map<String, dynamic>>[];
+  return null;
 }
 
 Map<String, dynamic> _coerceJsonMap(Object? value) {
diff --git a/lib/core/services/streaming_helper.dart b/lib/core/services/streaming_helper.dart
index c546181..dd267ba 100644
--- a/lib/core/services/streaming_helper.dart
+++ b/lib/core/services/streaming_helper.dart
@@ -617,6 +617,19 @@ ActiveSocketStream attachUnifiedChunkedStreaming({
 
       if (type == 'chat:completion' && payload != null) {
         if (payload is Map<String, dynamic>) {
+          final rawSources = payload['sources'] ?? payload['citations'];
+          final normalizedSources = _normalizeSourcesPayload(rawSources);
+          if (normalizedSources != null && normalizedSources.isNotEmpty) {
+            final parsedSources = parseOpenWebUISourceList(normalizedSources);
+            if (parsedSources.isNotEmpty) {
+              final targetId = _resolveTargetMessageId(messageId, getMessages);
+              if (targetId != null) {
+                for (final source in parsedSources) {
+                  appendSourceReference(targetId, source);
+                }
+              }
+            }
+          }
           if (payload.containsKey('tool_calls')) {
             final tc = payload['tool_calls'];
             if (tc is List) {
@@ -1387,6 +1400,33 @@ Map<String, dynamic>? _asStringMap(dynamic value) {
   return null;
 }
 
+List<dynamic>? _normalizeSourcesPayload(dynamic raw) {
+  if (raw == null) {
+    return null;
+  }
+  if (raw is List) {
+    return raw;
+  }
+  if (raw is Iterable) {
+    return raw.toList(growable: false);
+  }
+  if (raw is Map) {
+    return [raw];
+  }
+  if (raw is String && raw.isNotEmpty) {
+    try {
+      final decoded = jsonDecode(raw);
+      if (decoded is List) {
+        return decoded;
+      }
+      if (decoded is Map) {
+        return [decoded];
+      }
+    } catch (_) {}
+  }
+  return null;
+}
+
 String? _resolveTargetMessageId(
   String? messageId,
   List<ChatMessage> Function() getMessages,