def test_agent_handles_large_context(client):
    """Test agent handles large amount of context efficiently."""
    # Create a document with substantial content
    large_content = "Philosophy " * 2000  # ~16K chars
    doc_id = client.documents.create(raw_text=large_content).results.document_id

    start_time = time.time()
    response = client.retrieval.agent(
        message={"role": "user", "content": "Summarize everything you know about philosophy."},
        search_settings={"filters": {"document_id": {"$eq": str(doc_id)}}},
        rag_generation_config={"stream": False, "max_tokens_to_sample": 200},
    )
    end_time = time.time()

    response_time = end_time - start_time
    assert response.results.messages[-1].content, "Agent should produce a summary with large context"
    assert response_time < 20, f"Large context processing should complete in reasonable time, took {response_time:.2f}s"

    # Clean up
    client.documents.delete(id=doc_id)
Test agent handles large amount of context efficiently.
py/tests/integration/test_agent.py