Update langchain-python example to LangChain 1.x

examples/langchain-python/.env.template

@@ -1,5 +1,5 @@
-# TODO: Get your E2B API key from https://e2b.dev/docs
+# Get your E2B API key from https://e2b.dev/dashboard?tab=keys
 E2B_API_KEY=""
 
-# TODO: Get your OpenAI API key from https://platform.openai.com
+# Get your OpenAI API key from https://platform.openai.com/settings/organization/api-keys
 OPENAI_API_KEY=""

examples/langchain-python/langchain_e2b_python/code_interpreter_tool.py

@@ -1,80 +1,53 @@
 import os
-import json
+from typing import Any
 
-from typing import Any, List
-from langchain_core.tools import Tool
-from pydantic.v1 import BaseModel, Field
 from e2b_code_interpreter import Sandbox
-from langchain_core.messages import BaseMessage, ToolMessage
-from langchain.agents.output_parsers.tools import (
-    ToolAgentAction,
-)
+from langchain_core.tools import tool
 
 
-class LangchainCodeInterpreterToolInput(BaseModel):
-    code: str = Field(description="Python code to execute.")
-
-
-class CodeInterpreterFunctionTool:
+class CodeInterpreterTool:
     """
     This class calls arbitrary code against a Python Jupyter notebook.
     It requires an E2B_API_KEY to create a sandbox.
     """
 
-    tool_name: str = "code_interpreter"
-
     def __init__(self):
-        # Instantiate the E2B sandbox - this is a long lived object
-        # that's pinging E2B cloud to keep the sandbox alive.
         if "E2B_API_KEY" not in os.environ:
             raise Exception(
-                "Code Interpreter tool called while E2B_API_KEY environment variable is not set. Please get your E2B api key here https://e2b.dev/docs and set the E2B_API_KEY environment variable."
+                "Code Interpreter tool called while E2B_API_KEY environment variable is not set. "
+                "Please get your E2B api key here https://e2b.dev/dashboard?tab=keys and set the E2B_API_KEY environment variable."
             )
-        self.code_interpreter = Sandbox()
+        self.sandbox = Sandbox.create()
+        self.last_results = []
 
     def close(self):
-        self.code_interpreter.kill()
+        self.sandbox.kill()
 
-    def call(self, parameters: dict, **kwargs: Any):
-        code = parameters.get("code", "")
+    def __enter__(self):
+        return self
+
+    def __exit__(self, *args):
+        self.close()
+
+    def run_code(self, code: str) -> dict[str, Any]:
+        """Execute Python code in a Jupyter notebook cell."""
         print(f"***Code Interpreting...\n{code}\n====")
-        execution = self.code_interpreter.run_code(code)
+        execution = self.sandbox.run_code(code)
+        self.last_results = execution.results
         return {
             "results": execution.results,
             "stdout": execution.logs.stdout,
             "stderr": execution.logs.stderr,
             "error": execution.error,
         }
 
-    # langchain does not return a dict as a parameter, only a code string
-    def langchain_call(self, code: str):
-        return self.call({"code": code})
-
-    def to_langchain_tool(self) -> Tool:
-        tool = Tool(
-            name=self.tool_name,
-            description="Execute python code in a Jupyter notebook cell and returns any rich data (eg charts), stdout, stderr, and error.",
-            func=self.langchain_call,
-        )
-        tool.args_schema = LangchainCodeInterpreterToolInput
-        return tool
 
-    @staticmethod
-    def format_to_tool_message(
-        agent_action: ToolAgentAction,
-        observation: dict,
-    ) -> List[BaseMessage]:
-        """
-        Format the output of the CodeInterpreter tool to be returned as a ToolMessage.
-        """
-        new_messages = list(agent_action.message_log)
+def create_code_interpreter_tool(interpreter: CodeInterpreterTool):
+    """Create a LangChain tool from the code interpreter."""
 
-        # TODO: Add info about the results for the LLM
-        content = json.dumps(
-            {k: v for k, v in observation.items() if k not in ("results")}, indent=2
-        )
-        new_messages.append(
-            ToolMessage(content=content, tool_call_id=agent_action.tool_call_id)
-        )
+    @tool
+    def code_interpreter(code: str) -> dict[str, Any]:
+        """Execute python code in a Jupyter notebook cell and returns any rich data (eg charts), stdout, stderr, and error."""
+        return interpreter.run_code(code)
 
-        return new_messages
+    return code_interpreter

examples/langchain-python/langchain_e2b_python/main.py

@@ -1,91 +1,45 @@
 import base64
 
-from typing import List, Sequence, Tuple
 from dotenv import load_dotenv
-from langchain_core.prompts import ChatPromptTemplate
-from langchain_groq import ChatGroq
-from langchain_e2b_python.code_interpreter_tool import CodeInterpreterFunctionTool
-from langchain.agents import AgentExecutor
+from langchain.agents import create_agent
 from langchain_openai import ChatOpenAI
-from langchain_core.messages import BaseMessage
-from langchain_core.runnables import RunnablePassthrough
-from langchain.agents.output_parsers.tools import (
-    ToolAgentAction,
-    ToolsAgentOutputParser,
-)
-from e2b_code_interpreter import Result
 
+from langchain_e2b_python.code_interpreter_tool import (
+    CodeInterpreterTool,
+    create_code_interpreter_tool,
+)
 
 load_dotenv()
 
 
-def format_to_tool_messages(
-    intermediate_steps: Sequence[Tuple[ToolAgentAction, dict]],
-) -> List[BaseMessage]:
-    messages = []
-    for agent_action, observation in intermediate_steps:
-        if agent_action.tool == CodeInterpreterFunctionTool.tool_name:
-            new_messages = CodeInterpreterFunctionTool.format_to_tool_message(
-                agent_action,
-                observation,
-            )
-            messages.extend([new for new in new_messages if new not in messages])
-        else:
-            # Handle other tools
-            print("Not handling tool: ", agent_action.tool)
-
-    return messages
-
-
 def main():
-    # 1. Pick your favorite llm
-    llm = ChatOpenAI(model="gpt-3.5-turbo-0125", temperature=0)
-    # llm = ChatGroq(temperature=0, model_name="llama3-70b-8192")
-
-    # 2. Initialize the code interpreter tool
-    code_interpreter = CodeInterpreterFunctionTool()
-    code_interpreter_tool = code_interpreter.to_langchain_tool()
-    tools = [code_interpreter_tool]
-
-    # 3. Define the prompt
-    prompt = ChatPromptTemplate.from_messages(
-        [("human", "{input}"), ("placeholder", "{agent_scratchpad}")]
-    )
-
-    # 4. Define the agent
-    agent = (
-        RunnablePassthrough.assign(
-            agent_scratchpad=lambda x: format_to_tool_messages(x["intermediate_steps"])
+    # Pick your favourite LLM
+    llm = ChatOpenAI(model="gpt-4o-mini", temperature=0)
+
+    # Initialize CodeInterpreterTool - defined in code_interpreter_tool.py
+    with CodeInterpreterTool() as interpreter:
+        # Create agent from LangChain
+        agent = create_agent(
+            model=llm,
+            tools=[create_code_interpreter_tool(interpreter)],
+            system_prompt="You are a helpful assistant that can execute Python code to help answer questions.",
         )
-        | prompt
-        | llm.bind_tools(tools)
-        | ToolsAgentOutputParser()
-    )
-
-    agent_executor = AgentExecutor(
-        agent=agent,
-        tools=tools,
-        verbose=True,
-        return_intermediate_steps=True,
-    )
 
-    # 5. Invoke the agent
-    result = agent_executor.invoke({"input": "plot and show sinus"})
-
-    code_interpreter.close()
-
-    print(result)
+        # Invoke agent to plot and show sinus
+        result = agent.invoke(
+            {"messages": [{"role": "user", "content": "plot and show sinus"}]}
+        )
+        print(result)
 
-    # Each intermediate step is a Tuple[ToolAgentAction, dict]
-    r: Result = result["intermediate_steps"][0][1]["results"][0]
+        # Save PNG chart
+        for r in interpreter.last_results:
+            if hasattr(r, "png") and r.png:
+                png_data = base64.b64decode(r.png)
+                with open("chart.png", "wb") as f:
+                    f.write(png_data)
+                print("Saved chart to chart.png")
+                break
 
-    # Save the PNG chart that was received
-    if r.png:
-        # Decode the base64 encoded PNG data
-        png_data = base64.b64decode(r.png)
 
-        # Save the decoded PNG data to a file
-        filename = f"chart.png"
-        with open(filename, "wb") as f:
-            f.write(png_data)
-        print(f"Saved chart to {filename}")
+if __name__ == "__main__":
+    main()