From c74f0c2b0f08adfdb39aabcf709c7ad1c9108807 Mon Sep 17 00:00:00 2001 From: Gabriel Luiz Freitas Almeida Date: Fri, 1 Mar 2024 09:32:44 -0300 Subject: [PATCH] Refactor CustomComponent data handling --- .../custom_component/custom_component.py | 24 +++++++++++++++---- 1 file changed, 19 insertions(+), 5 deletions(-) diff --git a/src/backend/langflow/interface/custom/custom_component/custom_component.py b/src/backend/langflow/interface/custom/custom_component/custom_component.py index c3b06f90a..707ab6611 100644 --- a/src/backend/langflow/interface/custom/custom_component/custom_component.py +++ b/src/backend/langflow/interface/custom/custom_component/custom_component.py @@ -16,6 +16,7 @@ import yaml from cachetools import TTLCache, cachedmethod from fastapi import HTTPException from langchain_core.documents import Document +from pydantic import BaseModel from sqlmodel import select from langflow.interface.custom.code_parser.utils import ( @@ -137,15 +138,28 @@ class CustomComponent(Component): if not isinstance(data, Sequence): data = [data] for item in data: - if isinstance(item, str): - records.append(Record(text=item)) + if isinstance(item, Document): + item = {"text": item.page_content, "data": item.metadata} + elif isinstance(item, BaseModel): + model_dump = item.model_dump() + if text_key not in model_dump: + raise ValueError(f"Key '{text_key}' not found in BaseModel item.") + if data_key not in model_dump: + raise ValueError(f"Key '{data_key}' not found in BaseModel item.") + item = {"text": model_dump[text_key], "data": model_dump[data_key]} + elif isinstance(item, str): + item = {"text": item, "data": {}} elif isinstance(item, dict): - records.append(Record(text=item.get(text_key), data=item.get(data_key))) - elif isinstance(item, Document): - records.append(Record(text=item.page_content, data=item.metadata)) + if text_key not in item: + raise ValueError(f"Key '{text_key}' not found in dictionary item.") + if data_key not in item: + raise ValueError(f"Key '{data_key}' not found in dictionary item.") + item = {"text": item[text_key], "data": item[data_key]} else: raise ValueError(f"Invalid data type: {type(item)}") + records.append(Record(**item)) + return records def create_references_from_records(