fix: Solve dependency issues

py/src/braintrust/cli/install/redshift.py

@@ -4,7 +4,7 @@
 from hashlib import md5
 
 # pylint: disable=no-name-in-module
-from ... import log_conn, login
+from ... import api_conn, login
 
 # pylint: disable=no-name-in-module
 from ...aws import iam, redshift_serverless
@@ -186,7 +186,7 @@ def main(args):
     login_kwargs = {"org_name": args.org_name} if args.org_name else {}
     login(**login_kwargs)
 
-    resp = log_conn().get(
+    resp = api_conn().get(
         "/dw-test",
         params={"iam_role": role["Role"]["Arn"], "msk_cluster_arn": msk_cluster_arn},
     )

py/src/braintrust/wrappers/langchain.py

@@ -9,18 +9,31 @@
 _logger = logging.getLogger("braintrust.wrappers.langchain")
 
 try:
-    from langchain.callbacks.base import BaseCallbackHandler
-    from langchain.schema import Document
-    from langchain.schema.agent import AgentAction
-    from langchain.schema.messages import BaseMessage
-    from langchain.schema.output import LLMResult
+    # Modern langchain
+    from langchain_core.callbacks.base import BaseCallbackHandler
+    from langchain_core.documents.base import Document
+    from langchain_core.messages.base import BaseMessage
+    from langchain_core.outputs.llm_result import LLMResult
 except ImportError:
-    _logger.warning("Failed to import langchain, using stubs")
-    BaseCallbackHandler = object
-    Document = object
-    AgentAction = object
-    BaseMessage = object
-    LLMResult = object
+    try:
+        # after the release of langchain v1, these submodules were also migrated to langchain_classic
+        from langchain_classic.callbacks.base import BaseCallbackHandler
+        from langchain_classic.schema import Document
+        from langchain_classic.schema.messages import BaseMessage
+        from langchain_classic.schema.output import LLMResult
+    except ImportError:
+        try:
+            # Old langchain from before the v1 version and the creation of langchain_classic
+            from langchain.callbacks.base import BaseCallbackHandler
+            from langchain.schema import Document
+            from langchain.schema.messages import BaseMessage
+            from langchain.schema.output import LLMResult
+        except ImportError:
+            raise ImportError(
+                "Could not import langchain callbacks and schema submodules. "
+                "Install one of: langchain_classic, langchain-core, or langchain<1."
+            )
+
 
 langchain_parent = contextvars.ContextVar("langchain_current_span", default=None)