examples/chat_completion.py

"""Run an AffixIO-authorised Hugging Face chat completion."""

# Credit: @paparichens

from affix_huggingface import AdmissionChecks, AffixHuggingFace, ModelRef


def main() -> None:
    model = ModelRef(
        model_id="meta-llama/Llama-3.1-8B-Instruct",
        provider="hf-inference",
    )
    checks = AdmissionChecks(
        entitled=True,
        data_route_allowed=True,
        controls_satisfied=True,
    )

    with AffixHuggingFace() as client:
        result = client.chat_completion(
            [{"role": "user", "content": "Summarise the supplied operational note."}],
            model=model,
            checks=checks,
            max_tokens=150,
        )

    print(result.output)
    print(f"AffixIO receipt: {result.receipt.gate.receipt_id}")


if __name__ == "__main__":
    main()