diff --git a/lumen/ai/export.py b/lumen/ai/export.py index d677e1ca5..534713bed 100644 --- a/lumen/ai/export.py +++ b/lumen/ai/export.py @@ -42,7 +42,9 @@ def format_output(msg: ChatMessage): output = msg.object code = [] with config.param.update(serializer='csv'): - spec = json.dumps(output.component.to_spec(), indent=2).replace('true', 'True').replace('false', 'False') + spec = json.dumps( + output.component.to_spec(), indent=2, ensure_ascii=False + ).replace('true', 'True').replace('false', 'False').replace('null', 'None').replace('\\\\', '\\') if isinstance(output.component, Pipeline): code.extend([ f'pipeline = lm.Pipeline.from_spec({spec})', diff --git a/lumen/ai/tools.py b/lumen/ai/tools.py index ee4349e86..89399348e 100644 --- a/lumen/ai/tools.py +++ b/lumen/ai/tools.py @@ -62,8 +62,9 @@ class DocumentLookup(VectorLookupTool): def __init__(self, **params): super().__init__(**params) - self._memory.on_change('document_sources', self._update_vector_store) - self._update_vector_store(None, None, self._memory.get("document_sources", [])) + self._memory.on_change("document_sources", self._update_vector_store) + if "document_sources" in self._memory: + self._update_vector_store(None, None, self._memory["document_sources"]) def _update_vector_store(self, _, __, sources): for source in sources: diff --git a/lumen/ai/views.py b/lumen/ai/views.py index d862188d8..65b8cebbb 100644 --- a/lumen/ai/views.py +++ b/lumen/ai/views.py @@ -93,7 +93,7 @@ def __init__(self, **params): code_col = Column(code_editor, icons, sizing_mode="stretch_both") if self.render_output: placeholder = Column( - ParamMethod(self.render, inplace=True), + ParamMethod(self.render, inplace=True, lazy=True), sizing_mode="stretch_width" ) self._main = Tabs( diff --git a/lumen/tests/ai/test_export.py b/lumen/tests/ai/test_export.py new file mode 100644 index 000000000..3673b0298 --- /dev/null +++ b/lumen/tests/ai/test_export.py @@ -0,0 +1,259 @@ +import datetime as dt +import json +import os +import pathlib + +import pytest + +from panel.chat import ChatInterface, ChatMessage + +import lumen + +from lumen.pipeline import Pipeline +from lumen.sources.intake import IntakeSource +from lumen.views import Table + +try: + from lumen.ai.export import ( + LumenOutput, export_notebook, format_markdown, format_output, + make_preamble, + ) +except ImportError: + pytest.skip("Skipping tests that require lumen.ai", allow_module_level=True) + + +TEST_DIR = pathlib.Path(lumen.__file__).parent / "tests" +CATALOG_URI = str(TEST_DIR / 'sources' / 'catalog.yml') + +@pytest.fixture +def source(): + return IntakeSource(uri=CATALOG_URI, root=TEST_DIR) + + +@pytest.fixture +def source_tables(mixed_df): + df_test = mixed_df.copy() + df_test_sql = mixed_df.copy() + df_test_sql_none = mixed_df.copy() + df_test_sql_none["C"] = ["foo1", None, "foo3", None, "foo5"] + tables = { + "test": df_test, + "test_sql": df_test_sql, + "test_sql_with_none": df_test_sql_none, + } + return tables + + +@pytest.fixture +def fixed_datetime(monkeypatch): + fixed_time = dt.datetime(2024, 4, 27, 12, 0, 0) + + class FixedDateTime: + @classmethod + def now(cls): + return fixed_time + + monkeypatch.setattr(dt, "datetime", FixedDateTime) + return fixed_time + + +def test_make_preamble(fixed_datetime): + preamble = "# Initial preamble content" + cells = make_preamble(preamble) + + assert len(cells) == 2 + + # Check the header cell + header_cell = cells[0] + assert header_cell.cell_type == "markdown" + expected_header = f"# Lumen.ai - Chat Logs {fixed_datetime}" + assert header_cell.source == expected_header + + # Check the imports cell + imports_cell = cells[1] + assert imports_cell.cell_type == "code" + expected_source = ( + "# Initial preamble content\nimport lumen as lm\n" + "import panel as pn\n\npn.extension('tabulator')" + ) + assert imports_cell.source == expected_source + + +def test_format_markdown_with_https_avatar(): + msg = ChatMessage( + object="This is a test message.", + user="User", + avatar="https://example.com/avatar.png", + ) + cell = format_markdown(msg)[0] + + assert cell.cell_type == "markdown" + expected_avatar = ( + '' + ) + expected_header = ( + f'
{expected_avatar}User
\n' + ) + expected_content = "\nThis is a test message." + expected_source = f"{expected_header}{expected_content}" + assert cell.source == expected_source + + +def test_format_markdown_with_text_avatar(): + msg = ChatMessage( + object="Bot response here.", + user="Bot", + avatar="B", + ) + cell = format_markdown(msg)[0] + + assert cell.cell_type == "markdown" + expected_avatar = "B" + expected_header = ( + f'
{expected_avatar}Bot
\n' + ) + expected_content = "\n> Bot response here." + expected_source = f"{expected_header}{expected_content}" + assert cell.source == expected_source + + +def test_format_output_pipeline(source): + pipeline = Pipeline(source=source, table="test") + msg = ChatMessage(object=LumenOutput(component=pipeline), user="User") + + cells = format_output(msg) + cells[0].pop("id") + uri = source.uri + if os.name == "nt": + uri = uri.replace("\\", "\\\\") + expected_source = ( + "pipeline = lm.Pipeline.from_spec({\n" + ' "source": {\n' + ' "uri": "'+CATALOG_URI.replace('\\\\', '\\')+'",\n' + ' "type": "intake"\n' + " },\n" + ' "table": "test"\n' + "})\n" + "pipeline" + ) + + assert cells == [ + { + "cell_type": "code", + "metadata": {}, + "execution_count": None, + "source": expected_source, + "outputs": [], + } + ] + + +def test_format_output_view(source): + pipeline = Pipeline(source=source, table="test") + table = Table(pipeline=pipeline) + msg = ChatMessage(object=LumenOutput(component=table), user="User") + + cells = format_output(msg) + cells[0].pop("id") + uri = source.uri + if os.name == "nt": + uri = uri.replace("\\", "\\\\") + assert cells == [ + { + "cell_type": "code", + "execution_count": None, + "metadata": {}, + "outputs": [], + "source": "view = lm.View.from_spec({\n" + ' "pipeline": {\n' + ' "source": {\n' + f' "uri": "{uri}",\n' + ' "type": "intake"\n' + " },\n" + ' "table": "test"\n' + " },\n" + ' "type": "table"\n' + "})\n" + "view", + }, + ] + + +def test_export_notebook(source): + interface = ChatInterface( + ChatMessage( + object="", + user="Help", + ), + ChatMessage( + object=LumenOutput(component=Pipeline(source=source, table="test")), + user="User", + ), + ChatMessage( + object=LumenOutput( + component=Table(pipeline=Pipeline(source=source, table="test")) + ), + user="Bot", + ), + ) + + cells = json.loads(export_notebook(interface)) + cells["cells"][0].pop("source") + cells["cells"][0].pop("id") + cells["cells"][1].pop("id") + cells["cells"][2].pop("id") + cells["cells"][3].pop("id") + assert cells == { + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + }, + { + "cell_type": "code", + "execution_count": None, + "metadata": {}, + "outputs": [], + "source": [ + "import lumen as lm\n", + "import panel as pn\n", + "\n", + "pn.extension('tabulator')", + ], + }, + { + 'cell_type': 'code', + 'execution_count': None, + 'metadata': {}, + 'outputs': [], + 'source': [ + 'pipeline = lm.Pipeline.from_spec({\n', + ' "source": {\n', ' "uri": "'+CATALOG_URI.replace('\\\\', '\\')+'",\n', + ' "type": "intake"\n', ' },\n', + ' "table": "test"\n', '})\n', + 'pipeline' + ] + }, + { + 'cell_type': 'code', + 'execution_count': None, + 'metadata': {}, + 'outputs': [], + 'source': [ + 'view = lm.View.from_spec({\n', + ' "pipeline": {\n', ' "source": {\n', + ' "uri": "'+CATALOG_URI.replace('\\\\', '\\')+'",\n', + ' "type": "intake"\n', ' },\n', + ' "table": "test"\n', ' },\n', + ' "type": "table"\n', + '})\n', + 'view' + ] + } + ], + "metadata": {}, + "nbformat": 4, + "nbformat_minor": 5, + }