Improve passing context (#619)

Co-authored-by: Philipp Rudiger <[email protected]>

lumen/ai/assistant.py

@@ -8,7 +8,7 @@
 import param
 
 from panel import bind
-from panel.chat import ChatInterface
+from panel.chat import ChatInterface, ChatStep
 from panel.layout import Column, FlexBox, Tabs
 from panel.pane import HTML, Markdown
 from panel.viewable import Viewer
@@ -338,7 +338,6 @@ async def _get_agent(self, messages: list | str):
                 step.stream(output.chain_of_thought, replace=True)
                 agent = output.agent
                 step.success_title = f"Selected {agent}"
-            messages.append({"role": "assistant", "content": output.chain_of_thought})
 
         if agent is None:
             return None
@@ -391,6 +390,9 @@ async def _get_agent(self, messages: list | str):
     def _serialize(self, obj):
         if isinstance(obj, (Tabs, Column)):
             for o in obj:
+                if isinstance(obj, ChatStep) and not obj.title.startswith("Selected"):
+                    # only want the chain of thoughts from the selected agent
+                    continue
                 if hasattr(o, "visible") and o.visible:
                     break
             return self._serialize(o)
@@ -422,7 +424,8 @@ async def invoke(self, messages: list | str) -> str:
             print(f"{message['role']!r}: {message['content']}")
             print("ENTRY" + "-" * 10)
 
-        await agent.invoke(messages[-2:])
+        print("\n\033[95mAGENT:\033[0m", agent, messages[-3:])
+        await agent.invoke(messages[-3:])
         self._current_agent.object = "## No agent active"
 
         if "current_pipeline" in agent.provides:

lumen/ai/views.py

@@ -121,6 +121,9 @@ def __panel__(self):
     def __repr__(self):
         return self.spec
 
+    def __str__(self):
+        return self.spec
+
 
 class SQLOutput(LumenOutput):