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 cachetools import TTLCache, cachedmethod
from fastapi import HTTPException from fastapi import HTTPException
from langchain_core.documents import Document from langchain_core.documents import Document
from pydantic import BaseModel
from sqlmodel import select from sqlmodel import select
from langflow.interface.custom.code_parser.utils import ( from langflow.interface.custom.code_parser.utils import (
@ -137,15 +138,28 @@ class CustomComponent(Component):
if not isinstance(data, Sequence): if not isinstance(data, Sequence):
data = [data] data = [data]
for item in data: for item in data:
if isinstance(item, str): if isinstance(item, Document):
records.append(Record(text=item)) 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): elif isinstance(item, dict):
records.append(Record(text=item.get(text_key), data=item.get(data_key))) if text_key not in item:
elif isinstance(item, Document): raise ValueError(f"Key '{text_key}' not found in dictionary item.")
records.append(Record(text=item.page_content, data=item.metadata)) 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: else:
raise ValueError(f"Invalid data type: {type(item)}") raise ValueError(f"Invalid data type: {type(item)}")
records.append(Record(**item))
return records return records
def create_references_from_records( def create_references_from_records(