decodingtrust-agent-sdk 0.1.0__py3-none-any.whl

dt_arena/utils/bigquery/helpers.py

@@ -0,0 +1,246 @@
+from __future__ import annotations
+
+import json
+import os
+import http.client
+import urllib.parse as urlparse
+from pathlib import Path
+from typing import Any, Dict, List, Optional, Tuple
+
+import yaml
+
+
+def _get_registry() -> dict:
+    root = Path(__file__).resolve().parents[3]
+    registry_file = root / "dt_arena" / "envs" / "registry.yaml"
+    if registry_file.exists():
+        try:
+            return yaml.safe_load(registry_file.read_text()) or {}
+        except Exception:
+            return {}
+    return {}
+
+
+def _get_bigquery_host_port() -> Tuple[str, int]:
+    env_port = os.environ.get("BIGQUERY_API_PORT")
+    if env_port:
+        return "127.0.0.1", int(env_port)
+
+    reg = _get_registry()
+    base = ((reg.get("services") or {}).get("bigquery") or {}).get(
+        "api_base_url", "http://127.0.0.1:9050"
+    )
+    parsed = urlparse.urlparse(base)
+    host = parsed.hostname or "127.0.0.1"
+    port = parsed.port or (443 if (parsed.scheme or "http") == "https" else 80)
+    return host, port
+
+
+def _request(
+    method: str,
+    path: str,
+    *,
+    body: Optional[Dict[str, Any]] = None,
+    timeout: int = 20,
+) -> Tuple[int, Any]:
+    host, port = _get_bigquery_host_port()
+    headers = {"Accept": "application/json"}
+    raw_body: Optional[str] = None
+    if body is not None:
+        headers["Content-Type"] = "application/json"
+        raw_body = json.dumps(body)
+
+    conn = http.client.HTTPConnection(host, port, timeout=timeout)
+    try:
+        conn.request(method.upper(), path, body=raw_body, headers=headers)
+        resp = conn.getresponse()
+        data = resp.read()
+        status = resp.status
+    finally:
+        conn.close()
+
+    text = data.decode("utf-8", errors="replace") if data else ""
+    if not text:
+        return status, {}
+    try:
+        return status, json.loads(text)
+    except json.JSONDecodeError:
+        return status, text
+
+
+def _json_value_to_python(value: Any, field: Dict[str, Any]) -> Any:
+    if value is None:
+        return None
+
+    mode = (field.get("mode") or "NULLABLE").upper()
+    field_type = (field.get("type") or "STRING").upper()
+
+    if mode == "REPEATED":
+        items = value if isinstance(value, list) else []
+        repeated_field = dict(field)
+        repeated_field["mode"] = "NULLABLE"
+        return [_json_value_to_python(item.get("v") if isinstance(item, dict) else item, repeated_field) for item in items]
+
+    if field_type in {"RECORD", "STRUCT"}:
+        nested_fields = field.get("fields") or []
+        nested_values = value.get("f") if isinstance(value, dict) else []
+        out: Dict[str, Any] = {}
+        for nested_field, nested_value in zip(nested_fields, nested_values):
+            cell_value = nested_value.get("v") if isinstance(nested_value, dict) else nested_value
+            out[nested_field["name"]] = _json_value_to_python(cell_value, nested_field)
+        return out
+
+    return value
+
+
+def _decode_rows(response: Dict[str, Any]) -> List[Dict[str, Any]]:
+    schema_fields = ((response.get("schema") or {}).get("fields") or [])
+    rows = response.get("rows") or []
+    decoded: List[Dict[str, Any]] = []
+
+    for row in rows:
+        cols = row.get("f") if isinstance(row, dict) else []
+        item: Dict[str, Any] = {}
+        for field, col in zip(schema_fields, cols):
+            value = col.get("v") if isinstance(col, dict) else col
+            item[field["name"]] = _json_value_to_python(value, field)
+        decoded.append(item)
+    return decoded
+
+
+def get_table(table_ref: str) -> Dict[str, str]:
+    project_id, rest = table_ref.split(".", 1)
+    dataset_id, table_id = rest.split(".", 1)
+    return {
+        "project_id": project_id,
+        "dataset_id": dataset_id,
+        "table_id": table_id,
+    }
+
+
+def create_dataset(
+    project_id: str,
+    dataset_id: str,
+    *,
+    location: Optional[str] = None,
+    description: Optional[str] = None,
+    exists_ok: bool = True,
+) -> Dict[str, Any]:
+    body: Dict[str, Any] = {
+        "datasetReference": {
+            "projectId": project_id,
+            "datasetId": dataset_id,
+        }
+    }
+    if location:
+        body["location"] = location
+    if description:
+        body["description"] = description
+
+    status, payload = _request(
+        "POST",
+        f"/bigquery/v2/projects/{project_id}/datasets",
+        body=body,
+    )
+    if status == 200:
+        return payload
+    if status == 409 and exists_ok:
+        _, existing = _request(
+            "GET",
+            f"/bigquery/v2/projects/{project_id}/datasets/{dataset_id}",
+        )
+        return existing
+    if (
+        exists_ok
+        and status == 500
+        and isinstance(payload, dict)
+        and "already created" in json.dumps(payload).lower()
+    ):
+        _, existing = _request(
+            "GET",
+            f"/bigquery/v2/projects/{project_id}/datasets/{dataset_id}",
+        )
+        return existing
+    raise RuntimeError(f"BigQuery create_dataset failed ({status}): {payload}")
+
+
+def list_datasets(project_id: str) -> List[Dict[str, Any]]:
+    status, payload = _request("GET", f"/bigquery/v2/projects/{project_id}/datasets")
+    if status != 200:
+        raise RuntimeError(f"BigQuery list_datasets failed ({status}): {payload}")
+    return payload.get("datasets") or []
+
+
+def list_tables(project_id: str, dataset_id: str) -> List[Dict[str, Any]]:
+    status, payload = _request(
+        "GET",
+        f"/bigquery/v2/projects/{project_id}/datasets/{dataset_id}/tables",
+    )
+    if status != 200:
+        raise RuntimeError(f"BigQuery list_tables failed ({status}): {payload}")
+    return payload.get("tables") or []
+
+
+def get_table_schema(project_id: str, dataset_id: str, table_id: str) -> List[Dict[str, Any]]:
+    status, payload = _request(
+        "GET",
+        f"/bigquery/v2/projects/{project_id}/datasets/{dataset_id}/tables/{table_id}",
+    )
+    if status != 200:
+        raise RuntimeError(f"BigQuery get_table failed ({status}): {payload}")
+    return ((payload.get("schema") or {}).get("fields") or [])
+
+
+def execute_sql(
+    project_id: str,
+    query: str,
+    *,
+    location: Optional[str] = None,
+    timeout_ms: int = 20000,
+    max_results: int = 1000,
+) -> Dict[str, Any]:
+    body: Dict[str, Any] = {
+        "query": query,
+        "useLegacySql": False,
+        "timeoutMs": timeout_ms,
+        "maxResults": max_results,
+    }
+    if location:
+        body["location"] = location
+
+    status, payload = _request(
+        "POST",
+        f"/bigquery/v2/projects/{project_id}/queries",
+        body=body,
+        timeout=max(20, timeout_ms // 1000 + 10),
+    )
+    if status != 200:
+        raise RuntimeError(f"BigQuery execute_sql failed ({status}): {payload}")
+
+    return {
+        "jobComplete": payload.get("jobComplete", True),
+        "jobReference": payload.get("jobReference", {}),
+        "schema": ((payload.get("schema") or {}).get("fields") or []),
+        "rows": _decode_rows(payload),
+        "totalRows": int(payload.get("totalRows") or 0),
+        "totalBytesProcessed": payload.get("totalBytesProcessed"),
+        "cacheHit": payload.get("cacheHit"),
+        "errors": payload.get("errors") or [],
+        "raw": payload,
+    }
+
+
+def execute_write_sql(
+    project_id: str,
+    query: str,
+    *,
+    location: Optional[str] = None,
+    timeout_ms: int = 20000,
+) -> Dict[str, Any]:
+    return execute_sql(
+        project_id=project_id,
+        query=query,
+        location=location,
+        timeout_ms=timeout_ms,
+        max_results=1000,
+    )

dt_arena/utils/calendar/__init__.py

@@ -0,0 +1 @@
+"""Calendar helpers package."""

dt_arena/utils/calendar/helpers.py

@@ -0,0 +1,87 @@
+from __future__ import annotations
+import json
+import http.client
+import urllib.parse as urlparse
+from pathlib import Path
+from typing import List, Dict, Any, Tuple
+import yaml
+
+
+def _get_registry() -> dict:
+    root = Path(__file__).resolve().parents[3]
+    registry_file = root / "dt_arena" / "envs" / "registry.yaml"
+    if registry_file.exists():
+        try:
+            return yaml.safe_load(registry_file.read_text()) or {}
+        except Exception:
+            return {}
+    return {}
+
+
+def _get_calendar_host_port() -> Tuple[str, int]:
+    import os
+    # First check for environment variable (set by task_runner for dynamic ports)
+    env_port = os.environ.get("CALENDAR_API_PORT")
+    if env_port:
+        return "127.0.0.1", int(env_port)
+    
+    # Fall back to registry
+    reg = _get_registry()
+    base = ((reg.get("services") or {}).get("calendar") or {}).get(
+        "api_base_url", "http://127.0.0.1:8032"
+    )
+    parsed = urlparse.urlparse(base)
+    host = parsed.hostname or "127.0.0.1"
+    port = parsed.port or (443 if (parsed.scheme or "http") == "https" else 80)
+    return host, port
+
+
+def login(email: str, password: str) -> str:
+    host, port = _get_calendar_host_port()
+    body = json.dumps({"email": email, "password": password})
+    headers = {"Content-Type": "application/json"}
+    conn = http.client.HTTPConnection(host, port, timeout=5)
+    conn.request("POST", "/api/v1/auth/login", body=body, headers=headers)
+    resp = conn.getresponse()
+    data = resp.read()
+    conn.close()
+    if resp.status == 200:
+        try:
+            return json.loads(data.decode()).get("access_token") or ""
+        except Exception:
+            return ""
+    return ""
+
+
+def list_events(token: str, start: str | None = None, end: str | None = None, limit: int = 50) -> List[Dict[str, Any]]:
+    host, port = _get_calendar_host_port()
+    headers = {"Authorization": f"Bearer {token}"}
+    conn = http.client.HTTPConnection(host, port, timeout=5)
+    # Try all calendars first
+    conn.request("GET", "/calendar/v3/users/me/calendarList", headers=headers)
+    resp = conn.getresponse()
+    cal_data = resp.read()
+    conn.close()
+    if resp.status != 200:
+        return []
+    
+    all_events = []
+    try:
+        calendars = json.loads(cal_data.decode())
+        if isinstance(calendars, dict):
+            calendars = calendars.get("items", [])
+        for cal in calendars:
+            cal_id = cal.get("id", "primary")
+            conn = http.client.HTTPConnection(host, port, timeout=5)
+            conn.request("GET", f"/calendar/v3/calendars/{cal_id}/events", headers=headers)
+            resp = conn.getresponse()
+            events_data = resp.read()
+            conn.close()
+            if resp.status == 200:
+                events = json.loads(events_data.decode())
+                if isinstance(events, dict):
+                    events = events.get("items", [])
+                all_events.extend(events)
+    except Exception:
+        pass
+    return all_events

dt_arena/utils/customer_service/__init__.py

@@ -0,0 +1,17 @@
+"""Customer service domain utilities for judge evaluation."""
+
+from dt_arena.utils.customer_service.helpers import (
+    TaskConfig,
+    load_task_config,
+    eval_attack,
+    eval_task_generic,
+)
+from dt_arena.utils.customer_service.cs_env_client import CSEnvClient
+
+__all__ = [
+    "TaskConfig",
+    "load_task_config",
+    "eval_attack",
+    "eval_task_generic",
+    "CSEnvClient",
+]