Refactoring jinja code

app.py

@@ -1,60 +1,21 @@
 from time import perf_counter
 
 import gradio as gr
-from jinja2 import Template
+from jinja2 import Environment, FileSystemLoader
 
 from backend.query_llm import generate
 from backend.semantic_search import qd_retriever
+from pathlib import Path
 
-template_string = """
-Instructions: Use the following pieces of context to answer the question at the end. If you don't know the answer, just say that you don't know, don't try to make up an answer.
-Context:
----
-{% for doc in documents %}
-    {{ doc.content }}
-    ---
-{% endfor %}
-Query: {{ query }}
-"""
-
-md_template_string = """
-<b>Instructions</b>: 
-<span style="color: green;">Use the following pieces of context to answer the question at the end.<br>If you don't know the answer, just say that you don't know, <span style="color: green; font-weight: bold;">don't try to make up an answer.</span></span><br>
-
-<b>Context</b>:
-{% for doc in documents %}
-    <div id=\"box{{ loop.index }}\" style=\"border: 2px solid #aaa; padding: 10px; margin-top: 10px; border-radius: 5px; background-color: #1E90FF; position: relative; cursor: pointer;\">
-        <div style=\"font-size: 0.8em; position: absolute; top: 10px; left: 10px;\"><b>Doc {{ loop.index }}</b></div>
-        <span id=\"doc{{ loop.index }}-short\" style=\"color: white; display: block; margin-top: 20px;\">{{ doc.content[:50] }}...</span>
-        <span id=\"doc{{ loop.index }}-full\" style=\"color: white; display: none; margin-top: 20px;\">{{ doc.content }}</span>
-    </div>
-{% endfor %}
-<b>Query</b>: <span style=\"color: yellow;\">{{ query }}</span>
-<script>
-document.addEventListener("DOMContentLoaded", function() {
-    {% for doc in documents %}
-        document.getElementById("box{{ loop.index }}").addEventListener("click", function() {
-            toggleContent('doc{{ loop.index }}');
-        });
-    {% endfor %}
-});
-
-function toggleContent(docID) {
-    var shortContent = document.getElementById(docID + '-short');
-    var fullContent = document.getElementById(docID + '-full');
-    if (fullContent.style.display === 'none') {
-        shortContent.style.display = 'none';
-        fullContent.style.display = 'block';
-    } else {
-        shortContent.style.display = 'block';
-        fullContent.style.display = 'none';
-    }
-}
-</script>
-"""
-
-template = Template(template_string)
-md_template = Template(md_template_string)
+proj_dir = Path(__file__).parent
+
+
+# Set up the template environment with the templates directory
+env = Environment(loader=FileSystemLoader(proj_dir/'templates'))
+
+# Load the templates directly from the environment
+template = env.get_template('template.j2')
+template_html = env.get_template('template_html.j2')
 
 
 def add_text(history, text):
@@ -74,12 +35,12 @@ def bot(history, system_prompt=""):
 
     # Create Prompt
     prompt = template.render(documents=documents, query=query)
-    md_prompt = md_template.render(documents=documents, query=query)
+    prompt_html = template_html.render(documents=documents, query=query)
 
     history[-1][1] = ""
     for character in generate(prompt, history[:-1]):
         history[-1][1] = character
-        yield history, md_prompt
+        yield history, prompt_html
 
 
 with gr.Blocks() as demo:
@@ -102,17 +63,17 @@ with gr.Blocks() as demo:
                     )
             txt_btn = gr.Button(value="Submit text", scale=1)
 
-        prompt_md = gr.HTML()
+        prompt_html = gr.HTML()
         # Turn off interactivity while generating if you hit enter
         txt_msg = txt_btn.click(add_text, [chatbot, txt], [chatbot, txt], queue=False).then(
-                bot, chatbot, [chatbot, prompt_md])
+                bot, chatbot, [chatbot, prompt_html])
 
         # Turn it back on
         txt_msg.then(lambda: gr.Textbox(interactive=True), None, [txt], queue=False)
 
         # Turn off interactivity while generating if you hit enter
         txt_msg = txt.submit(add_text, [chatbot, txt], [chatbot, txt], queue=False).then(
-                bot, chatbot, [chatbot, prompt_md])
+                bot, chatbot, [chatbot, prompt_html])
 
         # Turn it back on
         txt_msg.then(lambda: gr.Textbox(interactive=True), None, [txt], queue=False)

templates/template.j2

@@ -0,0 +1,8 @@
+Instructions: Use the following pieces of context to answer the question at the end. If you don't know the answer, just say that you don't know, don't try to make up an answer.
+Context:
+---
+{% for doc in documents %}
+    {{ doc.content }}
+    ---
+{% endfor %}
+Query: {{ query }}

templates/template_html.j2

@@ -0,0 +1,60 @@
+<b>Instructions</b>:
+<span style="color: green;">Use the following pieces of context to answer the question at the end.<br>If you don't know the answer, just say that you don't know, <span style="color: green; font-weight: bold;">don't try to make up an answer.</span></span><br>
+
+<b>Context</b>:
+{% for doc in documents %}
+    <div id="box{{ loop.index }}" class="doc-box">
+        <div class="doc-label"><b>Doc {{ loop.index }}</b></div>
+        <span id="doc{{ loop.index }}-short" class="doc-short">{{ doc.content[:50] }}...</span>
+        <span id="doc{{ loop.index }}-full" class="doc-full">{{ doc.content }}</span>
+    </div>
+{% endfor %}
+<b>Query</b>: <span style="color: yellow;">{{ query }}</span>
+
+<style>
+    .doc-box {
+        border: 2px solid #aaa;
+        padding: 10px;
+        margin-top: 10px;
+        border-radius: 5px;
+        background-color: #1E90FF;
+        position: relative;
+        cursor: pointer;
+    }
+    .doc-label {
+        font-size: 0.8em;
+        position: absolute;
+        top: 10px;
+        left: 10px;
+    }
+    .doc-short, .doc-full {
+        color: white;
+        display: block;
+        margin-top: 20px;
+    }
+    .doc-full {
+        display: none;
+    }
+</style>
+
+<script>
+document.addEventListener("DOMContentLoaded", function() {
+    {% for doc in documents %}
+        document.getElementById("box{{ loop.index }}").addEventListener("click", function() {
+            toggleContent('doc{{ loop.index }}');
+        });
+    {% endfor %}
+});
+
+function toggleContent(docID) {
+    var shortContent = document.getElementById(docID + '-short');
+    var fullContent = document.getElementById(docID + '-full');
+    if (fullContent.style.display === 'none') {
+        shortContent.style.display = 'none';
+        fullContent.style.display = 'block';
+    } else {
+        shortContent.style.display = 'block';
+        fullContent.style.display = 'none';
+    }
+}
+</script>