From 74625bb3a18b7d09edc8cc87976c0b4a95d50fd5 Mon Sep 17 00:00:00 2001
From: emrgnt-cmplxty <owen@algofi.org>
Date: Thu, 13 Feb 2025 14:04:13 -0800
Subject: [PATCH] add additional agent tests

---
 py/tests/integration/test_agent.py | 65 ++++++++++++++++++++++++++++++
 1 file changed, 65 insertions(+)
 create mode 100644 py/tests/integration/test_agent.py

diff --git a/py/tests/integration/test_agent.py b/py/tests/integration/test_agent.py
new file mode 100644
index 000000000..b023340ab
--- /dev/null
+++ b/py/tests/integration/test_agent.py
@@ -0,0 +1,65 @@
+from r2r import R2RClient
+
+
+def test_agentic_citations_0(client: R2RClient, test_collection):
+
+    response = client.retrieval.rag(
+        query="Who was Aristotle? USE YOUR SOURCES.",
+    )
+    assert (
+        response.results.citations[0].document_id
+        == test_collection["document_ids"][0]
+    ), "Expected first citation to first doc"
+
+
+def test_agentic_citations_1(client: R2RClient, test_collection):
+
+    response = client.retrieval.rag(
+        query="Who was Socrates? USE YOUR SOURCES.",
+    )
+    assert (
+        response.results.citations[0].document_id
+        == test_collection["document_ids"][1]
+    ), "Expected first citation to second doc"
+
+
+def test_agentic_citations_2(client: R2RClient, test_collection):
+
+    response = client.retrieval.rag(
+        query="Who were Rene Descartes and Immanuel Kant? USE YOUR SOURCES.",
+    )
+    assert test_collection["document_ids"][2] in [
+        citation.document_id for citation in response.results.citations
+    ], "Expected a citation to third doc"
+    assert test_collection["document_ids"][3] in [
+        citation.document_id for citation in response.results.citations
+    ], "Expected a citation to fourth doc"
+
+
+def test_agentic_citations_3(client: R2RClient, test_collection):
+    """
+    Tests the agent endpoint in non-streaming mode with a single user message.
+    Verifies final text and citations.
+    """
+    response = client.retrieval.agent(
+        message={
+            "role": "user",
+            "content": "Tell me about Socrates in detail. USE YOUR SOURCES.",
+        },
+        rag_generation_config={
+            "stream": False,
+        },
+    )
+    # Typically returns a WrappedAgentResponse
+    assert response.results.messages, "No messages returned"
+    assistant_msg = response.results.messages[-1]
+    assert "Socrates" in assistant_msg.content
+
+    print("response.results.messages = ", response.results.messages)
+    # If your server includes citations in `metadata`, you can check them here:
+    if assistant_msg.metadata and "citations" in assistant_msg.metadata:
+        citations = assistant_msg.metadata["citations"]
+        doc_ids = [c["document_id"] for c in citations]
+        assert (
+            test_collection["document_ids"][1] in doc_ids
+        ), "Expected Socrates doc citation"