diff --git a/src/fuser/__init__.py b/src/fuser/__init__.py index 33a3a6fb53..0d0cab8321 100644 --- a/src/fuser/__init__.py +++ b/src/fuser/__init__.py @@ -39,10 +39,14 @@ def __init__(self, config: RuntimeConfig): self.knowledge_base = None self.kb_min_score = 0.0 + self.kb_high_score = 0.92 + self.kb_low_score = 0.75 if config.knowledge_base: try: kb_config = dict(config.knowledge_base) self.kb_min_score = kb_config.get("min_score", 0.0) + self.kb_high_score = kb_config.get("high_score", 0.92) + self.kb_low_score = kb_config.get("low_score", 0.75) if self.kb_min_score > 0: logging.info( f"KnowledgeBase min_score threshold: {self.kb_min_score}" @@ -108,12 +112,33 @@ async def fuse( results = await self.knowledge_base.query( query_text, top_k=3, min_score=self.kb_min_score ) - if results: - kb_context = self.knowledge_base.format_context( - results, max_chars=1500 + high = [ + r + for r in results + if r.score is not None and r.score >= self.kb_high_score + ] + low = [ + r + for r in results + if r.score is not None + and self.kb_low_score <= r.score < self.kb_high_score + ] + kb_parts = [] + if high: + kb_parts.append( + self.knowledge_base.format_context(high, max_chars=1500) ) + if low: + kb_parts.append( + "[Potentially relevant, low confidence]\n" + + self.knowledge_base.format_context(low, max_chars=1000) + ) + if kb_parts: + kb_context = "\n".join(kb_parts) logging.info( - f"Knowledge base: {len(results)} docs passed to LLM" + f"Knowledge base: {len(high)} high," + f" {len(low)} low confidence;" + f" {len(results)} docs passed to LLM" ) else: logging.info( @@ -124,7 +149,7 @@ async def fuse( # Add knowledge base context to inputs if available if kb_context: - inputs_fused += f"\n\nKNOWLEDGE BASE:\n{kb_context}" + inputs_fused += f"\n\nYOUR KNOWLEDGE BASE:\n{kb_context}" # if we provide laws from blockchain, these override the locally stored rules # the rules are not provided in the system prompt, but as a separate INPUT, diff --git a/tests/fuser/test_init.py b/tests/fuser/test_init.py index d58a3cc6b5..f64a3f2ff0 100644 --- a/tests/fuser/test_init.py +++ b/tests/fuser/test_init.py @@ -225,9 +225,11 @@ async def test_fuser_with_knowledge_base_and_voice_input(): mock_kb.query.assert_called_once_with( "What is the capital of France?", top_k=3, min_score=0.0 ) - mock_kb.format_context.assert_called_once_with( - [mock_doc1, mock_doc2], max_chars=1500 - ) + # Two-tier KB filtering: high confidence (>=0.92) and low (0.75-0.92) + # are formatted separately, so format_context is called twice + assert mock_kb.format_context.call_count == 2 + mock_kb.format_context.assert_any_call([mock_doc1], max_chars=1500) + mock_kb.format_context.assert_any_call([mock_doc2], max_chars=1000) assert result is not None assert "KNOWLEDGE BASE:" in result assert "Paris is the capital of France." in result