Prompts¶

async def load_system_prompt(
    db: graph.Graph | None,
    ref: str,
    variables: collections.abc.Mapping[str, str],
    default: DefaultPrompt,
    *,
    override: str | None = None,
) -> SystemPrompt:
    """Return the consumer's system prompt.

    Parameters:
        db: The graph, or ``None`` when it is not open yet (the
            packaged prompt is used).
        ref: The CMS reference, such as ``imbi-assistant/system@stable``.
        variables: The consumer's template variables. Only the variables
            the version declares are passed to it, so an editor may drop
            one without breaking the render.
        default: The packaged prompt used when the CMS cannot answer.
        override: An operator-supplied template that replaces both.

    """
    if override:
        try:
            return SystemPrompt(
                text=await _render(override, variables), source='override'
            )
        except rendering.RenderError:
            LOGGER.exception(
                'Failed to render the system prompt override for %s; '
                'using the packaged prompt',
                default.ref,
            )
            return await _fallback(default, variables)

    if db is None:
        return await _fallback(default, variables)

    try:
        prompt, version, _label = await resolve.resolve(db, ref)
        declared = {
            name: value
            for name, value in variables.items()
            if name in version.variable_schema
        }
        rendered = await rendering.render(version, declared, {})
        model_id = await _model_id(db, version, ref)
    except resolve.PromptError as e:
        LOGGER.info('Using the packaged prompt for %s: %s', ref, e)
        return await _fallback(default, variables)
    except Exception:  # noqa: BLE001 - a turn must never fail on this
        LOGGER.warning(
            'Failed to load prompt %s from the CMS; using the packaged prompt',
            ref,
            exc_info=True,
        )
        return await _fallback(default, variables)

    return SystemPrompt(
        text=rendered.system,
        source='cms',
        ref=f'{prompt.namespace}/{prompt.slug}@{version.n}',
        model_id=model_id,
        max_tokens=version.params.max_tokens,
        temperature=version.params.temperature,
    )

async def load_model_id(db: graph.Graph, ref: str) -> str | None:
    """Return the model id of ``ref``'s version, when a consumer can call it.

    ``None`` when the prompt is missing, names no model, or names one
    that is disabled or served by an unsupported driver. Errors are
    logged, never raised.
    """
    try:
        _prompt, version, _label = await resolve.resolve(db, ref)
        return await _model_id(db, version, ref)
    except resolve.PromptError:
        return None
    except Exception:  # noqa: BLE001 - creating a conversation must not fail
        LOGGER.warning('Failed to read the model of %s', ref, exc_info=True)
        return None

async def resolve(
    db: graph.Graph, ref: str
) -> tuple[models.Prompt, models.PromptVersion, str | None]:
    """Find the version a reference names.

    Returns:
        The prompt, the version, and the label used (``None`` when the
        reference named a version number).

    Raises:
        InvalidRef: The reference is not valid.
        PromptNotFound: The prompt, version, or a usable label does not
            exist.

    """
    namespace, slug, selector = parse_ref(ref)
    prompt, _latest = await fetch_prompt(db, namespace, slug)
    if selector is not None and selector.isascii() and selector.isdigit():
        return prompt, await fetch_version(db, prompt, int(selector)), None
    by_name = {label.name: label for label in prompt.labels}
    label = by_name.get(selector or '') or by_name.get(prompt.default_label)
    if label is None:
        raise PromptNotFound(
            f'Prompt {namespace}/{slug} has no label {selector!r} '
            f'and no default label {prompt.default_label!r}'
        )
    version = await fetch_version(db, prompt, label.version)
    return prompt, version, label.name

def parse_ref(ref: str) -> tuple[str, str, str | None]:
    """Split ``namespace/slug@label`` into its parts.

    Raises:
        InvalidRef: The reference is not valid.

    """
    match = REF_PATTERN.match(ref)
    if match is None:
        raise InvalidRef(
            f'Invalid prompt reference {ref!r}; expected '
            'namespace/slug, namespace/slug@label, or namespace/slug@n'
        )
    return match['namespace'], match['slug'], match['label']

async def render(
    version: models.PromptVersion,
    variables: collections.abc.Mapping[str, object],
    providers: collections.abc.Mapping[str, Provider],
) -> RenderedPrompt:
    """Render the system prompt and each message of ``version``.

    Raises:
        RenderError: The variables are not valid, the template fails or
            breaks a limit, the render takes too long, or the output is
            too large.

    """
    sources = {'system': version.system} | {
        f'messages.{i}': m.content for i, m in enumerate(version.messages)
    }
    rendered = await _render_sources(version, variables, providers, sources)
    return RenderedPrompt(
        system=rendered['system'],
        messages=[
            models.PromptMessage(
                role=m.role, content=rendered[f'messages.{i}']
            )
            for i, m in enumerate(version.messages)
        ],
    )

async def render_decision(
    version: models.PromptVersion,
    variables: collections.abc.Mapping[str, object],
    providers: collections.abc.Mapping[str, Provider],
) -> RenderedDecision:
    """Render the state and every question template of ``version``.

    Raises:
        RenderError: As :func:`render`.

    """
    rendered = await _render_sources(
        version,
        variables,
        providers,
        decision_sources(version.state, version.questions),
    )
    questions: dict[str, models.DecisionQuestion] = {}
    for qid, question in version.questions.items():
        base = f'questions.{qid}'
        instructions = rendered[f'{base}.instructions']
        match question:
            case models.NoulQuestion():
                criteria = (
                    None
                    if question.criteria is None
                    else models.NoulCriteria(
                        true=rendered[f'{base}.criteria.true'],
                        false=rendered[f'{base}.criteria.false'],
                    )
                )
                questions[qid] = models.NoulQuestion(
                    instructions=instructions, criteria=criteria
                )
            case models.ChoiceQuestion():
                questions[qid] = models.ChoiceQuestion(
                    instructions=instructions,
                    criteria={
                        key: (
                            None
                            if text is None
                            else rendered[f'{base}.criteria.{key}']
                        )
                        for key, text in question.criteria.items()
                    },
                )
            case models.ScoreQuestion():
                questions[qid] = models.ScoreQuestion(
                    instructions=instructions,
                    criteria=[
                        rendered[f'{base}.criteria.{i}']
                        for i in range(len(question.criteria))
                    ],
                )
    return RenderedDecision(
        state=_state_value(rendered['state']), questions=questions
    )

def check_syntax(
    version_fields: collections.abc.Mapping[str, str],
) -> None:
    """Parse each template so a broken version cannot be saved.

    Parameters:
        version_fields: Field label (``system``, ``messages[0]``) to
            template source.

    Raises:
        RenderError: A template does not parse, uses a macro, or nests
            loops too deep.

    """
    for label, source in version_fields.items():
        try:
            _parse(source)
        except jinja2.TemplateSyntaxError as e:
            raise RenderError(f'{label}: line {e.lineno}: {e.message}') from e
        except RenderError as e:
            raise RenderError(f'{label}: {e}') from e

def validate_variables(
    schema: collections.abc.Mapping[str, models.PromptVariable],
    variables: collections.abc.Mapping[str, object],
) -> None:
    """Check the caller's variables against the version's schema.

    Raises:
        RenderError: A variable is unknown, missing, or the wrong type.

    """
    unknown = sorted(set(variables) - set(schema))
    if unknown:
        raise RenderError(f'Unknown variables: {", ".join(unknown)}')
    for name, spec in schema.items():
        if name not in variables:
            if spec.required:
                raise RenderError(f'Missing required variable {name!r}')
            continue
        value = variables[name]
        allowed = _TYPES[spec.type]
        # ``bool`` is an ``int`` in Python; do not accept it as a number.
        is_bool = isinstance(value, bool)
        if not isinstance(value, allowed) or (is_bool and spec.type != 'bool'):
            raise RenderError(f'Variable {name!r} must be of type {spec.type}')

@contextlib.asynccontextmanager
async def pool_lifespan() -> collections.abc.AsyncIterator[None]:
    """Start warm render workers with a service, and stop them after."""
    if not in_process():
        try:
            await get_pool().start()
        except Exception:  # rendering starts workers on demand instead
            LOGGER.exception('Could not start the prompt render workers')
    try:
        yield
    finally:
        await shutdown()

def get_pool() -> RenderPool:
    """Return the pool for the running event loop.

    Workers are tied to the loop that started them. A pool left by a
    loop that has ended (tests, ``TestClient`` portals) is killed and
    replaced.
    """
    global _current
    loop = asyncio.get_running_loop()
    if _current is not None and _current.loop is not loop:
        _current.abandon()
        _current = None
    if _current is None:
        _current = RenderPool()
    return _current