feat: allow caller-supplied Snowpark session in ChatSnowflake

CHANGELOG.md

@@ -14,6 +14,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 * `Chat` gains `close()` and `close_async()` methods (plus context-manager support) for releasing resources held by the provider -- HTTP connection pools, the Snowflake Snowpark session/connection, and (via `close_async()`) MCP server sessions. This is useful in long-lived applications like Shiny that create a chat per user session: `session.on_ended(chat.close)`. Providers only close resources they created themselves; caller-supplied clients are left open.
 
+* `ChatSnowflake()` gains a `session` parameter for supplying an existing `snowflake.snowpark.Session`, mirroring `ChatDatabricks()`'s `workspace_client`. This lets one session be shared across multiple chats; `Chat.close()` only closes sessions that chatlas created itself, leaving caller-supplied sessions open.
+
 * When running on Posit Connect, chatlas now forwards the Shiny viewer's session token to Connect's LLM gateway (as a `Posit-Connect-User-Session-Token` header) so gateway usage can be attributed to the viewer. This happens automatically for Shiny content and only affects requests to the gateway.
 
 ### Changes

chatlas/_provider_snowflake.py

@@ -31,6 +31,7 @@
 if TYPE_CHECKING:
     import snowflake.core.cortex.inference_service._generated.models as models
     from snowflake.core.rest import Event, SSEClient
+    from snowflake.snowpark import Session
 
     Completion = models.NonStreamingCompleteResponse
     CompletionChunk = models.StreamingCompleteResponseDataEvent
@@ -69,6 +70,7 @@ def ChatSnowflake(
     password: Optional[str] = None,
     private_key_file: Optional[str] = None,
     private_key_file_pwd: Optional[str] = None,
+    session: Optional["Session"] = None,
     kwargs: Optional[dict[str, "str | int"]] = None,
 ) -> Chat["CompleteRequest", "Completion"]:
     """
@@ -136,6 +138,13 @@ def ChatSnowflake(
     private_key_file_pwd
         The password for your private key file. Required if you are using key pair authentication.
         https://docs.snowflake.com/en/user-guide/key-pair-auth
+    session
+        A `snowflake.snowpark.Session` to use for the connection. If not
+        provided, a new session will be created from the other connection
+        parameters. Note that calling
+        [`Chat.close()`](`chatlas.Chat.close`) does not close a
+        caller-supplied `Session` -- it only closes a session that chatlas
+        created itself.
     kwargs
         Additional keyword arguments passed along to the Snowflake connection builder. These can
         include any parameters supported by the `snowflake-ml-python` package.
@@ -154,6 +163,7 @@ def ChatSnowflake(
             password=password,
             private_key_file=private_key_file,
             private_key_file_pwd=private_key_file_pwd,
+            session=session,
             kwargs=kwargs,
         ),
         system_prompt=system_prompt,
@@ -173,6 +183,7 @@ def __init__(
         password: Optional[str],
         private_key_file: Optional[str],
         private_key_file_pwd: Optional[str],
+        session: Optional["Session"],
         name: str = "Snowflake",
         kwargs: Optional[dict[str, "str | int"]],
     ):
@@ -193,27 +204,34 @@ def __init__(
         # https://docs.snowflake.com/en/developer-guide/python-connector/python-connector-api#functions
         application = os.environ.get("SF_PARTNER", "py_chatlas")
 
-        configs: dict[str, str | int] = drop_none(
-            {
-                "connection_name": connection_name,
-                "account": account,
-                "user": user,
-                "password": password,
-                "private_key_file": private_key_file,
-                "private_key_file_pwd": private_key_file_pwd,
-                "application": application,
-                **(kwargs or {}),
-            }
-        )
+        if session is None:
+            configs: dict[str, str | int] = drop_none(
+                {
+                    "connection_name": connection_name,
+                    "account": account,
+                    "user": user,
+                    "password": password,
+                    "private_key_file": private_key_file,
+                    "private_key_file_pwd": private_key_file_pwd,
+                    "application": application,
+                    **(kwargs or {}),
+                }
+            )
+            session = Session.builder.configs(configs).create()
+            self._owns_session = True
+        else:
+            self._owns_session = False
 
-        self._session = Session.builder.configs(configs).create()
+        self._session = session
         self._cortex_service = Root(self._session).cortex_inference_service
 
     def close(self) -> None:
         """
-        Close the underlying Snowpark session (and its Snowflake connection).
+        Close the underlying Snowpark session (and its Snowflake connection),
+        unless the session was supplied by the caller.
         """
-        self._session.close()
+        if self._owns_session:
+            self._session.close()
 
     def list_models(self):
         raise NotImplementedError(

tests/test_close.py

@@ -166,6 +166,17 @@ def test_snowflake_provider_close():
     assert session.close.call_count == 2
 
 
+def test_snowflake_provider_close_ownership():
+    from chatlas import ChatSnowflake
+
+    snowflake, session = _make_mock_snowflake_modules()
+    with patch.dict(sys.modules, {"snowflake": snowflake, "snowflake.snowpark": snowflake.snowpark, "snowflake.core": snowflake.core}):
+        chat = ChatSnowflake(model="llama3.1-70b", session=session)
+    chat.close()
+    # A caller-supplied session is not closed.
+    session.close.assert_not_called()
+
+
 def test_databricks_provider_close_ownership():
     from chatlas import ChatDatabricks