From 03d7bce95dc7448a3c90555e9a00cf3d2eec45b3 Mon Sep 17 00:00:00 2001
From: Gabriel Luiz Freitas Almeida <gabriel@logspace.ai>
Date: Wed, 6 Mar 2024 18:45:51 -0300
Subject: [PATCH] Refactor PromptComponent to convert dictionary values to
 strings

---
 .../langflow/components/prompts/Prompt.py     | 14 ++---
 .../components/prompts/base/__init__.py       |  0
 .../langflow/components/prompts/base/utils.py | 55 +++++++++++++++++++
 3 files changed, 61 insertions(+), 8 deletions(-)
 create mode 100644 src/backend/langflow/components/prompts/base/__init__.py
 create mode 100644 src/backend/langflow/components/prompts/base/utils.py

diff --git a/src/backend/langflow/components/prompts/Prompt.py b/src/backend/langflow/components/prompts/Prompt.py
index 975998919..7a6da307d 100644
--- a/src/backend/langflow/components/prompts/Prompt.py
+++ b/src/backend/langflow/components/prompts/Prompt.py
@@ -1,6 +1,7 @@
 from langchain_core.prompts import PromptTemplate
 
 from langflow import CustomComponent
+from langflow.components.prompts.base.utils import dict_values_to_string
 from langflow.field_typing import Prompt, TemplateField, Text
 
 
@@ -21,16 +22,13 @@ class PromptComponent(CustomComponent):
         **kwargs,
     ) -> Text:
         prompt_template = PromptTemplate.from_template(Text(template))
-
-        attributes_to_check = ["text", "page_content"]
-        for key, value in kwargs.copy().items():
-            for attribute in attributes_to_check:
-                if hasattr(value, attribute):
-                    kwargs[key] = getattr(value, attribute)
-
+        kwargs = dict_values_to_string(kwargs)
+        kwargs = {
+            k: "\n".join(v) if isinstance(v, list) else v for k, v in kwargs.items()
+        }
         try:
             formated_prompt = prompt_template.format(**kwargs)
         except Exception as exc:
             raise ValueError(f"Error formatting prompt: {exc}") from exc
-        self.status = f'Prompt: "{formated_prompt}"'
+        self.status = f'Prompt:\n"{formated_prompt}"'
         return formated_prompt
diff --git a/src/backend/langflow/components/prompts/base/__init__.py b/src/backend/langflow/components/prompts/base/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/src/backend/langflow/components/prompts/base/utils.py b/src/backend/langflow/components/prompts/base/utils.py
new file mode 100644
index 000000000..1f41ebda1
--- /dev/null
+++ b/src/backend/langflow/components/prompts/base/utils.py
@@ -0,0 +1,55 @@
+from langchain_core.documents import Document
+
+from langflow.schema import Record
+
+
+def dict_values_to_string(d: dict) -> dict:
+    """
+    Converts the values of a dictionary to strings.
+
+    Args:
+        d (dict): The dictionary whose values need to be converted.
+
+    Returns:
+        dict: The dictionary with values converted to strings.
+    """
+    # Do something similar to the above
+    for key, value in d.items():
+        # it could be a list of records or documents or strings
+        if isinstance(value, list):
+            for i, item in enumerate(value):
+                if isinstance(item, Record):
+                    d[key][i] = record_to_string(item)
+                elif isinstance(item, Document):
+                    d[key][i] = document_to_string(item)
+        elif isinstance(value, Record):
+            d[key] = record_to_string(value)
+        elif isinstance(value, Document):
+            d[key] = document_to_string(value)
+    return d
+
+
+def record_to_string(record: Record) -> str:
+    """
+    Convert a record to a string.
+
+    Args:
+        record (Record): The record to convert.
+
+    Returns:
+        str: The record as a string.
+    """
+    return record.text
+
+
+def document_to_string(document: Document) -> str:
+    """
+    Convert a document to a string.
+
+    Args:
+        document (Document): The document to convert.
+
+    Returns:
+        str: The document as a string.
+    """
+    return document.page_content