dataface-playground 0.0.1__py3-none-any.whl

dataface_playground/__init__.py

@@ -0,0 +1 @@
+"""Dataface interactive playground — OSS dashboard editor package."""

dataface_playground/ai_service.py

@@ -0,0 +1,535 @@
+"""AI service for playground chat functionality.
+
+This module provides OpenAI integration for generating dashboards, reports,
+and providing database exploration assistance. Tool calls are dispatched
+through the canonical Dataface AI tool surface.
+"""
+
+import json
+import os
+from pathlib import Path
+from typing import Any
+
+from dataface.agent_api.docs import read_full_text as _read_dataface_syntax
+from dataface.ai.context import DatafaceAIContext
+from dataface.ai.generate_sql import get_sql_generation_guidance
+from dataface.ai.prompts import build_context_section, load_prompt, load_shared_prompt
+from dataface.ai.tool_schemas import to_openai_tool
+from dataface.ai.tools import dispatch_tool_call
+from dataface.ai.yaml_utils import extract_yaml
+from dataface_playground.ai_tools import PLAYGROUND_TOOLS, _handle_render_dashboard
+from dataface_playground.error_display import error_summary as format_playground_error
+
+
+class AIService:
+    """Service for interacting with OpenAI API."""
+
+    _TOOL_INSTRUCTIONS = """
+
+## IMPORTANT: Rendering and Testing
+
+You have access to tools for working with dashboards:
+
+1. **ALWAYS render** your generated YAML using `render_dashboard` before returning it — this checks structure, executes queries, and verifies render output in one step
+2. **Use `execute_query`** to run SQL directly for data exploration (e.g. discovering tables, checking column names)
+3. **Use `schema`** to browse the data hierarchy: sources, schemas, tables, and columns
+4. **Fix any errors** reported by `render_dashboard` and render again until it succeeds
+5. **Only return YAML** that renders successfully
+6. If `schema(table=...)` says a table has **no cached profile**, that means the table may still exist live. Use the returned columns and/or `execute_query` to confirm instead of claiming the table is missing.
+
+If rendering fails, analyze the errors (including tips), fix them, and render again. Each render attempt is visible to the user in the playground editor, so keep iterating with complete YAML.
+"""
+
+    def __init__(
+        self,
+        api_key: str | None = None,
+        ai_context: DatafaceAIContext | None = None,
+    ) -> None:
+        self.api_key = api_key or os.getenv("OPENAI_API_KEY")
+        self._client: Any = None
+        self.prompts_dir = Path(__file__).parent / "prompts"
+        self._ai_context = ai_context
+
+    @property
+    def client(self) -> Any:
+        """Lazy-load OpenAI client."""
+        if self._client is None and self.api_key:
+            from openai import OpenAI
+
+            self._client = OpenAI(api_key=self.api_key)
+        return self._client
+
+    @property
+    def is_available(self) -> bool:
+        """Lazily initializes the client if API key is now available in environment."""
+        if self._client is None:
+            api_key = self.api_key or os.getenv("OPENAI_API_KEY")
+            if api_key:
+                self.api_key = api_key
+                self._client = None  # Reset so client property re-creates
+        return self.client is not None
+
+    def _load_prompt(self, prompt_name: str) -> str:
+        """Load a prompt template from markdown file.
+
+        Args:
+            prompt_name: Name of prompt file (without .md extension)
+
+        Returns:
+            Prompt content as string
+        """
+        # Use core utility for prompt loading
+        return load_prompt(prompt_name, self.prompts_dir)
+
+    def _build_system_prompt(
+        self,
+        prompt_type: str,
+        database_context: str | None = None,
+        yaml_context: str | None = None,
+        chart_context: str | None = None,
+    ) -> str:
+        """Build system prompt with context.
+
+        Args:
+            prompt_type: Type of prompt (database_exploration, yaml_generation, etc.)
+            database_context: Optional database schema information
+            yaml_context: Optional selected YAML code
+            chart_context: Optional chart information
+
+        Returns:
+            Complete system prompt
+        """
+        base_prompt = self._load_prompt(prompt_type)
+
+        # Prepend shared design guide for dashboard/report prompts.
+        # Shared guides under dataface.ai.skills carry universal design
+        # principles; the app-specific prompts add tool instructions.
+        shared_guide_map = {
+            "dashboard_design": "dashboard-design",
+            "report_generation": "report-design",
+        }
+        if prompt_type in shared_guide_map:
+            shared = load_shared_prompt(shared_guide_map[prompt_type])
+            if shared:
+                base_prompt = shared + "\n\n---\n\n" + base_prompt
+
+        if prompt_type in ("yaml_generation", "dashboard_design", "report_generation"):
+            base_prompt += "\n\n" + get_sql_generation_guidance()
+
+        schema_section = ""
+        if prompt_type in ("yaml_generation", "dashboard_design", "report_generation"):
+            schema_section = f"\n\n{_read_dataface_syntax()}\n\n"
+
+        # Use core utility for building context section
+        context_section = build_context_section(
+            database_context=database_context,
+            yaml_context=yaml_context,
+            chart_context=chart_context,
+        )
+
+        # Build final prompt: context first, then base prompt, then schema
+        if context_section:
+            return context_section + "\n\n" + base_prompt + schema_section
+
+        return base_prompt + schema_section
+
+    def _get_tools(self) -> list[dict[str, Any]]:
+        """Get playground tool schemas in OpenAI function-calling format.
+
+        Returns:
+            List of tool definitions for OpenAI function calling
+        """
+        return [to_openai_tool(tool) for tool in PLAYGROUND_TOOLS]
+
+    def _dispatch_tool(
+        self, function_name: str, function_args: dict[str, Any]
+    ) -> dict[str, Any]:
+        """Dispatch a tool call via the shared canonical tool surface.
+
+        When the AI omits a ``source`` on query/schema calls, injects the
+        context's ``default_source`` so the tool hits the same database the
+        schema preview discovered and the rendered YAML will use.
+        """
+        ctx = self._ai_context
+        if ctx is None:
+            from dataface_playground.routes import get_playground_ai_context
+
+            ctx = get_playground_ai_context()
+
+        if ctx.default_source and function_name in (
+            "execute_query",
+            "schema",
+        ):
+            function_args.setdefault("source", ctx.default_source)
+
+        if function_name == "render_dashboard":
+            return _handle_render_dashboard(function_args, ctx)
+
+        return dispatch_tool_call(function_name, function_args, context=ctx)
+
+    def _tool_instructions(self) -> str:
+        """Return the shared tool instructions for playground prompts."""
+        return self._TOOL_INSTRUCTIONS
+
+    def chat(
+        self,
+        user_message: str,
+        prompt_type: str = "yaml_generation",
+        database_context: str | None = None,
+        yaml_context: str | None = None,
+        chart_context: str | None = None,
+        conversation_history: list[dict[str, Any]] | None = None,
+        enable_validation: bool = True,
+    ) -> dict[str, Any]:
+        """Send a chat message to OpenAI with optional render tools.
+
+        Args:
+            user_message: User's message
+            prompt_type: Type of prompt to use (database_exploration, yaml_generation, report_generation, dashboard_design)
+            database_context: Optional database schema information
+            yaml_context: Optional selected YAML code
+            chart_context: Optional chart information
+            conversation_history: Optional list of previous messages [{role: str, content: str}]
+            enable_validation: If True, enable function calling for render checks (default: True)
+
+        Returns:
+            Dict with 'response' (str) and 'yaml' (optional str if YAML was generated)
+
+        Raises:
+            ValueError: If AI service is not available
+        """
+        if not self.is_available:
+            raise ValueError(
+                "OpenAI API key not configured. Set OPENAI_API_KEY environment variable."
+            )
+
+        # Build system prompt
+        system_prompt = self._build_system_prompt(
+            prompt_type=prompt_type,
+            database_context=database_context,
+            yaml_context=yaml_context,
+            chart_context=chart_context,
+        )
+
+        # Add instructions about playground tools
+        if enable_validation and prompt_type in (
+            "yaml_generation",
+            "dashboard_design",
+            "report_generation",
+        ):
+            system_prompt += self._tool_instructions()
+
+        # Build messages
+        messages: list[dict[str, Any]] = [{"role": "system", "content": system_prompt}]
+
+        # Add conversation history
+        if conversation_history:
+            messages.extend(conversation_history)
+
+        # Add current user message
+        messages.append({"role": "user", "content": user_message})
+
+        # Call OpenAI API with function calling if render checks are enabled
+        max_iterations = 10  # Increased to allow more iteration for complex dashboards
+        allow_return_with_errors_after = (
+            7  # After this many iterations, allow returning with errors
+        )
+        iteration = 0
+        last_yaml_content = None  # Track last generated YAML
+
+        try:
+            while iteration < max_iterations:
+                iteration += 1
+
+                # After threshold iterations, tell agent it can return with explanation if stuck
+                if iteration == allow_return_with_errors_after:
+                    messages.append(
+                        {
+                            "role": "user",
+                            "content": "You've tried several times to fix the errors. If you're hitting a real issue that you cannot fix (like a configuration problem, missing database connection, or something outside your control), you can return the YAML you've generated along with an explanation of what the issue is. The user will be able to see what you've created and understand the problem. However, if you think you can fix it with one more attempt, please try again.",
+                        }
+                    )
+
+                # Prepare API call
+                # Try GPT-5 Thinking (gpt-5.1) first, fallback to gpt-4o if not available
+                # GPT-5 supports advanced reasoning and function calling
+                model = "gpt-5.1"
+
+                api_kwargs: dict[str, Any] = {
+                    "model": model,
+                    "messages": messages,
+                    "temperature": 0.7,
+                    "max_tokens": 4000,
+                    "stream": False,  # Set to True for streaming (handled separately)
+                }
+
+                # Add reasoning parameter if supported (GPT-5 models)
+                # This may not be supported in all SDK versions, so we'll handle errors gracefully
+                api_kwargs["reasoning"] = {
+                    "summary": "auto",  # Enable reasoning summaries for status indicators
+                    "effort": "medium",  # Balance between quality and latency
+                }
+
+                # Add tools if render checks are enabled
+                if enable_validation and prompt_type in (
+                    "yaml_generation",
+                    "dashboard_design",
+                    "report_generation",
+                ):
+                    api_kwargs["tools"] = self._get_tools()
+                    api_kwargs["tool_choice"] = "auto"
+
+                # Try API call with reasoning, fallback without if not supported
+                try:
+                    response = self.client.chat.completions.create(**api_kwargs)
+                except (TypeError, ValueError) as e:
+                    # If reasoning parameter is not supported, retry without it
+                    error_str = str(e).lower()
+                    if "reasoning" in error_str or "unexpected keyword" in error_str:
+                        api_kwargs.pop("reasoning", None)
+                        # Also try fallback model if gpt-5.1 doesn't exist
+                        if model == "gpt-5.1":
+                            api_kwargs["model"] = "gpt-4o"
+                        response = self.client.chat.completions.create(**api_kwargs)
+                    else:
+                        raise
+
+                assistant_message = response.choices[0].message
+
+                # Check if function calling was used
+                if assistant_message.tool_calls:
+                    # Add assistant message to conversation
+                    # Convert tool_calls to dict format
+                    tool_calls_dict = [
+                        {
+                            "id": tc.id,
+                            "type": tc.type,
+                            "function": {
+                                "name": tc.function.name,
+                                "arguments": tc.function.arguments,
+                            },
+                        }
+                        for tc in assistant_message.tool_calls
+                    ]
+                    messages.append(
+                        {
+                            "role": "assistant",
+                            "content": assistant_message.content,
+                            "tool_calls": tool_calls_dict,
+                        }
+                    )
+
+                    # Execute tool calls
+                    for tool_call in assistant_message.tool_calls:
+                        # Handle both object and dict access
+                        if hasattr(tool_call, "function"):
+                            function_name = tool_call.function.name
+                            function_args_str = tool_call.function.arguments
+                            tool_call_id = tool_call.id
+                        else:
+                            function_name = tool_call["function"]["name"]
+                            function_args_str = tool_call["function"]["arguments"]
+                            tool_call_id = tool_call["id"]
+
+                        function_args = (
+                            json.loads(function_args_str) if function_args_str else {}
+                        )
+
+                        result = self._dispatch_tool(function_name, function_args)
+
+                        # Truncate large data fields for the AI context window
+                        if "data" in result and isinstance(result["data"], list):
+                            result = {
+                                **result,
+                                "data": result["data"][:20],
+                            }
+
+                        messages.append(
+                            {
+                                "role": "tool",
+                                "tool_call_id": tool_call_id,
+                                "name": function_name,
+                                "content": json.dumps(result, default=str),
+                            }
+                        )
+
+                    # Continue loop to get AI's response after function calls
+                    continue
+
+                # No function calls - we have the final response
+                assistant_content = assistant_message.content or ""
+
+                # Try to extract YAML from response
+                yaml_content = self._extract_yaml(assistant_content)
+
+                # Track the last YAML content we've seen
+                if yaml_content:
+                    last_yaml_content = yaml_content
+
+                # If we have YAML but render checks were enabled, do final checks
+                if (
+                    yaml_content
+                    and enable_validation
+                    and prompt_type
+                    in (
+                        "yaml_generation",
+                        "dashboard_design",
+                        "report_generation",
+                    )
+                ):
+                    render_result = self._dispatch_tool(
+                        "render_dashboard",
+                        {"yaml_content": yaml_content},
+                    )
+
+                    if not render_result.get("success", False):
+                        # After threshold, allow agent to return with explanation if it wants
+                        if iteration >= allow_return_with_errors_after:
+                            if any(
+                                phrase in assistant_content.lower()
+                                for phrase in [
+                                    "cannot fix",
+                                    "unable to",
+                                    "configuration",
+                                    "missing",
+                                    "outside my control",
+                                    "requires",
+                                    "need to",
+                                    "issue is",
+                                ]
+                            ):
+                                return {
+                                    "response": assistant_content,
+                                    "yaml": yaml_content,
+                                }
+
+                        # Continue iterating — feed errors back to AI
+                        render_error_summary = render_result.get("error_summary", "")
+                        errors = render_result.get("errors", [])
+                        tips = render_result.get("tips", [])
+
+                        error_feedback = (
+                            "Dashboard rendering failed. Errors:\n"
+                            f"{render_error_summary}"
+                        )
+                        if errors:
+                            error_feedback += "\n\nDetailed errors:\n" + "\n".join(
+                                format_playground_error(e) for e in errors[:5]
+                            )
+                        if tips:
+                            error_feedback += "\n\nTips:\n" + "\n".join(tips)
+                        error_feedback += "\n\nPlease fix these errors and regenerate the complete YAML. Use `execute_query` to discover table names if needed. Do not ask questions - fix the errors and return complete, valid YAML."
+
+                        messages.append({"role": "user", "content": error_feedback})
+                        continue
+
+                # Render checks passed (or were disabled) - return the response
+                return {
+                    "response": assistant_content,
+                    "yaml": yaml_content,
+                }
+
+            # If we've done too many iterations, return the last YAML we saw with explanation
+            # Ask the agent to explain what happened
+            if last_yaml_content:
+                # One final attempt to get an explanation from the agent
+                messages.append(
+                    {
+                        "role": "user",
+                        "content": "You've reached the maximum number of render attempts. Please return the YAML you've generated along with an explanation of what issues remain and why they couldn't be fixed. The user needs to see what you've created.",
+                    }
+                )
+
+                try:
+                    # Get final response with explanation
+                    api_kwargs = {
+                        "model": "gpt-5.1",
+                        "messages": messages,
+                        "temperature": 0.7,
+                        "max_tokens": 4000,
+                    }
+                    # Try adding reasoning if supported
+                    api_kwargs["reasoning"] = {
+                        "summary": "auto",
+                        "effort": "medium",
+                    }
+
+                    try:
+                        response = self.client.chat.completions.create(**api_kwargs)
+                    except (TypeError, ValueError) as e:
+                        # If reasoning parameter is not supported, retry without it
+                        error_str = str(e).lower()
+                        if (
+                            "reasoning" in error_str
+                            or "unexpected keyword" in error_str
+                        ):
+                            api_kwargs.pop("reasoning", None)
+                            # Also try fallback model if gpt-5.1 doesn't exist
+                            if api_kwargs.get("model") == "gpt-5.1":
+                                api_kwargs["model"] = "gpt-4o"
+                            response = self.client.chat.completions.create(**api_kwargs)
+                        else:
+                            raise
+                    assistant_content = response.choices[0].message.content or ""
+
+                    # Extract YAML if present, otherwise use last known YAML
+                    final_yaml = (
+                        self._extract_yaml(assistant_content) or last_yaml_content
+                    )
+
+                    return {
+                        "response": assistant_content
+                        or "Reached maximum render iterations. Here's the YAML that was generated. Please review it manually.",
+                        "yaml": final_yaml,
+                    }
+                except Exception:  # noqa: BLE001
+                    # If final call fails, still return what we have
+                    return {
+                        "response": "Reached maximum render iterations. Here's the YAML that was generated. Please review it manually.",
+                        "yaml": last_yaml_content,
+                    }
+            else:
+                # No YAML was ever generated
+                return {
+                    "response": "Reached maximum render iterations but no YAML was generated. Please try again with a different request.",
+                    "yaml": None,
+                }
+
+        except Exception as e:  # noqa: BLE001
+            return {
+                "response": f"Error: {str(e)}",
+                "yaml": None,
+            }
+
+    def _extract_yaml(self, text: str) -> str | None:
+        """Extract YAML code block from text.
+
+        Args:
+            text: Text that may contain YAML code blocks
+
+        Returns:
+            Extracted YAML content or None
+        """
+        # Use core utility for YAML extraction
+        return extract_yaml(text)
+
+
+# Global instance
+_ai_service: AIService | None = None
+
+
+def get_ai_service() -> AIService:
+    """Get global AI service instance."""
+    global _ai_service
+    if _ai_service is None:
+        _ai_service = AIService()
+    return _ai_service
+
+
+def reset_ai_service() -> None:
+    """Reset the global AI service instance.
+
+    This is useful when environment variables change (e.g., after loading .env file).
+    """
+    global _ai_service
+    _ai_service = None