Replace Archytas with LangGraph

README.md

@@ -3,7 +3,7 @@ Beaker is a next generation coding notebook built for the AI era. Beaker seamles
 
 Beyond that, Beaker solves one of the major challenges presented by coding notebooks--it introduces a true _undo_ mechanism so that the user can roll back to any previous state in the notebook. Beaker also lets you swap effortlessly between a notebook style coding interface and a chat style interface, giving you the best of both worlds. Since everything is interoperable with Jupyter, you can always export your notebook and use it in any other Jupyter-compatible environment.
 
-Beaker is powered by [Archytas](https://github.com/jataware/archytas), our framework for building AI agents that can interact with code and advanced users can generate their own custom agents to meet their specific needs. These agents can have custom ReAct toolsets built in and can be extended to support any number of use cases.
+Beaker is powered by [LangGraph](https://github.com/langchain-ai/langgraph), the modern framework for building AI agents that can interact with code. Advanced users can generate their own custom agents to meet their specific needs. These agents can have custom ReAct toolsets built in and can be extended to support any number of use cases.
 
 We like to think of Beaker as a (much better!) drop in replacement for workflows where you'd normally rely on Jupyter notebooks and we hope you'll give it a try and let us know what you think!
 

beaker_kernel/contexts/default/agent.py

@@ -1,34 +1,26 @@
-import json
-import logging
-import re
+"""
+Default agent implementation for Beaker.
 
-from archytas.tool_utils import AgentRef, LoopControllerRef, tool
+Provides a default agent with basic functionality including a joke tool.
+"""
 
+import logging
 from beaker_kernel.lib.agent import BeakerAgent
 from beaker_kernel.lib.context import BeakerContext
-from beaker_kernel.lib.utils import ExecutionError
-
+from beaker_kernel.lib.utils import ExecutionError, get_beaker_kernel
+from beaker_kernel.lib.tools import tool
 
 logger = logging.getLogger(__name__)
 
 
-class DefaultAgent(BeakerAgent):
-    """
-    You are an programming assistant to aid a developer working in a Jupyter notebook by answering their questions and helping write code for them based on their
-    prompt.
-    """
+@tool
+async def tell_a_joke(topic: str = "any") -> str:
+    """Generates a joke for the user.
 
-    @tool()
-    async def tell_a_joke(self, topic: str, agent: AgentRef, loop: LoopControllerRef) -> str:
-        """
-        Generates a joke for the user.
-
-        Args:
-            topic (str): A topic that the joke should be about. If no topic is provided, use the default value of "any".
-        Returns:
-            str: The text of the joke, possibly with or without formatting.
-        """
-        code = """
+    Args:
+        topic: A topic that the joke should be about. If no topic is provided, use the default value of 'any'.
+    """
+    code = """
 import requests
 url = 'https://v2.jokeapi.dev/joke/Programming,Miscellaneous,Pun,Spooky?blacklistFlags=nsfw,religious,political,racist,sexist,explicit'
 result = requests.get(url).json()
@@ -44,11 +36,32 @@ async def tell_a_joke(self, topic: str, agent: AgentRef, loop: LoopControllerRef
 formatted_joke
 """.strip()
 
-        try:
-            context = await agent.context.evaluate(code)
+    try:
+        # Get the kernel to access the context
+        kernel = get_beaker_kernel()
+        if kernel and hasattr(kernel, 'context'):
+            context = await kernel.context.evaluate(code)
             joke = context["return"]
-        except ExecutionError as e:
-            logger.warning(f"Error fetching joke: {e}")
-            joke = """Have you ever seen an elephant hiding in a tree? Me neither. They must be VERY good at it!"""
+        else:
+            raise Exception("No kernel context available")
+    except Exception as e:
+        logger.warning(f"Error fetching joke: {e}")
+        joke = """Have you ever seen an elephant hiding in a tree? Me neither. They must be VERY good at it!"""
+
+    return joke
+
+
+class DefaultAgent(BeakerAgent):
+    """
+    Default agent for Beaker contexts.
+
+    Provides basic functionality with a joke tool as an example.
+    """
 
-        return joke
+    def __init__(self, context: BeakerContext = None, tools=None, **kwargs):
+        # Add our default tools
+        default_tools = [tell_a_joke]
+        all_tools = default_tools + (tools or [])
+
+        # Initialize the parent agent
+        super().__init__(context=context, tools=all_tools, **kwargs)

beaker_kernel/kernel.py

@@ -273,39 +273,18 @@ async def send_kernel_state_info(self, parent_header=None):
                     self.send_response("iopub", "kernel_state_info", state_payload, parent_header=parent_header)
 
     async def send_chat_history(self, parent_header=None):
-        if self.context.agent.chat_history:
-            from archytas.chat_history import OutboundChatHistory, OutboundModel, SummaryRecord, MessageRecord
-            from dataclasses import asdict
-            chat_history = self.context.agent.chat_history
-            model = self.context.agent.model
-            records = await chat_history.records(auto_update_context=False)
-            output = OutboundChatHistory(
-                records=[{
-                    "message": {
-                        "text": record.message.text(),
-                        "raw_content": record.message.content,
-                        **record.message.model_dump(),
-                    },
-                    "uuid": record.uuid,
-                    "token_count": record.token_count,
-                    "metadata": record.metadata,
-                    "react_loop_id": record.react_loop_id,
-                } for record in records],
-                system_message=chat_history.system_message.message.text(),
-                tool_token_usage_estimate=chat_history.tool_token_estimate,
-                model=OutboundModel(
-                    provider=model.__class__.__name__,
-                    model_name=model.model_name,
-                    context_window=model.contextsize()
-                ),
-                message_token_count=sum(record.token_count for record in records if isinstance(record, MessageRecord) and record.token_count),
-                summary_token_count=sum(record.token_count for record in records if isinstance(record, SummaryRecord) and record.token_count),
-                overhead_token_count=chat_history.token_overhead,
-                summarization_threshold=model.summarization_threshold,
-                token_estimate=chat_history._token_estimate or await chat_history.token_estimate(model),
-            )
-            output_dict = asdict(output)
-            self.send_response("iopub", "chat_history", output_dict, parent_header=parent_header)
+        # All agents now use BeakerChatHistory with to_outbound_format method
+        if self.context and self.context.agent and self.context.agent.chat_history:
+            try:
+                from dataclasses import asdict
+                model = getattr(self.context.agent, 'model', None)
+                outbound_history = self.context.agent.chat_history.to_outbound_format(model)
+                output_dict = asdict(outbound_history)
+
+                self.send_response("iopub", "chat_history", output_dict, parent_header=parent_header)
+                logger.debug(f"Sent chat history: {len(outbound_history.records)} records, {outbound_history.total_token_count} tokens")
+            except Exception as e:
+                logger.error(f"Error sending chat history: {e}")
 
     async def update_connection_file(self, **kwargs):
         try:
@@ -524,7 +503,7 @@ def stderr(self, text, parent_header=None):
 
     @message_handler
     async def llm_request(self, message: JupyterMessage):
-        from archytas.exceptions import AuthenticationError
+        # Handle LLM requests from frontend
         content: dict = message.content
         request = content.get("request", None)
         if not request:
@@ -548,7 +527,7 @@ async def llm_request(self, message: JupyterMessage):
                 task = asyncio.create_task(self.context.agent.react_async(request, react_context={"message": message}))
                 self.running_actions[request_key] = task
                 result = await task
-            except AuthenticationError as err:
+            except ValueError as err:  # Updated from AuthenticationError
                 self.send_response(
                     stream="iopub",
                     msg_or_type="llm_auth_failure",