fix: make langchain imports compatible with langchain v2

src/agentlab/analyze/agent_xray.py

@@ -14,7 +14,15 @@
 import pandas as pd
 from attr import dataclass
 from browsergym.experiments.loop import StepInfo as BGymStepInfo
-from langchain.schema import BaseMessage, HumanMessage
+
+try:
+    from langchain_core.messages import BaseMessage, HumanMessage
+except ImportError:
+    try:
+        from langchain.schema import BaseMessage, HumanMessage
+    except ImportError:
+        BaseMessage = None
+        HumanMessage = None
 from openai import OpenAI
 from openai.types.responses import ResponseFunctionToolCall
 from PIL import Image
@@ -746,7 +754,7 @@ def format_chat_message(message: BaseMessage | MessageBuilder | dict):
     """
     Format a message to markdown.
     """
-    if isinstance(message, BaseMessage):
+    if BaseMessage is not None and isinstance(message, BaseMessage):
         return message.content
     elif isinstance(message, MessageBuilder):
         return message.to_markdown()
@@ -847,8 +855,12 @@ 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
-        assert isinstance(chat_messages[1], HumanMessage), "Second message should be user"
+    if BaseMessage is not None and isinstance(
+        chat_messages[1], BaseMessage
+    ):  # TODO remove once langchain is deprecated
+        assert HumanMessage is not None and isinstance(
+            chat_messages[1], HumanMessage
+        ), "Second message should be user"
         chat_messages = [
             make_system_message(chat_messages[0].content),
             make_user_message(chat_messages[1].content),

src/agentlab/llm/llm_utils.py

@@ -22,8 +22,17 @@
 
 langchain_community = importlib.util.find_spec("langchain_community")
 if langchain_community is not None:
-    from langchain.schema import BaseMessage as LangchainBaseMessage
-    from langchain_community.adapters.openai import convert_message_to_dict
+    try:
+        from langchain_core.messages import BaseMessage as LangchainBaseMessage
+    except ImportError:
+        try:
+            from langchain.schema import BaseMessage as LangchainBaseMessage
+        except ImportError:
+            LangchainBaseMessage = None
+    try:
+        from langchain_community.adapters.openai import convert_message_to_dict
+    except ImportError:
+        convert_message_to_dict = None
 else:
     LangchainBaseMessage = None
     convert_message_to_dict = None
@@ -40,7 +49,13 @@ def messages_to_dict(messages: list[dict] | list[LangchainBaseMessage]) -> dict:
         elif isinstance(m, str):
             new_messages.add_message({"role": "<unknown role>", "content": m})
         elif LangchainBaseMessage is not None and isinstance(m, LangchainBaseMessage):
-            new_messages.add_message(convert_message_to_dict(m))
+            if convert_message_to_dict is not None:
+                new_messages.add_message(convert_message_to_dict(m))
+            else:
+                role_map = {"human": "user", "ai": "assistant", "system": "system"}
+                new_messages.add_message(
+                    {"role": role_map.get(m.type, m.type), "content": str(m.content)}
+                )
         else:
             raise ValueError(f"Unknown message type: {type(m)}")
     return new_messages

src/agentlab/llm/tracking.py

@@ -13,7 +13,11 @@
 
 langchain_community = importlib.util.find_spec("langchain_community")
 if langchain_community is not None:
-    from langchain_community.callbacks import bedrock_anthropic_callback, openai_info
+    try:
+        from langchain_community.callbacks import bedrock_anthropic_callback, openai_info
+    except ImportError:
+        bedrock_anthropic_callback = None
+        openai_info = None
 else:
     bedrock_anthropic_callback = None
     openai_info = None