Feat/custom data files

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,
                 )

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",
+]

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)

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)

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.

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)
 

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

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"):