From 1e99e0df53dcb5beb4721f7f44a4b0af7457ff34 Mon Sep 17 00:00:00 2001
From: Chester Curme <chester.curme@gmail.com>
Date: Wed, 8 Jan 2025 11:16:23 -0500
Subject: [PATCH] update langchain-xai

---
 libs/partners/xai/Makefile                                 | 3 +++
 libs/partners/xai/langchain_xai/chat_models.py             | 7 ++++---
 .../__snapshots__/test_chat_models_standard.ambr           | 1 -
 3 files changed, 7 insertions(+), 4 deletions(-)

diff --git a/libs/partners/xai/Makefile b/libs/partners/xai/Makefile
index 963563e2d7099..6859cc789a179 100644
--- a/libs/partners/xai/Makefile
+++ b/libs/partners/xai/Makefile
@@ -11,6 +11,9 @@ integration_test integration_tests: TEST_FILE=tests/integration_tests/
 test tests:
 	poetry run pytest --disable-socket --allow-unix-socket $(TEST_FILE)
 
+test_watch:
+	poetry run ptw --snapshot-update --now . -- -vv $(TEST_FILE)
+
 integration_test integration_tests:
 	poetry run pytest $(TEST_FILE)
 
diff --git a/libs/partners/xai/langchain_xai/chat_models.py b/libs/partners/xai/langchain_xai/chat_models.py
index 775d22740cd4e..a854be5487d4c 100644
--- a/libs/partners/xai/langchain_xai/chat_models.py
+++ b/libs/partners/xai/langchain_xai/chat_models.py
@@ -320,9 +320,9 @@ def _get_ls_params(
     @model_validator(mode="after")
     def validate_environment(self) -> Self:
         """Validate that api key and python package exists in environment."""
-        if self.n < 1:
+        if self.n is not None and self.n < 1:
             raise ValueError("n must be at least 1.")
-        if self.n > 1 and self.streaming:
+        if self.n is not None and self.n > 1 and self.streaming:
             raise ValueError("n must be 1 when streaming.")
 
         client_params: dict = {
@@ -331,10 +331,11 @@ def validate_environment(self) -> Self:
             ),
             "base_url": self.xai_api_base,
             "timeout": self.request_timeout,
-            "max_retries": self.max_retries,
             "default_headers": self.default_headers,
             "default_query": self.default_query,
         }
+        if self.max_retries is not None:
+            client_params["max_retries"] = self.max_retries
 
         if client_params["api_key"] is None:
             raise ValueError(
diff --git a/libs/partners/xai/tests/unit_tests/__snapshots__/test_chat_models_standard.ambr b/libs/partners/xai/tests/unit_tests/__snapshots__/test_chat_models_standard.ambr
index 5c6f113f2174a..4cd1261555c90 100644
--- a/libs/partners/xai/tests/unit_tests/__snapshots__/test_chat_models_standard.ambr
+++ b/libs/partners/xai/tests/unit_tests/__snapshots__/test_chat_models_standard.ambr
@@ -10,7 +10,6 @@
       'max_retries': 2,
       'max_tokens': 100,
       'model_name': 'grok-beta',
-      'n': 1,
       'request_timeout': 60.0,
       'stop': list([
       ]),