examples/openai_tool_flow.py

from __future__ import annotations

import json
import sys
from pathlib import Path

REPO_ROOT = Path(__file__).resolve().parents[1]
if str(REPO_ROOT) not in sys.path:
    sys.path.insert(0, str(REPO_ROOT))

from reframr.cli import compose_generation_context
from reframr.model import ReframrModel


def main() -> None:
    model = ReframrModel.load(REPO_ROOT / "model.safetensors")

    messages = [
        {
            "role": "system",
            "content": (
                "Answer from tool evidence when it is provided. "
                "If the question needs fresh information and no source is available, say what is missing."
            ),
        },
        {
            "role": "user",
            "content": "Who won the Rivergate mayoral runoff, and what number should I cite?",
        },
        {
            "role": "assistant",
            "tool_calls": [
                {
                    "id": "call_1",
                    "type": "function",
                    "function": {
                        "name": "web.search",
                        "arguments": json.dumps({"query": "Rivergate mayoral runoff result vote share"}),
                    },
                }
            ],
        },
        {
            "role": "tool",
            "tool_call_id": "call_1",
            "name": "web.search",
            "content": json.dumps(
                {
                    "ok": True,
                    "source": {
                        "title": "Local Civic Wire",
                        "url": "https://example.org/rivergate-runoff",
                        "snippet": "Mara Ibekwe won the Rivergate mayoral runoff with 52.4 percent of the vote.",
                    },
                }
            ),
        },
    ]

    context = compose_generation_context("", messages=messages)
    print(
        model.generate_text(
            context,
            max_tokens=90,
            temperature=0.58,
            top_k=64,
            top_p=0.92,
            repetition_penalty=1.25,
        )
    )


if __name__ == "__main__":
    main()