jyapayne/langflow
1
0
Fork 0
Code Issues Pull requests Projects Releases 1 Packages Wiki Activity Actions

refactor: Update AstraDBMessageWriterComponent to use AstraVectorStoreComponent

Browse source 
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.
This commit is contained in:
ogabrielluiz 2024-06-19 01:07:05 -03:00
commit b475e51c51
1 changed files with 12 additions and 17 deletions
Download patch file Download diff file Expand all files Collapse all files

29
tests/integration/astra/test_astra_component.py
View file

@ -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"