@@ -49,7 +49,7 @@ public async Task QueryAsync_WithValidQuery_ReturnsResponseWithSources()
4949 var query = new RagQuery { Query = "What is the capital of France?" } ;
5050 var queryEmbedding = new float [ ] { 0.1f , 0.2f , 0.3f } ;
5151 var searchResults = CreateSearchResults ( 2 ) ;
52- var llmResponse = "The capital of France is Paris. [Source 1] " ;
52+ var llmResponse = "The capital of France is Paris." ;
5353
5454 _embeddingProviderMock
5555 . Setup ( x => x . GenerateEmbeddingAsync ( It . IsAny < string > ( ) , It . IsAny < CancellationToken > ( ) ) )
@@ -100,7 +100,7 @@ public async Task QueryAsync_WithNoSearchResults_ReturnsNoResultsResponse()
100100
101101 // Assert
102102 Assert . NotNull ( response ) ;
103- Assert . Contains ( "couldn't find any relevant information " , response . Answer ) ;
103+ Assert . Contains ( "Sorry " , response . Answer ) ;
104104 Assert . Empty ( response . Sources ) ;
105105 Assert . Equal ( "test-model" , response . Meta . Model ) ;
106106
@@ -366,16 +366,16 @@ public async Task QueryAsync_BuildsPromptWithContextCorrectly()
366366
367367 // Assert
368368 Assert . NotNull ( capturedUserPrompt ) ;
369- Assert . Contains ( "CONTEXT " , capturedUserPrompt ) ;
370- Assert . Contains ( "Source 1 " , capturedUserPrompt ) ;
371- Assert . Contains ( "Source 2 " , capturedUserPrompt ) ;
372- Assert . Contains ( "test-file-0.txt " , capturedUserPrompt ) ;
373- Assert . Contains ( "test-file-1.txt " , capturedUserPrompt ) ;
369+ Assert . Contains ( "INFORMATION " , capturedUserPrompt ) ;
370+ Assert . Contains ( "Item 1: " , capturedUserPrompt ) ;
371+ Assert . Contains ( "Item 2: " , capturedUserPrompt ) ;
372+ Assert . Contains ( "Test content from chunk 0 " , capturedUserPrompt ) ;
373+ Assert . Contains ( "Test content from chunk 1 " , capturedUserPrompt ) ;
374374 Assert . Contains ( "What is the weather?" , capturedUserPrompt ) ;
375375 }
376376
377377 [ Fact ]
378- public async Task QueryAsync_WithPageNumbers_IncludesPageInfoInPrompt ( )
378+ public async Task QueryAsync_WithPageNumbers_IncludesPageInfoInSources ( )
379379 {
380380 // Arrange
381381 var query = new RagQuery { Query = "Test query" } ;
@@ -406,18 +406,16 @@ public async Task QueryAsync_WithPageNumbers_IncludesPageInfoInPrompt()
406406 . Setup ( x => x . SearchAsync ( It . IsAny < float [ ] > ( ) , It . IsAny < int > ( ) , It . IsAny < SearchFilters > ( ) , It . IsAny < CancellationToken > ( ) ) )
407407 . ReturnsAsync ( searchResults ) ;
408408
409- string ? capturedUserPrompt = null ;
410409 _llmProviderMock
411410 . Setup ( x => x . GenerateAsync ( It . IsAny < string > ( ) , It . IsAny < string > ( ) , It . IsAny < LlmGenerationOptions > ( ) , It . IsAny < CancellationToken > ( ) ) )
412- . Callback < string , string , LlmGenerationOptions , CancellationToken > ( ( s , u , o , c ) => capturedUserPrompt = u )
413411 . ReturnsAsync ( "Answer" ) ;
414412
415413 // Act
416- await _ragEngine . QueryAsync ( query ) ;
414+ var response = await _ragEngine . QueryAsync ( query ) ;
417415
418- // Assert
419- Assert . NotNull ( capturedUserPrompt ) ;
420- Assert . Contains ( "Page 5" , capturedUserPrompt ) ;
416+ // Assert - page number is in sources metadata, not in prompt
417+ Assert . Single ( response . Sources ) ;
418+ Assert . Equal ( 5 , response . Sources [ 0 ] . PageNumber ) ;
421419 }
422420
423421 [ Fact ]
0 commit comments