async def test_rag_response_with_citations(self, mock_providers, sample_chunk_results):
        """Test generating a RAG response with citations."""
        # Function to generate RAG response with citations
        async def generate_rag_response_with_citations(query, search_results):
            # Build prompt with citations
            context = ""
            citation_metadata = {}

            for i, result in enumerate(search_results):
                citation_id = result.get("metadata", {}).get("citation_id")
                if citation_id:
                    context += f"\n[{i+1}] {result['text']} [{citation_id}]"

                    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."

            # Generate response (mocked)
            # In real implementation, this would call the LLM
            mock_providers.llm.aget_completion.return_value = {
                "choices": [{
                    "message": {
                        "content": "The main concept is explained in [cite0] and further elaborated in [cite2]."
                    }
                }]
            }

            response = await mock_providers.llm.aget_completion(prompt=prompt)
            content = response["choices"][0]["message"]["content"]

            return content, citation_metadata

        # Generate response
        query = "What is the main concept?"
        response, citation_metadata = await generate_rag_response_with_citations(query, sample_chunk_results)

        # Verify response contains citations
        assert "[cite0]" in response
        assert "[cite2]" in response

        # Extract citations from response
        def extract_citations_from_response(text):
            citation_pattern = r'\[([\w\d]+)\]'
            citations = re.findall(citation_pattern, text)
            return citations

        citations = extract_citations_from_response(response)
        assert "cite0" in citations
        assert "cite2" in citations
Test generating a RAG response with citations.
py/tests/unit/retrieval/test_citations.py