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}"