async def test_consolidate_citations_in_final_answer(self, mock_providers):
        """Test consolidating citations in the final answer."""
        # Create a citation tracker with some spans
        tracker = CitationTracker()
        tracker.is_new_span("cite0", (10, 18))
        tracker.is_new_span("cite0", (30, 38))
        tracker.is_new_span("cite2", (50, 58))

        # Create citation metadata
        citation_metadata = {
            "cite0": {
                "document_id": "doc-0",
                "chunk_id": "chunk-0",
                "metadata": {"source": "source-0", "title": "Document 0"}
            },
            "cite2": {
                "document_id": "doc-1",
                "chunk_id": "chunk-2",
                "metadata": {"source": "source-2", "title": "Document 1"}
            }
        }

        # Function to consolidate citations
        def consolidate_citations(response_text, citation_tracker, citation_metadata):
            # Get all citations from the tracker
            all_citation_spans = citation_tracker.get_all_citation_spans()

            # Build consolidated citations
            consolidated_citations = {}
            for citation_id, spans in all_citation_spans.items():
                if citation_id in citation_metadata:
                    metadata = citation_metadata[citation_id]
                    consolidated_citations[citation_id] = {
                        "spans": spans,
                        "document_id": metadata["document_id"],
                        "chunk_id": metadata["chunk_id"],
                        "metadata": metadata["metadata"]
                    }

            # Return the response with consolidated citations
            return {
                "response": response_text,
                "citations": consolidated_citations
            }

        # Test response
        response_text = "The main concept is explained in [cite0] and further elaborated in [cite2]."

        # Consolidate citations
        result = consolidate_citations(response_text, tracker, citation_metadata)

        # Verify result
        assert "response" in result
        assert "citations" in result
        assert result["response"] == response_text

        # Verify consolidated citations
        assert "cite0" in result["citations"]
        assert "cite2" in result["citations"]
        assert len(result["citations"]["cite0"]["spans"]) == 2
        assert len(result["citations"]["cite2"]["spans"]) == 1
        assert result["citations"]["cite0"]["document_id"] == "doc-0"
        assert result["citations"]["cite2"]["document_id"] == "doc-1"
Test consolidating citations in the final answer.
py/tests/unit/retrieval/test_citations.py