diff --git a/src/models.py b/src/models.py new file mode 100644 index 0000000..57e9546 --- /dev/null +++ b/src/models.py @@ -0,0 +1,123 @@ +"""Data models for the prompt manager.""" + +from dataclasses import dataclass, field +from datetime import datetime +from typing import Any + + +@dataclass +class PromptVariable: + """Definition of a template variable.""" + name: str + description: str = "" + required: bool = True + default: str = "" + + +@dataclass +class Prompt: + """Represents a prompt template.""" + name: str + template: str + description: str = "" + tags: list[str] = field(default_factory=list) + variables: list[dict[str, Any]] = field(default_factory=list) + provider: str = "" + model: str = "" + created_at: str = "" + updated_at: str = "" + + def __post_init__(self): + if not self.created_at: + self.created_at = datetime.now().isoformat() + if not self.updated_at: + self.updated_at = datetime.now().isoformat() + + def to_dict(self) -> dict[str, Any]: + """Convert prompt to dictionary.""" + return { + "name": self.name, + "template": self.template, + "description": self.description, + "tags": self.tags, + "variables": self.variables, + "provider": self.provider, + "model": self.model, + "created_at": self.created_at, + "updated_at": self.updated_at, + } + + @classmethod + def from_dict(cls, data: dict[str, Any]) -> "Prompt": + """Create prompt from dictionary.""" + return cls( + name=data.get("name", ""), + template=data.get("template", ""), + description=data.get("description", ""), + tags=data.get("tags", []), + variables=data.get("variables", []), + provider=data.get("provider", ""), + model=data.get("model", ""), + created_at=data.get("created_at", ""), + updated_at=data.get("updated_at", ""), + ) + + def get_required_variables(self) -> list[str]: + """Get list of required variable names.""" + return [ + v["name"] for v in self.variables + if v.get("required", True) + ] + + +@dataclass +class Tag: + """Represents a tag with associated prompts.""" + name: str + prompts: list[str] = field(default_factory=list) + + def to_dict(self) -> dict[str, Any]: + """Convert tag to dictionary.""" + return { + "name": self.name, + "prompts": self.prompts, + } + + @classmethod + def from_dict(cls, data: dict[str, Any]) -> "Tag": + """Create tag from dictionary.""" + return cls( + name=data.get("name", ""), + prompts=data.get("prompts", []), + ) + + +@dataclass +class Config: + """Configuration settings.""" + prompt_dir: str = "~/.config/llm-prompt-manager/prompts" + ollama_url: str = "http://localhost:11434" + lmstudio_url: str = "http://localhost:1234" + default_model: str = "llama3.2" + default_provider: str = "ollama" + + def to_dict(self) -> dict[str, Any]: + """Convert config to dictionary.""" + return { + "prompt_dir": self.prompt_dir, + "ollama_url": self.ollama_url, + "lmstudio_url": self.lmstudio_url, + "default_model": self.default_model, + "default_provider": self.default_provider, + } + + @classmethod + def from_dict(cls, data: dict[str, Any]) -> "Config": + """Create config from dictionary.""" + return cls( + prompt_dir=data.get("prompt_dir", cls.prompt_dir), + ollama_url=data.get("ollama_url", cls.ollama_url), + lmstudio_url=data.get("lmstudio_url", cls.lmstudio_url), + default_model=data.get("default_model", cls.default_model), + default_provider=data.get("default_provider", cls.default_provider), + )