From b475e51c51abd984e6445598326b0eb6da9a9249 Mon Sep 17 00:00:00 2001
From: ogabrielluiz <gabriel@langflow.org>
Date: Wed, 19 Jun 2024 01:07:05 -0300
Subject: [PATCH] refactor: Update AstraDBMessageWriterComponent to use
 AstraVectorStoreComponent

This commit updates the AstraDBMessageWriterComponent class in the AstraDBMessageWriter.py file to use the AstraVectorStoreComponent for vector storage. By making this change, we ensure consistency and improve code organization. This update aligns with recent changes in the AstraDBVectorStoreComponent and promotes code reuse.
---
 .../integration/astra/test_astra_component.py | 29 ++++++++-----------
 1 file changed, 12 insertions(+), 17 deletions(-)

diff --git a/tests/integration/astra/test_astra_component.py b/tests/integration/astra/test_astra_component.py
index 161e6436c..2a67924c4 100644
--- a/tests/integration/astra/test_astra_component.py
+++ b/tests/integration/astra/test_astra_component.py
@@ -1,20 +1,15 @@
 import os
+
 import pytest
+from langchain_core.documents import Document
+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 AstraVectorStoreComponent
+from langflow.schema.data import Data
 
 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"
@@ -54,7 +49,7 @@ def test_astra_setup(astra_fixture):
     api_endpoint = os.getenv("ASTRA_DB_API_ENDPOINT")
     embedding = MockEmbeddings()
 
-    component = AstraDBVectorStoreComponent()
+    component = AstraVectorStoreComponent()
     component.build(
         token=application_token,
         api_endpoint=api_endpoint,
@@ -74,9 +69,9 @@ def test_astra_embeds_and_search(astra_fixture):
     embedding = MockEmbeddings()
 
     documents = [Document(page_content="test1"), Document(page_content="test2")]
-    records = [Record.from_document(d) for d in documents]
+    records = [Data.from_document(d) for d in documents]
 
-    component = AstraDBVectorStoreComponent()
+    component = AstraVectorStoreComponent()
     component.build(
         token=application_token,
         api_endpoint=api_endpoint,
@@ -109,7 +104,7 @@ def test_astra_memory():
     writer = AstraDBMessageWriterComponent()
     reader = AstraDBMessageReaderComponent()
 
-    input_value = Record.from_document(
+    input_value = Data.from_document(
         Document(
             page_content="memory1",
             metadata={"session_id": 1, "sender": "human", "sender_name": "Bob"},
@@ -131,7 +126,7 @@ def test_astra_memory():
         collection_name=MEMORY_COLLECTION,
     )
     assert len(records) == 1
-    assert isinstance(records[0], Record)
+    assert isinstance(records[0], Data)
     content = records[0].get_text()
     assert content == "memory1"