merge dev into two_edges

tests/integration/astra/test_astra_component.py

@@ -0,0 +1,160 @@
+import os
+import pytest
+
+from integration.utils import MockEmbeddings, check_env_vars
+
+from langflow.components.memories.AstraDBMessageReader import (
+    AstraDBMessageReaderComponent,
+)
+from langflow.components.memories.AstraDBMessageWriter import (
+    AstraDBMessageWriterComponent,
+)
+from langflow.components.vectorsearch.AstraDBSearch import AstraDBSearchComponent
+from langflow.components.vectorstores.AstraDB import AstraDBVectorStoreComponent
+from langflow.schema.record import Record
+
+from langchain_core.documents import Document
+
+COLLECTION = "test_basic"
+SEARCH_COLLECTION = "test_search"
+MEMORY_COLLECTION = "test_memory"
+
+
+@pytest.fixture()
+def astra_fixture(request):
+    """
+    Sets up the astra collection and cleans up after
+    """
+    try:
+        from langchain_astradb import AstraDBVectorStore
+    except ImportError:
+        raise ImportError(
+            "Could not import langchain Astra DB integration package. Please install it with `pip install langchain-astradb`."
+        )
+
+    store = AstraDBVectorStore(
+        collection_name=request.param,
+        embedding=MockEmbeddings(),
+        api_endpoint=os.getenv("ASTRA_DB_API_ENDPOINT"),
+        token=os.getenv("ASTRA_DB_APPLICATION_TOKEN"),
+    )
+
+    yield
+
+    store.delete_collection()
+
+
+@pytest.mark.skipif(
+    not check_env_vars("ASTRA_DB_APPLICATION_TOKEN", "ASTRA_DB_API_ENDPOINT"),
+    reason="missing astra env vars",
+)
+@pytest.mark.parametrize("astra_fixture", [COLLECTION], indirect=True)
+def test_astra_setup(astra_fixture):
+    application_token = os.getenv("ASTRA_DB_APPLICATION_TOKEN")
+    api_endpoint = os.getenv("ASTRA_DB_API_ENDPOINT")
+    embedding = MockEmbeddings()
+
+    component = AstraDBVectorStoreComponent()
+    component.build(
+        token=application_token,
+        api_endpoint=api_endpoint,
+        collection_name=COLLECTION,
+        embedding=embedding,
+    )
+
+
+@pytest.mark.skipif(
+    not check_env_vars("ASTRA_DB_APPLICATION_TOKEN", "ASTRA_DB_API_ENDPOINT"),
+    reason="missing astra env vars",
+)
+@pytest.mark.parametrize("astra_fixture", [SEARCH_COLLECTION], indirect=True)
+def test_astra_embeds_and_search(astra_fixture):
+    application_token = os.getenv("ASTRA_DB_APPLICATION_TOKEN")
+    api_endpoint = os.getenv("ASTRA_DB_API_ENDPOINT")
+    embedding = MockEmbeddings()
+
+    documents = [Document(page_content="test1"), Document(page_content="test2")]
+    records = [Record.from_document(d) for d in documents]
+
+    component = AstraDBVectorStoreComponent()
+    component.build(
+        token=application_token,
+        api_endpoint=api_endpoint,
+        collection_name=SEARCH_COLLECTION,
+        embedding=embedding,
+        inputs=records,
+    )
+
+    component = AstraDBSearchComponent()
+    records = component.build(
+        token=application_token,
+        api_endpoint=api_endpoint,
+        collection_name=SEARCH_COLLECTION,
+        embedding=embedding,
+        input_value="test1",
+        number_of_results=1,
+    )
+
+    assert len(records) == 1
+
+
+@pytest.mark.skipif(
+    not check_env_vars("ASTRA_DB_APPLICATION_TOKEN", "ASTRA_DB_API_ENDPOINT"),
+    reason="missing astra env vars",
+)
+def test_astra_memory():
+    application_token = os.getenv("ASTRA_DB_APPLICATION_TOKEN")
+    api_endpoint = os.getenv("ASTRA_DB_API_ENDPOINT")
+
+    writer = AstraDBMessageWriterComponent()
+    reader = AstraDBMessageReaderComponent()
+
+    input_value = Record.from_document(
+        Document(
+            page_content="memory1",
+            metadata={"session_id": 1, "sender": "human", "sender_name": "Bob"},
+        )
+    )
+    writer.build(
+        input_value=input_value,
+        session_id=1,
+        token=application_token,
+        api_endpoint=api_endpoint,
+        collection_name=MEMORY_COLLECTION,
+    )
+
+    # verify reading w/ same session id pulls the same record
+    records = reader.build(
+        session_id=1,
+        token=application_token,
+        api_endpoint=api_endpoint,
+        collection_name=MEMORY_COLLECTION,
+    )
+    assert len(records) == 1
+    assert isinstance(records[0], Record)
+    content = records[0].get_text()
+    assert content == "memory1"
+
+    # verify reading w/ different session id does not pull the same record
+    records = reader.build(
+        session_id=2,
+        token=application_token,
+        api_endpoint=api_endpoint,
+        collection_name=MEMORY_COLLECTION,
+    )
+    assert len(records) == 0
+
+    # Cleanup store - doing here rather than fixture (see https://github.com/langchain-ai/langchain-datastax/pull/36)
+    try:
+        from langchain_astradb import AstraDBVectorStore
+    except ImportError:
+        raise ImportError(
+            "Could not import langchain Astra DB integration package. Please install it with `pip install langchain-astradb`."
+        )
+    store = AstraDBVectorStore(
+        collection_name=MEMORY_COLLECTION,
+        embedding=MockEmbeddings(),
+        api_endpoint=api_endpoint,
+        token=application_token,
+    )
+    store.delete_collection()

tests/integration/utils.py

@@ -0,0 +1,35 @@
+import os
+from typing import List
+
+from langflow.field_typing import Embeddings
+
+
+def check_env_vars(*vars):
+    """
+    Check if all specified environment variables are set.
+
+    Args:
+    *vars (str): The environment variables to check.
+
+    Returns:
+    bool: True if all environment variables are set, False otherwise.
+    """
+    return all(os.getenv(var) for var in vars)
+
+
+class MockEmbeddings(Embeddings):
+    def __init__(self):
+        self.embedded_documents = None
+        self.embedded_query = None
+
+    @staticmethod
+    def mock_embedding(text: str):
+        return [len(text) / 2, len(text) / 5, len(text) / 10]
+
+    def embed_documents(self, texts: List[str]) -> List[List[float]]:
+        self.embedded_documents = texts
+        return [self.mock_embedding(text) for text in texts]
+
+    def embed_query(self, text: str) -> List[float]:
+        self.embedded_query = text
+        return self.mock_embedding(text)

tests/test_endpoints.py

@@ -707,36 +707,36 @@ def test_run_flow_with_caching_invalid_input_format(client: TestClient, starter_
     assert response.status_code == status.HTTP_422_UNPROCESSABLE_ENTITY
 
 
-@pytest.mark.api_key_required
-def test_run_flow_with_session_id(client, starter_project, created_api_key):
-    headers = {"x-api-key": created_api_key.api_key}
-    flow_id = starter_project["id"]
-    payload = {
-        "input_value": "value1",
-        "input_type": "text",
-        "output_type": "text",
-        "session_id": "test-session-id",
-    }
-    response = client.post(f"/api/v1/run/{flow_id}", json=payload, headers=headers)
-    assert response.status_code == status.HTTP_404_NOT_FOUND
-    data = response.json()
-    assert {"detail": "Session test-session-id not found"} == data
+# @pytest.mark.api_key_required
+# def test_run_flow_with_session_id(client, starter_project, created_api_key):
+#     headers = {"x-api-key": created_api_key.api_key}
+#     flow_id = starter_project["id"]
+#     payload = {
+#         "input_value": "value1",
+#         "input_type": "text",
+#         "output_type": "text",
+#         "session_id": "test-session-id",
+#     }
+#     response = client.post(f"/api/v1/run/{flow_id}", json=payload, headers=headers)
+#     assert response.status_code == status.HTTP_404_NOT_FOUND
+#     data = response.json()
+#     assert {"detail": "Session test-session-id not found"} == data
 
 
-def test_run_flow_with_invalid_session_id(client, starter_project, created_api_key):
-    headers = {"x-api-key": created_api_key.api_key}
-    flow_id = starter_project["id"]
-    payload = {
-        "input_value": "value1",
-        "input_type": "text",
-        "output_type": "text",
-        "session_id": "invalid-session-id",
-    }
-    response = client.post(f"/api/v1/run/{flow_id}", json=payload, headers=headers)
-    assert response.status_code == status.HTTP_404_NOT_FOUND
-    data = response.json()
-    assert "detail" in data
-    assert f"Session {payload['session_id']} not found" in data["detail"]
+# def test_run_flow_with_invalid_session_id(client, starter_project, created_api_key):
+#     headers = {"x-api-key": created_api_key.api_key}
+#     flow_id = starter_project["id"]
+#     payload = {
+#         "input_value": "value1",
+#         "input_type": "text",
+#         "output_type": "text",
+#         "session_id": "invalid-session-id",
+#     }
+#     response = client.post(f"/api/v1/run/{flow_id}", json=payload, headers=headers)
+#     assert response.status_code == status.HTTP_404_NOT_FOUND
+#     data = response.json()
+#     assert "detail" in data
+#     assert f"Session {payload['session_id']} not found" in data["detail"]
 
 
 @pytest.mark.api_key_required