Merge pull request #299 from ServiceNow/make-langchain-optional

make langchain optional in xray

Author: recursix

src/agentlab/analyze/agent_xray.py

@@ -14,7 +14,6 @@
 import pandas as pd
 from attr import dataclass
 from browsergym.experiments.loop import StepInfo as BGymStepInfo
-from langchain.schema import BaseMessage, HumanMessage
 from openai import OpenAI
 from openai.types.responses import ResponseFunctionToolCall
 from PIL import Image
@@ -30,6 +29,11 @@
 from agentlab.llm.llm_utils import Discussion
 from agentlab.llm.response_api import MessageBuilder, ToolCalls
 
+try:
+    from langchain.schema import BaseMessage, HumanMessage
+except ImportError:
+    BaseMessage, HumanMessage = None, None
+
 select_dir_instructions = "Select Experiment Directory"
 AGENT_NAME_KEY = "agent.agent_name"
 TASK_NAME_KEY = "env.task_name"
@@ -746,7 +750,7 @@ def format_chat_message(message: BaseMessage | MessageBuilder | dict):
     """
     Format a message to markdown.
     """
-    if isinstance(message, BaseMessage):
+    if BaseMessage and isinstance(message, BaseMessage):
         return message.content
     elif isinstance(message, MessageBuilder):
         return message.to_markdown()
@@ -859,7 +863,9 @@ def submit_action(input_text):
     global info
     agent_info = info.exp_result.steps_info[info.step].agent_info
     chat_messages = deepcopy(agent_info.get("chat_messages", ["No Chat Messages"])[:2])
-    if isinstance(chat_messages[1], BaseMessage):  # TODO remove once langchain is deprecated
+    if BaseMessage and isinstance(
+        chat_messages[1], BaseMessage
+    ):  # TODO remove once langchain is deprecated
         assert isinstance(chat_messages[1], HumanMessage), "Second message should be user"
         chat_messages = [
             make_system_message(chat_messages[0].content),