Skip to content

Commit 22544d6

Browse files
rcholicrcholic
authored
Merge pull request #94 from SentienceAPI/fix_screenshot2
upload screenshots along with traces
2 parents fc7e7ff + f45d1f1 commit 22544d6

File tree

10 files changed

+1268
-153
lines changed

10 files changed

+1268
-153
lines changed

‎sentience/agent.py‎

Lines changed: 93 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@
99
from typing import TYPE_CHECKING, Any, Optional
1010

1111
from .actions import click, click_async, press, press_async, type_text, type_text_async
12+
from .agent_config import AgentConfig
1213
from .base_agent import BaseAgent, BaseAgentAsync
1314
from .browser import AsyncSentienceBrowser, SentienceBrowser
1415
from .llm_provider import LLMProvider, LLMResponse
@@ -25,7 +26,6 @@
2526
from .snapshot import snapshot, snapshot_async
2627

2728
if TYPE_CHECKING:
28-
from .agent_config import AgentConfig
2929
from .tracing import Tracer
3030

3131

@@ -78,8 +78,9 @@ def __init__(
7878
self.default_snapshot_limit = default_snapshot_limit
7979
self.verbose = verbose
8080
self.tracer = tracer
81-
self.config = config
81+
self.config = config or AgentConfig()
8282

83+
# Screenshot sequence counter
8384
# Execution history
8485
self.history: list[dict[str, Any]] = []
8586

@@ -150,6 +151,21 @@ def act( # noqa: C901
150151
if snap_opts.goal is None:
151152
snap_opts.goal = goal
152153

154+
# Apply AgentConfig screenshot settings if not overridden by snapshot_options
155+
if snapshot_options is None and self.config:
156+
if self.config.capture_screenshots:
157+
# Create ScreenshotConfig from AgentConfig
158+
snap_opts.screenshot = ScreenshotConfig(
159+
format=self.config.screenshot_format,
160+
quality=(
161+
self.config.screenshot_quality
162+
if self.config.screenshot_format == "jpeg"
163+
else None
164+
),
165+
)
166+
else:
167+
snap_opts.screenshot = False
168+
153169
# Call snapshot with options object (matches TypeScript API)
154170
snap = snapshot(self.browser, snap_opts)
155171

@@ -178,14 +194,36 @@ def act( # noqa: C901
178194
for el in filtered_elements[:50] # Limit to first 50 for performance
179195
]
180196

197+
# Build snapshot event data
198+
snapshot_data = {
199+
"url": snap.url,
200+
"element_count": len(snap.elements),
201+
"timestamp": snap.timestamp,
202+
"elements": elements_data, # Add element data for overlay
203+
}
204+
205+
# Always include screenshot in trace event for studio viewer compatibility
206+
# CloudTraceSink will extract and upload screenshots separately, then remove
207+
# screenshot_base64 from events before uploading the trace file.
208+
if snap.screenshot:
209+
# Extract base64 string from data URL if needed
210+
if snap.screenshot.startswith("data:image"):
211+
# Format: "data:image/jpeg;base64,{base64_string}"
212+
screenshot_base64 = (
213+
snap.screenshot.split(",", 1)[1]
214+
if "," in snap.screenshot
215+
else snap.screenshot
216+
)
217+
else:
218+
screenshot_base64 = snap.screenshot
219+
220+
snapshot_data["screenshot_base64"] = screenshot_base64
221+
if snap.screenshot_format:
222+
snapshot_data["screenshot_format"] = snap.screenshot_format
223+
181224
self.tracer.emit(
182225
"snapshot",
183-
{
184-
"url": snap.url,
185-
"element_count": len(snap.elements),
186-
"timestamp": snap.timestamp,
187-
"elements": elements_data, # Add element data for overlay
188-
},
226+
snapshot_data,
189227
step_id=step_id,
190228
)
191229

@@ -721,8 +759,9 @@ def __init__(
721759
self.default_snapshot_limit = default_snapshot_limit
722760
self.verbose = verbose
723761
self.tracer = tracer
724-
self.config = config
762+
self.config = config or AgentConfig()
725763

764+
# Screenshot sequence counter
726765
# Execution history
727766
self.history: list[dict[str, Any]] = []
728767

@@ -790,6 +829,23 @@ async def act( # noqa: C901
790829
if snap_opts.goal is None:
791830
snap_opts.goal = goal
792831

832+
# Apply AgentConfig screenshot settings if not overridden by snapshot_options
833+
# Only apply if snapshot_options wasn't provided OR if screenshot wasn't explicitly set
834+
# (snapshot_options.screenshot defaults to False, so we check if it's still False)
835+
if self.config and (snapshot_options is None or snap_opts.screenshot is False):
836+
if self.config.capture_screenshots:
837+
# Create ScreenshotConfig from AgentConfig
838+
snap_opts.screenshot = ScreenshotConfig(
839+
format=self.config.screenshot_format,
840+
quality=(
841+
self.config.screenshot_quality
842+
if self.config.screenshot_format == "jpeg"
843+
else None
844+
),
845+
)
846+
else:
847+
snap_opts.screenshot = False
848+
793849
# Call snapshot with options object (matches TypeScript API)
794850
snap = await snapshot_async(self.browser, snap_opts)
795851

@@ -818,14 +874,36 @@ async def act( # noqa: C901
818874
for el in filtered_elements[:50] # Limit to first 50 for performance
819875
]
820876

877+
# Build snapshot event data
878+
snapshot_data = {
879+
"url": snap.url,
880+
"element_count": len(snap.elements),
881+
"timestamp": snap.timestamp,
882+
"elements": elements_data, # Add element data for overlay
883+
}
884+
885+
# Always include screenshot in trace event for studio viewer compatibility
886+
# CloudTraceSink will extract and upload screenshots separately, then remove
887+
# screenshot_base64 from events before uploading the trace file.
888+
if snap.screenshot:
889+
# Extract base64 string from data URL if needed
890+
if snap.screenshot.startswith("data:image"):
891+
# Format: "data:image/jpeg;base64,{base64_string}"
892+
screenshot_base64 = (
893+
snap.screenshot.split(",", 1)[1]
894+
if "," in snap.screenshot
895+
else snap.screenshot
896+
)
897+
else:
898+
screenshot_base64 = snap.screenshot
899+
900+
snapshot_data["screenshot_base64"] = screenshot_base64
901+
if snap.screenshot_format:
902+
snapshot_data["screenshot_format"] = snap.screenshot_format
903+
821904
self.tracer.emit(
822905
"snapshot",
823-
{
824-
"url": snap.url,
825-
"element_count": len(snap.elements),
826-
"timestamp": snap.timestamp,
827-
"elements": elements_data, # Add element data for overlay
828-
},
906+
snapshot_data,
829907
step_id=step_id,
830908
)
831909

0 commit comments

Comments
 (0)