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

Refactor CustomComponent data handling

Browse source
This commit is contained in:
Gabriel Luiz Freitas Almeida 2024-03-01 09:32:44 -03:00
commit c74f0c2b0f
1 changed files with 19 additions and 5 deletions
Download patch file Download diff file Expand all files Collapse all files

24
src/backend/langflow/interface/custom/custom_component/custom_component.py
View file

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