From 1863d463d0deab90f8418d0f84be31ba21d7e13c Mon Sep 17 00:00:00 2001 From: Gabriel Luiz Freitas Almeida Date: Fri, 14 Jul 2023 14:20:09 -0300 Subject: [PATCH] =?UTF-8?q?=F0=9F=94=A8=20refactor(conftest.py):=20change?= =?UTF-8?q?=20return=20type=20of=20`build`=20method=20in=20`CSVLoaderCompo?= =?UTF-8?q?nent`=20from=20List[Document]=20to=20Document=20to=20match=20th?= =?UTF-8?q?e=20actual=20return=20type=20=F0=9F=94=A7=20chore(conftest.py):?= =?UTF-8?q?=20add=20new=20fixture=20`filter=5Fdocs`=20for=20testing=20`Doc?= =?UTF-8?q?umentFilterByLengthComponent`?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- tests/conftest.py | 22 +++++++++++++++++++++- 1 file changed, 21 insertions(+), 1 deletion(-) diff --git a/tests/conftest.py b/tests/conftest.py index 328a168ad..79704e3b6 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -250,7 +250,7 @@ class CSVLoaderComponent(CustomComponent): "column_name": {"field_type": "str", "required": True}, } - def build(self, filename: str, column_name: str) -> List[Document]: + def build(self, filename: str, column_name: str) -> Document: # Load the CSV file df = pd.read_csv(filename) @@ -265,3 +265,23 @@ class CSVLoaderComponent(CustomComponent): documents.append(Document(page_content=str(content), metadata=metadata)) return documents""" + + +@pytest.fixture +def filter_docs(): + return """from langchain.schema import Document +from langflow.interface.custom.base import CustomComponent +from typing import List + +class DocumentFilterByLengthComponent(CustomComponent): + display_name: str = "Document Filter By Length" + field_config = { + "documents": {"field_type": "Document", "required": True}, + "max_length": {"field_type": "int", "required": True}, + } + + def build(self, documents: List[Document], max_length: int) -> List[Document]: + # Filter the documents by length + filtered_documents = [doc for doc in documents if len(doc.page_content) <= max_length] + + return filtered_documents"""