File size: 3,358 Bytes
abd4352 c5f9c5f abd4352 c5f9c5f abd4352 c5f9c5f abd4352 c5f9c5f abd4352 c5f9c5f abd4352 c5f9c5f abd4352 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 | """
agent/nodes/executor.py
Executes validated SQL or Pandas code.
Caches results in Upstash Redis (TTL 1 hour).
OPTIMIZATION: Singleton Redis client, removed per-call instantiation overhead.
"""
import hashlib
import json
import os
import time
from functools import lru_cache
from typing import Any, Dict, List, Optional
from agent.state import AgentState
from connectors.base import get_connector
from sandbox.python_sandbox import run_pandas
@lru_cache(maxsize=1)
def _get_redis():
"""Singleton Redis client — created once, reused across all requests."""
try:
from upstash_redis import Redis
return Redis(
url=os.environ["UPSTASH_REDIS_REST_URL"],
token=os.environ["UPSTASH_REDIS_REST_TOKEN"],
)
except Exception:
return None
def _cache_key(connector_id: str, code: str, code_type: str) -> str:
raw = f"{connector_id}:{code_type}:{code}"
return "exec:" + hashlib.sha256(raw.encode()).hexdigest()[:32]
def executor(state: AgentState) -> AgentState:
# Skip if safety block was triggered
if (state.get("execution_error") or "").startswith("SAFETY_BLOCK"):
return state
code = state["generated_code"]
code_type = state["code_type"]
connector_id = state["connector_id"]
# ── Cache lookup ──────────────────────────────────────────────────────────
redis = _get_redis()
cache_key = _cache_key(connector_id, code, code_type)
try:
if redis:
cached = redis.get(cache_key)
if cached:
result = json.loads(cached)
return {**state, "execution_result": result, "from_cache": True, "execution_error": None}
except Exception:
pass # Cache miss or Redis error — proceed to execution
# ── Execute ───────────────────────────────────────────────────────────────
start = time.time()
connector = get_connector(connector_id)
try:
if code_type == "sql":
rows = connector.execute_sql(code)
result = rows[:500] # hard cap
else:
# Pandas: load dataframe from connector then run sandboxed code
df = connector.load_dataframe()
result_df = run_pandas(code, df)
result = json.loads(result_df.to_json(orient="records", date_format="iso"))
latency_ms = int((time.time() - start) * 1000)
# ── Write to cache (non-blocking) ─────────────────────────────────────
try:
if redis:
redis.setex(cache_key, 3600, json.dumps(result))
except Exception:
pass # Non-fatal cache write failure
return {
**state,
"execution_result": result,
"execution_error": None,
"from_cache": False,
"latency_ms": latency_ms,
}
except Exception as exc:
return {
**state,
"execution_result": None,
"execution_error": str(exc),
"from_cache": False,
}
|