hexdag 0.5.0.dev1__py3-none-any.whl

hexdag/builtin/nodes/llm_node.py

@@ -0,0 +1,491 @@
+"""LLMNode - Unified LLM node for prompt building, API calls, and parsing.
+
+This is the primary LLM node in hexdag, providing an n8n-style unified interface
+for all LLM interactions. It combines:
+- Prompt templating (Jinja2-style variable substitution)
+- LLM API calls (via the llm port)
+- Optional structured output parsing (JSON/Pydantic validation)
+"""
+
+from __future__ import annotations
+
+import json
+import re
+import time
+from typing import TYPE_CHECKING, Any
+
+from pydantic import BaseModel, ValidationError
+
+from hexdag.core.context import get_port
+from hexdag.core.exceptions import ParseError
+from hexdag.core.logging import get_logger
+from hexdag.core.orchestration.prompt.template import PromptTemplate
+from hexdag.core.ports.llm import Message
+from hexdag.core.protocols import to_dict
+
+from .base_node_factory import BaseNodeFactory
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+    from hexdag.core.domain.dag import NodeSpec
+    from hexdag.core.orchestration.prompt import PromptInput, TemplateType
+
+logger = get_logger(__name__)
+
+
+def _convert_dicts_to_messages(message_dicts: list[dict[str, str]]) -> list[Message]:
+    """Convert list of message dicts to Message objects."""
+    return [Message(**msg) for msg in message_dicts]
+
+
+class LLMNode(BaseNodeFactory):
+    """Unified LLM node - prompt building, API calls, and optional parsing.
+
+    This is the primary node for LLM interactions in hexdag. It provides a simple,
+    n8n-style interface that handles the complete LLM workflow in a single node.
+
+    Capabilities
+    ------------
+    1. **Prompt Templating**: Jinja2-style variable substitution ({{variable}})
+    2. **LLM API Calls**: Calls the configured LLM port
+    3. **Structured Output**: Optional JSON parsing with Pydantic validation
+    4. **System Prompts**: Optional system message support
+    5. **Message History**: Support for conversation context
+
+    Examples
+    --------
+    Simple text generation::
+
+        llm = LLMNode()
+        spec = llm(
+            name="summarizer",
+            prompt_template="Summarize this text: {{text}}"
+        )
+
+    With system prompt::
+
+        spec = llm(
+            name="assistant",
+            prompt_template="Answer: {{question}}",
+            system_prompt="You are a helpful assistant."
+        )
+
+    Structured output with JSON parsing::
+
+        from pydantic import BaseModel
+
+        class Analysis(BaseModel):
+            sentiment: str
+            confidence: float
+            keywords: list[str]
+
+        spec = llm(
+            name="analyzer",
+            prompt_template="Analyze this text: {{text}}",
+            output_schema=Analysis,
+            parse_json=True
+        )
+
+    YAML Pipeline Usage::
+
+        - kind: llm_node
+          metadata:
+            name: analyzer
+          spec:
+            prompt_template: "Analyze: {{input}}"
+            system_prompt: "You are an analyst"
+            parse_json: true
+            output_schema:
+              summary: str
+              confidence: float
+    """
+
+    def __init__(self, **kwargs: Any) -> None:
+        """Initialize LLMNode."""
+        super().__init__()
+
+    def __call__(
+        self,
+        name: str,
+        prompt_template: PromptInput | str | None = None,
+        output_schema: dict[str, Any] | type[BaseModel] | None = None,
+        system_prompt: str | None = None,
+        parse_json: bool = False,
+        parse_strategy: str = "json",
+        deps: list[str] | None = None,
+        template: PromptInput | str | None = None,  # Alias for prompt_template (YAML compat)
+        **kwargs: Any,
+    ) -> NodeSpec:
+        """Create a unified LLM node specification.
+
+        Parameters
+        ----------
+        name : str
+            Node name (must be unique in the graph)
+        prompt_template : PromptInput | str
+            Template for the user prompt. Supports Jinja2-style {{variable}} syntax.
+            Can be a string or PromptTemplate/ChatPromptTemplate object.
+            Can also be provided as 'template' (alias for YAML compatibility).
+        output_schema : dict[str, Any] | type[BaseModel] | None, optional
+            Expected output schema for structured output. If provided with parse_json=True,
+            the LLM response will be parsed and validated against this schema.
+        system_prompt : str | None, optional
+            System message to prepend to the conversation.
+        parse_json : bool, optional
+            If True, parse the LLM response as JSON and validate against output_schema.
+            Default is False (returns raw text).
+        parse_strategy : str, optional
+            JSON parsing strategy: "json", "json_in_markdown", or "yaml".
+            Default is "json".
+        deps : list[str] | None, optional
+            List of dependency node names.
+        **kwargs : Any
+            Additional parameters passed to NodeSpec.
+
+        Returns
+        -------
+        NodeSpec
+            Complete node specification ready for execution.
+
+        Raises
+        ------
+        ValueError
+            If parse_json is True but no output_schema is provided.
+
+        Examples
+        --------
+        >>> llm = LLMNode()
+        >>> spec = llm(
+        ...     name="greeter",
+        ...     prompt_template="Say hello to {{name}}",
+        ...     system_prompt="You are friendly."
+        ... )
+        """
+        # Handle template alias for YAML compatibility
+        actual_template = prompt_template or template
+        if actual_template is None:
+            raise ValueError("prompt_template (or template) is required")
+
+        if parse_json and output_schema is None:
+            raise ValueError("output_schema is required when parse_json=True")
+
+        # Prepare template
+        prepared_template = self._prepare_template(actual_template)
+
+        # Infer input schema from template variables
+        input_schema = self.infer_input_schema_from_template(
+            prepared_template, special_params={"context_history", "system_prompt"}
+        )
+
+        # Create output model if schema provided
+        output_model: type[BaseModel] | None = None
+        if output_schema is not None:
+            output_model = self.create_pydantic_model(f"{name}Output", output_schema)
+
+        # Create the LLM wrapper function
+        llm_wrapper = self._create_llm_wrapper(
+            name=name,
+            template=prepared_template,
+            output_model=output_model,
+            system_prompt=system_prompt,
+            parse_json=parse_json,
+            parse_strategy=parse_strategy,
+        )
+
+        return self.create_node_with_mapping(
+            name=name,
+            wrapped_fn=llm_wrapper,
+            input_schema=input_schema,
+            output_schema=output_schema if parse_json else None,
+            deps=deps,
+            **kwargs,
+        )
+
+    # Template Processing
+    # -------------------
+
+    @staticmethod
+    def _prepare_template(template: PromptInput | str) -> TemplateType:
+        """Convert string input to PromptTemplate if needed."""
+        if isinstance(template, str):
+            return PromptTemplate(template)
+        return template
+
+    # LLM Wrapper Creation
+    # --------------------
+
+    def _create_llm_wrapper(
+        self,
+        name: str,
+        template: TemplateType,
+        output_model: type[BaseModel] | None,
+        system_prompt: str | None,
+        parse_json: bool,
+        parse_strategy: str,
+    ) -> Callable[..., Any]:
+        """Create an async LLM wrapper function."""
+
+        async def llm_wrapper(validated_input: dict[str, Any]) -> Any:
+            """Execute LLM call with optional parsing."""
+            node_logger = logger.bind(node=name, node_type="llm_node")
+            start_time = time.perf_counter()
+
+            llm = get_port("llm")
+            if not llm:
+                raise RuntimeError("LLM port not available in execution context")
+
+            try:
+                # Convert input to dict if needed
+                try:
+                    input_dict = to_dict(validated_input)
+                except TypeError:
+                    input_dict = validated_input
+
+                # Log input variables at debug level
+                node_logger.debug(
+                    "Prompt variables",
+                    variables=list(input_dict.keys()),
+                    variable_count=len(input_dict),
+                )
+
+                # Log execution start
+                node_logger.info(
+                    "Calling LLM",
+                    has_system_prompt=system_prompt is not None,
+                    parse_json=parse_json,
+                    parse_strategy=parse_strategy if parse_json else None,
+                )
+
+                # Enhance template with schema instructions if using structured output
+                enhanced_template = template
+                if parse_json and output_model:
+                    enhanced_template = self._enhance_template_with_schema(template, output_model)
+
+                # Generate messages from template
+                messages = self._generate_messages(enhanced_template, input_dict, system_prompt)
+
+                # Call LLM
+                response = await llm.aresponse(messages)
+
+                # Log LLM response received
+                duration_ms = (time.perf_counter() - start_time) * 1000
+                node_logger.debug(
+                    "LLM response received",
+                    response_length=len(response) if isinstance(response, str) else None,
+                    duration_ms=f"{duration_ms:.2f}",
+                )
+
+                # Parse response if requested
+                if parse_json and output_model:
+                    result = self._parse_response(response, output_model, parse_strategy)
+                    node_logger.debug(
+                        "Response parsed",
+                        output_type=type(result).__name__,
+                    )
+                    return result
+
+                return response
+
+            except Exception as e:
+                duration_ms = (time.perf_counter() - start_time) * 1000
+                node_logger.error(
+                    "LLM call failed",
+                    duration_ms=f"{duration_ms:.2f}",
+                    error=str(e),
+                    error_type=type(e).__name__,
+                )
+                raise
+
+        return llm_wrapper
+
+    def _generate_messages(
+        self,
+        template: TemplateType,
+        input_data: dict[str, Any],
+        system_prompt: str | None,
+    ) -> list[Message]:
+        """Generate messages from template and input data."""
+        message_dicts = template.to_messages(**input_data)
+
+        # Add system prompt if provided and not already present
+        if system_prompt:
+            has_system = any(msg.get("role") == "system" for msg in message_dicts)
+            if not has_system:
+                message_dicts.insert(0, {"role": "system", "content": system_prompt})
+
+        return _convert_dicts_to_messages(message_dicts)
+
+    # Schema Enhancement
+    # ------------------
+
+    def _enhance_template_with_schema(
+        self, template: TemplateType, output_model: type[BaseModel]
+    ) -> TemplateType:
+        """Add schema instructions to template for structured output."""
+        schema_instruction = self._create_schema_instruction(output_model)
+        return template + schema_instruction
+
+    def _create_schema_instruction(self, output_model: type[BaseModel]) -> str:
+        """Create schema instruction for structured output."""
+        schema = output_model.model_json_schema()
+
+        fields_info = []
+        if "properties" in schema:
+            for field_name, field_schema in schema["properties"].items():
+                field_type = field_schema.get("type", "any")
+                field_desc = field_schema.get("description", "")
+                desc_part = f" - {field_desc}" if field_desc else ""
+                fields_info.append(f"  - {field_name}: {field_type}{desc_part}")
+
+        fields_text = "\n".join(fields_info) if fields_info else "  - (no specific fields defined)"
+
+        example_data = {field: f"<{field}_value>" for field in schema.get("properties", {})}
+        example_json = json.dumps(example_data, indent=2)
+
+        return f"""
+
+## Output Format
+Respond with valid JSON matching this schema:
+{fields_text}
+
+Example: {example_json}
+"""
+
+    # Response Parsing
+    # ----------------
+
+    def _parse_response(
+        self, response: str, output_model: type[BaseModel], strategy: str
+    ) -> BaseModel:
+        """Parse LLM response into structured output."""
+        try:
+            if strategy == "json":
+                parsed_data = self._parse_json(response)
+            elif strategy == "json_in_markdown":
+                parsed_data = self._parse_json_in_markdown(response)
+            elif strategy == "yaml":
+                parsed_data = self._parse_yaml(response)
+            else:
+                parsed_data = self._parse_json(response)
+
+        except (json.JSONDecodeError, ValueError, SyntaxError) as e:
+            error_msg = self._create_parse_error_message(response, str(e), strategy)
+            raise ParseError(error_msg) from e
+
+        # Validate against schema
+        try:
+            return output_model.model_validate(parsed_data)
+        except ValidationError as e:
+            error_msg = self._create_validation_error_message(
+                response, parsed_data, e, output_model
+            )
+            raise ParseError(error_msg) from e
+
+    def _parse_json(self, text: str) -> dict[str, Any]:
+        """Parse JSON from text."""
+        cleaned = text.strip()
+
+        try:
+            result: dict[str, Any] = json.loads(cleaned)
+            return result
+        except json.JSONDecodeError:
+            # Try to extract JSON from surrounding text
+            json_match = re.search(r"(\{.*\}|\[.*\])", cleaned, re.DOTALL)
+            if json_match:
+                return json.loads(json_match.group(1))  # type: ignore[no-any-return]
+            raise
+
+    def _parse_json_in_markdown(self, text: str) -> dict[str, Any]:
+        """Extract and parse JSON from markdown code blocks."""
+        code_block_pattern = r"```(?:json)?\s*(.*?)\s*```"
+        matches = re.findall(code_block_pattern, text, re.DOTALL)
+
+        if matches:
+            for block in matches:
+                try:
+                    result: dict[str, Any] = json.loads(block)
+                    return result
+                except json.JSONDecodeError:
+                    continue
+
+        return self._parse_json(text)
+
+    def _parse_yaml(self, text: str) -> dict[str, Any]:
+        """Parse YAML from text."""
+        import yaml
+
+        result: dict[str, Any] = yaml.safe_load(text)
+        return result
+
+    # Error Messages
+    # --------------
+
+    def _create_parse_error_message(self, text: str, error: str, strategy: str) -> str:
+        """Create helpful error message for parse failures."""
+        preview = text[:200] + ("..." if len(text) > 200 else "")
+
+        return f"""
+Failed to parse LLM output using strategy '{strategy}'.
+
+Error: {error}
+
+Output preview:
+{preview}
+
+Retry hints:
+1. Ensure the LLM output is valid {strategy.upper()} format
+2. Check for trailing commas, missing quotes, or malformed syntax
+3. Consider using 'json_in_markdown' strategy if JSON is in code blocks
+"""
+
+    def _create_validation_error_message(
+        self,
+        text: str,
+        parsed_data: Any,
+        error: ValidationError,
+        model: type[BaseModel],
+    ) -> str:
+        """Create helpful error message for validation failures."""
+        schema = model.model_json_schema()
+        required_fields = schema.get("required", [])
+
+        preview = str(parsed_data)[:200]
+
+        return f"""
+Parsed data does not match expected schema.
+
+Expected schema: {model.__name__}
+Required fields: {required_fields}
+
+Parsed data preview:
+{preview}
+
+Validation errors:
+{error}
+"""
+
+    # Legacy Compatibility
+    # --------------------
+
+    @classmethod
+    def from_template(
+        cls,
+        name: str,
+        template: PromptInput | str,
+        output_schema: dict[str, Any] | type[BaseModel] | None = None,
+        deps: list[str] | None = None,
+        **kwargs: Any,
+    ) -> NodeSpec:
+        """Create a NodeSpec from template (legacy compatibility method).
+
+        This method provides backward compatibility with the old LLMNode API.
+        """
+        return cls()(
+            name=name,
+            prompt_template=template,
+            output_schema=output_schema,
+            parse_json=output_schema is not None,
+            deps=deps,
+            **kwargs,
+        )