Skip to content

Commit acacf1d

Browse files
committed
feat: add Arraylake Code button to web UI, increase download limit to 15GB
- Show 📦 Arraylake Code button inline with Enlarge/Download/Show Code - Auto-generates ready-to-paste Python snippet with exact query params - Only shows snippet for current request (no accumulation across turns) - Raise max_download_size_gb from 5 to 15
1 parent ef352dc commit acacf1d

File tree

5 files changed

+219
-20
lines changed

5 files changed

+219
-20
lines changed

src/eurus/retrieval.py

Lines changed: 19 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -41,29 +41,32 @@ def _arraylake_snippet(
4141
min_lon: float,
4242
max_lon: float,
4343
) -> str:
44-
"""Generate a self-contained Python snippet for direct Arraylake data access."""
44+
"""Generate a ready-to-paste Python snippet for direct Arraylake access."""
45+
# Convert negative lons to 0-360 for ERA5
46+
era5_min = min_lon % 360 if min_lon < 0 else min_lon
47+
era5_max = max_lon % 360 if max_lon < 0 else max_lon
4548
return (
46-
"# ── Direct Arraylake Retrieval (copy-paste into any Python env) ──\n"
47-
"import os, xarray as xr\n"
48-
"from arraylake import Client\n"
49-
"\n"
50-
"client = Client(token=os.environ['ARRAYLAKE_API_KEY'])\n"
51-
f"repo = client.get_repo('{CONFIG.data_source}')\n"
52-
"session = repo.readonly_session('main')\n"
53-
"\n"
49+
f"\n📦 Reproduce this download yourself (copy-paste into Jupyter):\n"
50+
f"```python\n"
51+
f"import os, xarray as xr\n"
52+
f"from arraylake import Client\n"
53+
f"\n"
54+
f"client = Client(token=os.environ['ARRAYLAKE_API_KEY'])\n"
55+
f"repo = client.get_repo('{CONFIG.data_source}')\n"
56+
f"session = repo.readonly_session('main')\n"
57+
f"\n"
5458
f"ds = xr.open_dataset(session.store, engine='zarr',\n"
5559
f" consolidated=False, zarr_format=3,\n"
5660
f" chunks=None, group='{query_type}')\n"
57-
"\n"
61+
f"\n"
5862
f"subset = ds['{variable}'].sel(\n"
5963
f" time=slice('{start_date}', '{end_date}'),\n"
60-
f" latitude=slice({max_lat}, {min_lat}), # ERA5: descending lat\n"
61-
f" longitude=slice({min_lon}, {max_lon}),\n"
62-
")\n"
63-
"\n"
64-
"# Compute & save locally\n"
64+
f" latitude=slice({max_lat}, {min_lat}), # ERA5: descending\n"
65+
f" longitude=slice({era5_min}, {era5_max}),\n"
66+
f")\n"
67+
f"\n"
6568
f"subset.load().to_dataset(name='{variable}').to_zarr('my_data.zarr', mode='w')\n"
66-
"# ────────────────────────────────────────────────────────────────\n"
69+
f"```"
6770
)
6871

6972

web/agent_wrapper.py

Lines changed: 28 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,7 @@
2525

2626
# IMPORT FROM EURUS PACKAGE - SINGLE SOURCE OF TRUTH
2727
from eurus.config import CONFIG, AGENT_SYSTEM_PROMPT
28+
from eurus.retrieval import _arraylake_snippet
2829
from eurus.memory import get_memory, SmartConversationMemory # Singleton for datasets, per-session for chat
2930
from eurus.tools import get_all_tools
3031
from eurus.tools.repl import PythonREPLTool
@@ -181,12 +182,14 @@ async def process_message(
181182
lambda: self._agent.invoke({"messages": self._messages}, config=config)
182183
)
183184

184-
# Update messages
185+
# Only scan NEW messages from this turn
186+
prev_count = len(self._messages)
185187
self._messages = result["messages"]
188+
new_messages = self._messages[prev_count:]
186189

187-
# Parse messages to show tool calls made
190+
# Parse NEW messages to show tool calls made
188191
tool_calls_made = []
189-
for msg in self._messages:
192+
for msg in new_messages:
190193
if hasattr(msg, 'tool_calls') and msg.tool_calls:
191194
for tc in msg.tool_calls:
192195
tool_name = tc.get('name', 'unknown')
@@ -198,6 +201,24 @@ async def process_message(
198201
await stream_callback("status", f"🛠️ Used tools: {tools_str}")
199202
await asyncio.sleep(0.5)
200203

204+
# Collect Arraylake snippet from NEW messages only
205+
arraylake_snippets = []
206+
for msg in new_messages:
207+
if hasattr(msg, 'tool_calls') and msg.tool_calls:
208+
for tc in msg.tool_calls:
209+
if tc.get('name') == 'retrieve_era5_data':
210+
args = tc.get('args', {})
211+
arraylake_snippets.append(_arraylake_snippet(
212+
variable=args.get('variable_id', 'sst'),
213+
query_type=args.get('query_type', 'spatial'),
214+
start_date=args.get('start_date', ''),
215+
end_date=args.get('end_date', ''),
216+
min_lat=args.get('min_latitude', -90),
217+
max_lat=args.get('max_latitude', 90),
218+
min_lon=args.get('min_longitude', 0),
219+
max_lon=args.get('max_longitude', 360),
220+
))
221+
201222
# Extract response
202223
last_message = self._messages[-1]
203224

@@ -244,6 +265,10 @@ async def process_message(
244265
# Default to plot (png, jpg, etc.)
245266
await stream_callback("plot", "", data=base64_data, path=filepath, code=code)
246267

268+
# Send Arraylake snippets AFTER response + plots exist in DOM
269+
for snippet in arraylake_snippets:
270+
await stream_callback("arraylake_snippet", snippet)
271+
247272
# Save to memory
248273
self._conversation.add_message("assistant", response_text)
249274

web/static/css/style.css

Lines changed: 68 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -880,4 +880,72 @@ dialog .close-modal:hover {
880880
.save-keys-btn:disabled {
881881
opacity: 0.5;
882882
cursor: not-allowed;
883+
}
884+
885+
/* ===== ARRAYLAKE SNIPPET (in-message) ===== */
886+
.arraylake-snippet-section {
887+
margin-top: 0.75rem;
888+
}
889+
890+
.arraylake-actions {
891+
display: flex;
892+
gap: 0.5rem;
893+
}
894+
895+
.arraylake-btn {
896+
padding: 0.5rem 0.875rem;
897+
font-size: 0.75rem;
898+
font-weight: 500;
899+
border: 1px solid var(--glass-border);
900+
border-radius: 0.375rem;
901+
background: var(--bg-tertiary);
902+
color: var(--text-secondary);
903+
cursor: pointer;
904+
transition: all 0.15s ease;
905+
}
906+
907+
.arraylake-btn:hover {
908+
border-color: var(--accent-primary);
909+
color: var(--accent-primary);
910+
}
911+
912+
.arraylake-code {
913+
margin-top: 0.5rem;
914+
border-radius: 0.5rem;
915+
overflow: hidden;
916+
border: 1px solid var(--glass-border);
917+
}
918+
919+
.arraylake-code pre {
920+
margin: 0;
921+
padding: 0.875rem;
922+
background: var(--code-bg);
923+
overflow-x: auto;
924+
font-size: 0.8125rem;
925+
line-height: 1.5;
926+
}
927+
928+
.arraylake-code code {
929+
font-family: 'SF Mono', Monaco, Consolas, 'Liberation Mono', monospace;
930+
}
931+
932+
.copy-snippet-btn {
933+
display: block;
934+
margin: 0;
935+
padding: 0.5rem 0.875rem;
936+
width: 100%;
937+
font-size: 0.75rem;
938+
font-weight: 500;
939+
border: none;
940+
border-top: 1px solid var(--glass-border);
941+
background: var(--bg-tertiary);
942+
color: var(--text-secondary);
943+
cursor: pointer;
944+
transition: all 0.15s ease;
945+
text-align: center;
946+
}
947+
948+
.copy-snippet-btn:hover {
949+
color: var(--accent-primary);
950+
background: var(--bg-secondary);
883951
}

web/static/js/chat.js

Lines changed: 103 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -430,6 +430,10 @@ class EurusChat {
430430
this.sendBtn.disabled = false;
431431
break;
432432

433+
case 'arraylake_snippet':
434+
this.addArraylakeSnippet(data.content);
435+
break;
436+
433437
case 'error':
434438
this.showError(data.content);
435439
this.sendBtn.disabled = false;
@@ -677,6 +681,105 @@ class EurusChat {
677681
this.currentAssistantMessage = null;
678682
}
679683

684+
addArraylakeSnippet(snippetText) {
685+
// Find the latest assistant message
686+
const messages = this.messagesContainer.querySelectorAll('.assistant-message');
687+
const targetMsg = messages.length > 0 ? messages[messages.length - 1] : null;
688+
if (!targetMsg) return;
689+
690+
// Strip markdown fences for raw code display
691+
let cleanCode = snippetText
692+
.replace(/^\n?📦[^\n]*\n/, '')
693+
.replace(/^```python\n?/, '')
694+
.replace(/\n?```$/, '')
695+
.trim();
696+
697+
// Find the last plot figure's action bar to add button inline
698+
const figures = targetMsg.querySelectorAll('.plot-figure');
699+
const lastFigure = figures.length > 0 ? figures[figures.length - 1] : null;
700+
const actionsDiv = lastFigure ? lastFigure.querySelector('.plot-actions') : null;
701+
702+
if (actionsDiv) {
703+
// Add button inline with Enlarge/Download/Show Code
704+
const btn = document.createElement('button');
705+
btn.className = 'code-btn';
706+
btn.title = 'Arraylake Code';
707+
btn.textContent = '📦 Arraylake Code';
708+
actionsDiv.appendChild(btn);
709+
710+
// Add code block to figure (same pattern as Show Code)
711+
const codeDiv = document.createElement('div');
712+
codeDiv.className = 'plot-code';
713+
codeDiv.style.display = 'none';
714+
715+
const pre = document.createElement('pre');
716+
const codeEl = document.createElement('code');
717+
codeEl.className = 'language-python hljs';
718+
try {
719+
codeEl.innerHTML = hljs.highlight(cleanCode, { language: 'python' }).value;
720+
} catch (e) {
721+
codeEl.textContent = cleanCode;
722+
}
723+
pre.appendChild(codeEl);
724+
codeDiv.appendChild(pre);
725+
726+
// Copy button inside code block
727+
const copyBtn = document.createElement('button');
728+
copyBtn.className = 'copy-snippet-btn';
729+
copyBtn.textContent = 'Copy Code';
730+
copyBtn.addEventListener('click', () => {
731+
navigator.clipboard.writeText(cleanCode).then(() => {
732+
copyBtn.textContent = '✓ Copied!';
733+
setTimeout(() => copyBtn.textContent = 'Copy Code', 2000);
734+
});
735+
});
736+
codeDiv.appendChild(copyBtn);
737+
lastFigure.appendChild(codeDiv);
738+
739+
// Toggle
740+
btn.addEventListener('click', () => {
741+
if (codeDiv.style.display === 'none') {
742+
codeDiv.style.display = 'block';
743+
btn.textContent = '📦 Hide Arraylake';
744+
} else {
745+
codeDiv.style.display = 'none';
746+
btn.textContent = '📦 Arraylake Code';
747+
}
748+
});
749+
} else {
750+
// No plot figure — add as standalone section under the message
751+
const wrapper = document.createElement('div');
752+
wrapper.className = 'arraylake-snippet-section';
753+
wrapper.innerHTML = `
754+
<div class="plot-actions">
755+
<button class="code-btn" title="Arraylake Code">📦 Arraylake Code</button>
756+
</div>
757+
<div class="plot-code" style="display: none;">
758+
<pre><code class="language-python hljs"></code></pre>
759+
</div>
760+
`;
761+
const codeEl = wrapper.querySelector('code');
762+
try {
763+
codeEl.innerHTML = hljs.highlight(cleanCode, { language: 'python' }).value;
764+
} catch (e) {
765+
codeEl.textContent = cleanCode;
766+
}
767+
const btn = wrapper.querySelector('.code-btn');
768+
const codeDiv = wrapper.querySelector('.plot-code');
769+
btn.addEventListener('click', () => {
770+
if (codeDiv.style.display === 'none') {
771+
codeDiv.style.display = 'block';
772+
btn.textContent = '📦 Hide Arraylake';
773+
} else {
774+
codeDiv.style.display = 'none';
775+
btn.textContent = '📦 Arraylake Code';
776+
}
777+
});
778+
targetMsg.appendChild(wrapper);
779+
}
780+
this.scrollToBottom();
781+
}
782+
680783
showError(message) {
681784
this.removeThinkingIndicator();
682785

web/templates/index.html

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -60,5 +60,5 @@ <h3>Cached Datasets</h3>
6060
{% endblock %}
6161

6262
{% block scripts %}
63-
<script src="/static/js/chat.js?v=20260218c"></script>
63+
<script src="/static/js/chat.js?v=20260219b"></script>
6464
{% endblock %}

0 commit comments

Comments
 (0)