Update web session on use

Author: ajhai

llmstack/processors/providers/api_processor_interface.py

@@ -5,6 +5,7 @@
 
 import ujson as json
 from django import db
+from django.shortcuts import get_object_or_404
 from pydantic import BaseModel
 
 from llmstack.apps.app_session_utils import get_app_session_data, save_app_session_data
@@ -133,6 +134,13 @@ def _create_asset_stream(self, mime_type, file_name=None):
             return AssetStream(asset)
         return None
 
+    def update_connection(self, connection):
+        from llmstack.base.models import Profile
+
+        if self._request_user and self._request_user.is_authenticated:
+            profile = get_object_or_404(Profile, user=self._request_user)
+            profile.add_connection(connection)
+
     def __init__(
         self,
         input,

llmstack/processors/providers/promptly/static_web_browser.py

@@ -225,6 +225,8 @@ def process(self) -> dict:
         output_stream = self._output_stream
         browser_response = None
         session_videos = []
+        updated_session_data = None
+
         with WebBrowser(
             f"{settings.RUNNER_HOST}:{settings.RUNNER_PORT}",
             interactive=self._config.stream_video,
@@ -237,6 +239,7 @@ def process(self) -> dict:
                 else ""
             ),
             record_video=self._config.capture_session_video,
+            persist_session=bool(self._config.connection_id),
         ) as web_browser:
             if self._config.stream_video and web_browser.get_wss_url():
                 async_to_sync(
@@ -262,6 +265,9 @@ def process(self) -> dict:
             if self._config.capture_session_video:
                 session_videos = web_browser.get_videos()
 
+            if self._config.connection_id:
+                updated_session_data = web_browser.terminate()
+
         if browser_response:
             output_text = browser_response.text or "\n".join(
                 list(map(lambda entry: entry.output, browser_response.command_outputs))
@@ -315,4 +321,15 @@ def process(self) -> dict:
 
         output = output_stream.finalize()
 
+        if self._config.connection_id and updated_session_data:
+            self.update_connection(
+                {
+                    **self._env["connections"][self._config.connection_id],
+                    "configuration": {
+                        **self._env["connections"][self._config.connection_id]["configuration"],
+                        "_storage_state": updated_session_data,
+                    },
+                }
+            )
+
         return output

llmstack/processors/providers/promptly/web_browser.py

@@ -465,6 +465,8 @@ def _process_anthropic(self) -> dict:
         credentials_tool_schema["input_schema"] = credentials_tool_schema["parameters"]
         del credentials_tool_schema["parameters"]
 
+        updated_session_data = None
+
         with WebBrowserClient(
             f"{settings.RUNNER_HOST}:{settings.RUNNER_PORT}",
             interactive=self._config.stream_video,
@@ -477,6 +479,7 @@ def _process_anthropic(self) -> dict:
                 else ""
             ),
             record_video=self._config.record_session_video,
+            persist_session=bool(self._config.connection_id),
         ) as web_browser:
             browser_response = web_browser.run_commands(
                 commands=[
@@ -667,6 +670,9 @@ def _process_anthropic(self) -> dict:
             if self._config.record_session_video:
                 session_videos = web_browser.get_videos()
 
+            if self._config.connection_id:
+                updated_session_data = web_browser.terminate()
+
         if browser_response:
             output_text = "\n".join(list(map(lambda entry: entry.output, browser_response.command_outputs)))
             browser_content = WebBrowserContent(**browser_response.model_dump(exclude=("screenshot", "downloads")))
@@ -715,12 +721,26 @@ def _process_anthropic(self) -> dict:
             async_to_sync(self._output_stream.write)(WebBrowserOutput(videos=encoded_videos))
 
         output = self._output_stream.finalize()
+
+        # Update the session data if we have a connection ID
+        if self._config.connection_id and updated_session_data:
+            self.update_connection(
+                {
+                    **self._env["connections"][self._config.connection_id],
+                    "configuration": {
+                        **self._env["connections"][self._config.connection_id]["configuration"],
+                        "_storage_state": updated_session_data,
+                    },
+                }
+            )
         return output
 
     def process(self) -> dict:
         if self._config.provider_config.provider == "anthropic":
             return self._process_anthropic()
 
+        updated_session_data = None
+
         client = get_llm_client_from_provider_config(
             provider=self._config.provider_config.provider,
             model_slug=self._config.provider_config.model.value,
@@ -752,6 +772,7 @@ def process(self) -> dict:
                 if self._config.connection_id
                 else ""
             ),
+            persist_session=bool(self._config.connection_id),
         ) as web_browser:
             # Start streaming video if enabled
             if self._config.stream_video and web_browser.get_wss_url():
@@ -938,6 +959,9 @@ def process(self) -> dict:
                         )
                     )
 
+            if self._config.connection_id:
+                updated_session_data = web_browser.terminate()
+
         if browser_response:
             output_text = "\n".join(list(map(lambda entry: entry.output, browser_response.command_outputs)))
             browser_content = browser_response.model_dump(exclude=("screenshot",))
@@ -951,4 +975,17 @@ def process(self) -> dict:
             async_to_sync(self._output_stream.write)(WebBrowserOutput(text=output_text, content=browser_content))
 
         output = self._output_stream.finalize()
+
+        # Update the session data if we have a connection ID
+        if self._config.connection_id and updated_session_data:
+            self.update_connection(
+                {
+                    **self._env["connections"][self._config.connection_id],
+                    "configuration": {
+                        **self._env["connections"][self._config.connection_id]["configuration"],
+                        "_storage_state": updated_session_data,
+                    },
+                }
+            )
+
         return output