def test_extract_citations_with_spans(self):
        """Test citation extraction with text spans."""
        # Test function to extract citations with spans
        def extract_citations_with_spans(text):
            citation_pattern = r'\[([\w\d]+)\]'
            citations_with_spans = []

            for match in re.finditer(citation_pattern, text):
                citation_id = match.group(1)
                start = match.start()
                end = match.end()

                # Get the context (text before and after the citation)
                context_start = max(0, start - 50)
                context_end = min(len(text), end + 50)
                context = text[context_start:context_end]

                citations_with_spans.append({
                    "citation_id": citation_id,
                    "start": start,
                    "end": end,
                    "context": context
                })

            return citations_with_spans

        # Test text
        text = (
            "Aristotle discussed virtue ethics in his Nicomachean Ethics [abc1234]. "
            "According to Plato [xyz5678], the ideal state is described in The Republic. "
            "Socrates' method of questioning is demonstrated in many dialogues [ghi9012]."
        )

        # Extract citations with spans
        extracted = extract_citations_with_spans(text)

        # Verify the correct number of citations was extracted
        assert len(extracted) == 3

        # Verify citation IDs are correct
        assert extracted[0]["citation_id"] == "abc1234"
        assert extracted[1]["citation_id"] == "xyz5678"
        assert extracted[2]["citation_id"] == "ghi9012"

        # Verify spans and context
        for citation in extracted:
            assert citation["start"] < citation["end"]
            assert text[citation["start"]:citation["end"]] == f"[{citation['citation_id']}]"
            assert citation["citation_id"] in citation["context"]
Test citation extraction with text spans.
py/tests/unit/retrieval/test_citations.py