async def test_rag_with_citation_metadata(self, mock_providers, sample_chunk_results):
        """Test RAG with citation metadata in search results."""
        # Function to build a RAG prompt with citations
        def build_rag_prompt_with_citations(query, search_results):
            context = ""
            citation_metadata = {}

            for i, result in enumerate(search_results):
                # Extract citation information
                citation_id = result.get("metadata", {}).get("citation_id")
                if citation_id:
                    # Add to context with citation marker
                    context += f"\n[{i+1}] {result['text']} [{citation_id}]"

                    # Store metadata
                    citation_metadata[citation_id] = {
                        "document_id": result["document_id"],
                        "chunk_id": result["chunk_id"],
                        "metadata": result.get("metadata", {})
                    }
                else:
                    context += f"\n[{i+1}] {result['text']}"

            prompt = f"Question: {query}\n\nContext:{context}\n\nPlease answer the question based on the provided context."

            return prompt, citation_metadata

        # Build prompt
        query = "What is the main concept?"
        prompt, citation_metadata = build_rag_prompt_with_citations(query, sample_chunk_results)

        # Verify prompt contains citations
        for i in range(5):
            assert f"[cite{i}]" in prompt

        # Verify metadata is stored
        assert len(citation_metadata) == 5
        for i in range(5):
            assert f"cite{i}" in citation_metadata
            assert citation_metadata[f"cite{i}"]["document_id"] == f"doc-{i//2}"
            assert citation_metadata[f"cite{i}"]["chunk_id"] == f"chunk-{i}"
Test RAG with citation metadata in search results.
py/tests/unit/retrieval/test_citations.py