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