109 lines
3.9 KiB
Python
109 lines
3.9 KiB
Python
"""
|
|
Asks a local llama-swap model to suggest a topic, subtopic, and
|
|
destination folder for a file, given its filename, current folder, and
|
|
(if available) an extracted content excerpt.
|
|
|
|
Talks to llama-swap's OpenAI-compatible endpoint over plain HTTP using
|
|
only the stdlib — nothing about a file's name or content ever leaves the
|
|
machine this endpoint points at.
|
|
"""
|
|
|
|
import json
|
|
import os
|
|
import urllib.error
|
|
import urllib.request
|
|
|
|
DEFAULT_ENDPOINT = os.environ.get("VAULT_RENAMER_LLM_ENDPOINT", "http://localhost:8080/v1/chat/completions")
|
|
DEFAULT_MODEL = os.environ.get("VAULT_RENAMER_LLM_MODEL", "daily")
|
|
TIMEOUT_SECONDS = 30
|
|
|
|
PROMPT_TEMPLATE = """You are helping organise a personal file vault. Given a file's name, its current folder, an optional content excerpt, and the list of real destination folders that already exist, suggest a short topic, a short subtopic, and which existing folder it belongs in.
|
|
|
|
Rules:
|
|
- topic and subtopic: a few lowercase words each, no punctuation, describing what the file is actually about — not a restatement of the filename.
|
|
- dest_folder: MUST be exactly one of the folders listed below, copied verbatim. If genuinely unsure, pick the closest match rather than inventing a new folder.
|
|
- Reply with ONLY a JSON object, nothing else, no markdown fence: {{"topic": "...", "subtopic": "...", "dest_folder": "..."}}
|
|
|
|
Current filename: {filename}
|
|
Current folder: {current_folder}
|
|
Content excerpt (may be empty): {excerpt}
|
|
|
|
Real destination folders:
|
|
{folders}
|
|
"""
|
|
|
|
|
|
class LLMSuggestError(Exception):
|
|
"""Callers should catch this and treat it as 'no suggestion available'
|
|
for this file — never let a model/network hiccup crash the batch."""
|
|
|
|
|
|
def suggest(filename: str, current_folder: str, excerpt: str, folders: list,
|
|
endpoint: str = DEFAULT_ENDPOINT, model: str = DEFAULT_MODEL) -> dict:
|
|
prompt = PROMPT_TEMPLATE.format(
|
|
filename=filename,
|
|
current_folder=current_folder or "(root)",
|
|
excerpt=(excerpt or "")[:1500],
|
|
folders="\n".join(folders),
|
|
)
|
|
|
|
payload = {
|
|
"model": model,
|
|
"messages": [{"role": "user", "content": prompt}],
|
|
"temperature": 0.2,
|
|
"max_tokens": 200,
|
|
}
|
|
|
|
req = urllib.request.Request(
|
|
endpoint,
|
|
data=json.dumps(payload).encode("utf-8"),
|
|
headers={"Content-Type": "application/json"},
|
|
method="POST",
|
|
)
|
|
|
|
try:
|
|
with urllib.request.urlopen(req, timeout=TIMEOUT_SECONDS) as resp:
|
|
data = json.loads(resp.read().decode("utf-8"))
|
|
except urllib.error.URLError as e:
|
|
raise LLMSuggestError(f"Couldn't reach {endpoint} — is llama-swap running? ({e})")
|
|
except TimeoutError:
|
|
raise LLMSuggestError(f"No response within {TIMEOUT_SECONDS}s — model may still be loading.")
|
|
|
|
try:
|
|
content = data["choices"][0]["message"]["content"]
|
|
except (KeyError, IndexError):
|
|
raise LLMSuggestError(f"Unexpected response shape from {endpoint}: {data}")
|
|
|
|
parsed = _extract_json(content)
|
|
if parsed is None:
|
|
raise LLMSuggestError(f"Model didn't reply with valid JSON: {content[:200]!r}")
|
|
|
|
for key in ("topic", "subtopic", "dest_folder"):
|
|
parsed.setdefault(key, "")
|
|
|
|
if parsed["dest_folder"] not in folders:
|
|
# The model invented a folder that doesn't exist — never apply a
|
|
# destination that isn't real.
|
|
parsed["dest_folder"] = ""
|
|
|
|
return parsed
|
|
|
|
|
|
def _extract_json(text: str):
|
|
text = text.strip()
|
|
if text.startswith("```"):
|
|
text = text.strip("`")
|
|
if text.lower().startswith("json"):
|
|
text = text[4:]
|
|
text = text.strip()
|
|
try:
|
|
return json.loads(text)
|
|
except json.JSONDecodeError:
|
|
pass
|
|
start, end = text.find("{"), text.rfind("}")
|
|
if start != -1 and end != -1 and end > start:
|
|
try:
|
|
return json.loads(text[start:end + 1])
|
|
except json.JSONDecodeError:
|
|
return None
|
|
return None
|