hexdag 0.5.0.dev1__py3-none-any.whl

hexdag/builtin/prompts/base.py

@@ -0,0 +1,422 @@
+"""Base prompt template classes for building composable prompts.
+
+These are the foundational prompt types that all builtin prompts inherit from.
+They provide the core functionality for chat-style prompts, few-shot learning,
+and combinations thereof.
+
+Moved from core.orchestration.prompt.template to keep core minimal and make
+these templates part of the builtin library.
+"""
+
+from __future__ import annotations
+
+import re
+from collections.abc import Callable
+from typing import Any
+
+from hexdag.core.orchestration.prompt.template import PromptTemplate, _extract_variables_cached
+
+
+class FewShotPromptTemplate(PromptTemplate):
+    r"""PromptTemplate with built-in support for few-shot examples.
+
+    Simplified implementation that builds the template once during initialization.
+
+    Examples
+    --------
+    Basic few-shot template::
+
+        template = FewShotPromptTemplate(
+            template="Classify: {{text}}",
+            examples=[
+                {"input": "Great product!", "output": "positive"},
+                {"input": "Terrible", "output": "negative"}
+            ]
+        )
+
+        result = template.format(text="Amazing!")
+        # Output includes examples + prompt
+
+    Custom example formatting::
+
+        def custom_formatter(ex):
+            return f"Q: {ex['input']}\nA: {ex['output']}"
+
+        template = FewShotPromptTemplate(
+            template="Q: {{question}}\nA:",
+            examples=[...],
+            format_example=custom_formatter
+        )
+    """
+
+    def __init__(
+        self,
+        template: str,
+        examples: list[dict[str, Any]] | None = None,
+        *,
+        format_example: Callable[[dict[str, Any]], str] | None = None,
+        example_separator: str = "\n\n",
+        **kwargs: Any,
+    ) -> None:
+        r"""Initialize FewShotPromptTemplate with examples.
+
+        Args
+        ----
+            template: Base template string with {{variables}}
+            examples: List of example dicts (typically with 'input' and 'output' keys)
+            format_example: Function to format each example (optional)
+            example_separator: String to separate examples (default: "\n\n")
+            **kwargs: Additional arguments passed to PromptTemplate
+        """
+        self.examples: list[dict[str, Any]] = examples or []
+        self.example_separator = example_separator
+        self.base_template = template
+
+        if format_example is None:
+
+            def _default_formatter(ex: dict[str, Any]) -> str:
+                inp = ex.get("input", "")
+                out = ex.get("output", "")
+                return f"Input: {inp}\nOutput: {out}"
+
+            format_example = _default_formatter
+
+        self.format_example = format_example
+
+        full_template = self._build_template()
+
+        super().__init__(full_template, **kwargs)
+
+    def _build_template(self) -> str:
+        """Build the complete template with examples prepended.
+
+        Returns
+        -------
+        str
+            Complete template string with examples and base template
+        """
+        if not self.examples:
+            return self.base_template
+
+        examples_text = self.example_separator.join(
+            self.format_example(example) for example in self.examples
+        )
+
+        return f"{examples_text}{self.example_separator}{self.base_template}"
+
+    def add_example(self, example: dict[str, Any]) -> None:
+        """Add an example and rebuild the template.
+
+        Args
+        ----
+            example: Example dict to add
+        """
+        self.examples.append(example)
+
+        # Rebuild the complete template
+        self.template = self._build_template()
+
+        # Re-extract variables from the new template using cached function
+        self.input_vars = list(_extract_variables_cached(self.template))
+
+    def __add__(self, other: str | PromptTemplate) -> FewShotPromptTemplate:
+        """Add text or template using + operator.
+
+        Returns
+        -------
+        FewShotPromptTemplate
+            New FewShotPromptTemplate instance with enhanced content
+        """
+        if isinstance(other, PromptTemplate):
+            enhanced_base_template = self.base_template + other.template
+        else:
+            enhanced_base_template = self.base_template + other
+
+        return FewShotPromptTemplate(
+            enhanced_base_template,
+            self.examples,
+            format_example=self.format_example,
+            example_separator=self.example_separator,
+        )
+
+    def add(self, text: str | PromptTemplate) -> FewShotPromptTemplate:
+        """Add text to template (alias for + operator).
+
+        Returns
+        -------
+        FewShotPromptTemplate
+            New FewShotPromptTemplate instance with enhanced content
+        """
+        return self + text
+
+
+class ChatPromptTemplate(PromptTemplate):
+    """Enhanced prompt template supporting multi-message conversations like LangChain.
+
+    Supports:
+    - System messages
+    - Conversation history from context
+    - Multi-message templates
+    - Role-based message construction
+
+    Examples
+    --------
+    Simple system + user::
+
+        template = ChatPromptTemplate(
+            system_message="You are an expert {{domain}} analyst.",
+            human_message="Analyze this {{data_type}}: {{data}}"
+        )
+
+        messages = template.to_messages(
+            domain="financial",
+            data_type="report",
+            data="Q4 earnings..."
+        )
+
+    Multi-message template::
+
+        template = ChatPromptTemplate(messages=[
+            {"role": "system", "content": "You are an expert analyst."},
+            {"role": "user", "content": "Context: {{context}}"},
+            {"role": "assistant", "content": "I understand."},
+            {"role": "user", "content": "Analyze: {{data}}"}
+        ])
+    """
+
+    def __init__(
+        self,
+        messages: list[dict[str, str]] | None = None,
+        system_message: str | None = None,
+        human_message: str | None = None,
+        **kwargs: Any,
+    ) -> None:
+        """Initialize ChatPromptTemplate with multiple message types.
+
+        Args
+        ----
+            messages: List of message templates with role and content
+            system_message: System message template (optional)
+            human_message: Human/user message template (optional)
+            **kwargs: Additional arguments passed to PromptTemplate
+        """
+        self.system_message = system_message
+        self.human_message = human_message
+        self.message_templates = messages or []
+
+        if not self.message_templates:
+            if system_message and human_message:
+                combined_template = f"{system_message}\n{human_message}"
+            elif human_message:
+                combined_template = human_message
+            elif system_message:
+                combined_template = system_message
+            else:
+                combined_template = ""
+        else:
+            combined_template = "\n".join(msg.get("content", "") for msg in self.message_templates)
+
+        super().__init__(combined_template, **kwargs)
+
+    def to_messages(
+        self,
+        system_prompt: str | None = None,
+        context_history: list[dict[str, str]] | None = None,
+        **kwargs: Any,
+    ) -> list[dict[str, str]]:
+        """Convert to role-based messages with conversation history support.
+
+        Args
+        ----
+            system_prompt: Optional system prompt to include
+            context_history: Previous conversation messages from context
+            **kwargs: Variable values for template substitution
+
+        Returns
+        -------
+        list[dict[str, str]]
+            List of message dictionaries ready for LLM
+        """
+        messages = []
+
+        if system_prompt:
+            system_content = self._render_template(system_prompt, **kwargs)
+            messages.append({"role": "system", "content": system_content})
+        elif self.system_message:
+            system_content = self._render_template(self.system_message, **kwargs)
+            messages.append({"role": "system", "content": system_content})
+
+        if context_history:
+            messages.extend(context_history)
+
+        if self.message_templates:
+            for msg_template in self.message_templates:
+                role = msg_template.get("role", "user")
+                content_template = msg_template.get("content", "")
+                rendered_content = self._render_template(content_template, **kwargs)
+                messages.append({"role": role, "content": rendered_content})
+        elif self.human_message:
+            human_content = self._render_template(self.human_message, **kwargs)
+            messages.append({"role": "user", "content": human_content})
+
+        return messages
+
+    def _render_template(self, template: str, **kwargs: Any) -> str:
+        """Render a single template string with variables.
+
+        Returns
+        -------
+        str
+            Rendered template string with variables substituted
+        """
+        # Use the parent class variable substitution logic
+        pattern = r"\{\{\s*([a-zA-Z_][a-zA-Z0-9_]*(?:\.[a-zA-Z_][a-zA-Z0-9_]*)*)\s*\}\}"
+
+        def replace_var(match: Any) -> str:
+            var_path = match.group(1).strip()
+            try:
+                value = self._get_nested_value(kwargs, var_path)
+                return str(value)
+            except (KeyError, AttributeError):
+                return f"{{{{{var_path}}}}}"  # Keep unreplaced variables
+
+        return re.sub(pattern, replace_var, template)
+
+    @classmethod
+    def from_messages(cls, messages: list[dict[str, str]]) -> ChatPromptTemplate:
+        """Create ChatPromptTemplate from a list of message dictionaries.
+
+        Args
+        ----
+            messages: List of messages with 'role' and 'content' keys
+
+        Returns
+        -------
+        ChatPromptTemplate
+            ChatPromptTemplate instance
+        """
+        return cls(messages=messages)
+
+    def __add__(self, other: str | PromptTemplate) -> ChatPromptTemplate:
+        """Add text or template to system message using + operator.
+
+        Returns
+        -------
+        ChatPromptTemplate
+            New ChatPromptTemplate instance with enhanced system message
+        """
+        if isinstance(other, PromptTemplate):
+            enhanced_system = (self.system_message or "") + other.template
+        else:
+            enhanced_system = (self.system_message or "") + other
+
+        return ChatPromptTemplate(
+            messages=self.message_templates,
+            system_message=enhanced_system,
+            human_message=self.human_message,
+        )
+
+    def add(self, text: str | PromptTemplate) -> ChatPromptTemplate:
+        """Add text to system message (alias for + operator).
+
+        Returns
+        -------
+        ChatPromptTemplate
+            New ChatPromptTemplate instance with enhanced system message
+        """
+        return self + text
+
+
+class ChatFewShotTemplate(ChatPromptTemplate):
+    """Chat template with few-shot examples support.
+
+    Combines the role-based messaging of ChatPromptTemplate with the example formatting of
+    FewShotPromptTemplate. Has the EXACT same API as ChatPromptTemplate, just enhanced with
+    examples.
+
+    Examples
+    --------
+    Chat with examples::
+
+        template = ChatFewShotTemplate(
+            system_message="You are a classifier",
+            human_message="Classify: {{text}}",
+            examples=[
+                {"input": "Great!", "output": "positive"},
+                {"input": "Bad", "output": "negative"}
+            ]
+        )
+
+        messages = template.to_messages(text="Amazing!")
+        # System message includes examples
+    """
+
+    def __init__(
+        self,
+        messages: list[dict[str, str]] | None = None,
+        system_message: str | None = None,
+        human_message: str | None = None,
+        examples: list[dict[str, Any]] | None = None,
+        example_separator: str = "\n\n",
+        format_example: Callable[[dict[str, Any]], str] | None = None,
+        **kwargs: Any,
+    ) -> None:
+        """Initialize ChatFewShotTemplate with same API as ChatPromptTemplate plus examples.
+
+        Args
+        ----
+            messages: List of message templates with role and content (same as ChatPromptTemplate)
+            system_message: System message template (same as ChatPromptTemplate)
+            human_message: Human/user message template (same as ChatPromptTemplate)
+            examples: List of example dicts for few-shot learning (NEW)
+            example_separator: String to separate examples (NEW)
+            format_example: Function to format each example (NEW)
+            **kwargs: Additional arguments passed to ChatPromptTemplate
+        """
+        self.examples = examples or []
+        self.example_separator = example_separator
+
+        if format_example is None:
+
+            def _default_formatter(ex: dict[str, Any]) -> str:
+                inp = ex.get("input", "")
+                out = ex.get("output", "")
+                return f"Input: {inp}\nOutput: {out}"
+
+            format_example = _default_formatter
+        self.format_example = format_example
+
+        # If we have examples, enhance the system message
+        enhanced_system_message = system_message
+        if self.examples:
+            examples_text = self.example_separator.join(
+                self.format_example(example) for example in self.examples
+            )
+
+            if system_message:
+                enhanced_system_message = f"{examples_text}{self.example_separator}{system_message}"
+            else:
+                enhanced_system_message = examples_text
+
+        super().__init__(
+            messages=messages,
+            system_message=enhanced_system_message,
+            human_message=human_message,
+            **kwargs,
+        )
+
+        self._original_system_message = system_message
+
+    def add_example(self, example: dict[str, Any]) -> None:
+        """Add an example and rebuild system message."""
+        self.examples.append(example)
+
+        # Rebuild system message with new examples
+        examples_text = self.example_separator.join(self.format_example(ex) for ex in self.examples)
+
+        if self._original_system_message:
+            self.system_message = (
+                f"{examples_text}{self.example_separator}{self._original_system_message}"
+            )
+        else:
+            self.system_message = examples_text