Skip to content

Commit 260c19c

Browse files
priyavratuniyalpriyavratuniyal
committed
fix(ai): address PR #4 review findings for safety, security, and routing
- Separate ObjectBox lookup from query routing so storage errors no longer bypass emergency/off-topic detection (safety-critical) - Guard Bearer token to trusted Hugging Face hosts only, preventing token leakage to arbitrary download URLs - Tighten personal-lab regex patterns to require medical terms after possessive pronouns, fixing false positives like "is my code correct" - Make QueryRouter resolve the intent classifier lazily via getter so Stage 2 embedding classification activates once loaded, not only if ready at screen init - Wrap HNSW query in try/finally to prevent resource leak on error - Fix emergency number to list both 911/112 for multi-region support - Correct "HuggingFace" brand spelling to "Hugging Face" - Remove unused routingFailedLog constant
1 parent 21f6aed commit 260c19c

8 files changed

Lines changed: 90 additions & 60 deletions

File tree

‎lib/constants/llm_strings.dart‎

Lines changed: 1 addition & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -69,7 +69,7 @@ abstract final class LlmStrings {
6969
/// Info banner text explaining common issues.
7070
static const String hfTokenInfoMessage =
7171
'Having a token does not guarantee access. You must also:\n'
72-
' 1. Accept the model\'s license on its HuggingFace page\n'
72+
' 1. Accept the model\'s license on its Hugging Face page\n'
7373
' 2. Ensure your token has "Read" permission\n\n'
7474
'If downloads fail with 401/403 errors after adding a token, visit '
7575
'the model page and accept the license agreement, then retry.';
@@ -98,10 +98,6 @@ abstract final class LlmStrings {
9898
static const String persistenceWarning =
9999
'Could not save chat history for one or more messages.';
100100

101-
/// Debug log when query routing fails.
102-
static const String routingFailedLog =
103-
'Query routing failed, falling through to LLM: ';
104-
105101
/// Debug log when retrying after empty/failed generation.
106102
static const String retryingGenerationLog =
107103
'Generation produced empty/failed output, retrying once';

‎lib/constants/response_templates.dart‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,7 @@ abstract final class ResponseTemplates {
66
/// Shown when the user describes potentially urgent symptoms.
77
static const String emergencyEscalation =
88
'Your message may describe a medical emergency. '
9-
'Please contact emergency services (112) or visit your nearest '
9+
'Please contact your local emergency services (e.g., 911/112) or visit your nearest '
1010
'emergency room immediately.\n\n'
1111
'Koshika is a lab-report assistant and cannot provide emergency '
1212
'medical guidance.';

‎lib/screens/chat_screen.dart‎

Lines changed: 14 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -60,7 +60,7 @@ class _ChatScreenState extends State<ChatScreen> {
6060
super.initState();
6161
_contextBuilder = ChatContextBuilder(vectorStore: vectorStoreService);
6262
_initIntentClassifier();
63-
_queryRouter = QueryRouter(classifier: _intentClassifier);
63+
_queryRouter = QueryRouter(classifierGetter: () => _intentClassifier);
6464
_statusSubscription = llmService.modelStatusStream.listen((info) {
6565
if (mounted) {
6666
setState(() => _modelInfo = info);
@@ -330,23 +330,24 @@ class _ChatScreenState extends State<ChatScreen> {
330330
// history-aware routing and prompt injection.
331331
final history = _buildConversationHistory();
332332

333-
// Route the query — may short-circuit with a deterministic response
334-
late final QueryRouteResult routeResult;
333+
// Route the query — may short-circuit with a deterministic response.
334+
// Lab-data lookup is separated from routing so that an ObjectBox error
335+
// doesn't bypass safety-critical emergency/off-topic detection.
336+
var hasLabData = false;
335337
try {
336-
final hasLabData = objectbox.getLatestResults().isNotEmpty;
337-
routeResult = await _queryRouter.route(
338-
text,
339-
hasLabData: hasLabData,
340-
conversationHistory: history,
341-
);
338+
hasLabData = objectbox.getLatestResults().isNotEmpty;
342339
} catch (e) {
343-
// If routing fails (e.g. ObjectBox error), fall through to LLM
344-
debugPrint('${LlmStrings.routingFailedLog}$e');
345-
routeResult = const QueryRouteResult(
346-
decision: QueryDecision.answerGeneralHealth,
340+
debugPrint(
341+
'Lab data lookup failed; assuming no lab data for routing: $e',
347342
);
348343
}
349344

345+
final routeResult = await _queryRouter.route(
346+
text,
347+
hasLabData: hasLabData,
348+
conversationHistory: history,
349+
);
350+
350351
if (!routeResult.requiresLlm) {
351352
if (mounted) {
352353
_addDeterministicResponse(

‎lib/services/intent_prefilter.dart‎

Lines changed: 23 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -67,11 +67,12 @@ abstract final class IntentPrefilter {
6767

6868
/// Phrases that explicitly reference the user's own reports or results.
6969
static final _personalLabPatterns = [
70-
// Explicit report/result references (English)
71-
RegExp(r'lab\s*report', caseSensitive: false),
72-
RegExp(r'test\s*result', caseSensitive: false),
73-
RegExp(r'blood\s*test', caseSensitive: false),
74-
RegExp(r'blood\s*work', caseSensitive: false),
70+
// Possessive + report/result references (English + Hinglish)
71+
// "my lab report", "mera test result", "my blood test"
72+
RegExp(
73+
'$_possessive\\s+(lab\\s*report|test\\s*results?|blood\\s*test|blood\\s*work)\\b',
74+
caseSensitive: false,
75+
),
7576

7677
// Possessive + generic result words (English + Hinglish)
7778
RegExp(
@@ -83,15 +84,24 @@ abstract final class IntentPrefilter {
8384
// "my cholesterol", "mera sugar", "meri thyroid"
8485
RegExp('$_possessive\\s+$_biomarkerAlt\\b', caseSensitive: false),
8586

86-
// "is my ... high/low/normal" (English)
87-
RegExp(r'is\s+my\b', caseSensitive: false),
88-
// "why is my ... high/low"
89-
RegExp(r'why\s+(is|are)\s+my\b', caseSensitive: false),
90-
// "show me my", "check my"
91-
RegExp(r'(show|check|tell)\s*(me\s+)?my\b', caseSensitive: false),
87+
// "is my <medical term> high/low/normal" (English)
88+
RegExp(
89+
'is\\s+$_possessive\\s+($_biomarkerAlt|report|results?|levels?|values?|numbers?|test)\\b',
90+
caseSensitive: false,
91+
),
92+
// "why is my <medical term> high/low"
93+
RegExp(
94+
'why\\s+(is|are)\\s+$_possessive\\s+($_biomarkerAlt|report|results?|levels?|values?|numbers?|test)\\b',
95+
caseSensitive: false,
96+
),
97+
// "show me my <medical term>", "check my <medical term>"
98+
RegExp(
99+
'(show|check|tell)\\s*(me\\s+)?$_possessive\\s+($_biomarkerAlt|report|results?|levels?|values?|numbers?|test)\\b',
100+
caseSensitive: false,
101+
),
92102

93103
// Hinglish report phrases
94-
// "meri report dikhao", "report dikhao", "mera report"
104+
// "meri report dikhao", "mera report"
95105
RegExp('$_possessive\\s*report', caseSensitive: false),
96106
// "kitna hai" / "kitni hai" (how much is) — common Hinglish question form
97107
RegExp('$_biomarkerAlt\\s*(kitna|kitni|kitne)\\b', caseSensitive: false),
@@ -102,7 +112,7 @@ abstract final class IntentPrefilter {
102112
// Only match medical compound units — excludes bare "g", "%", "ml" to
103113
// avoid false positives like "5g network", "100% sure", "500ml water".
104114
RegExp(
105-
r'\d+\.?\d*\s*(mg/dl|mg/l|mmol/l|mmol|miu/l|miu/ml|iu/l|iu/ml|µg/dl|µg/l|ng/ml|ng/dl|pg/ml|g/dl|g/l|meq/l|u/l|cells/cumm|lakh|thousand)',
115+
r'\d+\.?\d*\s*(mg/dl|mg/l|mmol/l|mmol|miu/l|miu/ml|iu/l|iu/ml|µg/dl|µg/l|ng/ml|ng/dl|pg/ml|g/dl|g/l|meq/l|u/l|cells/cumm)',
106116
caseSensitive: false,
107117
),
108118
];

‎lib/services/model_downloader.dart‎

Lines changed: 20 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -88,8 +88,14 @@ class ModelDownloader {
8888
startByte = tempFile.lengthSync();
8989
}
9090

91-
final request = await _client!.getUrl(Uri.parse(url));
92-
if (authToken != null && authToken.isNotEmpty) {
91+
final uri = Uri.parse(url);
92+
final request = await _client!.getUrl(uri);
93+
final shouldAttachAuth =
94+
authToken != null &&
95+
authToken.isNotEmpty &&
96+
uri.scheme == 'https' &&
97+
_isTrustedHfHost(uri.host);
98+
if (shouldAttachAuth) {
9399
request.headers.add(
94100
HttpHeaders.authorizationHeader,
95101
'Bearer $authToken',
@@ -184,4 +190,16 @@ class ModelDownloader {
184190
_cancelled = true;
185191
_client?.close(force: true);
186192
}
193+
194+
/// Whether [host] belongs to a known Hugging Face domain.
195+
static bool _isTrustedHfHost(String host) {
196+
final normalized = host.toLowerCase();
197+
const trustedHosts = {
198+
'huggingface.co',
199+
'hf.co',
200+
'cdn-lfs.huggingface.co',
201+
'cdn-lfs-us-1.huggingface.co',
202+
};
203+
return trustedHosts.contains(normalized);
204+
}
187205
}

‎lib/services/query_router.dart‎

Lines changed: 10 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -13,13 +13,16 @@ import 'llm_service.dart';
1313
/// This is a class (not static) because the embedding classifier is injected
1414
/// as an optional constructor dependency.
1515
class QueryRouter {
16-
final IntentClassifier? _classifier;
16+
final IntentClassifier? Function() _classifierGetter;
1717

18-
/// Create a router with an optional Stage 2 embedding classifier.
18+
/// Create a router with a lazy Stage 2 embedding classifier.
1919
///
20-
/// When [classifier] is null or not ready, falls back to Stage 1
21-
/// (regex/keyword) only.
22-
QueryRouter({IntentClassifier? classifier}) : _classifier = classifier;
20+
/// [classifierGetter] is called at route time so the classifier can become
21+
/// available after the router is constructed (e.g. once embeddings load).
22+
/// When the getter returns null or the classifier isn't ready, falls back
23+
/// to Stage 1 (regex/keyword) only.
24+
QueryRouter({required IntentClassifier? Function() classifierGetter})
25+
: _classifierGetter = classifierGetter;
2326

2427
/// Classify the [message] and decide how to respond.
2528
///
@@ -123,8 +126,8 @@ class QueryRouter {
123126
}
124127
}
125128

126-
// Fall through to Stage 2 classifier
127-
final classifier = _classifier;
129+
// Fall through to Stage 2 classifier (resolved lazily)
130+
final classifier = _classifierGetter();
128131
if (classifier == null || !classifier.isReady) {
129132
// No Stage 2 — default to general health (safe fallback)
130133
return _applyStrictnessMode(

‎lib/services/vector_store_service.dart‎

Lines changed: 20 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -75,24 +75,26 @@ class VectorStoreService {
7575
)
7676
.build();
7777

78-
final results = hnswQuery.find();
79-
hnswQuery.close();
80-
81-
return results.map((r) {
82-
final docId = '${r.biomarkerKey}_${r.report.targetId}';
83-
return RetrievalResult(
84-
id: docId,
85-
content: _buildChunkText(r),
86-
metadata: jsonEncode({
87-
'biomarkerKey': r.biomarkerKey,
88-
'reportId': r.report.targetId,
89-
'flag': r.flag.name,
90-
'category': r.category,
91-
'testDate': r.testDate.toIso8601String(),
92-
'labName': r.report.target?.labName,
93-
}),
94-
);
95-
}).toList();
78+
try {
79+
final results = hnswQuery.find();
80+
return results.map((r) {
81+
final docId = '${r.biomarkerKey}_${r.report.targetId}';
82+
return RetrievalResult(
83+
id: docId,
84+
content: _buildChunkText(r),
85+
metadata: jsonEncode({
86+
'biomarkerKey': r.biomarkerKey,
87+
'reportId': r.report.targetId,
88+
'flag': r.flag.name,
89+
'category': r.category,
90+
'testDate': r.testDate.toIso8601String(),
91+
'labName': r.report.target?.labName,
92+
}),
93+
);
94+
}).toList();
95+
} finally {
96+
hnswQuery.close();
97+
}
9698
} catch (e) {
9799
debugPrint('VectorStoreService.search failed: $e');
98100
return [];

‎test/services/query_router_test.dart‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,7 @@ void main() {
99

1010
setUp(() {
1111
// No classifier — tests Stage 1 (regex-only) routing
12-
router = QueryRouter();
12+
router = QueryRouter(classifierGetter: () => null);
1313
});
1414

1515
group('QueryRouter.route', () {

0 commit comments

Comments
 (0)