Custom provider: minimal implementation and usage#

This example shows how to write a small custom provider compatible with the coord2region.ai_model_interface.AIModelInterface.

The provider below simply echoes the supplied prompt. In real usage you would wrap an external API (e.g., your organization’s internal model endpoint).

What you’ll see here: - Define a ModelProvider subclass - Register it with AIModelInterface - Call sync, async and streaming generation methods

from coord2region.ai_model_interface import AIModelInterface, ModelProvider


class EchoProvider(ModelProvider):
    """Provider that returns the prompt verbatim."""

    def __init__(self) -> None:
        super().__init__({"echo-1": "echo-1"})

    def generate_text(self, model: str, prompt, max_tokens: int) -> str:
        if isinstance(prompt, str):
            return prompt
        return " ".join(m["content"] for m in prompt)


def main() -> None:
    ai = AIModelInterface()
    ai.register_provider(EchoProvider())
    # Basic generation
    print(ai.generate_text("echo-1", "Hello from EchoProvider"))

    # Streaming (yields chunks as they arrive)
    print("Streaming:", end=" ")
    for chunk in ai.stream_generate_text("echo-1", "streaming demo", max_tokens=8):
        print(chunk, end="")
    print()

    # Async version
    import asyncio

    async def _demo_async():
        res = await ai.generate_text_async("echo-1", "Hello async", max_tokens=8)
        print("Async:", res)

    asyncio.run(_demo_async())


if __name__ == "__main__":
    main()

Gallery generated by Sphinx-Gallery