Add core modules (prompt, template)

2026-02-04 12:30:55 +00:00
parent ea913096f8
commit 40667771a8
1 changed files with 127 additions and 0 deletions
--- a/src/promptforge/core/prompt.py
+++ b/src/promptforge/core/prompt.py
@@ -0,0 +1,127 @@
+import hashlib
+import uuid
+from datetime import datetime
+from pathlib import Path
+from typing import Any, Dict, List, Optional
+from enum import Enum
+
+import yaml
+from pydantic import BaseModel, Field, field_validator
+
+
+class VariableType(str, Enum):
+    STRING = "string"
+    INTEGER = "integer"
+    FLOAT = "float"
+    BOOLEAN = "boolean"
+    CHOICE = "choice"
+
+
+class PromptVariable(BaseModel):
+    name: str
+    type: VariableType = VariableType.STRING
+    description: Optional[str] = None
+    required: bool = True
+    default: Optional[Any] = None
+    choices: Optional[List[str]] = None
+
+
+class ValidationRule(BaseModel):
+    type: str
+    pattern: Optional[str] = None
+    json_schema: Optional[Dict[str, Any]] = None
+    message: Optional[str] = None
+
+
+class Prompt(BaseModel):
+    id: str = Field(default_factory=lambda: str(uuid.uuid4()))
+    name: str
+    description: Optional[str] = None
+    content: str
+    variables: List[PromptVariable] = Field(default_factory=list)
+    validation_rules: List[ValidationRule] = Field(default_factory=list)
+    provider: Optional[str] = None
+    tags: List[str] = Field(default_factory=list)
+    version: str = "1.0.0"
+    created_at: datetime = Field(default_factory=datetime.utcnow)
+    updated_at: datetime = Field(default_factory=datetime.utcnow)
+    hash: str = ""
+
+    @field_validator('hash', mode='before')
+    @classmethod
+    def compute_hash(cls, v, info):
+        if not v:
+            content = info.data.get('content', '')
+            return hashlib.md5(content.encode()).hexdigest()
+        return v
+
+    @classmethod
+    def from_yaml(cls, yaml_content: str) -> "Prompt":
+        content = yaml_content.strip()
+        
+        if not content.startswith('---'):
+            metadata = {}
+            prompt_content = content
+        else:
+            parts = content[4:].split('\n---', 1)
+            metadata = yaml.safe_load(parts[0]) or {}
+            prompt_content = parts[1].strip() if len(parts) > 1 else ''
+
+        data = {
+            'name': metadata.get('name', 'Untitled'),
+            'description': metadata.get('description'),
+            'content': prompt_content,
+            'variables': [PromptVariable(**v) for v in metadata.get('variables', [])],
+            'validation_rules': [ValidationRule(**r) for r in metadata.get('validation', [])],
+            'provider': metadata.get('provider'),
+            'tags': metadata.get('tags', []),
+            'version': metadata.get('version', '1.0.0'),
+        }
+        return cls(**data)
+
+    def to_yaml(self) -> str:
+        def var_to_dict(v):
+            d = v.model_dump()
+            d['type'] = v.type.value
+            return d
+        
+        def rule_to_dict(r):
+            return r.model_dump()
+        
+        metadata = {
+            'name': self.name,
+            'description': self.description,
+            'provider': self.provider,
+            'tags': self.tags,
+            'version': self.version,
+            'variables': [var_to_dict(v) for v in self.variables],
+            'validation': [rule_to_dict(r) for r in self.validation_rules],
+        }
+        yaml_str = yaml.dump(metadata, default_flow_style=False, allow_unicode=True)
+        return f"---\n{yaml_str}---\n{self.content}"
+
+    def save(self, prompts_dir: Path) -> Path:
+        prompts_dir.mkdir(parents=True, exist_ok=True)
+        filename = self.name.lower().replace(' ', '_').replace('/', '_') + '.yaml'
+        filepath = prompts_dir / filename
+        with open(filepath, 'w') as f:
+            f.write(self.to_yaml())
+        return filepath
+
+    @classmethod
+    def load(cls, filepath: Path) -> "Prompt":
+        with open(filepath, 'r') as f:
+            content = f.read()
+        return cls.from_yaml(content)
+
+    @classmethod
+    def list(cls, prompts_dir: Path) -> List["Prompt"]:
+        prompts = []
+        if not prompts_dir.exists():
+            return prompts
+        for filepath in prompts_dir.glob('*.yaml'):
+            try:
+                prompts.append(cls.load(filepath))
+            except Exception:
+                continue
+        return sorted(prompts, key=lambda p: p.name)