pixie-prompts 0.0.0__py3-none-any.whl

pixie/prompts/server.py

@@ -0,0 +1,263 @@
+"""FastAPI server for SDK."""
+
+import argparse
+import os
+import colorlog
+import logging
+import sys
+import importlib.util
+from pathlib import Path
+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.storage import initialize_prompt_storage
+from pixie.prompts.graphql import schema
+
+
+logger = logging.getLogger(__name__)
+
+# Global logging mode
+_logging_mode: str = "default"
+
+
+def discover_and_load_prompts():
+    """Discover and load all Python files that use pixie.prompts.create_prompt, or pixie.create_prompt.
+
+    This function recursively searches the current working directory for Python files
+    """
+    cwd = Path.cwd()
+    # Recursively find all Python files
+    python_files = list(cwd.rglob("*.py"))
+
+    if not python_files:
+        return
+
+    # Add current directory to Python path if not already there
+    if str(cwd) not in sys.path:
+        sys.path.insert(0, str(cwd))
+
+    loaded_count = 0
+    for py_file in python_files:
+        # Skip __init__.py, private files, and anything in site-packages/venv
+        if py_file.name.startswith("_") or any(
+            part in py_file.parts
+            for part in ["site-packages", ".venv", "venv", "__pycache__"]
+        ):
+            continue
+
+        # Load the module with a unique name based on path
+        relative_path = py_file.relative_to(cwd)
+        module_name = str(relative_path.with_suffix("")).replace("/", ".")
+        spec = importlib.util.spec_from_file_location(module_name, py_file)
+        if spec and spec.loader:
+            module = importlib.util.module_from_spec(spec)
+            sys.modules[module_name] = module
+            spec.loader.exec_module(module)
+            loaded_count += 1
+
+
+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_prompts()
+
+    dotenv.load_dotenv(os.getcwd() + "/.env")
+    storage_directory = os.getenv("PIXIE_PROMPT_STORAGE_DIR", ".pixie/prompts")
+    initialize_prompt_storage(storage_directory)
+
+    app = FastAPI(
+        title="Pixie SDK Server",
+        description="Server for running AI applications and agents",
+        version="0.1.0",
+    )
+
+    # Add CORS middleware
+    app.add_middleware(
+        CORSMiddleware,
+        allow_origins=["*"],  # Allows all origins
+        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 SDK 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()

pixie/prompts/storage.py

@@ -0,0 +1,228 @@
+import json
+import logging
+import os
+from types import NoneType
+from typing import Any, Dict, Protocol, Self, TypedDict
+
+from jsonsubschema import isSubschema
+
+from .prompt import (
+    BasePrompt,
+    BaseUntypedPrompt,
+    Prompt,
+    TPromptVar,
+    variables_definition_to_schema,
+)
+
+
+logger = logging.getLogger(__name__)
+
+
+class PromptStorage(Protocol):
+
+    def load(self) -> None: ...
+
+    def exists(self, prompt_id: str) -> bool: ...
+
+    def save(self, prompt: BaseUntypedPrompt) -> None: ...
+
+    def get(self, prompt_id: str) -> BaseUntypedPrompt: ...
+
+
+class _BasePromptJson(TypedDict):
+    versions: Dict[str, str]
+    defaultVersionId: str
+    variablesSchema: Dict[str, Any]
+
+
+class _FilePromptStorage(PromptStorage):
+
+    def __init__(self, directory: str) -> None:
+        self._directory = directory
+        self._prompts: Dict[str, BaseUntypedPrompt] = {}
+        self.load()
+
+    def load(self) -> None:
+        """prompts that are in storage"""
+        if not os.path.exists(self._directory):
+            os.makedirs(self._directory)
+        for filename in os.listdir(self._directory):
+            if filename.endswith(".json"):
+                prompt_id = filename[:-5]  # remove .json
+                filepath = os.path.join(self._directory, filename)
+                with open(filepath, "r") as f:
+                    data: _BasePromptJson = json.load(f)
+                versions = data["versions"]
+                default_version_id = data["defaultVersionId"]
+                variables_schema = data["variablesSchema"]
+                prompt = BaseUntypedPrompt(
+                    id=prompt_id,
+                    versions=versions,
+                    default_version_id=default_version_id,
+                    variables_schema=variables_schema,
+                )
+                self._prompts[prompt_id] = prompt
+
+    def exists(self, prompt_id: str) -> bool:
+        return prompt_id in self._prompts
+
+    def save(self, prompt: BaseUntypedPrompt) -> bool:
+        prompt_id = prompt.id
+        original = self._prompts.get(prompt_id)
+        new_schema = prompt.get_variables_schema()
+        if original:
+            original_schema = original.get_variables_schema()
+            if not isSubschema(original_schema, new_schema):
+                raise TypeError(
+                    "Original schema must be a subschema of the new schema."
+                )
+        data: _BasePromptJson = {
+            "versions": prompt.get_versions(),
+            "defaultVersionId": prompt.get_default_version_id(),
+            "variablesSchema": prompt.get_variables_schema(),
+        }
+        filepath = os.path.join(self._directory, f"{prompt_id}.json")
+        with open(filepath, "w") as f:
+            json.dump(data, f, indent=2)
+        try:
+            BasePrompt.update_prompt_registry(prompt)
+        except KeyError:
+            # Prompt not in type prompt registry yet, meaning there's no usage in code
+            # thus this untyped prompt would just be stored but not used in code
+            pass
+        self._prompts[prompt_id] = prompt
+        return original is None
+
+    def get(self, prompt_id: str) -> BaseUntypedPrompt:
+        return self._prompts[prompt_id]
+
+
+_storage_instance: PromptStorage | None = None
+
+
+# TODO allow other storage types later
+def initialize_prompt_storage(directory: str) -> None:
+    global _storage_instance
+    if _storage_instance is not None:
+        raise RuntimeError("Prompt storage has already been initialized.")
+    _storage_instance = _FilePromptStorage(directory)
+    logger.info(f"Initialized prompt storage at directory: {directory}")
+
+
+class StorageBackedPrompt(Prompt[TPromptVar]):
+
+    def __init__(
+        self,
+        id: str,
+        *,
+        variables_definition: type[TPromptVar] = NoneType,
+    ) -> None:
+        self._id = id
+        self._variables_definition = variables_definition
+        self._prompt: BasePrompt[TPromptVar] | None = None
+
+    @property
+    def id(self) -> str:
+        return self._id
+
+    @property
+    def variables_definition(self) -> type[TPromptVar]:
+        return self._variables_definition
+
+    def get_variables_schema(self) -> dict[str, Any]:
+        return variables_definition_to_schema(self._variables_definition)
+
+    def _get_prompt(self) -> BasePrompt[TPromptVar]:
+        if _storage_instance is None:
+            raise RuntimeError("Prompt storage has not been initialized.")
+        if self._prompt is None:
+            untyped_prompt = _storage_instance.get(self.id)
+            self._prompt = BasePrompt.from_untyped(
+                untyped_prompt,
+                variables_definition=self.variables_definition,
+            )
+            schema_from_storage = untyped_prompt.get_variables_schema()
+            schema_from_definition = self.get_variables_schema()
+            if not isSubschema(schema_from_definition, schema_from_storage):
+                raise TypeError(
+                    "Schema from definition is not a subschema of the schema from storage."
+                )
+        return self._prompt
+
+    def actualize(self) -> Self:
+        self._get_prompt()
+        return self
+
+    def exists_in_storage(self) -> bool:
+        if _storage_instance is None:
+            raise RuntimeError("Prompt storage has not been initialized.")
+        try:
+            self.actualize()
+            return True
+        except KeyError:
+            return False
+
+    def get_versions(self) -> dict[str, str]:
+        prompt = self._get_prompt()
+        return prompt.get_versions()
+
+    def get_version_count(self) -> int:
+        try:
+            prompt = self._get_prompt()
+            versions_dict = prompt.get_versions()
+            return len(versions_dict)
+        except KeyError:
+            return 0
+
+    def get_default_version_id(self) -> str:
+        prompt = self._get_prompt()
+        return prompt.get_default_version_id()
+
+    def compile(
+        self,
+        variables: TPromptVar = None,
+        *,
+        version_id: str | None = None,
+    ) -> str:
+        prompt = self._get_prompt()
+        return prompt.compile(variables=variables, version_id=version_id)
+
+    def append_version(
+        self,
+        version_id: str,
+        content: str,
+        set_as_default: bool = False,
+    ) -> BasePrompt[TPromptVar]:
+        if _storage_instance is None:
+            raise RuntimeError("Prompt storage has not been initialized.")
+        if self.exists_in_storage():
+            prompt = self._get_prompt()
+            prompt.append_version(
+                version_id=version_id,
+                content=content,
+                set_as_default=set_as_default,
+            )
+            _storage_instance.save(prompt)
+            return prompt
+        else:
+            # it should be safe to assume there's no actualized prompt for this id
+            # thus it should be same to create a new instance of BasePrompt
+            new_prompt = BasePrompt(
+                id=self.id,
+                versions={version_id: content},
+                variables_definition=self.variables_definition,
+                default_version_id=version_id,
+            )
+            _storage_instance.save(new_prompt)
+            return new_prompt
+
+    def update_default_version_id(
+        self,
+        version_id: str,
+    ) -> BasePrompt[TPromptVar]:
+        if _storage_instance is None:
+            raise RuntimeError("Prompt storage has not been initialized.")
+        prompt = self._get_prompt()
+        prompt.update_default_version_id(version_id)
+        _storage_instance.save(prompt)
+        return prompt