diff --git a/libs/ktem/ktem/reasoning/simple.py b/libs/ktem/ktem/reasoning/simple.py
index 3397250de..29f48fed0 100644
--- a/libs/ktem/ktem/reasoning/simple.py
+++ b/libs/ktem/ktem/reasoning/simple.py
@@ -8,6 +8,7 @@
import tiktoken
from ktem.llms.manager import llms
+from ktem.utils.render import Render
from kotaemon.base import (
BaseComponent,
@@ -559,8 +560,6 @@ async def ainvoke( # type: ignore
def stream( # type: ignore
self, message: str, conv_id: str, history: list, **kwargs # type: ignore
) -> Generator[Document, None, Document]:
- import markdown
-
docs = []
doc_ids = []
if self.use_rewrite:
@@ -571,19 +570,15 @@ def stream( # type: ignore
if doc.doc_id not in doc_ids:
docs.append(doc)
doc_ids.append(doc.doc_id)
+
for doc in docs:
- # TODO: a better approach to show the information
- text = markdown.markdown(
- doc.text, extensions=["markdown.extensions.tables"]
- )
yield Document(
- content=(
- ""
- f"{doc.metadata['file_name']}
"
- f"{text}"
- "
"
- ),
channel="info",
+ content=Render.collapsible(
+ header=doc.metadata["file_name"],
+ content=Render.table(doc.text),
+ open=True,
+ ),
)
evidence_mode, evidence = self.evidence_pipeline(docs).content
@@ -638,23 +633,17 @@ def stream( # type: ignore
ss = sorted(ss, key=lambda x: x["start"])
text = id2docs[id].text[: ss[0]["start"]]
for idx, span in enumerate(ss):
- text += (
- "" + id2docs[id].text[span["start"] : span["end"]] + ""
- )
+ text += Render.highlight(id2docs[id].text[span["start"] : span["end"]])
if idx < len(ss) - 1:
text += id2docs[id].text[span["end"] : ss[idx + 1]["start"]]
text += id2docs[id].text[ss[-1]["end"] :]
- text_out = markdown.markdown(
- text, extensions=["markdown.extensions.tables"]
- )
yield Document(
- content=(
- ""
- f"{id2docs[id].metadata['file_name']}
"
- f"{text_out}"
- "
"
- ),
channel="info",
+ content=Render.collapsible(
+ header=id2docs[id].metadata["file_name"],
+ content=Render.table(text),
+ open=True,
+ ),
)
lack_evidence = False
@@ -667,17 +656,13 @@ def stream( # type: ignore
content="Retrieved segments without matching evidence:\n",
)
for id in list(not_detected):
- text_out = markdown.markdown(
- id2docs[id].text, extensions=["markdown.extensions.tables"]
- )
yield Document(
- content=(
- ""
- f"{id2docs[id].metadata['file_name']}
"
- f"{text_out}"
- "
"
- ),
channel="info",
+ content=Render.collapsible(
+ header=id2docs[id].metadata["file_name"],
+ content=Render.table(id2docs[id].text),
+ open=False,
+ ),
)
return answer
diff --git a/libs/ktem/ktem/utils/__init__.py b/libs/ktem/ktem/utils/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/libs/ktem/ktem/utils/render.py b/libs/ktem/ktem/utils/render.py
new file mode 100644
index 000000000..5890d3327
--- /dev/null
+++ b/libs/ktem/ktem/utils/render.py
@@ -0,0 +1,21 @@
+import markdown
+
+
+class Render:
+ """Default text rendering into HTML for the UI"""
+
+ @staticmethod
+ def collapsible(header, content, open: bool = False) -> str:
+ """Render an HTML friendly collapsible section"""
+ o = " open" if open else ""
+ return f"{header}
{content}
"
+
+ @staticmethod
+ def table(text: str) -> str:
+ """Render table from markdown format into HTML"""
+ return markdown.markdown(text, extensions=["markdown.extensions.tables"])
+
+ @staticmethod
+ def highlight(text: str) -> str:
+ """Highlight text"""
+ return f"{text}"