r2r/tests/integration/test_retrieval_advanced.py

import pytest

from r2r import Message, R2RException, SearchMode


# Semantic Search Tests
def test_semantic_search_with_near_duplicates(client):
    """Test semantic search can handle and differentiate near-duplicate content"""
    # Create two similar but distinct documents
    doc1 = client.documents.create(
        raw_text="Aristotle was a Greek philosopher who studied logic."
    )["results"]["document_id"]
    doc2 = client.documents.create(
        raw_text="Aristotle, the Greek philosopher, studied formal logic."
    )["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": 5},
    )
    results = resp["results"]["chunk_search_results"]

    # Both documents should be returned but with different scores
    scores = [r["score"] for r in results if r["document_id"] in [doc1, 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):
    """Test semantic search handles multilingual content"""
    # Create documents in different languages
    docs = [
        ("Aristotle was a philosopher", "English"),
        ("Aristóteles fue un filósofo", "Spanish"),
        ("アリストテレスは哲学者でした", "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):
#     """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):
    """Test RAG handles documents at or near context window limits"""
    # Create a document that approaches the context window limit
    large_text = (
        "Aristotle " * 1000
    )  # Adjust multiplier based on your context window
    doc_id = client.documents.create(raw_text=large_text)["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 "results" in resp, "RAG should handle large context gracefully"


# UNCOMMENT LATER
# def test_rag_empty_chunk_handling(client):
#     """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):
#     """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):
#     """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):
#     """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):
#     """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)