From 894610cd1606c840b3ef2e5e8964f204e079efb3 Mon Sep 17 00:00:00 2001
From: Daiyi Peng <daiyip@google.com>
Date: Wed, 18 Sep 2024 09:14:30 -0700
Subject: [PATCH] Add `lf.Template.raw_str` and `lf.Template.from_raw_str`.

This makes it easier for users to include raw strings in the prompt.

PiperOrigin-RevId: 676015977
---
 langfun/core/template.py      | 10 ++++++++++
 langfun/core/template_test.py | 15 +++++++++++++++
 2 files changed, 25 insertions(+)

diff --git a/langfun/core/template.py b/langfun/core/template.py
index a88a830..c3dd426 100644
--- a/langfun/core/template.py
+++ b/langfun/core/template.py
@@ -231,6 +231,16 @@ def missing_vars(self) -> Set[str]:
     """Returns the missing variable names."""
     return self.vars(closure=True, specified=False)
 
+  @classmethod
+  def raw_str(cls, text: str) -> str:
+    """Returns a template string that preserve the text as original."""
+    return '{% raw %}' + text + '{% endraw %}'
+
+  @classmethod
+  def from_raw_str(cls, text: str) -> 'Template':
+    """Returns a template that preserve the text as original."""
+    return cls(cls.raw_str(text), clean=False)
+
   def render(
       self,
       *,
diff --git a/langfun/core/template_test.py b/langfun/core/template_test.py
index 93d3f86..30082a2 100644
--- a/langfun/core/template_test.py
+++ b/langfun/core/template_test.py
@@ -98,6 +98,21 @@ class Foo(component.Component):
     self.assertEqual(d.z.render(), 'Bye, 1')
     self.assertEqual(d.p.render(), 'Again Hello, 1')
 
+  def test_raw_text(self):
+    self.assertEqual(
+        Template(
+            '{{a}}' + Template.raw_str('\n{{d}}, {%x%}\n') + '{{b}}',
+            a='hi', b=1
+        ).render().text,
+        'hi\n{{d}}, {%x%}\n1'
+    )
+
+  def test_from_raw_str(self):
+    self.assertEqual(
+        Template.from_raw_str('\n{{d}}, {%x%}\n').render().text,
+        '\n{{d}}, {%x%}\n'
+    )
+
 
 class DefinitionTest(unittest.TestCase):