Skip to content

Latest commit

 

History

History
317 lines (254 loc) · 11.3 KB

File metadata and controls

317 lines (254 loc) · 11.3 KB

Adding LLM Source Attribution

Problem Statement

The current Java Chat application cannot reliably distinguish between content sourced from RAG (Retrieval-Augmented Generation) and content generated by the LLM's inherent knowledge. This makes it difficult to:

  • Verify factual accuracy against documentation
  • Build user trust through transparent sourcing
  • Debug hallucinations or incorrect information
  • Comply with potential attribution requirements

Current Architecture Analysis

Data Flow

User Query → RetrievalService → Qdrant/LocalSearch → Documents
     ↓
Documents + Query → ChatService → System Context Creation
     ↓
System Context (with [CTX n] markers) + User Message → LLM
     ↓
LLM Response (mixed content) → Stream to User

Why Attribution Is Currently Impossible

  1. Context Injection: Retrieved documents are injected as system context with [CTX n] markers
  2. Single LLM Call: Everything goes to the LLM in one unified prompt
  3. Unstructured Response: LLM returns a single text stream mixing both RAG and inherent knowledge
  4. Lost Attribution: Once the LLM processes the context, source attribution is lost in the response

Implementation Strategies

Strategy 1: Structured Response with Source Markers (Recommended - Short Term)

Modify the LLM prompt to enforce structured output with explicit source tagging.

Implementation in ChatService.java:

public Flux<String> streamAnswer(List<Message> history, String latestUserMessage) {
    List<Document> contextDocs = retrievalService.retrieve(latestUserMessage);
    StringBuilder systemContext = new StringBuilder(
        "You are a Java learning assistant focused on Java 24 and current stable JDK releases.\n" +
        "CRITICAL: Structure your response with explicit source markers:\n" +
        "- {{rag:text}} for information directly from provided context\n" +
        "- {{llm:text}} for your own knowledge/reasoning\n" +
        "- {{hybrid:text}} for combined RAG+LLM insights\n" +
        "- Always include [CTX n] citations immediately after {{rag:}} content\n" +
        "- Be explicit about which information comes from which source\n\n" +
        "Example: {{rag:Java records are immutable data carriers[CTX 1]}} {{llm:which makes them ideal for DTOs}}\n"
    );
    // ... rest of implementation
}

Post-Processing for Structured Output:

public class AnnotatedResponse {
    private List<ContentSegment> segments;
    
    public static class ContentSegment {
        private String text;
        private SourceType source; // enum: RAG, LLM, HYBRID
        private List<Integer> contextRefs; // [1,2] for [CTX 1], [CTX 2]
        private float confidence; // 0.0-1.0
    }
}

public AnnotatedResponse parseAnnotatedResponse(String rawResponse) {
    Pattern pattern = Pattern.compile("\\{\\{(rag|llm|hybrid):([^}]*)\\}\\}");
    Matcher matcher = pattern.matcher(rawResponse);
    List<ContentSegment> segments = new ArrayList<>();
    
    while (matcher.find()) {
        String sourceType = matcher.group(1);
        String content = matcher.group(2);
        
        // Extract context references
        List<Integer> refs = extractContextRefs(content);
        
        segments.add(new ContentSegment(
            content.replaceAll("\\[CTX \\d+\\]", "").trim(),
            SourceType.valueOf(sourceType.toUpperCase()),
            refs,
            calculateConfidence(sourceType, refs)
        ));
    }
    return new AnnotatedResponse(segments);
}

Strategy 2: Dual-Pass Processing

Implement a two-phase approach for clear source separation.

public class DualPassChatService {
    
    public Flux<AnnotatedChunk> streamWithAttribution(String query) {
        // Phase 1: Extract facts from RAG
        List<Document> docs = retrievalService.retrieve(query);
        String ragFacts = extractFactsFromDocuments(docs);
        
        // Phase 2: Enhance with LLM knowledge
        String prompt = String.format(
            "Given these facts from documentation:\n%s\n\n" +
            "User question: %s\n\n" +
            "Provide additional context and explanations, clearly marking what's from the facts above vs your knowledge.",
            ragFacts, query
        );
        
        return llmClient.stream(prompt)
            .map(chunk -> annotateChunkSource(chunk, ragFacts));
    }
    
    private AnnotatedChunk annotateChunkSource(String chunk, String ragFacts) {
        // Use similarity matching to determine source
        double similarity = calculateSimilarity(chunk, ragFacts);
        return new AnnotatedChunk(
            chunk,
            similarity > 0.7 ? SourceType.RAG : SourceType.LLM,
            similarity
        );
    }
}

Strategy 3: Confidence Scoring with Provenance

Add metadata to track content confidence and source.

import com.fasterxml.jackson.databind.JsonNode;

public class EnrichedContent {
    private String text;
    private float ragConfidence;  // 0.0-1.0 based on similarity score
    private float llmConfidence;  // Based on temperature/certainty
    private SourceType primarySource; // RAG, LLM, or HYBRID
    private List<Citation> citations;
    private Map<String, JsonNode> metadata;
    
    public static EnrichedContent fromRetrievalContext(
            String text, 
            List<Document> sourceeDocs,
            float similarityScore) {
        EnrichedContent content = new EnrichedContent();
        content.text = text;
        content.ragConfidence = similarityScore;
        content.llmConfidence = 0.0f;
        content.primarySource = SourceType.RAG;
        content.citations = extractCitations(sourceDocs);
        return content;
    }
}

Strategy 4: Real-time Annotation Pipeline

Stream processing with live source tracking.

public class StreamAnnotationService {
    private final TextSimilarityService similarityService;
    
    public Flux<AnnotatedChunk> annotateStream(
            Flux<String> rawStream, 
            List<Document> contextDocs) {
        
        // Build context fingerprint for matching
        String contextFingerprint = buildContextFingerprint(contextDocs);
        
        return rawStream
            .bufferTimeout(10, Duration.ofMillis(100))
            .map(chunks -> {
                String combined = String.join("", chunks);
                boolean fromRag = matchesContext(combined, contextFingerprint);
                float confidence = calculateMatchConfidence(combined, contextDocs);
                
                return new AnnotatedChunk(
                    combined,
                    fromRag ? SourceType.RAG : SourceType.LLM,
                    confidence
                );
            });
    }
    
    private boolean matchesContext(String text, String contextFingerprint) {
        // Use fuzzy matching, n-gram similarity, or embedding similarity
        return similarityService.cosineSimilarity(text, contextFingerprint) > 0.65;
    }
}

Strategy 5: Hybrid Approach with Complete Metadata

Return responses with comprehensive metadata headers.

import com.fasterxml.jackson.databind.JsonNode;

public class AttributedResponse {
    private String response;
    private ResponseMetadata metadata;
    
    public static class ResponseMetadata {
        private float ragContribution;  // 0.0-1.0 (e.g., 0.65 = 65% from RAG)
        private float llmContribution;  // 0.0-1.0 (e.g., 0.35 = 35% from LLM)
        private List<Citation> citations;
        private List<ContentSegment> segments;
        private Map<String, JsonNode> debugInfo;
    }
    
    public static class ContentSegment {
        private int startIndex;
        private int endIndex;
        private SourceType source;
        private Integer contextRef; // Reference to context document
        private float confidence;
    }
}

// Modified controller endpoint
@PostMapping("/chat/attributed")
public Flux<AttributedResponse> chatWithAttribution(@RequestBody ChatRequest request) {
    return chatService.processWithAttribution(request)
        .map(this::enrichWithMetadata);
}

Implementation Roadmap

Phase 1: Immediate (1-2 days)

  • Fix configuration inconsistencies in README
  • Add basic source markers to LLM prompts
  • Implement simple regex-based parsing for {{rag:}} and {{llm:}} markers

Phase 2: Short-term (1 week)

  • Implement Strategy 1 fully with post-processing
  • Add AnnotatedResponse model
  • Create new /api/chat/attributed endpoint
  • Update frontend to display source indicators

Phase 3: Medium-term (2-3 weeks)

  • Add confidence scoring (Strategy 3)
  • Implement similarity-based source detection
  • Add metrics collection for attribution accuracy
  • Create attribution debugging tools

Phase 4: Long-term (1-2 months)

  • Full metadata implementation (Strategy 5)
  • Machine learning model for source classification
  • A/B testing framework for attribution strategies
  • Performance optimization for real-time annotation

Technical Considerations

Performance Impact

  • Strategy 1: Minimal (~5% overhead for parsing)
  • Strategy 2: Moderate (requires two LLM calls)
  • Strategy 3-4: Low-Moderate (similarity calculations)
  • Strategy 5: Moderate (comprehensive metadata generation)

Storage Requirements

  • Annotated responses: ~20% larger than plain text
  • Metadata storage: Additional 1-2KB per response
  • Context fingerprints: ~500 bytes per document

Compatibility

  • All strategies maintain backward compatibility
  • New endpoints alongside existing ones
  • Gradual migration path for clients

Testing Strategy

Unit Tests

@Test
public void testSourceAttribution() {
    String response = "{{rag:Java records are immutable[CTX 1]}} {{llm:similar to Kotlin data classes}}";
    AnnotatedResponse parsed = parser.parse(response);
    
    assertEquals(2, parsed.getSegments().size());
    assertEquals(SourceType.RAG, parsed.getSegments().get(0).getSource());
    assertEquals(SourceType.LLM, parsed.getSegments().get(1).getSource());
    assertEquals(List.of(1), parsed.getSegments().get(0).getContextRefs());
}

Integration Tests

  • Verify attribution accuracy against known documents
  • Test streaming performance with annotation
  • Validate confidence scores against ground truth

Metrics to Track

  • Attribution accuracy rate
  • False positive/negative rates for source detection
  • User trust scores (via feedback)
  • Performance impact on response time

Security & Privacy Considerations

  • No sensitive data in attribution metadata
  • Citations should not expose internal document structure
  • Confidence scores should not reveal model internals
  • Rate limiting on attributed endpoints to prevent abuse

Future Enhancements

  1. ML-Based Classification: Train a classifier to identify source based on writing style
  2. Interactive Attribution: Allow users to query source for specific claims
  3. Versioned Attribution: Track how attribution changes across model updates
  4. Cross-Reference Validation: Verify RAG content against multiple sources
  5. Attribution Analytics: Dashboard showing attribution patterns and accuracy

Conclusion

Implementing source attribution requires architectural changes to the current single-pass, unstructured approach. Strategy 1 (Structured Response with Source Markers) offers the best balance of implementation simplicity and value delivery for the short term, while Strategy 5 (Complete Metadata) represents the ideal long-term solution for comprehensive provenance tracking.