r2rpy/py/tests/integration/test_retrieval_advanced.py

import uuid

from r2r import R2RClient


# Semantic Search Tests
def test_semantic_search_with_near_duplicates(client: R2RClient):
    """Test semantic search can handle and differentiate near-duplicate
    content."""
    random_1 = str(uuid.uuid4())
    random_2 = str(uuid.uuid4())
    # Create two similar but distinct documents
    doc1 = client.documents.create(
        raw_text=
        f"Aristotle was a Greek philosopher who studied logic {random_1}."
    ).results.document_id
    doc2 = client.documents.create(
        raw_text=
        f"Aristotle, the Greek philosopher, studied formal logic {random_2}."
    ).results.document_id

    resp = client.retrieval.search(
        query="Tell me about Aristotle's work in logic",
        search_mode="custom",
        search_settings={
            "use_semantic_search": True,
            "limit": 25
        },
    )
    results = resp.results.chunk_search_results

    # Both documents should be returned but with different scores
    scores = [
        r.score for r in results
        if str(r.document_id) in [str(doc1), str(doc2)]
    ]
    assert len(scores) == 2, "Expected both similar documents"
    assert len(
        set(scores)) == 2, ("Expected different scores for similar documents")


def test_semantic_search_multilingual(client: R2RClient):
    """Test semantic search handles multilingual content."""
    # Create documents in different languages
    random_1 = str(uuid.uuid4())
    random_2 = str(uuid.uuid4())
    random_3 = str(uuid.uuid4())

    docs = [
        (f"Aristotle was a philosopher {random_1}", "English"),
        (f"Aristóteles fue un filósofo {random_2}", "Spanish"),
        (f"アリストテレスは哲学者でした {random_3}", "Japanese"),
    ]
    doc_ids = []
    for text, lang in docs:
        doc_id = client.documents.create(raw_text=text,
                                         metadata={
                                             "language": lang
                                         }).results.document_id
        doc_ids.append(doc_id)

    # Query in different languages
    queries = [
        "Who was Aristotle?",
        "¿Quién fue Aristóteles?",
        "アリストテレスとは誰でしたか？",
    ]

    for query in queries:
        resp = client.retrieval.search(
            query=query,
            search_mode="custom",
            search_settings={
                "use_semantic_search": True,
                "limit": len(doc_ids),
            },
        )
        results = resp.results.chunk_search_results
        assert len(results) > 0, f"No results found for query: {query}"


# UNCOMMENT LATER
# # Hybrid Search Tests
# def test_hybrid_search_weight_balance(client: R2RClient):
#     """Test hybrid search balances semantic and full-text scores appropriately"""
#     # Create a document with high semantic relevance but low keyword match
#     semantic_doc = client.documents.create(
#         raw_text="The ancient Greek thinker who studied under Plato made significant contributions to logic."
#     ).results.document_id

#     # Create a document with high keyword match but low semantic relevance
#     keyword_doc = client.documents.create(
#         raw_text="Aristotle is a common name in certain regions. This text mentions Aristotle but is not about philosophy."
#     ).results.document_id

#     resp = client.retrieval.search(
#         query="What were Aristotle's philosophical contributions?",
#         search_mode="custom",
#         search_settings={
#             "use_hybrid_search": True,
#             "hybrid_settings": {
#                 "semantic_weight": 0.7,
#                 "full_text_weight": 0.3,
#             },
#         },
#     )
#     results = resp["results"]["chunk_search_results"]

#     # The semantic document should rank higher
#     semantic_rank = next(
#         i for i, r in enumerate(results) if r["document_id"] == semantic_doc
#     )
#     keyword_rank = next(
#         i for i, r in enumerate(results) if r["document_id"] == keyword_doc
#     )
#     assert (
#         semantic_rank < keyword_rank
#     ), "Semantic relevance should outweigh keyword matches"


# RAG Tests
def test_rag_context_window_limits(client: R2RClient):
    """Test RAG handles documents at or near context window limits."""
    # Create a document that approaches the context window limit
    random_1 = str(uuid.uuid4())
    large_text = ("Aristotle " * 1000
                  )  # Adjust multiplier based on your context window
    doc_id = client.documents.create(
        raw_text=f"{large_text} {random_1}").results.document_id

    resp = client.retrieval.rag(
        query="Summarize this text about Aristotle",
        search_settings={"filters": {
            "document_id": {
                "$eq": str(doc_id)
            }
        }},
        rag_generation_config={"max_tokens": 100},
    )
    assert resp.results is not None, (
        "RAG should handle large context gracefully")


# UNCOMMENT LATER
# def test_rag_empty_chunk_handling(client: R2RClient):
#     """Test RAG properly handles empty or whitespace-only chunks"""
#     doc_id = client.documents.create(chunks=["", " ", "\n", "Valid content"])[
#         "results"
#     ]["document_id"]

#     resp = client.retrieval.rag(
#         query="What is the content?",
#         search_settings={"filters": {"document_id": {"$eq": str(doc_id)}}},
#     )
#     assert "results" in resp, "RAG should handle empty chunks gracefully"

# # Agent Tests
# def test_agent_clarification_requests(client: R2RClient):
#     """Test agent's ability to request clarification for ambiguous queries"""
#     msg = Message(role="user", content="Compare them")
#     resp = client.retrieval.agent(
#         message=msg,
#         search_settings={"use_semantic_search": True},
#     )
#     content = resp["results"]["messages"][-1]["content"]
#     assert any(
#         phrase in content.lower()
#         for phrase in [
#             "could you clarify",
#             "who do you",
#             "what would you",
#             "please specify",
#         ]
#     ), "Agent should request clarification for ambiguous queries"

## TODO - uncomment later
# def test_agent_source_citation_consistency(client: R2RClient):
#     """Test agent consistently cites sources across conversation turns"""
#     conversation_id = client.conversations.create()["results"]["id"]

#     # First turn - asking about a specific topic
#     msg1 = Message(role="user", content="What did Aristotle say about ethics?")
#     resp1 = client.retrieval.agent(
#         message=msg1,
#         conversation_id=conversation_id,
#         include_title_if_available=True,
#     )

#     # Second turn - asking for more details
#     msg2 = Message(role="user", content="Can you elaborate on that point?")
#     resp2 = client.retrieval.agent(
#         message=msg2,
#         conversation_id=conversation_id,
#         include_title_if_available=True,
#     )

#     # Check that sources are consistently cited across turns
#     sources1 = _extract_sources(resp1["results"]["messages"][-1]["content"])
#     sources2 = _extract_sources(resp2["results"]["messages"][-1]["content"])
#     assert (
#         len(sources1) > 0 and len(sources2) > 0
#     ), "Both responses should cite sources"
#     assert any(
#         s in sources2 for s in sources1
#     ), "Follow-up should reference some original sources"

## TODO - uncomment later
# # Error Handling Tests
# def test_malformed_filter_handling(client: R2RClient):
#     """Test system properly handles malformed filter conditions"""
#     invalid_filters = [
#         {"$invalid": {"$eq": "value"}},
#         {"field": {"$unsupported": "value"}},
#         {"$and": [{"field": "incomplete_operator"}]},
#         {"$or": []},  # Empty OR condition
#         {"$and": [{}]},  # Empty filter in AND
#     ]

#     for invalid_filter in invalid_filters:
#         with pytest.raises(R2RException) as exc_info:
#             client.retrieval.search(
#                 query="test", search_settings={"filters": invalid_filter}
#             )
#         assert exc_info.value.status_code in [
#             400,
#             422,
#         ], f"Expected validation error for filter: {invalid_filter}"

## TODO - Uncomment later
# def test_concurrent_search_stability(client: R2RClient):
#     """Test system handles concurrent search requests properly"""
#     import asyncio

#     async def concurrent_searches():
#         tasks = []
#         for i in range(10):  # Adjust number based on system capabilities
#             task = asyncio.create_task(
#                 client.retrieval.search_async(
#                     query=f"Concurrent test query {i}", search_mode="basic"
#                 )
#             )
#             tasks.append(task)

#         results = await asyncio.gather(*tasks, return_exceptions=True)
#         return results

#     results = asyncio.run(concurrent_searches())
#     assert all(
#         not isinstance(r, Exception) for r in results
#     ), "Concurrent searches should complete without errors"


# Helper function for source extraction
def _extract_sources(content: str) -> list[str]:
    """Extract source citations from response content."""
    # This is a simplified version - implement based on your citation format
    import re

    return re.findall(r'"([^"]*)"', content)