This commit is contained in:
211
src/storage.py
Normal file
211
src/storage.py
Normal file
@@ -0,0 +1,211 @@
|
|||||||
|
"""Storage management for prompts and tags."""
|
||||||
|
|
||||||
|
from pathlib import Path
|
||||||
|
from typing import Any
|
||||||
|
|
||||||
|
import yaml
|
||||||
|
|
||||||
|
from .models import Config, Prompt, Tag
|
||||||
|
|
||||||
|
|
||||||
|
class PromptStorage:
|
||||||
|
"""Manages prompt storage in YAML files."""
|
||||||
|
|
||||||
|
def __init__(self, prompt_dir: str = None):
|
||||||
|
if prompt_dir is None:
|
||||||
|
config = ConfigManager().load()
|
||||||
|
prompt_dir = config.prompt_dir
|
||||||
|
self.prompt_dir = Path(prompt_dir)
|
||||||
|
self._ensure_dir(str(self.prompt_dir))
|
||||||
|
self.tags_index = self.prompt_dir / "tags.yaml"
|
||||||
|
|
||||||
|
def _ensure_dir(self, path: str) -> Path:
|
||||||
|
"""Create directory if it doesn't exist."""
|
||||||
|
path = Path(path)
|
||||||
|
path.mkdir(parents=True, exist_ok=True)
|
||||||
|
return path
|
||||||
|
|
||||||
|
def get_prompt_path(self, name: str) -> Path:
|
||||||
|
"""Get the file path for a prompt."""
|
||||||
|
safe_name = "".join(c if c.isalnum() or c in "-_" else "_" for c in name)
|
||||||
|
return self.prompt_dir / f"{safe_name}.yaml"
|
||||||
|
|
||||||
|
def list_prompts(self) -> list[str]:
|
||||||
|
"""List all prompt names."""
|
||||||
|
prompts = []
|
||||||
|
if self.prompt_dir.exists():
|
||||||
|
for f in self.prompt_dir.glob("*.yaml"):
|
||||||
|
if f.name != "tags.yaml":
|
||||||
|
prompts.append(f.stem)
|
||||||
|
return sorted(prompts)
|
||||||
|
|
||||||
|
def get_prompt(self, name: str) -> Prompt | None:
|
||||||
|
"""Load a prompt by name."""
|
||||||
|
path = self.get_prompt_path(name)
|
||||||
|
if path.exists():
|
||||||
|
with open(path) as f:
|
||||||
|
data = yaml.safe_load(f)
|
||||||
|
if data:
|
||||||
|
return Prompt.from_dict(data)
|
||||||
|
return None
|
||||||
|
|
||||||
|
def save_prompt(self, prompt: Prompt) -> None:
|
||||||
|
"""Save a prompt to file."""
|
||||||
|
path = self.get_prompt_path(prompt.name)
|
||||||
|
from datetime import datetime
|
||||||
|
if not prompt.created_at:
|
||||||
|
prompt.created_at = datetime.now().isoformat()
|
||||||
|
prompt.updated_at = datetime.now().isoformat()
|
||||||
|
with open(path, "w") as f:
|
||||||
|
yaml.dump(prompt.to_dict(), f, default_flow_style=False, sort_keys=False)
|
||||||
|
|
||||||
|
def delete_prompt(self, name: str) -> bool:
|
||||||
|
"""Delete a prompt file."""
|
||||||
|
path = self.get_prompt_path(name)
|
||||||
|
if path.exists():
|
||||||
|
path.unlink()
|
||||||
|
self._remove_from_tags(name)
|
||||||
|
return True
|
||||||
|
return False
|
||||||
|
|
||||||
|
def prompt_exists(self, name: str) -> bool:
|
||||||
|
"""Check if a prompt exists."""
|
||||||
|
return self.get_prompt_path(name).exists()
|
||||||
|
|
||||||
|
def list_tags(self) -> list[str]:
|
||||||
|
"""List all tags."""
|
||||||
|
if self.tags_index.exists():
|
||||||
|
with open(self.tags_index) as f:
|
||||||
|
data = yaml.safe_load(f) or {}
|
||||||
|
return list(data.keys())
|
||||||
|
return []
|
||||||
|
|
||||||
|
def get_tag(self, name: str) -> Tag | None:
|
||||||
|
"""Get a tag with its associated prompts."""
|
||||||
|
if self.tags_index.exists():
|
||||||
|
with open(self.tags_index) as f:
|
||||||
|
data = yaml.safe_load(f) or {}
|
||||||
|
if name in data:
|
||||||
|
return Tag(name=name, prompts=data[name])
|
||||||
|
return None
|
||||||
|
|
||||||
|
def add_tag_to_prompt(self, prompt_name: str, tag: str) -> None:
|
||||||
|
"""Add a tag to a prompt in the index."""
|
||||||
|
data = {}
|
||||||
|
if self.tags_index.exists():
|
||||||
|
with open(self.tags_index) as f:
|
||||||
|
data = yaml.safe_load(f) or {}
|
||||||
|
if tag not in data:
|
||||||
|
data[tag] = []
|
||||||
|
if prompt_name not in data[tag]:
|
||||||
|
data[tag].append(prompt_name)
|
||||||
|
with open(self.tags_index, "w") as f:
|
||||||
|
yaml.dump(data, f)
|
||||||
|
|
||||||
|
def remove_tag_from_prompt(self, prompt_name: str, tag: str) -> bool:
|
||||||
|
"""Remove a tag from a prompt in the index."""
|
||||||
|
if not self.tags_index.exists():
|
||||||
|
return False
|
||||||
|
with open(self.tags_index) as f:
|
||||||
|
data = yaml.safe_load(f) or {}
|
||||||
|
if tag in data and prompt_name in data[tag]:
|
||||||
|
data[tag].remove(prompt_name)
|
||||||
|
if not data[tag]:
|
||||||
|
del data[tag]
|
||||||
|
with open(self.tags_index, "w") as f:
|
||||||
|
yaml.dump(data, f)
|
||||||
|
return True
|
||||||
|
return False
|
||||||
|
|
||||||
|
def _remove_from_tags(self, prompt_name: str) -> None:
|
||||||
|
"""Remove prompt from all tags."""
|
||||||
|
if not self.tags_index.exists():
|
||||||
|
return
|
||||||
|
with open(self.tags_index) as f:
|
||||||
|
data = yaml.safe_load(f) or {}
|
||||||
|
modified = False
|
||||||
|
for tag in list(data.keys()):
|
||||||
|
if prompt_name in data[tag]:
|
||||||
|
data[tag].remove(prompt_name)
|
||||||
|
modified = True
|
||||||
|
if not data[tag]:
|
||||||
|
del data[tag]
|
||||||
|
if modified:
|
||||||
|
with open(self.tags_index, "w") as f:
|
||||||
|
yaml.dump(data, f)
|
||||||
|
|
||||||
|
def search_prompts(
|
||||||
|
self,
|
||||||
|
name: str = None,
|
||||||
|
content: str = None,
|
||||||
|
tag: str = None
|
||||||
|
) -> list[Prompt]:
|
||||||
|
"""Search prompts by name, content, or tag."""
|
||||||
|
results = []
|
||||||
|
for prompt_name in self.list_prompts():
|
||||||
|
prompt = self.get_prompt(prompt_name)
|
||||||
|
if not prompt:
|
||||||
|
continue
|
||||||
|
if name and name.lower() not in prompt.name.lower():
|
||||||
|
continue
|
||||||
|
if content and content.lower() not in prompt.template.lower():
|
||||||
|
continue
|
||||||
|
if tag and tag.lower() not in [t.lower() for t in prompt.tags]:
|
||||||
|
continue
|
||||||
|
results.append(prompt)
|
||||||
|
return results
|
||||||
|
|
||||||
|
def get_prompts_by_tag(self, tag: str) -> list[Prompt]:
|
||||||
|
"""Get all prompts with a specific tag."""
|
||||||
|
tag_data = self.get_tag(tag)
|
||||||
|
if not tag_data:
|
||||||
|
return []
|
||||||
|
prompts = []
|
||||||
|
for prompt_name in tag_data.prompts:
|
||||||
|
prompt = self.get_prompt(prompt_name)
|
||||||
|
if prompt:
|
||||||
|
prompts.append(prompt)
|
||||||
|
return prompts
|
||||||
|
|
||||||
|
|
||||||
|
class ConfigManager:
|
||||||
|
"""Manages user configuration."""
|
||||||
|
|
||||||
|
def __init__(self, config_path: str = None):
|
||||||
|
if config_path is None:
|
||||||
|
home = Path.home()
|
||||||
|
self.config_path = home / ".config" / "llm-prompt-manager" / "config.yaml"
|
||||||
|
else:
|
||||||
|
self.config_path = Path(config_path)
|
||||||
|
|
||||||
|
def _ensure_dir(self, path: str) -> Path:
|
||||||
|
"""Create directory if it doesn't exist."""
|
||||||
|
path = Path(path)
|
||||||
|
path.mkdir(parents=True, exist_ok=True)
|
||||||
|
return path
|
||||||
|
|
||||||
|
def load(self) -> Config:
|
||||||
|
"""Load configuration from file."""
|
||||||
|
self._ensure_dir(str(self.config_path.parent))
|
||||||
|
if self.config_path.exists():
|
||||||
|
with open(self.config_path) as f:
|
||||||
|
data = yaml.safe_load(f) or {}
|
||||||
|
return Config.from_dict(data)
|
||||||
|
return Config()
|
||||||
|
|
||||||
|
def save(self, config: Config) -> None:
|
||||||
|
"""Save configuration to file."""
|
||||||
|
self._ensure_dir(str(self.config_path.parent))
|
||||||
|
with open(self.config_path, "w") as f:
|
||||||
|
yaml.dump(config.to_dict(), f, default_flow_style=False)
|
||||||
|
|
||||||
|
def get(self, key: str) -> Any:
|
||||||
|
"""Get a configuration value."""
|
||||||
|
config = self.load()
|
||||||
|
return getattr(config, key, None)
|
||||||
|
|
||||||
|
def set(self, key: str, value: Any) -> None:
|
||||||
|
"""Set a configuration value."""
|
||||||
|
config = self.load()
|
||||||
|
setattr(config, key, value)
|
||||||
|
self.save(config)
|
||||||
Reference in New Issue
Block a user