mirror of
https://github.com/hansjone/oclaw.git
synced 2026-10-16 20:10:45 +08:00
refactor: root-package imports (svc/runtime/interfaces) and fix PYTHONPATH
- Rename platform/ to svc/ to avoid shadowing stdlib platform. - Replace from oclaw.* with from svc/runtime/interfaces; update -m CLI paths. - tests/conftest: prepend repo root to sys.path (no parent-folder package name). - CI: paths and offline_eval script under repo root. - Ops scripts: PYTHONPATH must be repo root for python -m runtime.* (fixes gateway/WhatsApp sidecar startup). - Fix default oclaw.json path in tabular/file attachment limits; stabilize attachment test config. Co-authored-by: Cursor <cursoragent@cursor.com>
This commit is contained in:
parent
7941b27f58
commit
420abac9f1
405 changed files with 1316 additions and 1345 deletions
0
svc/files/__init__.py
Normal file
0
svc/files/__init__.py
Normal file
202
svc/files/archive_processor.py
Normal file
202
svc/files/archive_processor.py
Normal file
|
|
@ -0,0 +1,202 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import gzip
|
||||
import io
|
||||
import tarfile
|
||||
import zipfile
|
||||
from pathlib import PurePosixPath
|
||||
from typing import Any, Callable
|
||||
|
||||
MAX_ARCHIVE_DEPTH = 2
|
||||
MAX_ARCHIVE_FILE_COUNT = 200
|
||||
MAX_ARCHIVE_ENTRY_BYTES = 10 * 1024 * 1024
|
||||
MAX_ARCHIVE_TOTAL_UNCOMPRESSED_BYTES = 50 * 1024 * 1024
|
||||
MAX_ARCHIVE_MEMBER_NAME_LENGTH = 240
|
||||
|
||||
|
||||
def is_safe_archive_member_name(name: str) -> bool:
|
||||
n = str(name or "").replace("\\", "/").strip()
|
||||
if not n:
|
||||
return False
|
||||
if len(n) > MAX_ARCHIVE_MEMBER_NAME_LENGTH:
|
||||
return False
|
||||
if any(ord(ch) < 32 for ch in n):
|
||||
return False
|
||||
p = PurePosixPath(n)
|
||||
if p.is_absolute():
|
||||
return False
|
||||
if any(part in ("..", "") for part in p.parts):
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
def detect_archive_kind(name: str) -> str | None:
|
||||
n = str(name or "").strip().lower()
|
||||
if n.endswith(".tar.gz") or n.endswith(".tgz"):
|
||||
return "tgz"
|
||||
if n.endswith(".zip"):
|
||||
return "zip"
|
||||
if n.endswith(".tar"):
|
||||
return "tar"
|
||||
if n.endswith(".gz"):
|
||||
return "gz"
|
||||
if n.endswith(".rar"):
|
||||
return "rar"
|
||||
if n.endswith(".7z"):
|
||||
return "7z"
|
||||
return None
|
||||
|
||||
|
||||
def sniff_archive_kind(data: bytes) -> str | None:
|
||||
b = bytes(data or b"")
|
||||
if len(b) >= 4 and b[:4] == b"PK\x03\x04":
|
||||
return "zip"
|
||||
if len(b) >= 2 and b[:2] == b"\x1f\x8b":
|
||||
return "gz"
|
||||
if len(b) >= 6 and b[:6] == b"Rar!\x1a\x07":
|
||||
return "rar"
|
||||
if len(b) >= 6 and b[:6] == b"7z\xbc\xaf\x27\x1c":
|
||||
return "7z"
|
||||
if len(b) >= 262 and b[257:262] == b"ustar":
|
||||
return "tar"
|
||||
return None
|
||||
|
||||
|
||||
def process_archive(
|
||||
*,
|
||||
archive_name: str,
|
||||
archive_data: bytes,
|
||||
process_member: Callable[[str, bytes, int], list[dict[str, Any]]],
|
||||
depth: int = 0,
|
||||
limits: dict[str, int] | None = None,
|
||||
) -> list[dict[str, Any]]:
|
||||
lim = limits or {}
|
||||
max_depth = int(lim.get("max_depth") or MAX_ARCHIVE_DEPTH)
|
||||
max_file_count = int(lim.get("max_file_count") or MAX_ARCHIVE_FILE_COUNT)
|
||||
max_entry_bytes = int(lim.get("max_entry_bytes") or MAX_ARCHIVE_ENTRY_BYTES)
|
||||
max_total_uncompressed_bytes = int(
|
||||
lim.get("max_total_uncompressed_bytes") or MAX_ARCHIVE_TOTAL_UNCOMPRESSED_BYTES
|
||||
)
|
||||
kind_from_sig = sniff_archive_kind(archive_data)
|
||||
kind_from_name = detect_archive_kind(archive_name)
|
||||
# For .tgz/.tar.gz, signature looks like gzip; keep richer name hint.
|
||||
if kind_from_name == "tgz" and kind_from_sig == "gz":
|
||||
kind = "tgz"
|
||||
else:
|
||||
kind = kind_from_sig or kind_from_name
|
||||
if kind is None:
|
||||
return []
|
||||
if depth > max_depth:
|
||||
return [{"type": "text", "name": "zip-error", "content": f"ZIP nesting too deep (>{max_depth})"}]
|
||||
out: list[dict[str, Any]] = []
|
||||
|
||||
def _err(msg: str, *, error_code: str) -> list[dict[str, Any]]:
|
||||
base = {
|
||||
"type": "text",
|
||||
"name": ("zip-error" if kind == "zip" else "archive-error"),
|
||||
"error_code": str(error_code or "archive_error"),
|
||||
"content": (
|
||||
f"[ArchiveError]\n- error_code: {error_code}\n- format: {kind}\n- detail: {msg}"
|
||||
),
|
||||
}
|
||||
return [base]
|
||||
|
||||
try:
|
||||
if kind in {"rar", "7z"}:
|
||||
raise ValueError(
|
||||
f"unsupported archive format: {kind}; currently supported formats are zip/tar/tgz/gz"
|
||||
)
|
||||
if kind == "zip":
|
||||
with zipfile.ZipFile(io.BytesIO(archive_data)) as zf:
|
||||
infos = [x for x in zf.infolist() if not x.is_dir()]
|
||||
if len(infos) > max_file_count:
|
||||
raise ValueError(f"ZIP has too many files ({len(infos)} > {max_file_count})")
|
||||
total_uncompressed = 0
|
||||
for info in infos:
|
||||
if not is_safe_archive_member_name(info.filename):
|
||||
raise ValueError(f"ZIP contains unsafe path: {info.filename}")
|
||||
fsize = int(info.file_size or 0)
|
||||
total_uncompressed += fsize
|
||||
if fsize > max_entry_bytes:
|
||||
raise ValueError(f"ZIP entry too large: {info.filename} ({fsize} > {max_entry_bytes})")
|
||||
if total_uncompressed > max_total_uncompressed_bytes:
|
||||
raise ValueError(
|
||||
f"ZIP total uncompressed size too large ({total_uncompressed} > {max_total_uncompressed_bytes})"
|
||||
)
|
||||
with zf.open(info) as f:
|
||||
out.extend(process_member(info.filename, f.read(), depth + 1))
|
||||
return out
|
||||
|
||||
if kind in {"tar", "tgz"}:
|
||||
mode = "r:gz" if kind == "tgz" else "r:"
|
||||
with tarfile.open(fileobj=io.BytesIO(archive_data), mode=mode) as tf:
|
||||
all_members = list(tf.getmembers())
|
||||
for m in all_members:
|
||||
if m.issym() or m.islnk():
|
||||
raise ValueError(f"TAR contains link entry: {m.name}")
|
||||
if m.ischr() or m.isblk() or m.isfifo():
|
||||
raise ValueError(f"TAR contains device/fifo entry: {m.name}")
|
||||
members = [m for m in all_members if m.isfile()]
|
||||
if len(members) > max_file_count:
|
||||
raise ValueError(f"TAR has too many files ({len(members)} > {max_file_count})")
|
||||
total_uncompressed = 0
|
||||
for m in members:
|
||||
if not is_safe_archive_member_name(m.name):
|
||||
raise ValueError(f"TAR contains unsafe path: {m.name}")
|
||||
fsize = int(m.size or 0)
|
||||
total_uncompressed += fsize
|
||||
if fsize > max_entry_bytes:
|
||||
raise ValueError(f"TAR entry too large: {m.name} ({fsize} > {max_entry_bytes})")
|
||||
if total_uncompressed > max_total_uncompressed_bytes:
|
||||
raise ValueError(
|
||||
f"TAR total uncompressed size too large ({total_uncompressed} > {max_total_uncompressed_bytes})"
|
||||
)
|
||||
ext = tf.extractfile(m)
|
||||
if ext is None:
|
||||
continue
|
||||
out.extend(process_member(m.name, ext.read(), depth + 1))
|
||||
return out
|
||||
|
||||
if kind == "gz":
|
||||
raw = gzip.decompress(archive_data)
|
||||
if len(raw) > max_entry_bytes:
|
||||
raise ValueError(f"GZ entry too large ({len(raw)} > {max_entry_bytes})")
|
||||
name = str(archive_name or "file.gz")
|
||||
inner_name = name[:-3] if name.lower().endswith(".gz") else f"{name}.out"
|
||||
out.extend(process_member(inner_name, raw, depth + 1))
|
||||
return out
|
||||
except Exception as e:
|
||||
s = str(e).lower()
|
||||
code = "archive_parse_failed"
|
||||
if "unsupported archive format" in s:
|
||||
code = "archive_unsupported_format"
|
||||
elif "unsafe path" in s:
|
||||
code = "archive_path_traversal"
|
||||
elif "nesting too deep" in s:
|
||||
code = "archive_max_depth_exceeded"
|
||||
elif "too many files" in s:
|
||||
code = "archive_max_file_count_exceeded"
|
||||
elif "entry too large" in s:
|
||||
code = "archive_max_entry_bytes_exceeded"
|
||||
elif "total uncompressed size too large" in s:
|
||||
code = "archive_max_total_uncompressed_bytes_exceeded"
|
||||
elif "link entry" in s:
|
||||
code = "archive_link_entry_forbidden"
|
||||
elif "device/fifo entry" in s:
|
||||
code = "archive_special_entry_forbidden"
|
||||
return _err(str(e), error_code=code)
|
||||
return out
|
||||
|
||||
|
||||
__all__ = [
|
||||
"MAX_ARCHIVE_DEPTH",
|
||||
"MAX_ARCHIVE_ENTRY_BYTES",
|
||||
"MAX_ARCHIVE_FILE_COUNT",
|
||||
"MAX_ARCHIVE_MEMBER_NAME_LENGTH",
|
||||
"MAX_ARCHIVE_TOTAL_UNCOMPRESSED_BYTES",
|
||||
"detect_archive_kind",
|
||||
"sniff_archive_kind",
|
||||
"is_safe_archive_member_name",
|
||||
"process_archive",
|
||||
]
|
||||
|
||||
378
svc/files/attachment_assets.py
Normal file
378
svc/files/attachment_assets.py
Normal file
|
|
@ -0,0 +1,378 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import base64
|
||||
import hashlib
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
import time
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Any, Final, Optional
|
||||
|
||||
from svc.config.paths import attachments_dir
|
||||
_META_SUFFIX: Final[str] = ".meta.json"
|
||||
_ATTACHMENT_ID_RE: Final[re.Pattern[str]] = re.compile(r"^[a-f0-9]{64}$")
|
||||
# Must match extensions produced by ``save_bytes`` so ``load_bytes`` / ``get_local_path`` / ``gc`` can find blobs.
|
||||
_ATTACHMENT_BLOB_EXTENSIONS: Final[tuple[str, ...]] = (
|
||||
".png",
|
||||
".jpg",
|
||||
".jpeg",
|
||||
".webp",
|
||||
".gif",
|
||||
".mp4",
|
||||
".webm",
|
||||
".mov",
|
||||
"", # extensionless fallback
|
||||
)
|
||||
|
||||
|
||||
def _utc_ts() -> int:
|
||||
return int(time.time())
|
||||
|
||||
|
||||
def _safe_rel_name(name: str) -> str:
|
||||
# keep it simple: drop path parts, trim, and avoid empty
|
||||
raw = (name or "").replace("\\", "/").split("/")[-1].strip()
|
||||
return raw or "image"
|
||||
|
||||
|
||||
def _ext_from_mime(mime: str | None) -> str:
|
||||
m = (mime or "").strip().lower()
|
||||
if m == "image/png":
|
||||
return ".png"
|
||||
if m in ("image/jpg", "image/jpeg"):
|
||||
return ".jpg"
|
||||
if m == "image/webp":
|
||||
return ".webp"
|
||||
if m == "image/gif":
|
||||
return ".gif"
|
||||
if m == "video/mp4":
|
||||
return ".mp4"
|
||||
if m == "video/webm":
|
||||
return ".webm"
|
||||
if m in ("video/quicktime", "video/mov"):
|
||||
return ".mov"
|
||||
return ""
|
||||
|
||||
|
||||
def _guess_mime_from_ext(ext: str) -> str:
|
||||
e = (ext or "").lower().lstrip(".")
|
||||
if e == "png":
|
||||
return "image/png"
|
||||
if e in ("jpg", "jpeg"):
|
||||
return "image/jpeg"
|
||||
if e == "webp":
|
||||
return "image/webp"
|
||||
if e == "gif":
|
||||
return "image/gif"
|
||||
if e == "mp4":
|
||||
return "video/mp4"
|
||||
if e == "webm":
|
||||
return "video/webm"
|
||||
if e == "mov":
|
||||
return "video/quicktime"
|
||||
return "application/octet-stream"
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class AttachmentMeta:
|
||||
attachment_id: str
|
||||
name: str
|
||||
mime: str
|
||||
bytes: int
|
||||
created_at: int
|
||||
last_access: int
|
||||
width: int | None = None
|
||||
height: int | None = None
|
||||
|
||||
def to_dict(self) -> dict[str, Any]:
|
||||
d: dict[str, Any] = {
|
||||
"attachment_id": self.attachment_id,
|
||||
"name": self.name,
|
||||
"mime": self.mime,
|
||||
"bytes": int(self.bytes),
|
||||
"created_at": int(self.created_at),
|
||||
"last_access": int(self.last_access),
|
||||
}
|
||||
if self.width is not None:
|
||||
d["width"] = int(self.width)
|
||||
if self.height is not None:
|
||||
d["height"] = int(self.height)
|
||||
return d
|
||||
|
||||
@staticmethod
|
||||
def from_dict(d: dict[str, Any]) -> "AttachmentMeta":
|
||||
return AttachmentMeta(
|
||||
attachment_id=str(d.get("attachment_id") or ""),
|
||||
name=str(d.get("name") or ""),
|
||||
mime=str(d.get("mime") or "application/octet-stream"),
|
||||
bytes=int(d.get("bytes") or 0),
|
||||
created_at=int(d.get("created_at") or 0),
|
||||
last_access=int(d.get("last_access") or 0),
|
||||
width=int(d["width"]) if d.get("width") is not None else None,
|
||||
height=int(d["height"]) if d.get("height") is not None else None,
|
||||
)
|
||||
|
||||
|
||||
class AttachmentAssetStore:
|
||||
"""Disk-backed attachment store (project-local by default)."""
|
||||
|
||||
def __init__(self, root_dir: str | Path | None = None):
|
||||
self.root = Path(root_dir) if root_dir is not None else attachments_dir()
|
||||
self.root.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
@staticmethod
|
||||
def _normalize_attachment_id(attachment_id: str) -> str:
|
||||
aid = str(attachment_id or "").strip().lower()
|
||||
if not _ATTACHMENT_ID_RE.fullmatch(aid):
|
||||
raise ValueError("attachment_id_invalid")
|
||||
return aid
|
||||
|
||||
def _data_path(self, attachment_id: str, *, ext: str) -> Path:
|
||||
# bucket to avoid huge single dir
|
||||
aid = self._normalize_attachment_id(attachment_id)
|
||||
p1, p2 = (aid[:2] or "xx"), (aid[2:4] or "yy")
|
||||
d = self.root / p1 / p2
|
||||
d.mkdir(parents=True, exist_ok=True)
|
||||
return d / f"{aid}{ext}"
|
||||
|
||||
def _meta_path(self, attachment_id: str) -> Path:
|
||||
aid = self._normalize_attachment_id(attachment_id)
|
||||
p1, p2 = (aid[:2] or "xx"), (aid[2:4] or "yy")
|
||||
d = self.root / p1 / p2
|
||||
d.mkdir(parents=True, exist_ok=True)
|
||||
return d / f"{aid}{_META_SUFFIX}"
|
||||
|
||||
def save_bytes(
|
||||
self,
|
||||
data: bytes,
|
||||
*,
|
||||
filename: str,
|
||||
mime: str | None,
|
||||
width: int | None = None,
|
||||
height: int | None = None,
|
||||
) -> AttachmentMeta:
|
||||
blob = bytes(data or b"")
|
||||
h = hashlib.sha256(blob).hexdigest()
|
||||
ext = _ext_from_mime(mime) or (("." + filename.split(".")[-1].lower()) if "." in filename else "")
|
||||
ext = ext if ext in _ATTACHMENT_BLOB_EXTENSIONS[:-1] else _ext_from_mime(mime) or ""
|
||||
data_path = self._data_path(h, ext=ext)
|
||||
meta_path = self._meta_path(h)
|
||||
|
||||
now = _utc_ts()
|
||||
if not data_path.exists():
|
||||
tmp = str(data_path) + ".tmp"
|
||||
with open(tmp, "wb") as f:
|
||||
f.write(blob)
|
||||
os.replace(tmp, data_path)
|
||||
|
||||
meta = AttachmentMeta(
|
||||
attachment_id=h,
|
||||
name=_safe_rel_name(filename),
|
||||
mime=(mime or _guess_mime_from_ext(ext)).strip() or _guess_mime_from_ext(ext),
|
||||
bytes=len(blob),
|
||||
created_at=now,
|
||||
last_access=now,
|
||||
width=width,
|
||||
height=height,
|
||||
)
|
||||
# best-effort: keep existing created_at if present
|
||||
if meta_path.exists():
|
||||
try:
|
||||
prev = json.loads(meta_path.read_text("utf-8"))
|
||||
if isinstance(prev, dict) and prev.get("created_at"):
|
||||
meta = AttachmentMeta(
|
||||
**{**meta.to_dict(), "created_at": int(prev.get("created_at") or now)} # type: ignore[arg-type]
|
||||
)
|
||||
except Exception:
|
||||
pass
|
||||
meta_path.write_text(json.dumps(meta.to_dict(), ensure_ascii=False), encoding="utf-8")
|
||||
return meta
|
||||
|
||||
def get_meta(self, attachment_id: str) -> Optional[AttachmentMeta]:
|
||||
try:
|
||||
mp = self._meta_path(attachment_id)
|
||||
except Exception:
|
||||
return None
|
||||
if not mp.exists():
|
||||
return None
|
||||
try:
|
||||
obj = json.loads(mp.read_text("utf-8"))
|
||||
if not isinstance(obj, dict):
|
||||
return None
|
||||
return AttachmentMeta.from_dict(obj)
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
def touch(self, attachment_id: str) -> None:
|
||||
try:
|
||||
mp = self._meta_path(attachment_id)
|
||||
except Exception:
|
||||
return
|
||||
if not mp.exists():
|
||||
return
|
||||
try:
|
||||
obj = json.loads(mp.read_text("utf-8"))
|
||||
if not isinstance(obj, dict):
|
||||
return
|
||||
obj["last_access"] = _utc_ts()
|
||||
mp.write_text(json.dumps(obj, ensure_ascii=False), encoding="utf-8")
|
||||
except Exception:
|
||||
return
|
||||
|
||||
def load_bytes(self, attachment_id: str) -> tuple[bytes, Optional[AttachmentMeta]]:
|
||||
try:
|
||||
aid = self._normalize_attachment_id(attachment_id)
|
||||
except Exception:
|
||||
return b"", None
|
||||
meta = self.get_meta(aid)
|
||||
# try find data file by scanning known extensions
|
||||
data_path = None
|
||||
for ext in _ATTACHMENT_BLOB_EXTENSIONS:
|
||||
p = self._data_path(aid, ext=ext)
|
||||
if p.exists():
|
||||
data_path = p
|
||||
break
|
||||
if data_path is None:
|
||||
return b"", meta
|
||||
try:
|
||||
blob = data_path.read_bytes()
|
||||
except Exception:
|
||||
return b"", meta
|
||||
self.touch(aid)
|
||||
return blob, meta
|
||||
|
||||
def get_local_path(self, attachment_id: str) -> Path | None:
|
||||
try:
|
||||
aid = self._normalize_attachment_id(attachment_id)
|
||||
except Exception:
|
||||
return None
|
||||
for ext in _ATTACHMENT_BLOB_EXTENSIONS:
|
||||
p = self._data_path(aid, ext=ext)
|
||||
if p.exists():
|
||||
self.touch(aid)
|
||||
return p
|
||||
return None
|
||||
|
||||
def inspect_asset(self, attachment_id: str) -> dict[str, Any]:
|
||||
meta = self.get_meta(attachment_id)
|
||||
p = self.get_local_path(attachment_id)
|
||||
if meta is None and p is None:
|
||||
return {"state": "missing", "attachment_id": attachment_id}
|
||||
if meta is not None and p is None:
|
||||
return {
|
||||
"state": "meta_exists_blob_missing",
|
||||
"attachment_id": attachment_id,
|
||||
"meta": meta.to_dict(),
|
||||
}
|
||||
return {
|
||||
"state": "ok",
|
||||
"attachment_id": attachment_id,
|
||||
"meta": meta.to_dict() if meta else None,
|
||||
"path": str(p) if p is not None else "",
|
||||
}
|
||||
|
||||
def get_render_cache_key(self, attachment_id: str, *, mime_hint: str | None = None) -> tuple[str, str, int, int]:
|
||||
meta = self.get_meta(attachment_id)
|
||||
return (
|
||||
str(attachment_id or "").strip(),
|
||||
str(mime_hint or (meta.mime if meta else "") or "image/jpeg"),
|
||||
int(meta.bytes if meta else 0),
|
||||
int(meta.created_at if meta else 0),
|
||||
)
|
||||
|
||||
def gc(self, *, ttl_days: int = 7) -> dict[str, Any]:
|
||||
ttl_sec = max(1, int(ttl_days)) * 86400
|
||||
now = _utc_ts()
|
||||
removed = 0
|
||||
kept = 0
|
||||
|
||||
if not self.root.exists():
|
||||
return {"ok": True, "removed": 0, "kept": 0}
|
||||
|
||||
for mp in self.root.rglob(f"*{_META_SUFFIX}"):
|
||||
try:
|
||||
obj = json.loads(mp.read_text("utf-8"))
|
||||
if not isinstance(obj, dict):
|
||||
continue
|
||||
last_access = int(obj.get("last_access") or obj.get("created_at") or 0)
|
||||
aid = str(obj.get("attachment_id") or mp.name.replace(_META_SUFFIX, ""))
|
||||
except Exception:
|
||||
continue
|
||||
if last_access and (now - last_access) <= ttl_sec:
|
||||
kept += 1
|
||||
continue
|
||||
|
||||
# delete meta + possible data files
|
||||
try:
|
||||
mp.unlink(missing_ok=True) # py3.8+; on 3.14 ok
|
||||
except Exception:
|
||||
pass
|
||||
for ext in _ATTACHMENT_BLOB_EXTENSIONS:
|
||||
try:
|
||||
self._data_path(aid, ext=ext).unlink(missing_ok=True)
|
||||
except Exception:
|
||||
pass
|
||||
removed += 1
|
||||
|
||||
return {"ok": True, "removed": removed, "kept": kept}
|
||||
|
||||
|
||||
_LAST_GC_TS: int = 0
|
||||
|
||||
|
||||
def gc_if_due(*, ttl_days: int = 7, min_interval_sec: int = 3600, root_dir: str | Path | None = None) -> None:
|
||||
"""Best-effort throttled GC to be called in UI loop."""
|
||||
global _LAST_GC_TS
|
||||
now = _utc_ts()
|
||||
if _LAST_GC_TS and (now - _LAST_GC_TS) < max(10, int(min_interval_sec)):
|
||||
return
|
||||
_LAST_GC_TS = now
|
||||
try:
|
||||
AttachmentAssetStore(root_dir).gc(ttl_days=ttl_days)
|
||||
except Exception:
|
||||
return
|
||||
|
||||
|
||||
def image_ref_to_data_url(att: dict[str, Any], *, root_dir: str | Path | None = None) -> tuple[str, str]:
|
||||
"""Helper for UI: convert image_ref to (mime, data_url_base64_payload)."""
|
||||
aid = str(att.get("attachment_id") or "").strip()
|
||||
if not aid:
|
||||
return "image/jpeg", ""
|
||||
store = AttachmentAssetStore(root_dir)
|
||||
blob, meta = store.load_bytes(aid)
|
||||
mime = str(att.get("mime") or (meta.mime if meta else "image/jpeg") or "image/jpeg")
|
||||
if not blob:
|
||||
return mime, ""
|
||||
b64 = base64.b64encode(blob).decode("ascii")
|
||||
return mime, b64
|
||||
|
||||
|
||||
def attachment_id_to_data_url(
|
||||
attachment_id: str,
|
||||
*,
|
||||
mime: str | None = None,
|
||||
root_dir: str | Path | None = None,
|
||||
) -> str:
|
||||
"""Convert stored attachment bytes into data URL form."""
|
||||
aid = str(attachment_id or "").strip()
|
||||
if not aid:
|
||||
return ""
|
||||
store = AttachmentAssetStore(root_dir)
|
||||
blob, meta = store.load_bytes(aid)
|
||||
if not blob:
|
||||
return ""
|
||||
resolved_mime = str(mime or (meta.mime if meta else "") or "image/jpeg")
|
||||
b64 = base64.b64encode(blob).decode("ascii")
|
||||
return f"data:{resolved_mime};base64,{b64}"
|
||||
|
||||
|
||||
__all__ = [
|
||||
"AttachmentAssetStore",
|
||||
"AttachmentMeta",
|
||||
"attachment_id_to_data_url",
|
||||
"gc_if_due",
|
||||
"image_ref_to_data_url",
|
||||
]
|
||||
|
||||
776
svc/files/file_attachments.py
Normal file
776
svc/files/file_attachments.py
Normal file
|
|
@ -0,0 +1,776 @@
|
|||
from __future__ import annotations
|
||||
|
||||
"""将上传文件解析为附件字典(不依赖 Streamlit)。"""
|
||||
|
||||
import io
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
import subprocess
|
||||
import zipfile
|
||||
from functools import lru_cache
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import pandas as pd
|
||||
from docx import Document
|
||||
from PIL import Image
|
||||
from PyPDF2 import PdfReader
|
||||
|
||||
from svc.config.paths import PROJECT_ROOT
|
||||
from svc.files.archive_processor import (
|
||||
MAX_ARCHIVE_DEPTH,
|
||||
MAX_ARCHIVE_ENTRY_BYTES,
|
||||
MAX_ARCHIVE_FILE_COUNT,
|
||||
MAX_ARCHIVE_MEMBER_NAME_LENGTH,
|
||||
MAX_ARCHIVE_TOTAL_UNCOMPRESSED_BYTES,
|
||||
detect_archive_kind,
|
||||
is_safe_archive_member_name,
|
||||
process_archive,
|
||||
)
|
||||
from svc.files.attachment_assets import AttachmentAssetStore
|
||||
from svc.files.tabular_attachment_store import save_dataframe, save_workbook
|
||||
from svc.files.text_attachment_store import (
|
||||
DEFAULT_TEXT_CHUNK_OVERLAP,
|
||||
DEFAULT_TEXT_CHUNK_SIZE,
|
||||
DEFAULT_TEXT_INLINE_MAX_CHARS,
|
||||
save_text_document,
|
||||
)
|
||||
|
||||
|
||||
def _sniff_image_mime(data: bytes) -> str | None:
|
||||
if len(data) < 12:
|
||||
return None
|
||||
if data[:8] == b"\x89PNG\r\n\x1a\n":
|
||||
return "image/png"
|
||||
if data[:3] == b"\xff\xd8\xff":
|
||||
return "image/jpeg"
|
||||
if data[:6] in (b"GIF87a", b"GIF89a"):
|
||||
return "image/gif"
|
||||
if len(data) >= 12 and data[:4] == b"RIFF" and data[8:12] == b"WEBP":
|
||||
return "image/webp"
|
||||
return None
|
||||
|
||||
|
||||
def _ffprobe_exists() -> bool:
|
||||
try:
|
||||
p = subprocess.run(["ffprobe", "-version"], capture_output=True, text=True, encoding="utf-8", errors="replace", timeout=3)
|
||||
return p.returncode == 0
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
|
||||
def _extract_video_meta_from_ffprobe(obj: dict[str, Any]) -> dict[str, Any]:
|
||||
fmt = obj.get("format") if isinstance(obj.get("format"), dict) else {}
|
||||
streams = obj.get("streams") if isinstance(obj.get("streams"), list) else []
|
||||
duration_sec = None
|
||||
if isinstance(fmt, dict) and fmt.get("duration") is not None:
|
||||
try:
|
||||
duration_sec = float(fmt.get("duration"))
|
||||
except Exception:
|
||||
duration_sec = None
|
||||
width = None
|
||||
height = None
|
||||
fps = None
|
||||
for s in streams:
|
||||
if not isinstance(s, dict):
|
||||
continue
|
||||
if str(s.get("codec_type") or "") != "video":
|
||||
continue
|
||||
if s.get("width") is not None and s.get("height") is not None:
|
||||
try:
|
||||
width = int(s.get("width") or 0) or None
|
||||
height = int(s.get("height") or 0) or None
|
||||
except Exception:
|
||||
width = width
|
||||
height = height
|
||||
fr = str(s.get("avg_frame_rate") or s.get("r_frame_rate") or "").strip()
|
||||
if fr and fr != "0/0" and "/" in fr:
|
||||
try:
|
||||
a, b = fr.split("/", 1)
|
||||
fa = float(a)
|
||||
fb = float(b)
|
||||
if fb:
|
||||
fps = fa / fb
|
||||
except Exception:
|
||||
fps = fps
|
||||
break
|
||||
out: dict[str, Any] = {}
|
||||
if duration_sec is not None:
|
||||
out["duration_sec"] = duration_sec
|
||||
if width is not None:
|
||||
out["width"] = width
|
||||
if height is not None:
|
||||
out["height"] = height
|
||||
if fps is not None:
|
||||
out["fps"] = fps
|
||||
return out
|
||||
|
||||
|
||||
def _ffprobe_video_meta(path: Path) -> dict[str, Any]:
|
||||
if not _ffprobe_exists():
|
||||
return {}
|
||||
try:
|
||||
p = subprocess.run(
|
||||
["ffprobe", "-v", "error", "-print_format", "json", "-show_format", "-show_streams", str(path)],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
encoding="utf-8",
|
||||
errors="replace",
|
||||
timeout=8,
|
||||
)
|
||||
if p.returncode != 0:
|
||||
return {}
|
||||
obj = json.loads(p.stdout or "{}")
|
||||
if not isinstance(obj, dict):
|
||||
return {}
|
||||
return _extract_video_meta_from_ffprobe(obj)
|
||||
except Exception:
|
||||
return {}
|
||||
|
||||
|
||||
MAX_TEXT_CHARS = 12_000
|
||||
EXCEL_PREVIEW_ROWS = 80
|
||||
LARGE_TABLE_PREVIEW_ROWS = 20
|
||||
MAX_ZIP_DEPTH = MAX_ARCHIVE_DEPTH
|
||||
MAX_ZIP_FILE_COUNT = MAX_ARCHIVE_FILE_COUNT
|
||||
MAX_ZIP_ENTRY_BYTES = MAX_ARCHIVE_ENTRY_BYTES
|
||||
MAX_ZIP_TOTAL_UNCOMPRESSED_BYTES = MAX_ARCHIVE_TOTAL_UNCOMPRESSED_BYTES
|
||||
DEFAULT_TABULAR_ROWS_READ = 5000
|
||||
DEFAULT_TABULAR_COLUMNS = 200
|
||||
DEFAULT_TABULAR_CELL_CHARS = 500
|
||||
DEFAULT_TABULAR_TOOL_MODE_ENABLED = True
|
||||
DEFAULT_TABULAR_TOOL_MODE_MIN_ROWS = 5_000
|
||||
DEFAULT_TABULAR_TOOL_MODE_MAX_BYTES = 30 * 1024 * 1024
|
||||
DEFAULT_LARGE_TABLE_PREVIEW_ROWS = LARGE_TABLE_PREVIEW_ROWS
|
||||
MAX_ZIP_MEMBER_NAME_LENGTH = MAX_ARCHIVE_MEMBER_NAME_LENGTH
|
||||
MAX_EXCEL_ZIP_FILE_COUNT = 500
|
||||
MAX_EXCEL_ZIP_ENTRY_BYTES = 20 * 1024 * 1024
|
||||
MAX_EXCEL_ZIP_TOTAL_UNCOMPRESSED_BYTES = 120 * 1024 * 1024
|
||||
DEFAULT_MAX_EXCEL_SHEETS = 50
|
||||
DEFAULT_TEXT_QUERY_TOP_K = 5
|
||||
|
||||
|
||||
def _truncate_text(text: str, limit: int = MAX_TEXT_CHARS) -> str:
|
||||
t = str(text or "")
|
||||
if len(t) <= limit:
|
||||
return t
|
||||
return t[:limit] + f"\n\n...[truncated to {limit} chars]"
|
||||
|
||||
|
||||
def _text_ref_summary(*, name: str, chars: int, chunks: int, source_kind: str, preview: str, top_k: int) -> str:
|
||||
return _truncate_text(
|
||||
(
|
||||
f"# Text Document Summary\n"
|
||||
f"- name: {name}\n"
|
||||
f"- source_kind: {source_kind}\n"
|
||||
f"- chars: {chars}\n"
|
||||
f"- chunks: {chunks}\n"
|
||||
f"- mode: text_ref\n"
|
||||
f"- query_tool: query_text_attachment\n"
|
||||
f"- note: use `text_id` from accompanying `text_ref` attachment; recommended top_k <= {top_k}\n\n"
|
||||
f"## Preview (first {min(len(preview), 1500)} chars)\n"
|
||||
f"{preview[:1500]}"
|
||||
),
|
||||
limit=MAX_TEXT_CHARS,
|
||||
)
|
||||
|
||||
|
||||
def _decode_text_bytes(data: bytes) -> str:
|
||||
try:
|
||||
return data.decode("utf-8")
|
||||
except UnicodeDecodeError:
|
||||
try:
|
||||
return data.decode("gbk")
|
||||
except Exception:
|
||||
return "Error decoding text file (unknown encoding)"
|
||||
|
||||
|
||||
def _clean_html_to_text(raw: str) -> str:
|
||||
text = re.sub(r"(?is)<script[^>]*>.*?</script>", " ", str(raw or ""))
|
||||
text = re.sub(r"(?is)<style[^>]*>.*?</style>", " ", text)
|
||||
text = re.sub(r"(?is)<[^>]+>", " ", text)
|
||||
text = re.sub(r"[ \t\r\f\v]+", " ", text)
|
||||
text = re.sub(r"\n{3,}", "\n\n", text)
|
||||
return text.strip()
|
||||
|
||||
|
||||
def _safe_int(raw: Any, default: int, *, min_value: int = 1, max_value: int = 2_000_000) -> int:
|
||||
try:
|
||||
value = int(raw)
|
||||
except Exception:
|
||||
return default
|
||||
if value < min_value:
|
||||
return default
|
||||
return min(value, max_value)
|
||||
|
||||
|
||||
@lru_cache(maxsize=1)
|
||||
def _attachments_limits() -> dict[str, Any]:
|
||||
cfg_path_raw = str(os.getenv("AIA_OCLAW_CONFIG_PATH") or "").strip()
|
||||
cfg_path = Path(cfg_path_raw).expanduser() if cfg_path_raw else (Path(PROJECT_ROOT) / "oclaw.json").resolve()
|
||||
if not cfg_path.is_absolute():
|
||||
cfg_path = (Path(PROJECT_ROOT) / cfg_path).resolve()
|
||||
tabular_cfg: dict[str, Any] = {}
|
||||
try:
|
||||
obj = json.loads(cfg_path.read_text(encoding="utf-8"))
|
||||
tabular_cfg = (
|
||||
(((obj or {}).get("plugins") or {}).get("entries") or {}).get("memory-wiki", {}).get("auto", {}).get("attachments", {}).get("tabular", {}) or {}
|
||||
)
|
||||
except Exception:
|
||||
tabular_cfg = {}
|
||||
return {
|
||||
"rows_read": _safe_int(tabular_cfg.get("max_rows_read"), DEFAULT_TABULAR_ROWS_READ),
|
||||
"columns": _safe_int(tabular_cfg.get("max_columns"), DEFAULT_TABULAR_COLUMNS),
|
||||
"cell_chars": _safe_int(tabular_cfg.get("max_cell_chars"), DEFAULT_TABULAR_CELL_CHARS),
|
||||
"max_excel_sheets": _safe_int(tabular_cfg.get("max_excel_sheets"), DEFAULT_MAX_EXCEL_SHEETS, max_value=500),
|
||||
"large_table_preview_rows": _safe_int(
|
||||
tabular_cfg.get("large_table_preview_rows"), DEFAULT_LARGE_TABLE_PREVIEW_ROWS, max_value=500
|
||||
),
|
||||
"tool_mode_enabled": bool(tabular_cfg.get("tool_mode_enabled", DEFAULT_TABULAR_TOOL_MODE_ENABLED)),
|
||||
"tool_mode_min_rows": _safe_int(tabular_cfg.get("tool_mode_min_rows"), DEFAULT_TABULAR_TOOL_MODE_MIN_ROWS),
|
||||
"tool_mode_max_bytes": _safe_int(
|
||||
tabular_cfg.get("tool_mode_max_bytes"), DEFAULT_TABULAR_TOOL_MODE_MAX_BYTES, max_value=500 * 1024 * 1024
|
||||
),
|
||||
"text_inline_max_chars": _safe_int(
|
||||
tabular_cfg.get("text_inline_max_chars"), DEFAULT_TEXT_INLINE_MAX_CHARS, max_value=200_000
|
||||
),
|
||||
"text_chunk_size": _safe_int(tabular_cfg.get("text_chunk_size"), DEFAULT_TEXT_CHUNK_SIZE, max_value=8_000),
|
||||
"text_chunk_overlap": _safe_int(
|
||||
tabular_cfg.get("text_chunk_overlap"), DEFAULT_TEXT_CHUNK_OVERLAP, max_value=4_000
|
||||
),
|
||||
"text_query_top_k": _safe_int(tabular_cfg.get("text_query_top_k"), DEFAULT_TEXT_QUERY_TOP_K, max_value=50),
|
||||
"archive_max_depth": _safe_int(tabular_cfg.get("archive_max_depth"), MAX_ARCHIVE_DEPTH, max_value=10),
|
||||
"archive_max_file_count": _safe_int(
|
||||
tabular_cfg.get("archive_max_file_count"), MAX_ARCHIVE_FILE_COUNT, max_value=20_000
|
||||
),
|
||||
"archive_max_entry_bytes": _safe_int(
|
||||
tabular_cfg.get("archive_max_entry_bytes"), MAX_ARCHIVE_ENTRY_BYTES, max_value=2_000_000_000
|
||||
),
|
||||
"archive_max_total_uncompressed_bytes": _safe_int(
|
||||
tabular_cfg.get("archive_max_total_uncompressed_bytes"),
|
||||
MAX_ARCHIVE_TOTAL_UNCOMPRESSED_BYTES,
|
||||
max_value=5_000_000_000,
|
||||
),
|
||||
}
|
||||
|
||||
|
||||
def clear_attachment_limits_cache() -> None:
|
||||
_attachments_limits.cache_clear()
|
||||
|
||||
|
||||
def _clip_tabular_cells(df: pd.DataFrame, *, max_cell_chars: int) -> tuple[pd.DataFrame, bool]:
|
||||
clipped_any = False
|
||||
|
||||
def _clip_cell(value: Any) -> Any:
|
||||
nonlocal clipped_any
|
||||
if value is None:
|
||||
return value
|
||||
s = str(value)
|
||||
if len(s) <= max_cell_chars:
|
||||
return value
|
||||
clipped_any = True
|
||||
return s[:max_cell_chars] + "...[cell-truncated]"
|
||||
|
||||
out = df.copy()
|
||||
for col in out.columns:
|
||||
out[col] = out[col].map(_clip_cell)
|
||||
return out, clipped_any
|
||||
|
||||
|
||||
def _dataframe_summary_text(
|
||||
df: pd.DataFrame,
|
||||
*,
|
||||
sampled: bool = False,
|
||||
clipped_columns: bool = False,
|
||||
clipped_cells: bool = False,
|
||||
preview_rows: int = EXCEL_PREVIEW_ROWS,
|
||||
) -> str:
|
||||
rows = int(len(df.index))
|
||||
cols = int(len(df.columns))
|
||||
header = [str(x) for x in list(df.columns)]
|
||||
dtypes = [f"{str(k)}:{str(v)}" for k, v in df.dtypes.items()]
|
||||
show_rows = max(1, int(preview_rows or EXCEL_PREVIEW_ROWS))
|
||||
sample = df.head(show_rows).to_string(index=False)
|
||||
rows_scope_line = (
|
||||
f"## Preview (first {min(rows, show_rows)} rows)"
|
||||
if rows > show_rows
|
||||
else f"## Full table included ({rows} rows)"
|
||||
)
|
||||
body = (
|
||||
f"# Table Summary\n"
|
||||
f"- rows: {rows}\n"
|
||||
f"- cols: {cols}\n"
|
||||
f"- sampled: {'yes' if sampled else 'no'}\n"
|
||||
f"- clipped_columns: {'yes' if clipped_columns else 'no'}\n"
|
||||
f"- clipped_cells: {'yes' if clipped_cells else 'no'}\n"
|
||||
f"- columns: {', '.join(header)}\n"
|
||||
f"- dtypes: {', '.join(dtypes)}\n\n"
|
||||
f"{rows_scope_line}\n"
|
||||
f"{sample}"
|
||||
)
|
||||
return _truncate_text(body)
|
||||
|
||||
|
||||
def _pdf_summary_text(reader: PdfReader) -> str:
|
||||
pages = list(getattr(reader, "pages", []) or [])
|
||||
chunks: list[str] = [f"# PDF Summary\n- pages: {len(pages)}\n"]
|
||||
for idx, page in enumerate(pages, start=1):
|
||||
txt = ""
|
||||
try:
|
||||
txt = str(page.extract_text() or "").strip()
|
||||
except Exception:
|
||||
txt = ""
|
||||
if not txt:
|
||||
txt = "(empty page text)"
|
||||
chunks.append(f"## Page {idx}\n{txt}")
|
||||
return _truncate_text("\n\n".join(chunks))
|
||||
|
||||
|
||||
def _is_safe_zip_member_name(name: str) -> bool:
|
||||
return is_safe_archive_member_name(name)
|
||||
|
||||
|
||||
def _validate_excel_zip_payload(data: bytes) -> str | None:
|
||||
try:
|
||||
with zipfile.ZipFile(io.BytesIO(data)) as zf:
|
||||
infos = [x for x in zf.infolist() if not x.is_dir()]
|
||||
if len(infos) > MAX_EXCEL_ZIP_FILE_COUNT:
|
||||
return f"excel zip has too many entries ({len(infos)} > {MAX_EXCEL_ZIP_FILE_COUNT})"
|
||||
total_uncompressed = 0
|
||||
for info in infos:
|
||||
if not _is_safe_zip_member_name(info.filename):
|
||||
return f"excel zip contains unsafe path: {info.filename}"
|
||||
fsize = int(info.file_size or 0)
|
||||
total_uncompressed += fsize
|
||||
if fsize > MAX_EXCEL_ZIP_ENTRY_BYTES:
|
||||
return f"excel zip entry too large: {info.filename} ({fsize} > {MAX_EXCEL_ZIP_ENTRY_BYTES})"
|
||||
if total_uncompressed > MAX_EXCEL_ZIP_TOTAL_UNCOMPRESSED_BYTES:
|
||||
return (
|
||||
"excel zip total uncompressed size too large "
|
||||
f"({total_uncompressed} > {MAX_EXCEL_ZIP_TOTAL_UNCOMPRESSED_BYTES})"
|
||||
)
|
||||
return None
|
||||
except Exception:
|
||||
return "invalid excel zip payload"
|
||||
|
||||
|
||||
def process_zip(zip_data: bytes, *, _depth: int = 0) -> list[dict[str, Any]]:
|
||||
limits = _attachments_limits()
|
||||
return process_archive(
|
||||
archive_name="bundle.zip",
|
||||
archive_data=zip_data,
|
||||
process_member=lambda n, b, d: process_file_data(n, b, _zip_depth=d),
|
||||
depth=_depth,
|
||||
limits={
|
||||
"max_depth": int(limits.get("archive_max_depth") or MAX_ARCHIVE_DEPTH),
|
||||
"max_file_count": int(limits.get("archive_max_file_count") or MAX_ARCHIVE_FILE_COUNT),
|
||||
"max_entry_bytes": int(limits.get("archive_max_entry_bytes") or MAX_ARCHIVE_ENTRY_BYTES),
|
||||
"max_total_uncompressed_bytes": int(
|
||||
limits.get("archive_max_total_uncompressed_bytes") or MAX_ARCHIVE_TOTAL_UNCOMPRESSED_BYTES
|
||||
),
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
def process_file_data(name: str, data: bytes, *, _zip_depth: int = 0) -> list[dict[str, Any]]:
|
||||
ext = name.split(".")[-1].lower() if "." in name else ""
|
||||
attachments: list[dict[str, Any]] = []
|
||||
|
||||
if ext in ("png", "jpg", "jpeg", "webp", "gif", "jfif", "pjpeg", "bmp", "tif", "tiff"):
|
||||
if ext == "png":
|
||||
mime = "image/png"
|
||||
elif ext in ("jpg", "jpeg", "jfif", "pjpeg"):
|
||||
mime = "image/jpeg"
|
||||
elif ext == "webp":
|
||||
mime = "image/webp"
|
||||
elif ext == "gif":
|
||||
mime = "image/gif"
|
||||
elif ext == "bmp":
|
||||
mime = "image/bmp"
|
||||
elif ext in ("tif", "tiff"):
|
||||
mime = "image/tiff"
|
||||
width = None
|
||||
height = None
|
||||
try:
|
||||
with Image.open(io.BytesIO(data)) as im:
|
||||
width, height = im.size
|
||||
except Exception:
|
||||
pass
|
||||
meta = AttachmentAssetStore().save_bytes(data, filename=name, mime=mime, width=width, height=height)
|
||||
attachments.append(
|
||||
{
|
||||
"type": "image_ref",
|
||||
"name": meta.name,
|
||||
"mime": meta.mime,
|
||||
"attachment_id": meta.attachment_id,
|
||||
"bytes": meta.bytes,
|
||||
"width": meta.width,
|
||||
"height": meta.height,
|
||||
}
|
||||
)
|
||||
return attachments
|
||||
|
||||
sniff_mime = _sniff_image_mime(data)
|
||||
if sniff_mime:
|
||||
width = None
|
||||
height = None
|
||||
try:
|
||||
with Image.open(io.BytesIO(data)) as im:
|
||||
width, height = im.size
|
||||
except Exception:
|
||||
pass
|
||||
meta = AttachmentAssetStore().save_bytes(data, filename=name, mime=sniff_mime, width=width, height=height)
|
||||
attachments.append(
|
||||
{
|
||||
"type": "image_ref",
|
||||
"name": meta.name,
|
||||
"mime": meta.mime,
|
||||
"attachment_id": meta.attachment_id,
|
||||
"bytes": meta.bytes,
|
||||
"width": meta.width,
|
||||
"height": meta.height,
|
||||
}
|
||||
)
|
||||
return attachments
|
||||
|
||||
if ext in ("mp4", "mov", "mkv", "webm", "avi", "m4v"):
|
||||
mime = "video/mp4"
|
||||
if ext == "webm":
|
||||
mime = "video/webm"
|
||||
elif ext == "mov":
|
||||
mime = "video/quicktime"
|
||||
elif ext == "mkv":
|
||||
mime = "video/x-matroska"
|
||||
elif ext == "avi":
|
||||
mime = "video/x-msvideo"
|
||||
meta = AttachmentAssetStore().save_bytes(data, filename=name, mime=mime)
|
||||
probe_meta: dict[str, Any] = {}
|
||||
try:
|
||||
p = AttachmentAssetStore().get_local_path(str(meta.attachment_id))
|
||||
if p is not None:
|
||||
probe_meta = _ffprobe_video_meta(p)
|
||||
except Exception:
|
||||
probe_meta = {}
|
||||
attachments.append(
|
||||
{
|
||||
"type": "video_ref",
|
||||
"name": meta.name,
|
||||
"mime": meta.mime,
|
||||
"attachment_id": meta.attachment_id,
|
||||
"bytes": meta.bytes,
|
||||
**probe_meta,
|
||||
}
|
||||
)
|
||||
attachments.append(
|
||||
{
|
||||
"type": "text",
|
||||
"name": name,
|
||||
"content": _truncate_text(
|
||||
"".join(
|
||||
[
|
||||
"# Video Attachment\n",
|
||||
f"- name: {meta.name}\n",
|
||||
f"- mime: {meta.mime}\n",
|
||||
f"- bytes: {meta.bytes}\n",
|
||||
(
|
||||
f"- duration_sec: {probe_meta.get('duration_sec')}\n"
|
||||
if probe_meta.get("duration_sec") is not None
|
||||
else ""
|
||||
),
|
||||
(
|
||||
f"- size: {probe_meta.get('width')}x{probe_meta.get('height')}\n"
|
||||
if probe_meta.get("width") and probe_meta.get("height")
|
||||
else ""
|
||||
),
|
||||
(f"- fps: {probe_meta.get('fps')}\n" if probe_meta.get("fps") is not None else ""),
|
||||
"- mode: video_ref\n",
|
||||
"- query_tool: query_video_attachment\n",
|
||||
"- note: use `attachment_id` from accompanying `video_ref` attachment.\n",
|
||||
]
|
||||
),
|
||||
limit=MAX_TEXT_CHARS,
|
||||
),
|
||||
}
|
||||
)
|
||||
return attachments
|
||||
|
||||
if ext in ("xlsx", "xls", "csv"):
|
||||
try:
|
||||
if ext == "xlsx":
|
||||
zip_err = _validate_excel_zip_payload(data)
|
||||
if zip_err:
|
||||
attachments.append({"type": "text", "name": name, "content": _truncate_text(f"Error parsing excel/csv: {zip_err}")})
|
||||
return attachments
|
||||
limits = _attachments_limits()
|
||||
rows_read = int(limits.get("rows_read") or DEFAULT_TABULAR_ROWS_READ)
|
||||
max_columns = int(limits.get("columns") or DEFAULT_TABULAR_COLUMNS)
|
||||
max_cell_chars = int(limits.get("cell_chars") or DEFAULT_TABULAR_CELL_CHARS)
|
||||
max_excel_sheets = int(limits.get("max_excel_sheets") or DEFAULT_MAX_EXCEL_SHEETS)
|
||||
large_table_preview_rows = int(
|
||||
limits.get("large_table_preview_rows") or DEFAULT_LARGE_TABLE_PREVIEW_ROWS
|
||||
)
|
||||
tool_mode_enabled = bool(limits.get("tool_mode_enabled", DEFAULT_TABULAR_TOOL_MODE_ENABLED))
|
||||
tool_mode_min_rows = int(limits.get("tool_mode_min_rows") or DEFAULT_TABULAR_TOOL_MODE_MIN_ROWS)
|
||||
tool_mode_max_bytes = int(limits.get("tool_mode_max_bytes") or DEFAULT_TABULAR_TOOL_MODE_MAX_BYTES)
|
||||
if ext == "csv":
|
||||
df = pd.read_csv(io.BytesIO(data), nrows=rows_read)
|
||||
else:
|
||||
df = pd.read_excel(io.BytesIO(data), nrows=rows_read)
|
||||
sampled = int(len(df.index)) >= rows_read
|
||||
estimated_rows = 0
|
||||
if ext == "csv":
|
||||
estimated_rows = max(0, int(data.count(b"\n")) - 1)
|
||||
else:
|
||||
estimated_rows = int(len(df.index))
|
||||
clipped_columns = int(len(df.columns)) > max_columns
|
||||
if clipped_columns:
|
||||
df = df.iloc[:, :max_columns]
|
||||
df, clipped_cells = _clip_tabular_cells(df, max_cell_chars=max_cell_chars)
|
||||
tool_mode = bool(
|
||||
tool_mode_enabled
|
||||
and int(len(data or b"")) <= tool_mode_max_bytes
|
||||
and (estimated_rows >= tool_mode_min_rows or (ext in ("xlsx", "xls") and sampled))
|
||||
)
|
||||
attachments.append(
|
||||
{
|
||||
"type": "text",
|
||||
"name": name,
|
||||
"content": _dataframe_summary_text(
|
||||
df,
|
||||
sampled=sampled,
|
||||
clipped_columns=clipped_columns,
|
||||
clipped_cells=clipped_cells,
|
||||
preview_rows=(large_table_preview_rows if tool_mode else EXCEL_PREVIEW_ROWS),
|
||||
)
|
||||
+ (
|
||||
(
|
||||
f"\n\n# Large Table Mode\n- estimated_rows: {estimated_rows}\n- query_tool: query_tabular_attachment\n"
|
||||
"- note: use the `table_id` from the accompanying `tabular_ref` attachment to query full data."
|
||||
)
|
||||
if tool_mode
|
||||
else ""
|
||||
),
|
||||
}
|
||||
)
|
||||
if tool_mode:
|
||||
workbook_sheets: dict[str, pd.DataFrame]
|
||||
clipped_sheets = False
|
||||
if ext == "csv":
|
||||
full_df = pd.read_csv(io.BytesIO(data), dtype=str)
|
||||
if int(len(full_df.columns)) > max_columns:
|
||||
full_df = full_df.iloc[:, :max_columns]
|
||||
full_df, _ = _clip_tabular_cells(full_df, max_cell_chars=max_cell_chars)
|
||||
workbook_sheets = {"Sheet1": full_df}
|
||||
else:
|
||||
excel_obj = pd.read_excel(io.BytesIO(data), dtype=str, sheet_name=None)
|
||||
workbook_sheets = {}
|
||||
sheet_items = list((excel_obj or {}).items())
|
||||
if len(sheet_items) > max_excel_sheets:
|
||||
clipped_sheets = True
|
||||
sheet_items = sheet_items[:max_excel_sheets]
|
||||
for sheet_name, sheet_df in sheet_items:
|
||||
sdf = sheet_df
|
||||
if int(len(sdf.columns)) > max_columns:
|
||||
sdf = sdf.iloc[:, :max_columns]
|
||||
sdf, _ = _clip_tabular_cells(sdf, max_cell_chars=max_cell_chars)
|
||||
workbook_sheets[str(sheet_name or f"Sheet{len(workbook_sheets)+1}")] = sdf
|
||||
mime = "text/csv" if ext == "csv" else "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet"
|
||||
meta = AttachmentAssetStore().save_bytes(data, filename=name, mime=mime)
|
||||
table_meta = (
|
||||
save_dataframe(attachment_id=str(meta.attachment_id), name=name, df=workbook_sheets.get("Sheet1", pd.DataFrame()))
|
||||
if ext == "csv"
|
||||
else save_workbook(attachment_id=str(meta.attachment_id), name=name, sheets=workbook_sheets)
|
||||
)
|
||||
attachments.append(
|
||||
{
|
||||
"type": "tabular_ref",
|
||||
"name": str(name),
|
||||
"attachment_id": str(meta.attachment_id),
|
||||
"table_id": str(table_meta.get("table_id") or ""),
|
||||
"rows": int(table_meta.get("rows") or 0),
|
||||
"cols": int(table_meta.get("cols") or 0),
|
||||
"columns": list(table_meta.get("columns") or []),
|
||||
"sheets": list(table_meta.get("sheets") or []),
|
||||
}
|
||||
)
|
||||
if clipped_sheets:
|
||||
attachments.append(
|
||||
{
|
||||
"type": "text",
|
||||
"name": f"{name}.sheet-limit",
|
||||
"content": (
|
||||
f"Excel workbook contains too many sheets; only the first {max_excel_sheets} sheets were stored "
|
||||
"for tool-mode querying. Adjust max_excel_sheets if full coverage is required."
|
||||
),
|
||||
}
|
||||
)
|
||||
except Exception as e:
|
||||
attachments.append({"type": "text", "name": name, "content": _truncate_text(f"Error parsing excel/csv: {e}")})
|
||||
|
||||
elif ext == "docx":
|
||||
try:
|
||||
doc = Document(io.BytesIO(data))
|
||||
text = "\n".join([p.text for p in doc.paragraphs])
|
||||
limits = _attachments_limits()
|
||||
inline_cap = int(limits.get("text_inline_max_chars") or DEFAULT_TEXT_INLINE_MAX_CHARS)
|
||||
if len(str(text or "")) <= inline_cap:
|
||||
attachments.append({"type": "text", "name": name, "content": _truncate_text(text, limit=inline_cap)})
|
||||
else:
|
||||
meta = AttachmentAssetStore().save_bytes(
|
||||
data,
|
||||
filename=name,
|
||||
mime="application/vnd.openxmlformats-officedocument.wordprocessingml.document",
|
||||
)
|
||||
text_meta = save_text_document(
|
||||
attachment_id=str(meta.attachment_id),
|
||||
name=name,
|
||||
text=str(text or ""),
|
||||
source_kind="docx",
|
||||
chunk_size=int(limits.get("text_chunk_size") or DEFAULT_TEXT_CHUNK_SIZE),
|
||||
chunk_overlap=int(limits.get("text_chunk_overlap") or DEFAULT_TEXT_CHUNK_OVERLAP),
|
||||
)
|
||||
top_k = int(limits.get("text_query_top_k") or DEFAULT_TEXT_QUERY_TOP_K)
|
||||
attachments.append(
|
||||
{
|
||||
"type": "text",
|
||||
"name": name,
|
||||
"content": _text_ref_summary(
|
||||
name=name,
|
||||
chars=int(text_meta.get("chars") or 0),
|
||||
chunks=int(text_meta.get("chunks") or 0),
|
||||
source_kind="docx",
|
||||
preview=str(text or ""),
|
||||
top_k=top_k,
|
||||
),
|
||||
}
|
||||
)
|
||||
attachments.append(
|
||||
{
|
||||
"type": "text_ref",
|
||||
"name": str(name),
|
||||
"attachment_id": str(meta.attachment_id),
|
||||
"text_id": str(text_meta.get("text_id") or ""),
|
||||
"chars": int(text_meta.get("chars") or 0),
|
||||
"chunks": int(text_meta.get("chunks") or 0),
|
||||
"source_kind": "docx",
|
||||
}
|
||||
)
|
||||
except Exception as e:
|
||||
attachments.append({"type": "text", "name": name, "content": _truncate_text(f"Error parsing docx: {e}")})
|
||||
|
||||
elif ext == "pdf":
|
||||
try:
|
||||
reader = PdfReader(io.BytesIO(data))
|
||||
text = _pdf_summary_text(reader)
|
||||
limits = _attachments_limits()
|
||||
inline_cap = int(limits.get("text_inline_max_chars") or DEFAULT_TEXT_INLINE_MAX_CHARS)
|
||||
if len(str(text or "")) <= inline_cap:
|
||||
attachments.append({"type": "text", "name": name, "content": _truncate_text(text, limit=inline_cap)})
|
||||
else:
|
||||
meta = AttachmentAssetStore().save_bytes(data, filename=name, mime="application/pdf")
|
||||
text_meta = save_text_document(
|
||||
attachment_id=str(meta.attachment_id),
|
||||
name=name,
|
||||
text=str(text or ""),
|
||||
source_kind="pdf",
|
||||
chunk_size=int(limits.get("text_chunk_size") or DEFAULT_TEXT_CHUNK_SIZE),
|
||||
chunk_overlap=int(limits.get("text_chunk_overlap") or DEFAULT_TEXT_CHUNK_OVERLAP),
|
||||
)
|
||||
top_k = int(limits.get("text_query_top_k") or DEFAULT_TEXT_QUERY_TOP_K)
|
||||
attachments.append(
|
||||
{
|
||||
"type": "text",
|
||||
"name": name,
|
||||
"content": _text_ref_summary(
|
||||
name=name,
|
||||
chars=int(text_meta.get("chars") or 0),
|
||||
chunks=int(text_meta.get("chunks") or 0),
|
||||
source_kind="pdf",
|
||||
preview=str(text or ""),
|
||||
top_k=top_k,
|
||||
),
|
||||
}
|
||||
)
|
||||
attachments.append(
|
||||
{
|
||||
"type": "text_ref",
|
||||
"name": str(name),
|
||||
"attachment_id": str(meta.attachment_id),
|
||||
"text_id": str(text_meta.get("text_id") or ""),
|
||||
"chars": int(text_meta.get("chars") or 0),
|
||||
"chunks": int(text_meta.get("chunks") or 0),
|
||||
"source_kind": "pdf",
|
||||
}
|
||||
)
|
||||
except Exception as e:
|
||||
attachments.append({"type": "text", "name": name, "content": _truncate_text(f"Error parsing pdf: {e}")})
|
||||
|
||||
elif ext in ("txt", "py", "js", "html", "css", "md", "json", "yaml", "yml", "sh", "log"):
|
||||
text = _decode_text_bytes(data)
|
||||
if ext == "html":
|
||||
text = _clean_html_to_text(text)
|
||||
limits = _attachments_limits()
|
||||
inline_cap = int(limits.get("text_inline_max_chars") or DEFAULT_TEXT_INLINE_MAX_CHARS)
|
||||
if len(str(text or "")) <= inline_cap:
|
||||
attachments.append({"type": "text", "name": name, "content": _truncate_text(text, limit=inline_cap)})
|
||||
else:
|
||||
meta = AttachmentAssetStore().save_bytes(data, filename=name, mime="text/plain")
|
||||
text_meta = save_text_document(
|
||||
attachment_id=str(meta.attachment_id),
|
||||
name=name,
|
||||
text=str(text or ""),
|
||||
source_kind=ext or "text",
|
||||
chunk_size=int(limits.get("text_chunk_size") or DEFAULT_TEXT_CHUNK_SIZE),
|
||||
chunk_overlap=int(limits.get("text_chunk_overlap") or DEFAULT_TEXT_CHUNK_OVERLAP),
|
||||
)
|
||||
top_k = int(limits.get("text_query_top_k") or DEFAULT_TEXT_QUERY_TOP_K)
|
||||
attachments.append(
|
||||
{
|
||||
"type": "text",
|
||||
"name": name,
|
||||
"content": _text_ref_summary(
|
||||
name=name,
|
||||
chars=int(text_meta.get("chars") or 0),
|
||||
chunks=int(text_meta.get("chunks") or 0),
|
||||
source_kind=str(text_meta.get("source_kind") or ext or "text"),
|
||||
preview=str(text or ""),
|
||||
top_k=top_k,
|
||||
),
|
||||
}
|
||||
)
|
||||
attachments.append(
|
||||
{
|
||||
"type": "text_ref",
|
||||
"name": str(name),
|
||||
"attachment_id": str(meta.attachment_id),
|
||||
"text_id": str(text_meta.get("text_id") or ""),
|
||||
"chars": int(text_meta.get("chars") or 0),
|
||||
"chunks": int(text_meta.get("chunks") or 0),
|
||||
"source_kind": str(text_meta.get("source_kind") or ext or "text"),
|
||||
}
|
||||
)
|
||||
|
||||
elif detect_archive_kind(name):
|
||||
limits = _attachments_limits()
|
||||
attachments.extend(
|
||||
process_archive(
|
||||
archive_name=name,
|
||||
archive_data=data,
|
||||
process_member=lambda n, b, d: process_file_data(n, b, _zip_depth=d),
|
||||
depth=_zip_depth,
|
||||
limits={
|
||||
"max_depth": int(limits.get("archive_max_depth") or MAX_ARCHIVE_DEPTH),
|
||||
"max_file_count": int(limits.get("archive_max_file_count") or MAX_ARCHIVE_FILE_COUNT),
|
||||
"max_entry_bytes": int(limits.get("archive_max_entry_bytes") or MAX_ARCHIVE_ENTRY_BYTES),
|
||||
"max_total_uncompressed_bytes": int(
|
||||
limits.get("archive_max_total_uncompressed_bytes") or MAX_ARCHIVE_TOTAL_UNCOMPRESSED_BYTES
|
||||
),
|
||||
},
|
||||
)
|
||||
)
|
||||
|
||||
return attachments
|
||||
|
||||
|
||||
__all__ = ["process_zip", "process_file_data", "clear_attachment_limits_cache"]
|
||||
66
svc/files/session_export.py
Normal file
66
svc/files/session_export.py
Normal file
|
|
@ -0,0 +1,66 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from typing import Any
|
||||
|
||||
from svc.persistence.sqlite_store import ChatMessage, SqliteStore
|
||||
|
||||
|
||||
def export_session_json(store: SqliteStore, session_id: str, message_limit: int = 50_000) -> str:
|
||||
title = ""
|
||||
s = store.get_session(session_id)
|
||||
if s:
|
||||
title = s.title
|
||||
msgs = store.get_messages(session_id=session_id, limit=message_limit)
|
||||
payload: dict[str, Any] = {
|
||||
"session_id": session_id,
|
||||
"title": title,
|
||||
"messages": [_message_to_dict(m) for m in msgs],
|
||||
}
|
||||
return json.dumps(payload, ensure_ascii=False, indent=2)
|
||||
|
||||
|
||||
def export_session_markdown(store: SqliteStore, session_id: str, message_limit: int = 50_000) -> str:
|
||||
s = store.get_session(session_id)
|
||||
title = (s.title if s else "") or session_id[:12]
|
||||
lines = [f"# {title}", "", f"Session: `{session_id}`", ""]
|
||||
for m in store.get_messages(session_id=session_id, limit=message_limit):
|
||||
lines.extend(_message_to_markdown_lines(m))
|
||||
lines.append("")
|
||||
return "\n".join(lines).rstrip() + "\n"
|
||||
|
||||
|
||||
def _message_to_dict(m: ChatMessage) -> dict[str, Any]:
|
||||
return {
|
||||
"id": m.id,
|
||||
"role": m.role,
|
||||
"content": m.content,
|
||||
"tool_calls": m.tool_calls,
|
||||
"attachments": m.attachments,
|
||||
"timestamp": m.timestamp,
|
||||
}
|
||||
|
||||
|
||||
def _message_to_markdown_lines(m: ChatMessage) -> list[str]:
|
||||
role = m.role.upper()
|
||||
ts = m.timestamp or ""
|
||||
head = f"## {role} · `{m.id}` · {ts}"
|
||||
body = m.content or ""
|
||||
if m.role == "tool" and m.tool_calls:
|
||||
head = f"## TOOL · `{m.id}` · {ts}"
|
||||
lines = [head, ""]
|
||||
if body.strip():
|
||||
lines.append(body)
|
||||
lines.append("")
|
||||
if m.tool_calls and m.role != "tool":
|
||||
lines.append("```json")
|
||||
lines.append(m.tool_calls)
|
||||
lines.append("```")
|
||||
elif m.tool_calls and m.role == "tool":
|
||||
lines.append(f"_meta_: `{m.tool_calls}`")
|
||||
if m.attachments:
|
||||
lines.append("_attachments_: see JSON export for full data.")
|
||||
return lines
|
||||
|
||||
|
||||
__all__ = ["export_session_json", "export_session_markdown"]
|
||||
732
svc/files/tabular_attachment_store.py
Normal file
732
svc/files/tabular_attachment_store.py
Normal file
|
|
@ -0,0 +1,732 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import hashlib
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
import sqlite3
|
||||
import subprocess
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import pandas as pd
|
||||
|
||||
from svc.config.paths import PROJECT_ROOT
|
||||
from svc.config.paths import attachments_dir
|
||||
|
||||
|
||||
def _tabular_root() -> Path:
|
||||
p = (attachments_dir() / "tabular").resolve()
|
||||
p.mkdir(parents=True, exist_ok=True)
|
||||
return p
|
||||
|
||||
|
||||
def _table_id(attachment_id: str, name: str) -> str:
|
||||
raw = f"{attachment_id}:{name}".encode("utf-8", errors="ignore")
|
||||
return hashlib.sha256(raw).hexdigest()
|
||||
|
||||
|
||||
def _db_path(table_id: str) -> Path:
|
||||
return _tabular_root() / f"{table_id}.sqlite"
|
||||
|
||||
|
||||
def _meta_path(table_id: str) -> Path:
|
||||
return _tabular_root() / f"{table_id}.meta.json"
|
||||
|
||||
|
||||
def _is_valid_table_id(table_id: str) -> bool:
|
||||
tid = str(table_id or "").strip().lower()
|
||||
return bool(re.fullmatch(r"[0-9a-f]{64}", tid))
|
||||
|
||||
|
||||
def _normalize_columns(raw_cols: list[Any]) -> tuple[list[str], list[dict[str, Any]]]:
|
||||
out: list[str] = []
|
||||
mapping: list[dict[str, Any]] = []
|
||||
seen: dict[str, int] = {}
|
||||
for idx, raw in enumerate(raw_cols):
|
||||
original = "" if raw is None else str(raw)
|
||||
base = original.strip()
|
||||
if not base or base.lower().startswith("unnamed:"):
|
||||
base = f"col_{idx + 1}"
|
||||
base = re.sub(r"\s+", "_", base)
|
||||
base = re.sub(r"\W+", "_", base, flags=re.UNICODE).strip("_")
|
||||
if not base:
|
||||
base = f"col_{idx + 1}"
|
||||
n = int(seen.get(base, 0)) + 1
|
||||
seen[base] = n
|
||||
name = base if n == 1 else f"{base}__{n}"
|
||||
out.append(name)
|
||||
mapping.append({"index": idx, "original": original, "normalized": name})
|
||||
return out, mapping
|
||||
|
||||
|
||||
def _prepare_writable_df(df: pd.DataFrame) -> tuple[pd.DataFrame, list[str], list[dict[str, Any]]]:
|
||||
raw_cols = [x for x in list(df.columns)]
|
||||
cols, col_map = _normalize_columns(raw_cols)
|
||||
writable = df.copy()
|
||||
writable.columns = cols
|
||||
for c in cols:
|
||||
writable[c] = writable[c].map(lambda x: "" if x is None else str(x))
|
||||
return writable, cols, col_map
|
||||
|
||||
|
||||
def _safe_sheet_table_name(sheet_name: str, idx: int) -> str:
|
||||
base = re.sub(r"[^0-9a-zA-Z_]+", "_", str(sheet_name or "").strip()).strip("_").lower()
|
||||
if not base:
|
||||
base = f"sheet_{idx + 1}"
|
||||
return f"sheet_{idx + 1}_{base}"
|
||||
|
||||
|
||||
def _mcp_sqlite_command() -> list[str]:
|
||||
raw = str(os.getenv("AIA_MCP_SQLITE_COMMAND") or "").strip()
|
||||
if not raw:
|
||||
return []
|
||||
return [x for x in raw.split(" ") if x]
|
||||
|
||||
|
||||
def _tabular_use_mcp() -> bool:
|
||||
raw = str(os.getenv("AIA_TABULAR_USE_MCP") or "").strip().lower()
|
||||
return raw in {"1", "true", "yes", "on"}
|
||||
|
||||
|
||||
def _tabular_sql_timeout_ms() -> int:
|
||||
def _clamp_timeout(raw: Any, default: int = 8_000) -> int:
|
||||
try:
|
||||
n = int(raw)
|
||||
except Exception:
|
||||
return default
|
||||
return max(100, min(n, 120_000))
|
||||
|
||||
raw = str(os.getenv("AIA_TABULAR_SQL_TIMEOUT_MS") or "").strip()
|
||||
if raw:
|
||||
return _clamp_timeout(raw)
|
||||
cfg_path_raw = str(os.getenv("AIA_OCLAW_CONFIG_PATH") or "").strip()
|
||||
cfg_path = Path(cfg_path_raw).expanduser() if cfg_path_raw else (Path(PROJECT_ROOT) / "oclaw.json").resolve()
|
||||
if not cfg_path.is_absolute():
|
||||
cfg_path = (Path(PROJECT_ROOT) / cfg_path).resolve()
|
||||
try:
|
||||
obj = json.loads(cfg_path.read_text(encoding="utf-8"))
|
||||
except Exception:
|
||||
return 8_000
|
||||
tabular_cfg = (
|
||||
(((obj or {}).get("plugins") or {}).get("entries") or {}).get("memory-wiki", {}).get("auto", {}).get("attachments", {}).get("tabular", {}) or {}
|
||||
)
|
||||
if not isinstance(tabular_cfg, dict):
|
||||
return 8_000
|
||||
return _clamp_timeout(tabular_cfg.get("sql_timeout_ms"), 8_000)
|
||||
|
||||
|
||||
def _extract_rows_from_mcp_result(res: dict[str, Any]) -> list[dict[str, Any]] | None:
|
||||
if not isinstance(res, dict) or not bool(res.get("ok")):
|
||||
return None
|
||||
data = res.get("data") if isinstance(res.get("data"), dict) else res.get("result")
|
||||
if isinstance(data, dict):
|
||||
direct_rows = data.get("rows")
|
||||
if isinstance(direct_rows, list):
|
||||
out: list[dict[str, Any]] = []
|
||||
for r in direct_rows:
|
||||
if isinstance(r, dict):
|
||||
out.append({str(k): r[k] for k in r.keys()})
|
||||
return out
|
||||
content = data.get("content")
|
||||
if isinstance(content, list):
|
||||
for it in content:
|
||||
if not isinstance(it, dict):
|
||||
continue
|
||||
txt = str(it.get("text") or "").strip()
|
||||
if not txt:
|
||||
continue
|
||||
try:
|
||||
obj = json.loads(txt)
|
||||
except Exception:
|
||||
continue
|
||||
if isinstance(obj, dict) and isinstance(obj.get("rows"), list):
|
||||
rows = obj.get("rows")
|
||||
out: list[dict[str, Any]] = []
|
||||
for r in rows:
|
||||
if isinstance(r, dict):
|
||||
out.append({str(k): r[k] for k in r.keys()})
|
||||
return out
|
||||
return None
|
||||
|
||||
|
||||
def _run_select_via_mcp(*, db_path: Path, sql: str, params: list[Any]) -> list[dict[str, Any]] | None:
|
||||
if not _tabular_use_mcp():
|
||||
return None
|
||||
cmd = _mcp_sqlite_command()
|
||||
if not cmd:
|
||||
return None
|
||||
|
||||
def _call_once(tool_name: str, arguments: dict[str, Any]) -> dict[str, Any]:
|
||||
try:
|
||||
p = subprocess.Popen(
|
||||
cmd,
|
||||
stdin=subprocess.PIPE,
|
||||
stdout=subprocess.PIPE,
|
||||
stderr=subprocess.PIPE,
|
||||
text=True,
|
||||
encoding="utf-8",
|
||||
)
|
||||
except Exception:
|
||||
return {"ok": False}
|
||||
try:
|
||||
if p.stdin is None or p.stdout is None:
|
||||
return {"ok": False}
|
||||
init_req = {
|
||||
"jsonrpc": "2.0",
|
||||
"id": 1,
|
||||
"method": "initialize",
|
||||
"params": {
|
||||
"protocolVersion": "2024-11-05",
|
||||
"capabilities": {},
|
||||
"clientInfo": {"name": "oclaw-tabular", "version": "0.1.0"},
|
||||
},
|
||||
}
|
||||
p.stdin.write(json.dumps(init_req, ensure_ascii=False) + "\n")
|
||||
p.stdin.flush()
|
||||
_ = p.stdout.readline()
|
||||
p.stdin.write(json.dumps({"jsonrpc": "2.0", "method": "notifications/initialized", "params": {}}, ensure_ascii=False) + "\n")
|
||||
p.stdin.flush()
|
||||
call_req = {
|
||||
"jsonrpc": "2.0",
|
||||
"id": 2,
|
||||
"method": "tools/call",
|
||||
"params": {"name": str(tool_name), "arguments": arguments},
|
||||
}
|
||||
p.stdin.write(json.dumps(call_req, ensure_ascii=False) + "\n")
|
||||
p.stdin.flush()
|
||||
line = p.stdout.readline()
|
||||
if not line:
|
||||
return {"ok": False}
|
||||
obj = json.loads(line)
|
||||
if not isinstance(obj, dict):
|
||||
return {"ok": False}
|
||||
if isinstance(obj.get("error"), dict):
|
||||
return {"ok": False}
|
||||
return {"ok": True, "result": obj.get("result")}
|
||||
except Exception:
|
||||
return {"ok": False}
|
||||
finally:
|
||||
try:
|
||||
p.terminate()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# Probe several common sqlite MCP tool names/argument shapes.
|
||||
attempts: list[tuple[str, dict[str, Any]]] = [
|
||||
("query", {"database": str(db_path), "sql": sql, "params": params}),
|
||||
("query", {"db_path": str(db_path), "sql": sql, "params": params}),
|
||||
("execute_sql", {"database": str(db_path), "sql": sql, "params": params}),
|
||||
("read_query", {"database": str(db_path), "query": sql, "params": params}),
|
||||
("sqlite_query", {"database": str(db_path), "sql": sql, "params": params}),
|
||||
]
|
||||
for tool_name, args in attempts:
|
||||
res = _call_once(tool_name=tool_name, arguments=args)
|
||||
rows = _extract_rows_from_mcp_result(res if isinstance(res, dict) else {})
|
||||
if rows is not None:
|
||||
return rows
|
||||
return None
|
||||
|
||||
|
||||
def _run_select_sql(
|
||||
*,
|
||||
db_path: Path,
|
||||
sql: str,
|
||||
params: list[Any],
|
||||
timeout_ms: int | None = None,
|
||||
) -> tuple[list[dict[str, Any]], str]:
|
||||
rows = _run_select_via_mcp(db_path=db_path, sql=sql, params=params)
|
||||
if rows is not None:
|
||||
return rows, "mcp_sqlite"
|
||||
out_rows: list[dict[str, Any]] = []
|
||||
timeout_ms_eff = max(100, int(timeout_ms or _tabular_sql_timeout_ms()))
|
||||
deadline = time.perf_counter() + (float(timeout_ms_eff) / 1000.0)
|
||||
timed_out = {"hit": False}
|
||||
with sqlite3.connect(str(db_path)) as conn:
|
||||
conn.row_factory = sqlite3.Row
|
||||
|
||||
def _progress_cb() -> int:
|
||||
if time.perf_counter() > deadline:
|
||||
timed_out["hit"] = True
|
||||
return 1
|
||||
return 0
|
||||
|
||||
conn.set_progress_handler(_progress_cb, 1000)
|
||||
try:
|
||||
for r in conn.execute(sql, params).fetchall():
|
||||
out_rows.append({k: r[k] for k in r.keys()})
|
||||
except sqlite3.OperationalError as exc:
|
||||
if timed_out["hit"] and "interrupted" in str(exc).lower():
|
||||
raise TimeoutError(f"tabular_sql_timeout_ms={timeout_ms_eff}") from exc
|
||||
raise
|
||||
finally:
|
||||
conn.set_progress_handler(None, 0)
|
||||
return out_rows, "builtin_sqlite"
|
||||
|
||||
|
||||
def save_dataframe(*, attachment_id: str, name: str, df: pd.DataFrame) -> dict[str, Any]:
|
||||
return save_workbook(attachment_id=attachment_id, name=name, sheets={"Sheet1": df})
|
||||
|
||||
|
||||
def save_workbook(*, attachment_id: str, name: str, sheets: dict[str, pd.DataFrame]) -> dict[str, Any]:
|
||||
tid = _table_id(attachment_id=attachment_id, name=name)
|
||||
db = _db_path(tid)
|
||||
meta = _meta_path(tid)
|
||||
sheet_rows: list[dict[str, Any]] = []
|
||||
default_table = "rows_data"
|
||||
default_columns: list[str] = []
|
||||
default_rows = 0
|
||||
default_cols = 0
|
||||
with sqlite3.connect(str(db)) as conn:
|
||||
for idx, (sheet_name, df) in enumerate((sheets or {}).items()):
|
||||
writable, cols, col_map = _prepare_writable_df(df)
|
||||
table_name = _safe_sheet_table_name(str(sheet_name or f"Sheet{idx + 1}"), idx)
|
||||
writable.to_sql(table_name, conn, if_exists="replace", index=False)
|
||||
if idx == 0:
|
||||
# backward compatibility: old queries use rows_data
|
||||
writable.to_sql("rows_data", conn, if_exists="replace", index=False)
|
||||
default_table = table_name
|
||||
default_columns = list(cols)
|
||||
default_rows = int(len(writable.index))
|
||||
default_cols = int(len(cols))
|
||||
sheet_rows.append(
|
||||
{
|
||||
"sheet_name": str(sheet_name or f"Sheet{idx + 1}"),
|
||||
"table_name": table_name,
|
||||
"rows": int(len(writable.index)),
|
||||
"cols": int(len(cols)),
|
||||
"columns": list(cols),
|
||||
"column_map": col_map,
|
||||
}
|
||||
)
|
||||
payload = {
|
||||
"table_id": tid,
|
||||
"attachment_id": str(attachment_id or ""),
|
||||
"name": str(name or ""),
|
||||
"rows": int(default_rows),
|
||||
"cols": int(default_cols),
|
||||
"columns": default_columns,
|
||||
"default_table": default_table,
|
||||
"sheets": sheet_rows,
|
||||
"db_path": str(db),
|
||||
}
|
||||
meta.write_text(json.dumps(payload, ensure_ascii=False), encoding="utf-8")
|
||||
return payload
|
||||
|
||||
|
||||
def get_table_meta(table_id: str) -> dict[str, Any] | None:
|
||||
tid = str(table_id or "").strip()
|
||||
if not _is_valid_table_id(tid):
|
||||
return None
|
||||
p = _meta_path(tid)
|
||||
if not p.exists():
|
||||
return None
|
||||
try:
|
||||
obj = json.loads(p.read_text(encoding="utf-8"))
|
||||
return obj if isinstance(obj, dict) else None
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
|
||||
def _resolve_sheet_table(meta: dict[str, Any], sheet: str | None) -> str:
|
||||
default_table = str(meta.get("default_table") or "rows_data")
|
||||
if not sheet:
|
||||
return default_table
|
||||
target = str(sheet).strip()
|
||||
if not target:
|
||||
return default_table
|
||||
sheets = meta.get("sheets") if isinstance(meta.get("sheets"), list) else []
|
||||
for row in sheets:
|
||||
if not isinstance(row, dict):
|
||||
continue
|
||||
if str(row.get("sheet_name") or "") == target or str(row.get("table_name") or "") == target:
|
||||
return str(row.get("table_name") or default_table)
|
||||
return ""
|
||||
|
||||
|
||||
def _columns_for_sheet(meta: dict[str, Any], table_name: str) -> list[str]:
|
||||
sheets = meta.get("sheets") if isinstance(meta.get("sheets"), list) else []
|
||||
for row in sheets:
|
||||
if not isinstance(row, dict):
|
||||
continue
|
||||
if str(row.get("table_name") or "") == str(table_name):
|
||||
cols = row.get("columns")
|
||||
if isinstance(cols, list):
|
||||
return [str(x) for x in cols]
|
||||
cols = meta.get("columns")
|
||||
if isinstance(cols, list):
|
||||
return [str(x) for x in cols]
|
||||
return []
|
||||
|
||||
|
||||
def _sheet_hint_payload(meta: dict[str, Any], *, selected_sheet: str | None, table_name: str) -> dict[str, Any] | None:
|
||||
sheets = meta.get("sheets") if isinstance(meta.get("sheets"), list) else []
|
||||
valid_rows = [x for x in sheets if isinstance(x, dict)]
|
||||
if len(valid_rows) <= 1:
|
||||
return None
|
||||
if str(selected_sheet or "").strip():
|
||||
return None
|
||||
names = [str((x or {}).get("sheet_name") or "") for x in valid_rows]
|
||||
names = [x for x in names if x]
|
||||
if len(names) <= 1:
|
||||
return None
|
||||
default_sheet_name = ""
|
||||
for row in valid_rows:
|
||||
if str(row.get("table_name") or "") == str(table_name):
|
||||
default_sheet_name = str(row.get("sheet_name") or "")
|
||||
break
|
||||
return {
|
||||
"available_sheets": names[:20],
|
||||
"default_sheet": default_sheet_name,
|
||||
"note": "Workbook contains multiple sheets; pass `sheet` to target a specific sheet.",
|
||||
}
|
||||
|
||||
|
||||
def query_table(
|
||||
*,
|
||||
table_id: str,
|
||||
columns: list[str] | None = None,
|
||||
limit: int = 50,
|
||||
offset: int = 0,
|
||||
where_contains: dict[str, str] | None = None,
|
||||
sheet: str | None = None,
|
||||
) -> dict[str, Any]:
|
||||
if not _is_valid_table_id(table_id):
|
||||
return {
|
||||
"ok": False,
|
||||
"error": "table_id_invalid_format",
|
||||
"hint": "Use table_id from tabular_ref attachment (64-char hex id).",
|
||||
}
|
||||
meta = get_table_meta(table_id)
|
||||
if not isinstance(meta, dict):
|
||||
return {"ok": False, "error": "table_not_found"}
|
||||
db = Path(str(meta.get("db_path") or ""))
|
||||
if not db.exists():
|
||||
return {"ok": False, "error": "table_store_missing"}
|
||||
table_name = _resolve_sheet_table(meta, sheet)
|
||||
if not table_name:
|
||||
return {"ok": False, "error": "sheet_not_found"}
|
||||
known_cols = _columns_for_sheet(meta, table_name)
|
||||
use_cols = [c for c in (columns or known_cols) if c in known_cols]
|
||||
if not use_cols:
|
||||
use_cols = known_cols[: min(10, len(known_cols))]
|
||||
lim = max(1, min(int(limit or 50), 200))
|
||||
off = max(0, int(offset or 0))
|
||||
select_sql = ", ".join([f'"{c}"' for c in use_cols]) if use_cols else "*"
|
||||
where_sql = ""
|
||||
params: list[Any] = []
|
||||
wc = where_contains if isinstance(where_contains, dict) else {}
|
||||
wc_col = str(wc.get("column") or "").strip()
|
||||
wc_kw = str(wc.get("keyword") or "").strip()
|
||||
if wc_col and wc_kw and wc_col in known_cols:
|
||||
where_sql = f' WHERE "{wc_col}" LIKE ? '
|
||||
params.append(f"%{wc_kw}%")
|
||||
sql = f'SELECT {select_sql} FROM "{table_name}"{where_sql} LIMIT ? OFFSET ?'
|
||||
params.extend([lim, off])
|
||||
out_rows: list[dict[str, Any]] = []
|
||||
total = 0
|
||||
out_rows, engine = _run_select_sql(db_path=db, sql=sql, params=params)
|
||||
with sqlite3.connect(str(db)) as conn:
|
||||
total = int(conn.execute(f'SELECT COUNT(1) AS n FROM "{table_name}"').fetchone()[0] or 0)
|
||||
hint = _sheet_hint_payload(meta, selected_sheet=sheet, table_name=table_name)
|
||||
return {
|
||||
"ok": True,
|
||||
"table_id": str(table_id),
|
||||
"name": str(meta.get("name") or ""),
|
||||
"sheet": str(sheet or ""),
|
||||
"table_name": table_name,
|
||||
"rows_total": total,
|
||||
"columns": use_cols,
|
||||
"rows": out_rows,
|
||||
"limit": lim,
|
||||
"offset": off,
|
||||
"engine": engine,
|
||||
**({"sheet_hint": hint} if isinstance(hint, dict) else {}),
|
||||
}
|
||||
|
||||
|
||||
def aggregate_table(
|
||||
*,
|
||||
table_id: str,
|
||||
metric: str,
|
||||
target_column: str | None = None,
|
||||
group_by: str | None = None,
|
||||
where_contains: dict[str, str] | None = None,
|
||||
top_n: int = 20,
|
||||
sheet: str | None = None,
|
||||
) -> dict[str, Any]:
|
||||
if not _is_valid_table_id(table_id):
|
||||
return {
|
||||
"ok": False,
|
||||
"error": "table_id_invalid_format",
|
||||
"hint": "Use table_id from tabular_ref attachment (64-char hex id).",
|
||||
}
|
||||
meta = get_table_meta(table_id)
|
||||
if not isinstance(meta, dict):
|
||||
return {"ok": False, "error": "table_not_found"}
|
||||
db = Path(str(meta.get("db_path") or ""))
|
||||
if not db.exists():
|
||||
return {"ok": False, "error": "table_store_missing"}
|
||||
table_name = _resolve_sheet_table(meta, sheet)
|
||||
if not table_name:
|
||||
return {"ok": False, "error": "sheet_not_found"}
|
||||
known_cols = _columns_for_sheet(meta, table_name)
|
||||
mt = str(metric or "").strip().lower()
|
||||
if mt not in {"count", "sum", "avg"}:
|
||||
return {"ok": False, "error": "invalid_metric"}
|
||||
grp = str(group_by or "").strip()
|
||||
if grp and grp not in known_cols:
|
||||
return {"ok": False, "error": "invalid_group_by"}
|
||||
tgt = str(target_column or "").strip()
|
||||
if mt in {"sum", "avg"} and tgt not in known_cols:
|
||||
return {"ok": False, "error": "target_column_required"}
|
||||
where_sql = ""
|
||||
params: list[Any] = []
|
||||
wc = where_contains if isinstance(where_contains, dict) else {}
|
||||
wc_col = str(wc.get("column") or "").strip()
|
||||
wc_kw = str(wc.get("keyword") or "").strip()
|
||||
if wc_col and wc_kw and wc_col in known_cols:
|
||||
where_sql = f' WHERE "{wc_col}" LIKE ? '
|
||||
params.append(f"%{wc_kw}%")
|
||||
lim = max(1, min(int(top_n or 20), 200))
|
||||
if mt == "count":
|
||||
metric_sql = "COUNT(1)"
|
||||
else:
|
||||
metric_sql = f'{"SUM" if mt == "sum" else "AVG"}(CAST(NULLIF("{tgt}", \'\') AS REAL))'
|
||||
if grp:
|
||||
sql = (
|
||||
f'SELECT "{grp}" AS group_key, {metric_sql} AS metric_value '
|
||||
f'FROM "{table_name}"{where_sql} GROUP BY "{grp}" '
|
||||
"ORDER BY metric_value DESC LIMIT ?"
|
||||
)
|
||||
params.append(lim)
|
||||
raw_rows, engine = _run_select_sql(db_path=db, sql=sql, params=params)
|
||||
rows: list[dict[str, Any]] = []
|
||||
for r in raw_rows:
|
||||
rows.append(
|
||||
{
|
||||
"group": r.get("group_key"),
|
||||
"value": float(r.get("metric_value") or 0.0),
|
||||
}
|
||||
)
|
||||
out = {
|
||||
"ok": True,
|
||||
"table_id": str(table_id),
|
||||
"sheet": str(sheet or ""),
|
||||
"table_name": table_name,
|
||||
"metric": mt,
|
||||
"target_column": tgt if mt in {"sum", "avg"} else None,
|
||||
"group_by": grp,
|
||||
"rows": rows,
|
||||
"top_n": lim,
|
||||
"engine": engine,
|
||||
}
|
||||
hint = _sheet_hint_payload(meta, selected_sheet=sheet, table_name=table_name)
|
||||
if isinstance(hint, dict):
|
||||
out["sheet_hint"] = hint
|
||||
return out
|
||||
sql = f'SELECT {metric_sql} AS metric_value FROM "{table_name}"{where_sql} '
|
||||
raw_rows, engine = _run_select_sql(db_path=db, sql=sql, params=params)
|
||||
row = raw_rows[0] if raw_rows else {}
|
||||
out = {
|
||||
"ok": True,
|
||||
"table_id": str(table_id),
|
||||
"sheet": str(sheet or ""),
|
||||
"table_name": table_name,
|
||||
"metric": mt,
|
||||
"target_column": tgt if mt in {"sum", "avg"} else None,
|
||||
"group_by": "",
|
||||
"value": float((row.get("metric_value") if isinstance(row, dict) else 0.0) or 0.0),
|
||||
"engine": engine,
|
||||
}
|
||||
hint = _sheet_hint_payload(meta, selected_sheet=sheet, table_name=table_name)
|
||||
if isinstance(hint, dict):
|
||||
out["sheet_hint"] = hint
|
||||
return out
|
||||
|
||||
|
||||
def run_table_sql(
|
||||
*,
|
||||
table_id: str,
|
||||
sql: str,
|
||||
limit: int = 200,
|
||||
sheet: str | None = None,
|
||||
) -> dict[str, Any]:
|
||||
if not _is_valid_table_id(table_id):
|
||||
return {
|
||||
"ok": False,
|
||||
"error": "table_id_invalid_format",
|
||||
"hint": "Use table_id from tabular_ref attachment (64-char hex id).",
|
||||
}
|
||||
meta = get_table_meta(table_id)
|
||||
if not isinstance(meta, dict):
|
||||
return {"ok": False, "error": "table_not_found"}
|
||||
db = Path(str(meta.get("db_path") or ""))
|
||||
if not db.exists():
|
||||
return {"ok": False, "error": "table_store_missing"}
|
||||
table_name = _resolve_sheet_table(meta, sheet)
|
||||
if not table_name:
|
||||
return {"ok": False, "error": "sheet_not_found"}
|
||||
raw = str(sql or "").strip()
|
||||
if not raw:
|
||||
return {"ok": False, "error": "sql_required"}
|
||||
lowered = raw.lower()
|
||||
deny = ("insert ", "update ", "delete ", "drop ", "alter ", "attach ", "detach ", "pragma ", "create ")
|
||||
if any(k in lowered for k in deny):
|
||||
return {"ok": False, "error": "sql_not_readonly"}
|
||||
if ";" in raw:
|
||||
return {"ok": False, "error": "sql_multiple_statements_forbidden"}
|
||||
if not (lowered.startswith("select ") or lowered.startswith("with ")):
|
||||
return {"ok": False, "error": "sql_must_start_with_select_or_with"}
|
||||
lim = max(1, min(int(limit or 200), 500))
|
||||
auto_limited = " limit " not in lowered
|
||||
sql_targeted = re.sub(r"(?i)\brows_data\b", f'"{table_name}"', raw)
|
||||
safe_sql = sql_targeted if not auto_limited else f"{sql_targeted}\nLIMIT {lim}"
|
||||
timeout_ms = _tabular_sql_timeout_ms()
|
||||
try:
|
||||
rows, engine = _run_select_sql(db_path=db, sql=safe_sql, params=[], timeout_ms=timeout_ms)
|
||||
except TimeoutError:
|
||||
return {
|
||||
"ok": False,
|
||||
"error": "sql_timeout",
|
||||
"table_id": str(table_id),
|
||||
"name": str(meta.get("name") or ""),
|
||||
"sheet": str(sheet or ""),
|
||||
"table_name": table_name,
|
||||
"input_sql": raw,
|
||||
"executed_sql": safe_sql,
|
||||
"sql_guard": {
|
||||
"readonly_enforced": True,
|
||||
"multi_statement_forbidden": True,
|
||||
"auto_limit_applied": bool(auto_limited),
|
||||
"result_row_cap": lim,
|
||||
"timeout_ms": int(timeout_ms),
|
||||
"timeout_hit": True,
|
||||
},
|
||||
}
|
||||
out_rows = rows[:lim]
|
||||
out = {
|
||||
"ok": True,
|
||||
"table_id": str(table_id),
|
||||
"name": str(meta.get("name") or ""),
|
||||
"sheet": str(sheet or ""),
|
||||
"table_name": table_name,
|
||||
"input_sql": raw,
|
||||
"executed_sql": safe_sql,
|
||||
"sql_guard": {
|
||||
"readonly_enforced": True,
|
||||
"multi_statement_forbidden": True,
|
||||
"auto_limit_applied": bool(auto_limited),
|
||||
"result_row_cap": lim,
|
||||
"timeout_ms": int(timeout_ms),
|
||||
"timeout_hit": False,
|
||||
},
|
||||
"rows": out_rows,
|
||||
"rows_returned": int(len(out_rows)),
|
||||
"limit": lim,
|
||||
"engine": engine,
|
||||
}
|
||||
hint = _sheet_hint_payload(meta, selected_sheet=sheet, table_name=table_name)
|
||||
if isinstance(hint, dict):
|
||||
out["sheet_hint"] = hint
|
||||
return out
|
||||
|
||||
|
||||
def analyze_table_full_scan(
|
||||
*,
|
||||
table_id: str,
|
||||
columns: list[str] | None = None,
|
||||
sheet: str | None = None,
|
||||
top_values_limit: int = 3,
|
||||
) -> dict[str, Any]:
|
||||
if not _is_valid_table_id(table_id):
|
||||
return {
|
||||
"ok": False,
|
||||
"error": "table_id_invalid_format",
|
||||
"hint": "Use table_id from tabular_ref attachment (64-char hex id).",
|
||||
}
|
||||
meta = get_table_meta(table_id)
|
||||
if not isinstance(meta, dict):
|
||||
return {"ok": False, "error": "table_not_found"}
|
||||
db = Path(str(meta.get("db_path") or ""))
|
||||
if not db.exists():
|
||||
return {"ok": False, "error": "table_store_missing"}
|
||||
table_name = _resolve_sheet_table(meta, sheet)
|
||||
if not table_name:
|
||||
return {"ok": False, "error": "sheet_not_found"}
|
||||
known_cols = _columns_for_sheet(meta, table_name)
|
||||
if not known_cols:
|
||||
return {"ok": False, "error": "table_columns_missing"}
|
||||
use_cols = [c for c in (columns or known_cols) if c in known_cols]
|
||||
if not use_cols:
|
||||
return {"ok": False, "error": "columns_not_found"}
|
||||
top_lim = max(0, min(int(top_values_limit or 0), 10))
|
||||
started = time.perf_counter()
|
||||
with sqlite3.connect(str(db)) as conn:
|
||||
total_rows = int(conn.execute(f'SELECT COUNT(1) FROM "{table_name}"').fetchone()[0] or 0)
|
||||
col_stats: list[dict[str, Any]] = []
|
||||
for col in use_cols:
|
||||
row = conn.execute(
|
||||
(
|
||||
f'SELECT '
|
||||
f'SUM(CASE WHEN NULLIF(TRIM("{col}"), \'\') IS NOT NULL THEN 1 ELSE 0 END) AS non_empty_rows, '
|
||||
f'COUNT(DISTINCT NULLIF(TRIM("{col}"), \'\')) AS distinct_values '
|
||||
f'FROM "{table_name}"'
|
||||
)
|
||||
).fetchone()
|
||||
non_empty_rows = int((row[0] if row else 0) or 0)
|
||||
distinct_values = int((row[1] if row else 0) or 0)
|
||||
item: dict[str, Any] = {
|
||||
"column": col,
|
||||
"non_empty_rows": non_empty_rows,
|
||||
"empty_rows": max(0, int(total_rows - non_empty_rows)),
|
||||
"distinct_values": distinct_values,
|
||||
"coverage_ratio": (float(non_empty_rows) / float(total_rows)) if total_rows > 0 else 0.0,
|
||||
}
|
||||
if top_lim > 0:
|
||||
tops = conn.execute(
|
||||
(
|
||||
f'SELECT "{col}" AS val, COUNT(1) AS n '
|
||||
f'FROM "{table_name}" '
|
||||
f'WHERE NULLIF(TRIM("{col}"), \'\') IS NOT NULL '
|
||||
f'GROUP BY "{col}" ORDER BY n DESC LIMIT ?'
|
||||
),
|
||||
[top_lim],
|
||||
).fetchall()
|
||||
item["top_values"] = [{"value": str(r[0]), "count": int(r[1] or 0)} for r in tops]
|
||||
col_stats.append(item)
|
||||
elapsed_ms = int((time.perf_counter() - started) * 1000)
|
||||
out = {
|
||||
"ok": True,
|
||||
"table_id": str(table_id),
|
||||
"name": str(meta.get("name") or ""),
|
||||
"sheet": str(sheet or ""),
|
||||
"table_name": table_name,
|
||||
"rows_scanned": int(total_rows),
|
||||
"columns_analyzed": list(use_cols),
|
||||
"column_stats": col_stats,
|
||||
"scan_audit": {
|
||||
"full_scan": True,
|
||||
"rows_scanned": int(total_rows),
|
||||
"columns_count": int(len(use_cols)),
|
||||
"elapsed_ms": int(elapsed_ms),
|
||||
},
|
||||
"engine": "builtin_sqlite_fullscan",
|
||||
}
|
||||
hint = _sheet_hint_payload(meta, selected_sheet=sheet, table_name=table_name)
|
||||
if isinstance(hint, dict):
|
||||
out["sheet_hint"] = hint
|
||||
return out
|
||||
|
||||
|
||||
__all__ = [
|
||||
"save_dataframe",
|
||||
"save_workbook",
|
||||
"get_table_meta",
|
||||
"query_table",
|
||||
"aggregate_table",
|
||||
"run_table_sql",
|
||||
"analyze_table_full_scan",
|
||||
]
|
||||
|
||||
197
svc/files/text_attachment_store.py
Normal file
197
svc/files/text_attachment_store.py
Normal file
|
|
@ -0,0 +1,197 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import hashlib
|
||||
import json
|
||||
import re
|
||||
import sqlite3
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
from svc.config.paths import attachments_dir
|
||||
|
||||
|
||||
DEFAULT_TEXT_INLINE_MAX_CHARS = 12_000
|
||||
DEFAULT_TEXT_CHUNK_SIZE = 1_600
|
||||
DEFAULT_TEXT_CHUNK_OVERLAP = 200
|
||||
|
||||
|
||||
def _text_root() -> Path:
|
||||
p = (attachments_dir() / "textual").resolve()
|
||||
p.mkdir(parents=True, exist_ok=True)
|
||||
return p
|
||||
|
||||
|
||||
def _text_id(attachment_id: str, name: str) -> str:
|
||||
raw = f"{attachment_id}:{name}".encode("utf-8", errors="ignore")
|
||||
return hashlib.sha256(raw).hexdigest()
|
||||
|
||||
|
||||
def _db_path(text_id: str) -> Path:
|
||||
return _text_root() / f"{text_id}.sqlite"
|
||||
|
||||
|
||||
def _meta_path(text_id: str) -> Path:
|
||||
return _text_root() / f"{text_id}.meta.json"
|
||||
|
||||
|
||||
def _is_valid_text_id(text_id: str) -> bool:
|
||||
tid = str(text_id or "").strip().lower()
|
||||
return bool(re.fullmatch(r"[0-9a-f]{64}", tid))
|
||||
|
||||
|
||||
def _chunk_text(*, text: str, chunk_size: int, overlap: int) -> list[dict[str, Any]]:
|
||||
body = str(text or "")
|
||||
if not body:
|
||||
return []
|
||||
size = max(200, min(int(chunk_size or DEFAULT_TEXT_CHUNK_SIZE), 8_000))
|
||||
ov = max(0, min(int(overlap or DEFAULT_TEXT_CHUNK_OVERLAP), size - 1))
|
||||
chunks: list[dict[str, Any]] = []
|
||||
i = 0
|
||||
idx = 0
|
||||
n = len(body)
|
||||
while i < n:
|
||||
end = min(n, i + size)
|
||||
chunk = body[i:end]
|
||||
chunks.append({"chunk_index": idx, "start_char": i, "end_char": end, "content": chunk})
|
||||
if end >= n:
|
||||
break
|
||||
i = end - ov
|
||||
idx += 1
|
||||
return chunks
|
||||
|
||||
|
||||
def save_text_document(
|
||||
*,
|
||||
attachment_id: str,
|
||||
name: str,
|
||||
text: str,
|
||||
source_kind: str,
|
||||
chunk_size: int = DEFAULT_TEXT_CHUNK_SIZE,
|
||||
chunk_overlap: int = DEFAULT_TEXT_CHUNK_OVERLAP,
|
||||
) -> dict[str, Any]:
|
||||
tid = _text_id(attachment_id=attachment_id, name=name)
|
||||
db = _db_path(tid)
|
||||
meta = _meta_path(tid)
|
||||
chunks = _chunk_text(text=str(text or ""), chunk_size=chunk_size, overlap=chunk_overlap)
|
||||
with sqlite3.connect(str(db)) as conn:
|
||||
conn.execute("DROP TABLE IF EXISTS text_chunks")
|
||||
conn.execute(
|
||||
(
|
||||
"CREATE TABLE text_chunks ("
|
||||
"chunk_index INTEGER NOT NULL, "
|
||||
"start_char INTEGER NOT NULL, "
|
||||
"end_char INTEGER NOT NULL, "
|
||||
"content TEXT NOT NULL)"
|
||||
)
|
||||
)
|
||||
for c in chunks:
|
||||
conn.execute(
|
||||
"INSERT INTO text_chunks(chunk_index,start_char,end_char,content) VALUES(?,?,?,?)",
|
||||
[
|
||||
int(c.get("chunk_index") or 0),
|
||||
int(c.get("start_char") or 0),
|
||||
int(c.get("end_char") or 0),
|
||||
str(c.get("content") or ""),
|
||||
],
|
||||
)
|
||||
conn.commit()
|
||||
payload = {
|
||||
"text_id": tid,
|
||||
"attachment_id": str(attachment_id or ""),
|
||||
"name": str(name or ""),
|
||||
"source_kind": str(source_kind or "text"),
|
||||
"chars": int(len(str(text or ""))),
|
||||
"chunks": int(len(chunks)),
|
||||
"chunk_size": int(chunk_size),
|
||||
"chunk_overlap": int(chunk_overlap),
|
||||
"db_path": str(db),
|
||||
}
|
||||
meta.write_text(json.dumps(payload, ensure_ascii=False), encoding="utf-8")
|
||||
return payload
|
||||
|
||||
|
||||
def get_text_meta(text_id: str) -> dict[str, Any] | None:
|
||||
tid = str(text_id or "").strip()
|
||||
if not _is_valid_text_id(tid):
|
||||
return None
|
||||
p = _meta_path(tid)
|
||||
if not p.exists():
|
||||
return None
|
||||
try:
|
||||
obj = json.loads(p.read_text(encoding="utf-8"))
|
||||
return obj if isinstance(obj, dict) else None
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
|
||||
def query_text_document(
|
||||
*,
|
||||
text_id: str,
|
||||
query: str | None = None,
|
||||
top_k: int = 5,
|
||||
offset: int = 0,
|
||||
) -> dict[str, Any]:
|
||||
tid = str(text_id or "").strip()
|
||||
if not _is_valid_text_id(tid):
|
||||
return {"ok": False, "error": "text_id_invalid_format"}
|
||||
meta = get_text_meta(tid)
|
||||
if not isinstance(meta, dict):
|
||||
return {"ok": False, "error": "text_not_found"}
|
||||
db = Path(str(meta.get("db_path") or ""))
|
||||
if not db.exists():
|
||||
return {"ok": False, "error": "text_store_missing"}
|
||||
lim = max(1, min(int(top_k or 5), 50))
|
||||
off = max(0, int(offset or 0))
|
||||
q = str(query or "").strip()
|
||||
rows: list[dict[str, Any]] = []
|
||||
with sqlite3.connect(str(db)) as conn:
|
||||
conn.row_factory = sqlite3.Row
|
||||
total = int(conn.execute("SELECT COUNT(1) AS n FROM text_chunks").fetchone()[0] or 0)
|
||||
if q:
|
||||
q_like = f"%{q.lower()}%"
|
||||
sql = (
|
||||
"SELECT chunk_index,start_char,end_char,content, "
|
||||
"(LENGTH(LOWER(content)) - LENGTH(REPLACE(LOWER(content), LOWER(?), ''))) / MAX(1, LENGTH(?)) AS score "
|
||||
"FROM text_chunks WHERE LOWER(content) LIKE ? "
|
||||
"ORDER BY score DESC, chunk_index ASC LIMIT ? OFFSET ?"
|
||||
)
|
||||
out = conn.execute(sql, [q, q, q_like, lim, off]).fetchall()
|
||||
else:
|
||||
out = conn.execute(
|
||||
"SELECT chunk_index,start_char,end_char,content, 0 AS score FROM text_chunks ORDER BY chunk_index ASC LIMIT ? OFFSET ?",
|
||||
[lim, off],
|
||||
).fetchall()
|
||||
for r in out:
|
||||
rows.append(
|
||||
{
|
||||
"chunk_index": int(r["chunk_index"] or 0),
|
||||
"start_char": int(r["start_char"] or 0),
|
||||
"end_char": int(r["end_char"] or 0),
|
||||
"content": str(r["content"] or ""),
|
||||
"score": float(r["score"] or 0.0),
|
||||
}
|
||||
)
|
||||
return {
|
||||
"ok": True,
|
||||
"text_id": tid,
|
||||
"name": str(meta.get("name") or ""),
|
||||
"source_kind": str(meta.get("source_kind") or "text"),
|
||||
"chars": int(meta.get("chars") or 0),
|
||||
"chunks_total": int(meta.get("chunks") or 0),
|
||||
"rows_total": int(total),
|
||||
"query": q,
|
||||
"rows": rows,
|
||||
"top_k": lim,
|
||||
"offset": off,
|
||||
}
|
||||
|
||||
|
||||
__all__ = [
|
||||
"DEFAULT_TEXT_INLINE_MAX_CHARS",
|
||||
"DEFAULT_TEXT_CHUNK_SIZE",
|
||||
"DEFAULT_TEXT_CHUNK_OVERLAP",
|
||||
"save_text_document",
|
||||
"get_text_meta",
|
||||
"query_text_document",
|
||||
]
|
||||
|
||||
Loading…
Add table
Add a link
Reference in a new issue