pixie-prompts 0.1.1__py3-none-any.whl

pixie/prompts/prompt.py

@@ -0,0 +1,373 @@
+from copy import deepcopy
+from dataclasses import dataclass
+import json
+from types import NoneType
+from typing import Any, Generic, Protocol, Self, TypeVar, cast, overload
+from uuid import uuid4
+
+import jinja2
+from jinja2 import StrictUndefined
+from jsonsubschema import isSubschema
+from pydantic import BaseModel
+
+
+class Variables(BaseModel):
+    # TODO add validation to prevent fields using reserved names
+    pass
+
+
+TPromptVar = TypeVar("TPromptVar", bound=Variables | None)
+
+
+_prompt_registry: dict[str, "BasePrompt"] = {}
+"""Registry of all actualized prompts.
+
+Purpose of the registry is to ensure there's single actualized prompt instance per prompt ID globally,
+so that every compiled prompt can track back to one single instance of the prompt it was compiled from.
+"""
+
+
+def get_prompt_by_id(prompt_id: str) -> "BasePrompt":
+    return _prompt_registry[prompt_id]
+
+
+@dataclass(frozen=True)
+class _CompiledPrompt:
+    value: str
+    prompt: "BasePrompt | OutdatedPrompt"
+    version_id: str
+    variables: Variables | None
+
+
+_compiled_prompt_registry: dict[int, _CompiledPrompt] = {}
+"""Registry of all compiled prompts.
+
+This is to keep track of every result string returned by BasePrompt.compile().
+key is the id() of the compiled string."""
+
+
+def _find_matching_prompt(obj):
+    if isinstance(obj, str):
+        for compiled in _compiled_prompt_registry.values():
+            if compiled.value == obj:
+                return compiled
+        return None
+    elif isinstance(obj, dict):
+        for value in obj.values():
+            result = _find_matching_prompt(value)
+            if result:
+                return result
+        return None
+    elif isinstance(obj, list):
+        for item in obj:
+            result = _find_matching_prompt(item)
+            if result:
+                return result
+        return None
+    else:
+        return None
+
+
+def get_compiled_prompt(text: str) -> _CompiledPrompt | None:
+    """Find the compiled prompt metadata for a given compiled prompt string."""
+    if not _compiled_prompt_registry:
+        return None
+    direct_match = _compiled_prompt_registry.get(id(text))
+    if direct_match:
+        return direct_match
+    for compiled in _compiled_prompt_registry.values():
+        if compiled.value == text:
+            return compiled
+    try:
+        obj = json.loads(text)
+        return _find_matching_prompt(obj)
+
+    except json.JSONDecodeError:
+        return None
+
+
+def _mark_compiled_prompts_outdated(
+    prompt_id: str, outdated_prompt: "OutdatedPrompt"
+) -> None:
+    for key in list(_compiled_prompt_registry.keys()):
+        compiled_prompt = _compiled_prompt_registry[key]
+        if compiled_prompt.prompt.id == prompt_id:
+            _compiled_prompt_registry[key] = _CompiledPrompt(
+                value=compiled_prompt.value,
+                version_id=compiled_prompt.version_id,
+                variables=compiled_prompt.variables,
+                prompt=outdated_prompt,
+            )
+
+
+DEFAULT_VERSION_ID = "v0"
+
+
+def _to_versions_dict(versions: str | dict[str, str]) -> dict[str, str]:
+    if isinstance(versions, str):
+        return {DEFAULT_VERSION_ID: versions}
+    return deepcopy(versions)
+
+
+class _UnTypedPrompt(Protocol):
+    @property
+    def id(self) -> str: ...
+
+    def get_default_version_id(self) -> str: ...
+
+    def get_versions(self) -> dict[str, str]: ...
+
+    def get_variables_schema(self) -> dict[str, Any]: ...
+
+
+EMPTY_VARIABLES_SCHEMA = {"type": "object", "properties": {}}
+
+
+class BaseUntypedPrompt(_UnTypedPrompt):
+
+    def __init__(
+        self,
+        *,
+        versions: str | dict[str, str],
+        default_version_id: str | None = None,
+        id: str | None = None,
+        variables_schema: dict[str, Any] | None = None,
+    ) -> None:
+        if not id:
+            id = uuid4().hex
+            while id in _prompt_registry:
+                id = uuid4().hex
+
+        self._id = id
+        if not versions:
+            raise ValueError("No versions provided for the prompt.")
+        self._versions: dict[str, str]
+        self._versions = _to_versions_dict(versions)
+        self._default_version = default_version_id or next(iter(self._versions))
+        self._variables_schema = variables_schema or EMPTY_VARIABLES_SCHEMA
+
+    @property
+    def id(self) -> str:
+        return self._id
+
+    def get_default_version_id(self) -> str:
+        return self._default_version
+
+    def get_versions(self) -> dict[str, str]:
+        return deepcopy(self._versions)
+
+    def get_variables_schema(self) -> dict[str, Any]:
+        return deepcopy(self._variables_schema)
+
+
+class Prompt(_UnTypedPrompt, Generic[TPromptVar]):
+    @property
+    def variables_definition(self) -> type[TPromptVar]: ...
+
+    @overload
+    def compile(
+        self: "BasePrompt[NoneType]", *, version_id: str | None = None
+    ) -> str: ...
+
+    @overload
+    def compile(
+        self, variables: TPromptVar, *, version_id: str | None = None
+    ) -> str: ...
+
+    def compile(
+        self,
+        variables: TPromptVar | None = None,
+        *,
+        version_id: str | None = None,
+    ) -> str: ...
+
+
+def variables_definition_to_schema(definition: type[TPromptVar]) -> dict[str, Any]:
+    if definition is NoneType:
+        return EMPTY_VARIABLES_SCHEMA
+
+    return cast(type[Variables], definition).model_json_schema()
+
+
+class BasePrompt(BaseUntypedPrompt, Generic[TPromptVar]):
+    @classmethod
+    def from_untyped(
+        cls,
+        untyped_prompt: "BaseUntypedPrompt",
+        variables_definition: type[TPromptVar] = NoneType,
+    ) -> "BasePrompt[TPromptVar]":
+        base_schema = untyped_prompt.get_variables_schema()
+        typed_schema = variables_definition_to_schema(variables_definition)
+        if not isSubschema(typed_schema, base_schema):
+            raise TypeError(
+                "The provided variables_definition is not compatible with the prompt's variables schema."
+            )
+        return cls(
+            variables_definition=variables_definition,
+            versions=untyped_prompt.get_versions(),
+            default_version_id=untyped_prompt.get_default_version_id(),
+            id=untyped_prompt.id,
+        )
+
+    def __init__(
+        self,
+        *,
+        versions: str | dict[str, str],
+        default_version_id: str | None = None,
+        variables_definition: type[TPromptVar] = NoneType,
+        id: str | None = None,
+    ) -> None:
+        super().__init__(
+            versions=versions,
+            default_version_id=default_version_id,
+            id=id,
+            variables_schema=variables_definition_to_schema(variables_definition),
+        )
+        self._variables_definition = variables_definition
+        _prompt_registry[self.id] = self
+
+    @property
+    def variables_definition(self) -> type[TPromptVar]:
+        return self._variables_definition
+
+    @overload
+    def compile(
+        self: "BasePrompt[NoneType]", *, version_id: str | None = None
+    ) -> str: ...
+
+    @overload
+    def compile(
+        self, variables: TPromptVar, *, version_id: str | None = None
+    ) -> str: ...
+
+    def compile(
+        self,
+        variables: TPromptVar | None = None,
+        *,
+        version_id: str | None = None,
+    ) -> str:
+        version_id = version_id or self._default_version
+        template_txt = self._versions[version_id]
+        if self._variables_definition is not NoneType:
+            if variables is None:
+                raise ValueError(
+                    f"Variables[{self._variables_definition}] are required for this prompt."
+                )
+            template = jinja2.Template(template_txt, undefined=StrictUndefined)
+            ret = template.render(**variables.model_dump(mode="json"))
+        else:
+            ret = template_txt
+        _compiled_prompt_registry[id(ret)] = _CompiledPrompt(
+            value=ret,
+            version_id=version_id,
+            prompt=self,
+            variables=variables,
+        )
+        return ret
+
+    def _update(
+        self,
+        *,
+        versions: str | dict[str, str] | None = None,
+        default_version_id: str | None = None,
+    ) -> "tuple[Self, OutdatedPrompt[TPromptVar]]":
+        outdated_prompt = OutdatedPrompt.from_prompt(self)
+        if versions is not None:
+            self._versions = _to_versions_dict(versions)
+        if default_version_id is not None:
+            self._default_version = default_version_id
+        _mark_compiled_prompts_outdated(self.id, outdated_prompt)
+        return self, outdated_prompt
+
+    @staticmethod
+    def update_prompt_registry(
+        untyped_prompt: "BaseUntypedPrompt",
+    ) -> "BasePrompt":
+        """IMPORTANT: should only be called from storage on storage load!
+
+        Update the matching entry in type prompt registry in-place.
+        DO NOT call other than from initial storage load, to keep immutability of prompts in code.
+        """
+        existing = get_prompt_by_id(untyped_prompt.id)
+        outdated_prompt = OutdatedPrompt.from_prompt(existing)
+        _mark_compiled_prompts_outdated(existing.id, outdated_prompt)
+        existing._update(
+            versions=untyped_prompt.get_versions(),
+            default_version_id=untyped_prompt.get_default_version_id(),
+        )
+        return existing
+
+    def append_version(
+        self,
+        *,
+        version_id: str,
+        content: str,
+        set_as_default: bool = False,
+    ) -> None:
+        if version_id in self._versions:
+            raise ValueError(f"Version ID '{version_id}' already exists.")
+        self._update(
+            versions={version_id: content, **self._versions},
+            default_version_id=version_id if set_as_default else None,
+        )
+
+    def update_default_version_id(
+        self,
+        default_version_id: str,
+    ) -> None:
+        if default_version_id not in self._versions:
+            raise ValueError(f"Version ID '{default_version_id}' does not exist.")
+        if self._default_version == default_version_id:
+            return
+        self._update(
+            default_version_id=default_version_id,
+        )
+
+
+class OutdatedPrompt(BasePrompt[TPromptVar]):
+
+    def __init__(
+        self,
+        *,
+        versions: str | dict[str, str],
+        default_version_id: str,
+        variables_definition: type[TPromptVar],
+        id: str,
+    ) -> None:
+        self._id = id
+        self._versions = _to_versions_dict(versions)
+        self._default_version = default_version_id
+        self._variables_definition = variables_definition
+
+    @classmethod
+    def from_prompt(
+        cls, prompt: BasePrompt[TPromptVar]
+    ) -> "OutdatedPrompt[TPromptVar]":
+        return cls(
+            variables_definition=prompt.variables_definition,
+            versions=prompt.get_versions(),
+            default_version_id=prompt.get_default_version_id(),
+            id=prompt.id,
+        )
+
+    def _update(
+        self,
+        *,
+        versions: str | dict[str, str] | None = None,
+        default_version_id: str | None = None,
+    ) -> "OutdatedPrompt[TPromptVar]":
+        raise ValueError("Cannot update an outdated prompt.")
+
+    def get_default_version_id(self) -> str:
+        return self._default_version
+
+    def get_versions(self) -> dict[str, str]:
+        return deepcopy(self._versions)
+
+    def compile(
+        self,
+        _variables: TPromptVar | None = None,
+        *,
+        _version_id: str | None = None,
+    ) -> str:
+        raise ValueError("This prompt is outdated and can no longer be used.")

pixie/prompts/prompt_management.py

@@ -0,0 +1,82 @@
+from dataclasses import dataclass
+import inspect
+import logging
+from types import NoneType
+
+from pixie.prompts.prompt import TPromptVar
+from pixie.prompts.storage import StorageBackedPrompt
+
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class StorageBackedPromptWithRegistration:
+    prompt: StorageBackedPrompt
+    description: str | None
+    module: str
+
+
+_registry: dict[str, StorageBackedPromptWithRegistration] = {}
+"""Registry for StorageBackedPrompts created by `create_prompt`.
+
+StorageBackedPrompt is different from BasePrompt because it can be imcomplete
+(when record is not yet fetched from storage, or record doesn't exist at all).
+Thus this registry could contain more entries than _prompt_registry in prompt.py."""
+
+
+def list_prompts() -> list[StorageBackedPromptWithRegistration]:
+    """List all StorageBackedPrompts created via `create_prompt`."""
+    return list(_registry.values())
+
+
+def get_prompt(id: str) -> StorageBackedPromptWithRegistration | None:
+    """Get a StorageBackedPrompt by id, if it was created via `create_prompt`."""
+    return _registry.get(id)
+
+
+def _get_calling_module_name():
+    """Find the name of the module that called this function."""
+    # Get the current frame and the frame above it (the caller's frame)
+    try:
+        # inspect.stack()[2] gets the caller's frame record
+        # frame[0] or frame.frame is the actual frame object
+        caller_frame_record = inspect.stack()[2]
+        caller_frame = caller_frame_record.frame
+    except IndexError:
+        # Handle cases where the stack might be shallower, though unlikely for normal calls
+        return "__main__"
+
+    # Get the module object from the frame object
+    module = inspect.getmodule(caller_frame)
+
+    if module is not None:
+        return module.__name__
+    else:
+        # If getmodule returns None (e.g., if called from the interactive prompt or __main__),
+        # try to get the name from the frame's globals
+        return caller_frame.f_globals.get("__name__", "__main__")
+
+
+def create_prompt(
+    id: str,
+    variables_definition: type[TPromptVar] = NoneType,
+    *,
+    description: str | None = None,
+) -> StorageBackedPrompt[TPromptVar]:
+    if id in _registry:
+        ret = _registry[id].prompt
+        if ret.variables_definition != variables_definition:
+            raise ValueError(
+                f"Prompt with id '{id}' already exists with a different variables definition."
+            )
+        return ret
+    ret = StorageBackedPrompt(id=id, variables_definition=variables_definition)
+    calling_module = _get_calling_module_name()
+    _registry[id] = StorageBackedPromptWithRegistration(
+        prompt=ret,
+        description=description,
+        module=calling_module,
+    )
+    logger.info(f"✅ Registered prompt: {id} ({calling_module})")
+    return ret

pixie/prompts/server.py

@@ -0,0 +1,231 @@
+"""FastAPI server for SDK."""
+
+import argparse
+import os
+import colorlog
+import logging
+from urllib.parse import quote
+
+import dotenv
+from fastapi import FastAPI
+from fastapi.middleware.cors import CORSMiddleware
+from strawberry.fastapi import GraphQLRouter
+import uvicorn
+
+from pixie.prompts.file_watcher import (
+    discover_and_load_modules,
+    init_prompt_storage,
+)
+from pixie.prompts.graphql import schema
+
+
+logger = logging.getLogger(__name__)
+
+# Global logging mode
+_logging_mode: str = "default"
+
+
+def setup_logging(mode: str = "default"):
+    """Configure logging for the entire application.
+
+    Sets up colored logging with consistent formatting for all loggers.
+
+    Args:
+        mode: Logging mode - "default", "verbose", or "debug"
+            - default: INFO for server events, WARNING+ for all modules
+            - verbose: INFO+ for all modules
+            - debug: DEBUG+ for all modules
+    """
+    global _logging_mode
+    _logging_mode = mode
+
+    # Determine log level based on mode
+    if mode == "debug":
+        level = logging.DEBUG
+    elif mode == "verbose":
+        level = logging.INFO
+    else:  # default
+        level = logging.INFO
+
+    colorlog.basicConfig(
+        level=level,
+        format="[%(log_color)s%(levelname)-8s%(reset)s][%(asctime)s]\t%(message)s",
+        datefmt="%H:%M:%S",
+        log_colors={
+            "DEBUG": "cyan",
+            "INFO": "green",
+            "WARNING": "yellow",
+            "ERROR": "red",
+            "CRITICAL": "red,bg_white",
+        },
+        force=True,
+    )
+
+    # Configure uvicorn loggers to use the same format
+    for logger_name in ["uvicorn", "uvicorn.access", "uvicorn.error"]:
+        uvicorn_logger = logging.getLogger(logger_name)
+        uvicorn_logger.handlers = []
+        uvicorn_logger.propagate = True
+
+    # In default mode, set most loggers to WARNING+ except specific modules
+    if mode == "default":
+        # Set root logger to WARNING
+        logging.getLogger().setLevel(logging.WARNING)
+        # Allow INFO for pixie modules
+        logging.getLogger("pixie").setLevel(logging.INFO)
+        # Suppress uvicorn access logs in default mode
+        logging.getLogger("uvicorn.access").setLevel(logging.WARNING)
+
+
+def create_app() -> FastAPI:
+    """Create and configure the FastAPI application.
+
+    Returns:
+        Configured FastAPI application instance with GraphQL router.
+    """
+    # Setup logging first (use global logging mode)
+    setup_logging(_logging_mode)
+
+    # Discover and load applications on every app creation (including reloads)
+    discover_and_load_modules()
+
+    dotenv.load_dotenv(os.getcwd() + "/.env")
+    lifespan = init_prompt_storage()
+
+    app = FastAPI(
+        title="Pixie Prompts Dev Server",
+        description="Server for managing prompts",
+        version="0.1.0",
+        lifespan=lifespan,
+    )
+    # Matches:
+    # 1. http://localhost followed by an optional port (:8080, :3000, etc.)
+    # 2. http://127.0.0.1 followed by an optional port
+    # 3. https://yourdomain.com (the production domain)
+    origins_regex = r"http://(localhost|127\.0\.0\.1)(:\d+)?|https://gopixie\.ai"
+    # Add CORS middleware
+    app.add_middleware(
+        CORSMiddleware,
+        allow_origin_regex=origins_regex,
+        allow_credentials=True,
+        allow_methods=["*"],  # Allows all methods
+        allow_headers=["*"],  # Allows all headers
+    )
+
+    # Add GraphQL router with GraphiQL enabled
+    graphql_app = GraphQLRouter(
+        schema,
+        graphiql=True,
+    )
+
+    app.include_router(graphql_app, prefix="/graphql")
+
+    @app.get("/")
+    async def root():
+        return {
+            "message": "Pixie Prompts Dev Server",
+            "graphiql": "/graphql",
+            "version": "0.1.0",
+        }
+
+    return app
+
+
+def start_server(
+    host: str = "0.0.0.0",
+    port: int = 8000,
+    reload: bool = False,
+    log_mode: str = "default",
+) -> None:
+    """Start the SDK server.
+
+    Args:
+        host: Host to bind to
+        port: Port to bind to
+        reload: Enable auto-reload for development
+        log_mode: Logging mode - "default", "verbose", or "debug"
+        storage_directory: Directory to store prompt definitions
+    """
+    global _logging_mode
+    _logging_mode = log_mode
+
+    # Setup logging (will be called again in create_app for reload scenarios)
+    setup_logging(log_mode)
+
+    # Determine server URL
+    server_url = f"http://{host}:{port}"
+    if host == "0.0.0.0":
+        server_url = f"http://127.0.0.1:{port}"
+
+    # Log server start info
+    logger.info("Starting Pixie SDK Server")
+    logger.info("Server: %s", server_url)
+    logger.info("GraphQL: %s/graphql", server_url)
+
+    # Display gopixie.ai web link
+    encoded_url = quote(f"{server_url}/graphql", safe="")
+    pixie_web_url = f"https://gopixie.ai?url={encoded_url}"
+    logger.info("")
+    logger.info("=" * 60)
+    logger.info("")
+    logger.info("🎨 Open Pixie Web UI:")
+    logger.info("")
+    logger.info("   %s", pixie_web_url)
+    logger.info("")
+    logger.info("=" * 60)
+    logger.info("")
+
+    uvicorn.run(
+        "pixie.prompts.server:create_app",
+        host=host,
+        port=port,
+        loop="asyncio",
+        reload=reload,
+        factory=True,
+        log_config=None,
+    )
+
+
+def main():
+    """Start the Pixie server.
+
+    Loads environment variables and starts the server with auto-reload enabled.
+    Supports --verbose and --debug flags for enhanced logging.
+    """
+    parser = argparse.ArgumentParser(description="Pixie Prompts development server")
+    parser.add_argument(
+        "--verbose",
+        "-v",
+        action="store_true",
+        help="Enable verbose logging (INFO+ for all modules)",
+    )
+    parser.add_argument(
+        "--debug",
+        "-d",
+        action="store_true",
+        help="Enable debug logging (DEBUG+ for all modules)",
+    )
+    parser.add_argument(
+        "--port",
+        "-p",
+        type=int,
+        default=None,
+        help="Port to run the server on (overrides PIXIE_SDK_PORT env var)",
+    )
+    args = parser.parse_args()
+
+    # Determine logging mode
+    log_mode = "default"
+    if args.debug:
+        log_mode = "debug"
+    elif args.verbose:
+        log_mode = "verbose"
+
+    dotenv.load_dotenv(os.getcwd() + "/.env")
+    port = args.port or int(os.getenv("PIXIE_SDK_PORT", "8000"))
+
+    start_server(port=port, reload=True, log_mode=log_mode)
+
+
+if __name__ == "__main__":
+    main()