a2a-source-edit: write agent.py
This commit is contained in:
a2a-cloud
680
agent.py
680
agent.py
@@ -1,189 +1,581 @@
|
||||
"""csv-answers-studio-v1 agent.
|
||||
"""CSVAnswers Studio full-stack A2A agent.
|
||||
|
||||
Starter stack:
|
||||
- DeepAgents for tool-calling orchestration
|
||||
- Caller-provided LLM credentials via ctx.llm
|
||||
- A tiny model-call middleware hook you can replace with tracing,
|
||||
routing, rate limits, or policy checks
|
||||
Deterministic, no-LLM CSV analysis with PlatformUserAuth, typed uploads,
|
||||
bounded browser base64 bridge, managed Postgres persistence, and receipt rows.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import base64
|
||||
import csv
|
||||
import hashlib
|
||||
import io
|
||||
import json
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
import os
|
||||
import re
|
||||
import uuid
|
||||
from datetime import datetime, timezone
|
||||
from decimal import Decimal, InvalidOperation
|
||||
from typing import Annotated, Any
|
||||
|
||||
from pydantic import BaseModel
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
import a2a_pack as a2a
|
||||
from a2a_pack import (
|
||||
A2AAgent,
|
||||
LLMProvisioning,
|
||||
{{ auth_type }},
|
||||
AgentDatabase,
|
||||
AgentDatabaseEnv,
|
||||
AgentDatabaseMigrations,
|
||||
AgentPlatformResources,
|
||||
FileUpload,
|
||||
PlatformUserAuth,
|
||||
Pricing,
|
||||
Resources,
|
||||
RunContext,
|
||||
WorkspaceAccess,
|
||||
WorkspaceMode,
|
||||
State,
|
||||
UploadedFile,
|
||||
)
|
||||
from a2a_pack.context import LLMCreds
|
||||
|
||||
MAX_CSV_BYTES = 64 * 1024
|
||||
MAX_ROWS = 500
|
||||
MAX_COLUMNS = 50
|
||||
MAX_CELL_CHARS = 2_000
|
||||
CSV_MEDIA_TYPES = {"text/csv", "application/csv", "text/plain", "application/vnd.ms-excel"}
|
||||
SAFE_ID_RE = re.compile(r"^[A-Za-z0-9][A-Za-z0-9_.:-]{0,127}$")
|
||||
|
||||
# Local-only fallback lets sandbox/unit tests prove behavior without receiving
|
||||
# managed DATABASE_URL. Hosted deployments use managed Postgres exclusively.
|
||||
_LOCAL_ANALYSES: dict[tuple[str, str], dict[str, Any]] = {}
|
||||
_LOCAL_RECEIPTS: list[dict[str, Any]] = []
|
||||
|
||||
|
||||
class CsvAnswersStudioV1Config(BaseModel):
|
||||
pass
|
||||
|
||||
|
||||
SYSTEM_PROMPT = """\
|
||||
You are a compact tool-calling agent.
|
||||
|
||||
Use the text_stats tool when the user asks about text, counts, summaries,
|
||||
or anything where exact length/word numbers would help. Mention tool results
|
||||
briefly instead of dumping raw JSON.
|
||||
"""
|
||||
|
||||
RUNTIME_SKILLS_DIR = "csv-answers-studio-v1/.deepagents/skills/"
|
||||
DEEPAGENTS_RECURSION_LIMIT = 500
|
||||
class BrowserCsvUpload(BaseModel):
|
||||
filename: str = Field(..., min_length=1, max_length=180)
|
||||
media_type: str = Field(..., min_length=1, max_length=120)
|
||||
data_base64: str = Field(..., min_length=1, max_length=((MAX_CSV_BYTES * 4) // 3) + 16)
|
||||
|
||||
|
||||
class CsvAnswersStudioV1(A2AAgent[CsvAnswersStudioV1Config, {{ auth_type }}]):
|
||||
class CsvAnalysisResult(BaseModel):
|
||||
ok: bool
|
||||
analysis_id: str
|
||||
answer: str | None = None
|
||||
row_count: int = 0
|
||||
column_count: int = 0
|
||||
question: str | None = None
|
||||
source_rows: list[dict[str, Any]] = Field(default_factory=list)
|
||||
chart_data: list[dict[str, Any]] = Field(default_factory=list)
|
||||
numeric_column: str | None = None
|
||||
row_label_column: str | None = None
|
||||
saved: bool = False
|
||||
receipt_id: str | None = None
|
||||
code: str | None = None
|
||||
message: str | None = None
|
||||
warnings: list[str] = Field(default_factory=list)
|
||||
|
||||
|
||||
class CsvAnalysisListResult(BaseModel):
|
||||
ok: bool
|
||||
analyses: list[dict[str, Any]] = Field(default_factory=list)
|
||||
code: str | None = None
|
||||
message: str | None = None
|
||||
|
||||
|
||||
class CsvAnswersStudioV1(A2AAgent[CsvAnswersStudioV1Config, PlatformUserAuth]):
|
||||
name = "csv-answers-studio-v1"
|
||||
description = "CSVAnswers: upload or paste a bounded CSV, ask grounded highest-value questions, persist user-scoped analyses, and reopen them in a polished one-page app."
|
||||
description = (
|
||||
"CSVAnswers: upload or paste a bounded CSV, ask which row has the "
|
||||
"highest numeric value, persist the grounded answer in user-scoped "
|
||||
"managed Postgres, and reopen it from a polished one-page app."
|
||||
)
|
||||
version = "0.1.0"
|
||||
|
||||
config_model = CsvAnswersStudioV1Config
|
||||
auth_model = {{ auth_type }}
|
||||
auth_model = PlatformUserAuth
|
||||
|
||||
# Hosted generated agents read the caller's saved LLM credential through
|
||||
# ctx.llm. The platform may proxy that credential through LiteLLM, but agent
|
||||
# code never reads provider keys, LiteLLM master keys, or OPENAI_API_KEY
|
||||
# directly.
|
||||
llm_provisioning = LLMProvisioning.PLATFORM
|
||||
state = State.DURABLE
|
||||
resources = Resources(cpu="500m", memory="512Mi", max_runtime_seconds=120)
|
||||
pricing = Pricing(
|
||||
price_per_call_usd=0.0,
|
||||
caller_pays_llm=True,
|
||||
notes="Starter agent uses the caller's saved LLM credential via ctx.llm.",
|
||||
caller_pays_llm=False,
|
||||
notes="Deterministic CSV parsing and managed Postgres persistence. No LLM credential required.",
|
||||
)
|
||||
workspace_access = WorkspaceAccess.dynamic(
|
||||
max_files=64,
|
||||
allowed_modes=(WorkspaceMode.READ_ONLY, WorkspaceMode.READ_WRITE_OVERLAY),
|
||||
require_reason=False,
|
||||
)
|
||||
tools_used = ("deepagents", "langchain")
|
||||
|
||||
@a2a.tool(description="Ask the starter DeepAgent to answer with tool calls when useful")
|
||||
async def ask(self, ctx: RunContext[{{ auth_type }}], prompt: str) -> str:
|
||||
creds = ctx.llm
|
||||
await ctx.emit_progress(f"llm: {creds.model} via {creds.source}")
|
||||
if not creds.api_key:
|
||||
return (
|
||||
"LLM key required. Add an LLM credential in Settings > LLM "
|
||||
"credentials before running this agent; for local --invoke "
|
||||
"runs set AGENT_LLM_KEY."
|
||||
)
|
||||
graph = self._build_deep_agent(ctx=ctx, creds=creds)
|
||||
state = await graph.ainvoke(
|
||||
{"messages": [{"role": "user", "content": prompt}]},
|
||||
config={"recursion_limit": DEEPAGENTS_RECURSION_LIMIT},
|
||||
tools_used = ("python-csv", "managed-postgres")
|
||||
platform_resources = AgentPlatformResources(
|
||||
databases=(
|
||||
AgentDatabase(
|
||||
name="csv-answers-studio-v1-data",
|
||||
scope="user",
|
||||
access_mode="read_write",
|
||||
env=AgentDatabaseEnv(url="DATABASE_URL"),
|
||||
migrations=AgentDatabaseMigrations(path="db/migrations"),
|
||||
),
|
||||
)
|
||||
await ctx.emit_progress("deepagent finished")
|
||||
return _last_message_text(state)
|
||||
)
|
||||
|
||||
def _build_deep_agent(
|
||||
@a2a.tool(
|
||||
description="Analyze bounded pasted CSV text, answer the highest numeric row question, save the result, and persist a receipt.",
|
||||
timeout_seconds=30,
|
||||
idempotent=True,
|
||||
cost_class="cheap",
|
||||
)
|
||||
async def analyze_csv(
|
||||
self,
|
||||
*,
|
||||
ctx: RunContext[{{ auth_type }}],
|
||||
creds: LLMCreds,
|
||||
) -> Any:
|
||||
# Lazy imports keep `a2a card` usable before local dependencies are
|
||||
# installed. `a2a deploy` installs requirements.txt during the build.
|
||||
from a2a_pack.deepagents import create_a2a_deep_agent
|
||||
from langchain.agents.middleware import wrap_model_call
|
||||
from langchain_core.tools import tool
|
||||
ctx: RunContext[PlatformUserAuth],
|
||||
analysis_id: str,
|
||||
csv_text: str,
|
||||
question: str,
|
||||
) -> CsvAnalysisResult:
|
||||
tenant = _tenant_key(ctx)
|
||||
result = _analyze_csv_text(analysis_id=analysis_id, csv_text=csv_text, question=question)
|
||||
if not result.ok:
|
||||
return result
|
||||
saved, receipt_id = _persist_analysis_and_receipt(
|
||||
tenant_key=tenant,
|
||||
analysis_id=analysis_id,
|
||||
question=question,
|
||||
csv_text=csv_text,
|
||||
result=result.model_dump(mode="json"),
|
||||
source="pasted_csv",
|
||||
)
|
||||
result.saved = saved
|
||||
result.receipt_id = receipt_id
|
||||
await ctx.emit_progress(f"saved analysis {analysis_id} for {tenant}")
|
||||
return result
|
||||
|
||||
@tool
|
||||
def text_stats(text: str) -> str:
|
||||
"""Return exact word, character, and line counts for text."""
|
||||
words = [part for part in text.split() if part.strip()]
|
||||
return json.dumps(
|
||||
@a2a.tool(
|
||||
description="Analyze a bounded browser CSV upload encoded as base64 JSON, then save the result and receipt.",
|
||||
timeout_seconds=30,
|
||||
idempotent=True,
|
||||
cost_class="cheap",
|
||||
)
|
||||
async def analyze_csv_upload_base64(
|
||||
self,
|
||||
ctx: RunContext[PlatformUserAuth],
|
||||
analysis_id: str,
|
||||
upload: BrowserCsvUpload,
|
||||
question: str,
|
||||
) -> CsvAnalysisResult:
|
||||
decoded = _decode_browser_upload(upload)
|
||||
if isinstance(decoded, CsvAnalysisResult):
|
||||
return decoded.model_copy(update={"analysis_id": analysis_id})
|
||||
return await self.analyze_csv(ctx, analysis_id=analysis_id, csv_text=decoded, question=question)
|
||||
|
||||
@a2a.tool(
|
||||
description="Analyze a typed external FileUpload CSV, then save the result and receipt.",
|
||||
timeout_seconds=30,
|
||||
idempotent=True,
|
||||
cost_class="cheap",
|
||||
)
|
||||
async def analyze_csv_file(
|
||||
self,
|
||||
ctx: RunContext[PlatformUserAuth],
|
||||
analysis_id: str,
|
||||
file: Annotated[
|
||||
UploadedFile,
|
||||
FileUpload(
|
||||
accept=("text/csv", "text/plain", "application/vnd.ms-excel"),
|
||||
max_bytes=MAX_CSV_BYTES,
|
||||
description="Bounded CSV file for CSVAnswers analysis.",
|
||||
),
|
||||
],
|
||||
question: str,
|
||||
) -> CsvAnalysisResult:
|
||||
if file.size_bytes and file.size_bytes > MAX_CSV_BYTES:
|
||||
return _error(analysis_id, "csv_too_large", f"CSV must be at most {MAX_CSV_BYTES} bytes.")
|
||||
if file.media_type not in CSV_MEDIA_TYPES:
|
||||
return _error(analysis_id, "unsupported_media_type", "Upload must be a CSV or plain text file.")
|
||||
try:
|
||||
raw = await _read_uploaded_file(ctx, file)
|
||||
except Exception:
|
||||
return _error(analysis_id, "upload_unreadable", "The uploaded CSV could not be read from the workspace.")
|
||||
try:
|
||||
csv_text = raw.decode("utf-8-sig")
|
||||
except UnicodeDecodeError:
|
||||
return _error(analysis_id, "invalid_encoding", "CSV must be valid UTF-8 text.")
|
||||
return await self.analyze_csv(ctx, analysis_id=analysis_id, csv_text=csv_text, question=question)
|
||||
|
||||
@a2a.tool(
|
||||
description="Reopen a saved user-scoped CSV analysis by id.",
|
||||
timeout_seconds=15,
|
||||
idempotent=True,
|
||||
cost_class="cheap",
|
||||
)
|
||||
async def get_analysis(
|
||||
self,
|
||||
ctx: RunContext[PlatformUserAuth],
|
||||
analysis_id: str,
|
||||
) -> CsvAnalysisResult:
|
||||
tenant = _tenant_key(ctx)
|
||||
record = _load_analysis(tenant, analysis_id)
|
||||
if record is None:
|
||||
return _error(analysis_id, "not_found", "No saved analysis with that id exists for this signed-in user.")
|
||||
return CsvAnalysisResult.model_validate(record)
|
||||
|
||||
@a2a.tool(
|
||||
description="List recent saved CSV analyses for the signed-in user.",
|
||||
timeout_seconds=15,
|
||||
idempotent=True,
|
||||
cost_class="cheap",
|
||||
)
|
||||
async def list_analyses(
|
||||
self,
|
||||
ctx: RunContext[PlatformUserAuth],
|
||||
limit: int = 10,
|
||||
) -> CsvAnalysisListResult:
|
||||
tenant = _tenant_key(ctx)
|
||||
safe_limit = min(max(int(limit), 1), 50)
|
||||
return CsvAnalysisListResult(ok=True, analyses=_list_analyses(tenant, safe_limit))
|
||||
|
||||
|
||||
def _tenant_key(ctx: RunContext[PlatformUserAuth]) -> str:
|
||||
stable_id = ctx.auth.user_id if ctx.auth.user_id is not None else ctx.auth.sub
|
||||
if not stable_id:
|
||||
raise PermissionError("stable platform identity required")
|
||||
return f"user:{stable_id}"
|
||||
|
||||
|
||||
def _validate_analysis_id(analysis_id: str) -> str | None:
|
||||
cleaned = (analysis_id or "").strip()
|
||||
if not SAFE_ID_RE.fullmatch(cleaned):
|
||||
return "analysis_id must be 1-128 characters and contain only letters, numbers, dots, underscores, colons, or hyphens."
|
||||
return None
|
||||
|
||||
|
||||
def _error(analysis_id: str, code: str, message: str) -> CsvAnalysisResult:
|
||||
return CsvAnalysisResult(ok=False, analysis_id=analysis_id, code=code, message=message, answer=None)
|
||||
|
||||
|
||||
def _decode_browser_upload(upload: BrowserCsvUpload) -> str | CsvAnalysisResult:
|
||||
filename = upload.filename.replace("\\", "/").split("/")[-1]
|
||||
if not filename or filename in {".", ".."}:
|
||||
return _error("", "invalid_filename", "Upload filename is invalid.")
|
||||
if upload.media_type not in CSV_MEDIA_TYPES:
|
||||
return _error("", "unsupported_media_type", "Upload must be a CSV or plain text file.")
|
||||
try:
|
||||
raw = base64.b64decode(upload.data_base64, validate=True)
|
||||
except Exception:
|
||||
return _error("", "invalid_base64", "Upload data_base64 must be valid base64.")
|
||||
if len(raw) > MAX_CSV_BYTES:
|
||||
return _error("", "csv_too_large", f"CSV must be at most {MAX_CSV_BYTES} bytes.")
|
||||
try:
|
||||
return raw.decode("utf-8-sig")
|
||||
except UnicodeDecodeError:
|
||||
return _error("", "invalid_encoding", "CSV must be valid UTF-8 text.")
|
||||
|
||||
|
||||
async def _read_uploaded_file(ctx: RunContext[PlatformUserAuth], file: UploadedFile) -> bytes:
|
||||
workspace = ctx.workspace
|
||||
reader = getattr(workspace, "read_bytes", None)
|
||||
if reader is not None:
|
||||
value = reader(file.path)
|
||||
if hasattr(value, "__await__"):
|
||||
return await value
|
||||
return value
|
||||
view = await workspace.open_view(
|
||||
purpose="Read uploaded CSV",
|
||||
hints=(file.path,),
|
||||
max_files=1,
|
||||
reason="Read the uploaded CSV staged for this invocation.",
|
||||
)
|
||||
return await view.read(file.path)
|
||||
|
||||
|
||||
def _analyze_csv_text(*, analysis_id: str, csv_text: str, question: str) -> CsvAnalysisResult:
|
||||
id_error = _validate_analysis_id(analysis_id)
|
||||
if id_error:
|
||||
return _error(analysis_id, "invalid_analysis_id", id_error)
|
||||
if not question or not question.strip():
|
||||
return _error(analysis_id, "missing_question", "Question is required.")
|
||||
try:
|
||||
encoded = csv_text.encode("utf-8")
|
||||
except UnicodeEncodeError:
|
||||
return _error(analysis_id, "invalid_encoding", "CSV must be valid UTF-8 text.")
|
||||
if not encoded:
|
||||
return _error(analysis_id, "empty_csv", "CSV text is required.")
|
||||
if len(encoded) > MAX_CSV_BYTES:
|
||||
return _error(analysis_id, "csv_too_large", f"CSV must be at most {MAX_CSV_BYTES} bytes.")
|
||||
|
||||
try:
|
||||
rows = _parse_csv_strict(csv_text)
|
||||
except csv.Error:
|
||||
return _error(analysis_id, "malformed_csv", "CSV is malformed: unterminated quoted field or invalid CSV structure.")
|
||||
if not rows:
|
||||
return _error(analysis_id, "empty_csv", "CSV must contain a header row and at least one data row.")
|
||||
header = rows[0]
|
||||
data_rows = rows[1:]
|
||||
if len(header) < 2 or not any(cell.strip() for cell in header):
|
||||
return _error(analysis_id, "missing_header", "CSV must include a header row with at least two columns.")
|
||||
if len(header) > MAX_COLUMNS:
|
||||
return _error(analysis_id, "too_many_columns", f"CSV may contain at most {MAX_COLUMNS} columns.")
|
||||
if not data_rows:
|
||||
return _error(analysis_id, "no_data_rows", "CSV must include at least one data row.")
|
||||
if len(data_rows) > MAX_ROWS:
|
||||
return _error(analysis_id, "too_many_rows", f"CSV may contain at most {MAX_ROWS} data rows.")
|
||||
if len(set(header)) != len(header):
|
||||
return _error(analysis_id, "duplicate_headers", "CSV headers must be unique.")
|
||||
for row in data_rows:
|
||||
if len(row) != len(header):
|
||||
return _error(analysis_id, "ragged_rows", "Every CSV row must have the same number of columns as the header.")
|
||||
if any(len(cell) > MAX_CELL_CHARS for cell in row):
|
||||
return _error(analysis_id, "cell_too_large", f"Each CSV cell must be at most {MAX_CELL_CHARS} characters.")
|
||||
|
||||
records = [dict(zip(header, row, strict=True)) for row in data_rows]
|
||||
numeric_columns = _numeric_columns(records, header)
|
||||
if not numeric_columns:
|
||||
return _error(analysis_id, "no_numeric_values", "CSV must contain at least one numeric column.")
|
||||
numeric_column = _choose_numeric_column(question, numeric_columns)
|
||||
row_label_column = _choose_label_column(header, numeric_columns)
|
||||
|
||||
best_record: dict[str, str] | None = None
|
||||
best_value: Decimal | None = None
|
||||
for record in records:
|
||||
value = _parse_number(record.get(numeric_column, ""))
|
||||
if value is None:
|
||||
continue
|
||||
if best_value is None or value > best_value:
|
||||
best_value = value
|
||||
best_record = record
|
||||
if best_record is None or best_value is None:
|
||||
return _error(analysis_id, "no_numeric_values", f"Column {numeric_column!r} does not contain numeric values.")
|
||||
|
||||
label = best_record.get(row_label_column) or f"row {records.index(best_record) + 1}"
|
||||
value_text = _format_decimal(best_value)
|
||||
answer = f"{label} has the highest {numeric_column} at {value_text}."
|
||||
chart_data = []
|
||||
for idx, record in enumerate(records):
|
||||
value = _parse_number(record.get(numeric_column, ""))
|
||||
if value is not None:
|
||||
chart_data.append(
|
||||
{
|
||||
"characters": len(text),
|
||||
"words": len(words),
|
||||
"lines": len(text.splitlines()) or 1,
|
||||
"label": record.get(row_label_column) or f"row {idx + 1}",
|
||||
"value": float(value),
|
||||
"source_row_number": idx + 1,
|
||||
}
|
||||
)
|
||||
source_rows = [
|
||||
{
|
||||
"row_number": records.index(best_record) + 1,
|
||||
"values": best_record,
|
||||
"matched_column": numeric_column,
|
||||
"matched_value": float(best_value),
|
||||
}
|
||||
]
|
||||
return CsvAnalysisResult(
|
||||
ok=True,
|
||||
analysis_id=analysis_id,
|
||||
answer=answer,
|
||||
row_count=len(records),
|
||||
column_count=len(header),
|
||||
question=question,
|
||||
source_rows=source_rows,
|
||||
chart_data=chart_data,
|
||||
numeric_column=numeric_column,
|
||||
row_label_column=row_label_column,
|
||||
saved=False,
|
||||
)
|
||||
|
||||
@wrap_model_call
|
||||
async def log_model_call(request: Any, handler: Any) -> Any:
|
||||
messages = request.state.get("messages", [])
|
||||
print(
|
||||
"[middleware] model_call "
|
||||
f"model={creds.model} source={creds.source} messages={len(messages)}"
|
||||
|
||||
def _parse_csv_strict(csv_text: str) -> list[list[str]]:
|
||||
stream = io.StringIO(csv_text, newline="")
|
||||
sample = csv_text[:2048]
|
||||
try:
|
||||
dialect = csv.Sniffer().sniff(sample) if sample.strip() else csv.excel
|
||||
except csv.Error:
|
||||
dialect = csv.excel
|
||||
reader = csv.reader(stream, dialect=dialect, strict=True)
|
||||
rows = [[cell.strip() for cell in row] for row in reader]
|
||||
if csv_text.count('"') % 2:
|
||||
raise csv.Error("unterminated quoted field")
|
||||
return [row for row in rows if any(cell != "" for cell in row)]
|
||||
|
||||
|
||||
def _parse_number(value: Any) -> Decimal | None:
|
||||
text = str(value).strip().replace(",", "")
|
||||
if not text:
|
||||
return None
|
||||
if text.endswith("%"):
|
||||
text = text[:-1]
|
||||
if text.startswith("$"):
|
||||
text = text[1:]
|
||||
try:
|
||||
number = Decimal(text)
|
||||
except (InvalidOperation, ValueError):
|
||||
return None
|
||||
return number if number.is_finite() else None
|
||||
|
||||
|
||||
def _numeric_columns(records: list[dict[str, str]], header: list[str]) -> list[str]:
|
||||
out: list[str] = []
|
||||
for column in header:
|
||||
if any(_parse_number(record.get(column, "")) is not None for record in records):
|
||||
out.append(column)
|
||||
return out
|
||||
|
||||
|
||||
def _choose_numeric_column(question: str, numeric_columns: list[str]) -> str:
|
||||
q = question.lower()
|
||||
for column in numeric_columns:
|
||||
normalized = column.lower().replace("_", " ").replace("-", " ")
|
||||
if normalized in q or column.lower() in q:
|
||||
return column
|
||||
return numeric_columns[0]
|
||||
|
||||
|
||||
def _choose_label_column(header: list[str], numeric_columns: list[str]) -> str:
|
||||
for column in header:
|
||||
if column not in numeric_columns:
|
||||
return column
|
||||
return header[0]
|
||||
|
||||
|
||||
def _format_decimal(value: Decimal) -> str:
|
||||
if value == value.to_integral_value():
|
||||
return str(value.quantize(Decimal(1)))
|
||||
return format(value.normalize(), "f")
|
||||
|
||||
|
||||
def _connect():
|
||||
url = os.environ.get("DATABASE_URL")
|
||||
if not url:
|
||||
return None
|
||||
import psycopg
|
||||
|
||||
return psycopg.connect(
|
||||
url,
|
||||
options="-c statement_timeout=10000 -c lock_timeout=5000 -c idle_in_transaction_session_timeout=30000",
|
||||
)
|
||||
|
||||
|
||||
def _persist_analysis_and_receipt(
|
||||
*,
|
||||
tenant_key: str,
|
||||
analysis_id: str,
|
||||
question: str,
|
||||
csv_text: str,
|
||||
result: dict[str, Any],
|
||||
source: str,
|
||||
) -> tuple[bool, str]:
|
||||
receipt_id = _receipt_id(tenant_key, analysis_id, csv_text, question)
|
||||
payload = dict(result)
|
||||
payload.update({"saved": True, "receipt_id": receipt_id})
|
||||
receipt = {
|
||||
"receipt_id": receipt_id,
|
||||
"analysis_id": analysis_id,
|
||||
"tenant_key_hash": hashlib.sha256(tenant_key.encode()).hexdigest(),
|
||||
"input_sha256": hashlib.sha256(csv_text.encode("utf-8")).hexdigest(),
|
||||
"question_sha256": hashlib.sha256(question.encode("utf-8")).hexdigest(),
|
||||
"source": source,
|
||||
"created_at": datetime.now(timezone.utc).isoformat(),
|
||||
"ok": True,
|
||||
}
|
||||
conn = _connect()
|
||||
if conn is None:
|
||||
_LOCAL_ANALYSES[(tenant_key, analysis_id)] = payload
|
||||
_LOCAL_RECEIPTS.append(receipt)
|
||||
return True, receipt_id
|
||||
with conn:
|
||||
with conn.cursor() as cur:
|
||||
cur.execute(
|
||||
"""
|
||||
INSERT INTO csv_analyses (tenant_key, analysis_id, question, answer, row_count, column_count, payload, updated_at)
|
||||
VALUES (%s, %s, %s, %s, %s, %s, %s::jsonb, NOW())
|
||||
ON CONFLICT (tenant_key, analysis_id) DO UPDATE SET
|
||||
question = EXCLUDED.question,
|
||||
answer = EXCLUDED.answer,
|
||||
row_count = EXCLUDED.row_count,
|
||||
column_count = EXCLUDED.column_count,
|
||||
payload = EXCLUDED.payload,
|
||||
updated_at = NOW()
|
||||
""",
|
||||
(
|
||||
tenant_key,
|
||||
analysis_id,
|
||||
question,
|
||||
payload.get("answer"),
|
||||
int(payload.get("row_count") or 0),
|
||||
int(payload.get("column_count") or 0),
|
||||
json.dumps(payload),
|
||||
),
|
||||
)
|
||||
return await handler(request)
|
||||
|
||||
backend = ctx.workspace_backend()
|
||||
skill_sources = _seed_runtime_skills(backend, ctx)
|
||||
# create_a2a_deep_agent resolves provider:model strings with
|
||||
# langchain.init_chat_model from ctx.llm, preserving LiteLLM routing,
|
||||
# provider-specific extra body, and runtime model overrides.
|
||||
return create_a2a_deep_agent(
|
||||
ctx,
|
||||
creds=creds,
|
||||
backend=backend,
|
||||
skills=skill_sources or None,
|
||||
tools=[text_stats],
|
||||
middleware=[log_model_call],
|
||||
system_prompt=SYSTEM_PROMPT,
|
||||
)
|
||||
cur.execute(
|
||||
"""
|
||||
INSERT INTO execution_receipts (tenant_key, receipt_id, analysis_id, skill_name, input_hash, payload)
|
||||
VALUES (%s, %s, %s, %s, %s, %s::jsonb)
|
||||
ON CONFLICT (tenant_key, receipt_id) DO UPDATE SET payload = EXCLUDED.payload
|
||||
""",
|
||||
(
|
||||
tenant_key,
|
||||
receipt_id,
|
||||
analysis_id,
|
||||
"analyze_csv",
|
||||
receipt["input_sha256"],
|
||||
json.dumps(receipt),
|
||||
),
|
||||
)
|
||||
return True, receipt_id
|
||||
|
||||
|
||||
def _runtime_skills_root(ctx: RunContext[Any]) -> str:
|
||||
workspace = getattr(ctx, "_workspace", None)
|
||||
prefixes = tuple(getattr(workspace, "write_prefixes", ()) or ())
|
||||
if not prefixes:
|
||||
outputs_prefix = getattr(workspace, "outputs_prefix", None)
|
||||
prefixes = (outputs_prefix or "outputs/",)
|
||||
prefix = str(prefixes[0]).strip("/")
|
||||
return f"/{prefix}/{RUNTIME_SKILLS_DIR}" if prefix else f"/{RUNTIME_SKILLS_DIR}"
|
||||
def _load_analysis(tenant_key: str, analysis_id: str) -> dict[str, Any] | None:
|
||||
conn = _connect()
|
||||
if conn is None:
|
||||
return _LOCAL_ANALYSES.get((tenant_key, analysis_id))
|
||||
with conn:
|
||||
with conn.cursor() as cur:
|
||||
cur.execute(
|
||||
"SELECT payload FROM csv_analyses WHERE tenant_key = %s AND analysis_id = %s",
|
||||
(tenant_key, analysis_id),
|
||||
)
|
||||
row = cur.fetchone()
|
||||
if not row:
|
||||
return None
|
||||
return row[0]
|
||||
|
||||
|
||||
def _seed_runtime_skills(backend: Any, ctx: RunContext[Any]) -> list[str]:
|
||||
"""Copy packaged DeepAgents skills into the invocation workspace.
|
||||
|
||||
DeepAgents loads skills from its backend, while source-controlled
|
||||
``skills/`` folders live in the image. This bridge lets generated agents
|
||||
ship reusable SKILL.md bundles without giving up durable A2A workspace
|
||||
files.
|
||||
"""
|
||||
root = Path(__file__).parent / "skills"
|
||||
if not root.exists():
|
||||
return []
|
||||
runtime_skills_root = _runtime_skills_root(ctx)
|
||||
uploads: list[tuple[str, bytes]] = []
|
||||
for path in root.rglob("*"):
|
||||
if path.is_file():
|
||||
rel = path.relative_to(root).as_posix()
|
||||
uploads.append((runtime_skills_root + rel, path.read_bytes()))
|
||||
if uploads:
|
||||
backend.upload_files(uploads)
|
||||
return [runtime_skills_root]
|
||||
return []
|
||||
def _list_analyses(tenant_key: str, limit: int) -> list[dict[str, Any]]:
|
||||
conn = _connect()
|
||||
if conn is None:
|
||||
rows = [payload for (tenant, _), payload in _LOCAL_ANALYSES.items() if tenant == tenant_key]
|
||||
return [
|
||||
{
|
||||
"analysis_id": item["analysis_id"],
|
||||
"answer": item.get("answer"),
|
||||
"row_count": item.get("row_count", 0),
|
||||
"updated_at": item.get("updated_at"),
|
||||
}
|
||||
for item in rows[:limit]
|
||||
]
|
||||
with conn:
|
||||
with conn.cursor() as cur:
|
||||
cur.execute(
|
||||
"""
|
||||
SELECT analysis_id, answer, row_count, updated_at
|
||||
FROM csv_analyses
|
||||
WHERE tenant_key = %s
|
||||
ORDER BY updated_at DESC
|
||||
LIMIT %s
|
||||
""",
|
||||
(tenant_key, limit),
|
||||
)
|
||||
rows = cur.fetchall()
|
||||
return [
|
||||
{
|
||||
"analysis_id": row[0],
|
||||
"answer": row[1],
|
||||
"row_count": row[2],
|
||||
"updated_at": row[3].isoformat() if hasattr(row[3], "isoformat") else str(row[3]),
|
||||
}
|
||||
for row in rows
|
||||
]
|
||||
|
||||
|
||||
def _last_message_text(state: dict[str, Any]) -> str:
|
||||
messages = state.get("messages") or []
|
||||
if not messages:
|
||||
return json.dumps(state, default=str)
|
||||
def _receipt_id(tenant_key: str, analysis_id: str, csv_text: str, question: str) -> str:
|
||||
digest = hashlib.sha256(f"{tenant_key}\0{analysis_id}\0{question}\0{csv_text}".encode("utf-8")).hexdigest()
|
||||
return f"rcpt_{digest[:32]}"
|
||||
|
||||
content = getattr(messages[-1], "content", None)
|
||||
if isinstance(content, str):
|
||||
return content
|
||||
if isinstance(content, list):
|
||||
parts: list[str] = []
|
||||
for item in content:
|
||||
if isinstance(item, dict):
|
||||
text = item.get("text") or item.get("content")
|
||||
if text:
|
||||
parts.append(str(text))
|
||||
elif item:
|
||||
parts.append(str(item))
|
||||
return "\n".join(parts) if parts else json.dumps(content, default=str)
|
||||
return str(content or messages[-1])
|
||||
|
||||
def _reset_local_state_for_tests() -> None:
|
||||
_LOCAL_ANALYSES.clear()
|
||||
_LOCAL_RECEIPTS.clear()
|
||||
|
||||
Reference in New Issue
Block a user