diff --git a/src/vtk_prompt/client.py b/src/vtk_prompt/client.py index 1f9b9d2..3177d6c 100644 --- a/src/vtk_prompt/client.py +++ b/src/vtk_prompt/client.py @@ -393,11 +393,14 @@ def query( if not yaml_messages: from .prompts import PYTHON_VERSION, VTK_VERSION + from .data import uploaded_names + prompt_data = assemble_vtk_prompt( request=message, ui_mode=ui_mode, context_snippets=context_snippets, mcp_active=bool(mcp_client), + uploaded_files=uploaded_names(), VTK_VERSION=VTK_VERSION, PYTHON_VERSION=PYTHON_VERSION, ) diff --git a/src/vtk_prompt/data/__init__.py b/src/vtk_prompt/data/__init__.py index dec1cba..f11cf43 100644 --- a/src/vtk_prompt/data/__init__.py +++ b/src/vtk_prompt/data/__init__.py @@ -1,5 +1,33 @@ -"""Sample-data resolution for example-style prompts that read files.""" +"""Data resolution for prompts that read files: sample data and user uploads.""" -from .resolver import artifacts, available_names, has_data_root, referenced, resolve +from .resolver import ( + artifacts, + available_names, + cached_names, + clear_cache, + has_data_root, + referenced, + resolve, +) +from .uploads import ( + add_upload, + clear_uploads, + remove_upload, + uploaded_names, + uploaded_path, +) -__all__ = ["artifacts", "available_names", "has_data_root", "referenced", "resolve"] +__all__ = [ + "add_upload", + "artifacts", + "available_names", + "cached_names", + "clear_cache", + "clear_uploads", + "has_data_root", + "referenced", + "remove_upload", + "resolve", + "uploaded_names", + "uploaded_path", +] diff --git a/src/vtk_prompt/data/resolver.py b/src/vtk_prompt/data/resolver.py index ac8c696..3685dca 100644 --- a/src/vtk_prompt/data/resolver.py +++ b/src/vtk_prompt/data/resolver.py @@ -18,6 +18,8 @@ from pathlib import Path from urllib.error import URLError +from . import uploads + logger = logging.getLogger(__name__) _STORE_URL = "https://data.kitware.com/api/v1/file/hashsum/sha512/{digest}/download" @@ -152,14 +154,17 @@ def stage_code(code: str) -> str: ``reader.SetFileName('cow.g')`` runs against the fetched file. Explicit paths and unrelated strings are left untouched. """ - index = _load_index() - if not index or not code: + if not code: return code + index = _load_index() def _replace(match: "re.Match[str]") -> str: quote, value = match.group(1), match.group(2) if ("/" in value) or ("\\" in value): return match.group(0) + uploaded = uploads.uploaded_path(value) + if uploaded: + return f"{quote}{uploaded}{quote}" if value in index: path = resolve(value) if path: @@ -171,15 +176,17 @@ def _replace(match: "re.Match[str]") -> str: def referenced(code: str) -> list[str]: """Return dataset names referenced as bare string literals in code (no fetch).""" - index = _load_index() - if not index or not code: + if not code: + return [] + names = set(_load_index()) | set(uploads.uploaded_names()) + if not names: return [] found: list[str] = [] for match in _LITERAL_RE.finditer(code): value = match.group(2) if ("/" in value) or ("\\" in value): continue - if value in index and value not in found: + if value in names and value not in found: found.append(value) return found @@ -193,6 +200,27 @@ def artifacts(code: str) -> list[dict]: cache = _cache_dir() result: list[dict] = [] for name in referenced(code): - path = cache / name - result.append({"name": name, "path": str(path), "cached": path.exists()}) + uploaded = uploads.uploaded_path(name) + if uploaded: + result.append( + {"name": name, "path": uploaded, "cached": True, "source": "upload"} + ) + else: + path = cache / name + result.append( + {"name": name, "path": str(path), "cached": path.exists(), "source": "sample"} + ) return result + + +def cached_names() -> list[str]: + """Sorted basenames of sample-data files already fetched to the local cache.""" + directory = _cache_dir() + return sorted(p.name for p in directory.iterdir() if p.is_file()) + + +def clear_cache() -> None: + """Delete all fetched sample-data files from the local cache.""" + for p in _cache_dir().iterdir(): + if p.is_file(): + p.unlink(missing_ok=True) diff --git a/src/vtk_prompt/data/uploads.py b/src/vtk_prompt/data/uploads.py new file mode 100644 index 0000000..4fbab0a --- /dev/null +++ b/src/vtk_prompt/data/uploads.py @@ -0,0 +1,87 @@ +"""Registry for user-uploaded custom data files. + +Complements ``resolver.py`` (which fetches known VTK sample data by name) by +letting users supply their own files. Uploaded files are written to a local +cache directory and can then be referenced by bare filename in generated code, +exactly like sample data. User uploads take precedence over sample data of the +same name. +""" + +import os +import re +from collections.abc import Iterable +from pathlib import Path +from typing import cast + +# Characters allowed in a stored filename; anything else is replaced. +_SAFE_NAME_RE = re.compile(r"[^A-Za-z0-9._-]") + + +def _uploads_dir() -> Path: + base = os.environ.get("XDG_CACHE_HOME") + root = Path(base) if base else Path.home() / ".cache" + directory = root / "vtk-prompt" / "uploads" + directory.mkdir(parents=True, exist_ok=True) + return directory + + +def _safe_name(filename: str) -> str: + """Reduce an arbitrary filename to a bare, filesystem-safe basename.""" + name = os.path.basename((filename or "").strip()) + name = _SAFE_NAME_RE.sub("_", name) + return name.lstrip(".") or "upload" + + +def _as_bytes(content: object) -> bytes: + """Coerce uploaded file content (bytes, bytearray, str, or ints) to bytes.""" + if isinstance(content, bytes): + return content + if isinstance(content, (bytearray, memoryview)): + return bytes(content) + if isinstance(content, str): + return content.encode("utf-8", "surrogateescape") + if isinstance(content, Iterable): + try: + return bytes(cast(Iterable[int], content)) # e.g. a list of byte values + except (TypeError, ValueError): + return b"" + return b"" + + +def add_upload(filename: str, content: object) -> str: + """Store uploaded content under its sanitized basename; return the path.""" + dest = _uploads_dir() / _safe_name(filename) + with open(dest, "wb") as handle: + handle.write(_as_bytes(content)) + return str(dest) + + +def uploaded_names() -> list[str]: + """Sorted basenames of all uploaded files currently available.""" + directory = _uploads_dir() + return sorted(p.name for p in directory.iterdir() if p.is_file()) + + +def uploaded_path(name: str) -> str | None: + """Local path for an uploaded file by basename, or None if absent.""" + name = os.path.basename((name or "").strip()) + if not name: + return None + path = _uploads_dir() / name + return str(path) if path.is_file() else None + + +def remove_upload(name: str) -> bool: + """Delete one uploaded file by basename. Returns True if it existed.""" + path = uploaded_path(name) + if not path: + return False + Path(path).unlink(missing_ok=True) + return True + + +def clear_uploads() -> None: + """Remove all uploaded files.""" + for p in _uploads_dir().iterdir(): + if p.is_file(): + p.unlink(missing_ok=True) diff --git a/src/vtk_prompt/prompts/components/user_data.yml b/src/vtk_prompt/prompts/components/user_data.yml new file mode 100644 index 0000000..905645a --- /dev/null +++ b/src/vtk_prompt/prompts/components/user_data.yml @@ -0,0 +1,5 @@ +role: user +content: | + I have uploaded the following data files. You can load each one directly by its + bare filename (do not invent a directory path): {{uploaded_files_list}}. + Choose the appropriate VTK reader for each file based on its extension. diff --git a/src/vtk_prompt/prompts/prompt_component_assembler.py b/src/vtk_prompt/prompts/prompt_component_assembler.py index f024f88..a0223d9 100644 --- a/src/vtk_prompt/prompts/prompt_component_assembler.py +++ b/src/vtk_prompt/prompts/prompt_component_assembler.py @@ -198,6 +198,7 @@ def assemble_vtk_prompt( ui_mode: bool = False, context_snippets: str | None = None, mcp_active: bool = False, + uploaded_files: list[str] | None = None, **variables: Any, ) -> PromptData: """Assemble VTK prompt from file-based components. @@ -207,6 +208,7 @@ def assemble_vtk_prompt( ui_mode: Whether to include UI-specific instructions context_snippets: Optional context snippets from vtk-mcp (enables rag_context component) mcp_active: Whether a vtk-mcp server is reachable (enables tool_use guidance) + uploaded_files: Names of user-uploaded data files (enables user_data component) **variables: Additional variables for substitution Returns: @@ -223,6 +225,7 @@ def assemble_vtk_prompt( assembler.add_if(mcp_active, "tool_use") assembler.add_if(bool(context_snippets), "rag_context") assembler.add_if(ui_mode, "ui_renderer") + assembler.add_if(bool(uploaded_files), "user_data") # Always add output format and request last assembler.add_component("output_format") @@ -233,6 +236,7 @@ def assemble_vtk_prompt( "VTK_VERSION": variables.get("VTK_VERSION", "9.6.1"), "PYTHON_VERSION": variables.get("PYTHON_VERSION", ">=3.10"), "context_snippets": context_snippets or "", + "uploaded_files_list": ", ".join(uploaded_files or []), } default_variables.update(variables) diff --git a/src/vtk_prompt/state/initializer.py b/src/vtk_prompt/state/initializer.py index 4f8c5ae..74ae8b7 100644 --- a/src/vtk_prompt/state/initializer.py +++ b/src/vtk_prompt/state/initializer.py @@ -28,6 +28,13 @@ def initialize_state(app: Any) -> None: app.state.generated_explanation = "" app.state.current_prompt = "" # the sent prompt shown inline with the explanation app.state.data_artifacts = [] # datasets referenced by the current code + from ..data.uploads import uploaded_names as _uploaded_names + + app.state.uploaded_data_files = _uploaded_names() # user-supplied data files + app.state.data_uploads = None # file-input model for new uploads + from ..data.resolver import cached_names as _cached_names + + app.state.cached_data_files = _cached_names() # sample data fetched to cache # Sample-data resolver root: defaults to the env var, overridable in Settings. import os as _os from ..data.resolver import set_data_root as _set_data_root diff --git a/src/vtk_prompt/ui/layout/content.py b/src/vtk_prompt/ui/layout/content.py index 58feb13..7b6c2c1 100644 --- a/src/vtk_prompt/ui/layout/content.py +++ b/src/vtk_prompt/ui/layout/content.py @@ -46,24 +46,52 @@ def build_content(layout: Any, app: Any) -> None: size="small", variant="tonal", classes="mr-1", - color=("a.cached ? 'success' : 'primary'", "primary"), + color=( + "a.source == 'upload' ? 'info'" + + " : (a.cached ? 'success' : 'primary')", + "primary", + ), prepend_icon=( - "a.cached ? 'mdi-file-check'" - + " : 'mdi-cloud-download-outline'", + "a.source == 'upload' ? 'mdi-tray-arrow-up'" + + " : (a.cached ? 'mdi-file-check'" + + " : 'mdi-cloud-download-outline')", "mdi-file-outline", ), ) with html.Div(): html.Div( - "{{ a.cached ? 'Fetched' :" - + " 'Will download on run' }}", + "{{ a.source == 'upload' ? 'Uploaded'" + + " : (a.cached ? 'Fetched'" + + " : 'Will download on run') }}", classes="font-weight-medium", ) html.Div("{{ a.path }}", classes="text-caption") html.Div( - "Source: VTK data store", + "{{ a.source == 'upload' ? 'Source: your upload'" + + " : 'Source: VTK data store' }}", classes="text-caption text-medium-emphasis", ) + # Upload a custom data file to reference by name in generated code. + with vuetify.VTooltip(text="Upload data file", location="bottom"): + with vuetify.Template(v_slot_activator="{ props }"): + vuetify.VChip( + "Add data", + v_bind="props", + size="small", + variant="outlined", + prepend_icon="mdi-tray-arrow-up", + classes="ml-2", + click=( + "window.document.querySelector(" + + "'#data-upload-wrap input').click()" + ), + ) + # Hidden input driven by the chip above; trame handles the bytes. + with html.Div(id="data-upload-wrap", classes="d-none"): + vuetify.VFileInput( + v_model=("data_uploads", None), + multiple=True, + ) vuetify.VSpacer() # Undo across code versions (generations, runs, edits) with vuetify.VTooltip(text="Undo code change", location="bottom"): diff --git a/src/vtk_prompt/ui/layout/settings_dialog.py b/src/vtk_prompt/ui/layout/settings_dialog.py index a51bc21..48cbdf0 100644 --- a/src/vtk_prompt/ui/layout/settings_dialog.py +++ b/src/vtk_prompt/ui/layout/settings_dialog.py @@ -33,11 +33,13 @@ def build_settings_dialog(layout: Any, app: Any) -> None: classes="px-2", ): vuetify.VTab("Config", value="files") + vuetify.VTab("Data", value="data") vuetify.VTab("Model", value="model") vuetify.VTab("Advanced", value="advanced") vuetify.VDivider() with vuetify.VTabsWindow(v_model=("active_settings_tab", "files")): _config_tab() + _data_tab() _model_tab() _advanced_tab() @@ -82,25 +84,6 @@ def _config_tab() -> None: hide_details=True, ) - vuetify.VDivider(classes="my-5") - _section("Sample data") - html.Div( - "Local VTK data tree used to resolve example datasets by name " - "(e.g. cow.g).", - classes=_DESC, - ) - vuetify.VTextField( - label="Sample data location", - v_model=("data_root", ""), - placeholder="/path/to/VTK/Testing/Data", - hint="Folder of .sha512 data pointers; blank uses the " - "VTK_PROMPT_DATA_ROOT environment variable", - persistent_hint=True, - density="compact", - variant="outlined", - clearable=True, - ) - def _model_tab() -> None: with vuetify.VTabsWindowItem(value="model"): @@ -243,3 +226,89 @@ def _advanced_tab() -> None: density="compact", variant="outlined", ) + + +def _data_tab() -> None: + with vuetify.VTabsWindowItem(value="data"): + with vuetify.VCardText(classes="pa-4"): + _section("Sample data location") + html.Div( + "Local VTK data tree used to resolve example datasets by name " + "(e.g. cow.g).", + classes=_DESC, + ) + vuetify.VTextField( + label="Sample data location", + v_model=("data_root", ""), + placeholder="/path/to/VTK/Testing/Data", + hint="Folder of .sha512 data pointers; blank uses the " + "VTK_PROMPT_DATA_ROOT environment variable", + persistent_hint=True, + density="compact", + variant="outlined", + clearable=True, + ) + + vuetify.VDivider(classes="my-5") + _section("Uploaded files") + html.Div( + "Your own data files, referenced by bare name in generated code.", + classes=_DESC, + ) + html.Div( + "No files uploaded yet.", + classes="text-caption text-disabled mb-2", + v_show="uploaded_data_files.length === 0", + ) + with html.Div( + classes="d-flex flex-wrap", + v_show="uploaded_data_files.length > 0", + ): + vuetify.VChip( + "{{ f }}", + v_for="f in uploaded_data_files", + key="f", + size="small", + variant="tonal", + color="info", + closable=True, + click_close="window.trame.trigger('remove_uploaded_file', [f])", + prepend_icon="mdi-file-outline", + classes="mr-1 mb-1", + ) + + vuetify.VDivider(classes="my-5") + with html.Div(classes="d-flex align-center justify-space-between"): + html.Div("Fetched sample data", classes=_LABEL) + vuetify.VBtn( + "Clear cache", + variant="text", + size="small", + color="primary", + prepend_icon="mdi-delete-outline", + click="window.trame.trigger('clear_data_cache')", + v_show="cached_data_files.length > 0", + ) + html.Div( + "Sample datasets downloaded to the local cache.", + classes=_DESC, + ) + html.Div( + "Nothing fetched yet.", + classes="text-caption text-disabled mb-2", + v_show="cached_data_files.length === 0", + ) + with html.Div( + classes="d-flex flex-wrap", + v_show="cached_data_files.length > 0", + ): + vuetify.VChip( + "{{ f }}", + v_for="f in cached_data_files", + key="f", + size="small", + variant="tonal", + color="success", + prepend_icon="mdi-file-check", + classes="mr-1 mb-1", + ) diff --git a/src/vtk_prompt/vtk_prompt_ui.py b/src/vtk_prompt/vtk_prompt_ui.py index 341fdec..e057e09 100644 --- a/src/vtk_prompt/vtk_prompt_ui.py +++ b/src/vtk_prompt/vtk_prompt_ui.py @@ -258,6 +258,49 @@ def _on_uploaded_files_change(self, uploaded_files, **kwargs): self.state.prompt_file = None self.state.conversation_file = None + @change("data_uploads") + def _on_data_uploads_change(self, data_uploads, **kwargs): + """Store user-uploaded custom data files so code can reference them.""" + if not data_uploads: + return + from .data import add_upload, artifacts, uploaded_names + + for file_obj in data_uploads: + name = file_obj.get("name") + if name: + add_upload(name, file_obj.get("content", b"")) + self.state.data_uploads = None + self.state.uploaded_data_files = uploaded_names() + # A newly uploaded file may already be named in the current code. + self.state.data_artifacts = artifacts(self.state.generated_code) + + @trigger("remove_uploaded_file") + def remove_uploaded_file(self, name): + """Remove one uploaded data file and refresh dependent state.""" + from .data import artifacts, remove_upload, uploaded_names + + remove_upload(name) + self.state.uploaded_data_files = uploaded_names() + self.state.data_artifacts = artifacts(self.state.generated_code) + + @change("advanced_settings_open") + def _on_settings_open(self, advanced_settings_open, **kwargs): + """Refresh the data lists whenever the settings dialog opens.""" + if advanced_settings_open: + from .data import cached_names, uploaded_names + + self.state.uploaded_data_files = uploaded_names() + self.state.cached_data_files = cached_names() + + @trigger("clear_data_cache") + def clear_data_cache(self): + """Delete fetched sample data and refresh dependent state.""" + from .data import artifacts, cached_names, clear_cache + + clear_cache() + self.state.cached_data_files = cached_names() + self.state.data_artifacts = artifacts(self.state.generated_code) + @change("conversation_object") def on_conversation_file_data_change( self, conversation_object: dict[str, Any] | None, **_: Any