Merge pull request #12 from Zsailer/agent-api

Simplify server extension to make agent development lightweight

Author: Zsailer

jupyterlab_magic_wand/agents/base.py

@@ -1,10 +1,31 @@
-from dataclasses import dataclass
-from langgraph.graph.graph import CompiledGraph
+import urllib.parse
+from typing import Type
+from pydantic import BaseModel, ConfigDict
+from langgraph.graph.state import CompiledStateGraph
 
 
-@dataclass
-class Agent:
+class Agent(BaseModel):
+    model_config = ConfigDict(arbitrary_types_allowed=True)
     name: str
     description: str
-    workflow: CompiledGraph
-    version: str
+    workflow: CompiledStateGraph
+    version: str
+    state: Type[BaseModel]
+    
+    @property
+    def state_schema_id(self):
+        name = urllib.parse.quote_plus(self.name)
+        return f"https://events.jupyter.org/jupyter-ai/agents/{name}/state"
+    
+    @property
+    def state_schema(self):
+        event_schema = {
+            "$id": self.state_schema_id,
+            "version": self.version,
+            "title": "",
+            "description": "",
+            "personal-data": True,
+            "type": "object",
+        }
+        event_schema.update(self.state.model_json_schema())
+        return event_schema

jupyterlab_magic_wand/agents/magic_agent.py

@@ -7,30 +7,50 @@
 """
 import json
 import uuid
-from typing import Sequence, Union
+from typing import Sequence, Union, List, Optional
 from langgraph.graph import StateGraph
 from langgraph.graph import END, START
 from langchain_core.runnables import RunnableConfig
 from langchain_core.messages import BaseMessage
-from jupyterlab_magic_wand.state import AIWorkflowState, ConfigSchema
-from jupyterlab_magic_wand.agents.lab_commands import (
+from jupyterlab_magic_wand.agents.tools import (
     update_cell_source, 
     show_diff, 
     insert_cell_below
 )
-from jupyterlab_magic_wand.agents.base import Agent
+from .base import Agent
 
 from langchain_core.messages import HumanMessage
+from pydantic import BaseModel
 
-graph = StateGraph(AIWorkflowState, config_schema=ConfigSchema)
+
+class LabCommand(BaseModel):
+    name: str
+    args: dict
+
+
+class Context(BaseModel):
+    cell_id: str
+    content: dict
+
+
+class State(BaseModel):
+    agent: Optional[str] = None
+    input: str
+    context: Context
+    messages: list = []
+    commands: Optional[List[LabCommand]] = None
+
+    
+graph = StateGraph(State)
 
 
 def get_jupyter_ai_model(jupyter_ai_config):
     lm_provider = jupyter_ai_config.lm_provider
     return lm_provider(**jupyter_ai_config.lm_provider_params)
 
-def get_cell(cell_id: str, state: AIWorkflowState) -> dict:
-    content = state["context"]["content"]
+
+def get_cell(cell_id: str, state: State) -> dict:
+    content = state.context.content
     cells = content["cells"]
 
     for cell in cells:
@@ -60,8 +80,8 @@ def sanitize_code(code: str) -> str:
     )
 
 
-async def router(state: AIWorkflowState) -> Sequence[str]:
-    cell_id = state["context"]["cell_id"]
+async def router(state: State) -> Sequence[str]:
+    cell_id = state.context.cell_id
     current = get_cell(cell_id, state)
     if current.get("cell_type") == "markdown":
         return ["route_markdown"]
@@ -103,29 +123,30 @@ def _cast_ai_response(response: Union[str, BaseMessage]):
         raise Exception("The response type must be 'str' or 'BaseMessage'.")
     return response
 
-async def route_markdown(state: AIWorkflowState, config: RunnableConfig) -> dict:
-    llm = get_jupyter_ai_model(config["configurable"]["jupyter_ai_config"])
-    cell_id = state["context"]["cell_id"]
+
+async def route_markdown(state: State, config: RunnableConfig) -> dict:
+    llm = get_jupyter_ai_model(config["configurable"]["jai_config_manager"])
+    cell_id = state.context.cell_id
     current = get_cell(cell_id, state)
     # Spell check
     if current["source"].strip() != "":
         response = _cast_ai_response(await llm.ainvoke(input=f"Does the following input look like a prompt to write code (answer 'code' only) or content to be editted (answer 'content' only)?\n Input: {current['source']}"))
         if "code" in response.lower():
             response = _cast_ai_response(await llm.ainvoke(input=f"Write code based on the prompt. Then, update the code to make it more efficient, add code comments, and respond with only the code and comments.\n Input: {current['source']}"))
             response = sanitize_code(response)
-            messages = state.get("messages", []) or []
+            messages = state.messages
             messages.append(response)
-            commands = state["commands"]
+            commands = state.commands
             new_cell_id = str(uuid.uuid4())
             commands.extend([
                 insert_cell_below(cell_id, source=response, type="code", new_cell_id=new_cell_id),
             ])
             return {"commands": commands, "messages": messages}
         prompt = SPELLCHECK_MARKDOWN.format(input=current["source"])
         response = _cast_ai_response(await llm.ainvoke(input=prompt))
-        messages = state.get("messages", []) or []
+        messages = state.messages
         messages.append(response)
-        commands = state["commands"]
+        commands = state.commands
         commands.extend([
             update_cell_source(cell_id, source=response),
             {
@@ -135,8 +156,8 @@ async def route_markdown(state: AIWorkflowState, config: RunnableConfig) -> dict
         ])
         return {"commands": commands, "messages": messages}
 
-    content = state["context"]["content"]
-    cell_id = state["context"]["cell_id"]
+    content = state.context.content
+    cell_id = state.context.cell_id
     cells = content["cells"]
 
 
@@ -151,7 +172,7 @@ async def route_markdown(state: AIWorkflowState, config: RunnableConfig) -> dict
             response = _cast_ai_response(await llm.ainvoke(input=prompt))
             messages = state.get("messages", []) or []
             messages.append(response)
-            commands = state["commands"]
+            commands = state.commands
             commands.extend([
                 update_cell_source(cell_id, source=response),
                 {
@@ -177,9 +198,9 @@ async def route_markdown(state: AIWorkflowState, config: RunnableConfig) -> dict
 {exception_value}
 """
 
-async def route_exception(state: AIWorkflowState, config: RunnableConfig) -> dict:
-    llm = get_jupyter_ai_model(config["configurable"]["jupyter_ai_config"])
-    cell_id = state["context"]["cell_id"]
+async def route_exception(state: State, config: RunnableConfig) -> dict:
+    llm = get_jupyter_ai_model(config["configurable"]["jai_config_manager"])
+    cell_id = state.context.cell_id
     current = get_cell(cell_id, state)
     exception = get_exception(current)
     prompt = exception_prompt.format(
@@ -189,9 +210,9 @@ async def route_exception(state: AIWorkflowState, config: RunnableConfig) -> dic
     )
     response = _cast_ai_response(await llm.ainvoke(input=prompt))
     response = sanitize_code(response)
-    messages = state.get("messages", []) or []
+    messages = state.messages
     messages.append(response)
-    commands = state["commands"]
+    commands = state.commands
     commands.extend([
         update_cell_source(cell_id, source=response),
         show_diff(cell_id, current["source"], response),
@@ -215,7 +236,7 @@ async def route_exception(state: AIWorkflowState, config: RunnableConfig) -> dic
 """
 
 def prompt_new_cell_using_context(cell_id, state):
-    content = state["context"]["content"]
+    content = state.context.content
     cells = content["cells"]
 
     for i, cell in enumerate(cells):
@@ -236,20 +257,20 @@ def prompt_new_cell_using_context(cell_id, state):
     return prompt
 
 
-async def route_code(state: AIWorkflowState, config: RunnableConfig):
-    llm = get_jupyter_ai_model(config["configurable"]["jupyter_ai_config"])
+async def route_code(state: State, config: RunnableConfig):
+    llm = get_jupyter_ai_model(config["configurable"]["jai_config_manager"])
 
-    cell_id = state["context"]["cell_id"]
+    cell_id = state.context.cell_id
     current = get_cell(cell_id, state)
     source = current["source"]
     source = source.strip()
     if source:
         prompt = IMPROVE_PROMPT.format(code=source)
         response = _cast_ai_response(await llm.ainvoke(prompt, stream=False))
         response = sanitize_code(response)
-        messages = state.get("messages", []) or []
+        messages = state.messages
         messages.append(response)
-        commands = state["commands"]
+        commands = state.commands
         commands.extend([
             update_cell_source(cell_id, source=response),
             show_diff(cell_id, current["source"], response),
@@ -259,9 +280,9 @@ async def route_code(state: AIWorkflowState, config: RunnableConfig):
     prompt = prompt_new_cell_using_context(cell_id, state)   
     response = _cast_ai_response(await llm.ainvoke(input=prompt, stream=False))
     response = sanitize_code(response)
-    messages = state.get("messages", []) or []
+    messages = state.messages
     messages.append(response)
-    commands = state["commands"]
+    commands = state.commands
     commands.extend([
         update_cell_source(cell_id, source=response),
     ])
@@ -281,9 +302,11 @@ async def route_code(state: AIWorkflowState, config: RunnableConfig):
 
 workflow = graph.compile()
 
+
 agent = Agent(
     name = "Magic Button Agent",
     description = "Magic Button Agent",
     workflow = workflow,
-    version = "0.0.1"
+    version = "0.0.1",
+    state=State
 )

…gic_wand/agents/lab_commands/__init__.py …rlab_magic_wand/agents/tools/__init__.pyjupyterlab_magic_wand/agents/lab_commands/__init__.py renamed to jupyterlab_magic_wand/agents/tools/__init__.py jupyterlab_magic_wand/agents/lab_commands/__init__.py renamed to jupyterlab_magic_wand/agents/tools/__init__.py

@@ -2,48 +2,40 @@
 from traitlets import Instance, Dict, Unicode, default
 from jupyter_server.extension.application import ExtensionApp
 
-from .magic_handler import MagicHandler
 from .rest_handlers import handlers
 from importlib_metadata import entry_points
-from .config import ConfigManager
+from .agents.base import Agent
 
 
-feedback_logger = logging.getLogger("jupyterlab_magic_wand_feedback")
-
 class AIMagicExtension(ExtensionApp):  
     name = "jupyterlab_magic_wand"
     handlers = handlers
 
-    magic_handler = Instance(MagicHandler, allow_none=True)
-    ai_config = Instance(ConfigManager, allow_none=True)
     agents = Dict(key_trait=Unicode, value_trait=Instance(object))
-    feedback = Instance(logging.Logger, allow_none=True)
-    
-    @default('feedback')
-    def _default_feedback(self):
-        return feedback_logger
+    default_agent = Unicode(
+        default_value="Magic Button Agent",
+        help=(
+            "The name of the default agent, if an agent is not "
+            "explicitly named when a request is made to the server."
+        )
+    ).tag(config=True)
 
     def initialize_settings(self):
         eps = entry_points()
         agents_eps = eps.select(group="jupyterlab_magic_wand.agents")
         for eps in agents_eps:
             try:
-                agent = eps.load()
+                agent: Agent = eps.load()
                 self.agents[agent.name] = agent
-                self.log.info(f"Successfully loaded workflow: {agent.name}")
+                import json
+                print(json.dumps(agent.state_schema, indent=2))
+                self.serverapp.event_logger.register_event_schema(agent.state_schema)
+                self.log.info(f"Successfully loaded agent: {agent.name}")
             except Exception as err:
                 self.log.error(err)
                 self.log.error(f"Unable to load {agent.name}")
-        
-        self.ai_config = ConfigManager(self.agents)
-        self.magic_handler = MagicHandler(
-            event_logger=self.serverapp.event_logger,
-            config=self.ai_config,
-            jupyter_ai_config=self.settings["jai_config_manager"]
-        )
+                
         self.settings.update({
-            "magic_handler": self.magic_handler,
             "agents": self.agents,
-            "ai_config": self.ai_config,
-            "feedback": self.feedback
+            "current_agent": self.default_agent
         })