From 521c1faa2838c7779d8975312487cb2c5d257523 Mon Sep 17 00:00:00 2001
From: ogabrielluiz <gabriel@langflow.org>
Date: Thu, 6 Jun 2024 14:20:46 -0300
Subject: [PATCH] Refactor utils.py to handle file paths and files in
 dict_values_to_string and add get_file_paths and get_files functions

---
 .../base/langflow/base/prompts/utils.py       | 36 ++++++++++++++++---
 .../base/langflow/components/inputs/Prompt.py | 25 ++++++-------
 2 files changed, 43 insertions(+), 18 deletions(-)

diff --git a/src/backend/base/langflow/base/prompts/utils.py b/src/backend/base/langflow/base/prompts/utils.py
index 2270035af..c4b4e08c2 100644
--- a/src/backend/base/langflow/base/prompts/utils.py
+++ b/src/backend/base/langflow/base/prompts/utils.py
@@ -1,9 +1,10 @@
+import base64
 from copy import deepcopy
 
-
 from langchain_core.documents import Document
 
 from langflow.schema import Record
+from langflow.services.deps import get_storage_service
 
 
 def record_to_string(record: Record) -> str:
@@ -19,7 +20,7 @@ def record_to_string(record: Record) -> str:
     return record.get_text()
 
 
-def dict_values_to_string(d: dict) -> dict:
+async def dict_values_to_string(d: dict) -> dict:
     """
     Converts the values of a dictionary to strings.
 
@@ -36,16 +37,43 @@ def dict_values_to_string(d: dict) -> dict:
         if isinstance(value, list):
             for i, item in enumerate(value):
                 if isinstance(item, Record):
-                    d_copy[key][i] = record_to_string(item)
+                    d_copy[key][i] = item.to_lc_message()
                 elif isinstance(item, Document):
                     d_copy[key][i] = document_to_string(item)
         elif isinstance(value, Record):
-            d_copy[key] = record_to_string(value)
+            if "files" in value and value.files:
+                files = await get_file_paths(value.files)
+                value.files = files
+            d_copy[key] = value.to_lc_message()
         elif isinstance(value, Document):
             d_copy[key] = document_to_string(value)
     return d_copy
 
 
+async def get_file_paths(files: list[str]):
+    storage_service = get_storage_service()
+    file_paths = []
+    for file in files:
+        flow_id, file_name = file.split("/")
+        file_paths.append(storage_service.build_full_path(flow_id=flow_id, file_name=file_name))
+    return file_paths
+
+
+async def get_files(
+    file_paths: str,
+    convert_to_base64: bool = False,
+):
+    storage_service = get_storage_service()
+    file_objects = []
+    for file_path in file_paths:
+        flow_id, file_name = file_path.split("/")
+        file_object = await storage_service.get_file(flow_id=flow_id, file_name=file_name)
+        if convert_to_base64:
+            file_object = base64.b64encode(file_object).decode("utf-8")
+        file_objects.append(file_object)
+    return file_objects
+
+
 def document_to_string(document: Document) -> str:
     """
     Convert a document to a string.
diff --git a/src/backend/base/langflow/components/inputs/Prompt.py b/src/backend/base/langflow/components/inputs/Prompt.py
index 2c76e6132..f51c1aaab 100644
--- a/src/backend/base/langflow/components/inputs/Prompt.py
+++ b/src/backend/base/langflow/components/inputs/Prompt.py
@@ -1,7 +1,9 @@
-from langchain_core.prompts import PromptTemplate
+from langchain_core.prompts import ChatPromptTemplate
 
+from langflow.base.prompts.utils import dict_values_to_string
 from langflow.custom import CustomComponent
 from langflow.field_typing import Prompt, TemplateField, Text
+from langflow.schema.schema import Record
 
 
 class PromptComponent(CustomComponent):
@@ -15,19 +17,14 @@ class PromptComponent(CustomComponent):
             "code": TemplateField(advanced=True),
         }
 
-    def build(
+    async def build(
         self,
         template: Prompt,
         **kwargs,
-    ) -> Text:
-        from langflow.base.prompts.utils import dict_values_to_string
-
-        prompt_template = PromptTemplate.from_template(Text(template))
-        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:\n"{formated_prompt}"'
-        return formated_prompt
+    ) -> Record:
+        prompt_template = ChatPromptTemplate.from_template(Text(template))
+        kwargs = await dict_values_to_string(kwargs)
+        messages = list(kwargs.values())
+        prompt = prompt_template + messages
+        self.status = f'Prompt:\n"{template}"'
+        return Record(data={"prompt": prompt.to_json()})