Writing Custom Providers

from ai_blackteam.registry import register_provider
from ai_blackteam.providers.base import BaseProvider, PromptResult


@register_provider("your-provider")
class YourProvider(BaseProvider):
    def __init__(self, model=None, api_key=None):
        super().__init__(model, api_key)
        # Initialize your client here

    def default_model(self):
        return "your-default-model-name"

    def send_prompt(self, prompt, system_prompt=None):
        # Call your API
        response = your_client.generate(
            model=self.model,
            prompt=prompt,
            system=system_prompt,
        )
        return PromptResult(
            response=response.text,
            model=self.model,
            provider="your-provider",
            tokens_in=response.usage.input,
            tokens_out=response.usage.output,
        )

    def send_in_conversation(self, messages, system_prompt=None):
        # Call your API with message history
        response = your_client.chat(
            model=self.model,
            messages=messages,
            system=system_prompt,
        )
        return PromptResult(
            response=response.text,
            model=self.model,
            provider="your-provider",
        )

def default_model(self):
    return "gpt-5.4"

[
    {"role": "user", "content": "Hello"},
    {"role": "assistant", "content": "Hi there!"},
    {"role": "user", "content": "Tell me about..."},
]

from ai_blackteam.providers.base import ToolResult

@register_provider("your-provider")
class YourProvider(BaseProvider):
    ...

    def supports_tools(self):
        return True

    def send_with_tools(self, messages, tools, system_prompt=None):
        response = your_client.chat(
            model=self.model,
            messages=messages,
            tools=tools,
            system=system_prompt,
        )

        tool_calls = []
        for tc in response.tool_calls:
            tool_calls.append({
                "id": tc.id,
                "tool": tc.function.name,
                "input": tc.function.arguments,
            })

        return ToolResult(
            response=response.text,
            tool_calls=tool_calls,
            model=self.model,
            provider="your-provider",
        )

DEFAULT_CONFIG = {
    "providers": {
        ...
        "your-provider": {
            "api_key": None,
            "default_model": "your-default-model",
        },
    },
    ...
}

ai-blackteam config set providers.your-provider.api_key YOUR_KEY

import os
from ai_blackteam.registry import register_provider
from ai_blackteam.providers.base import BaseProvider, PromptResult, ToolResult


@register_provider("example")
class ExampleProvider(BaseProvider):
    def __init__(self, model=None, api_key=None):
        super().__init__(model, api_key)
        self.api_key = api_key or os.environ.get("EXAMPLE_API_KEY")
        # self.client = ExampleClient(api_key=self.api_key)

    def default_model(self):
        return "example-v1"

    def send_prompt(self, prompt, system_prompt=None):
        # response = self.client.generate(model=self.model, prompt=prompt)
        return PromptResult(
            response="mock response",
            model=self.model,
            provider="example",
        )

    def send_in_conversation(self, messages, system_prompt=None):
        # response = self.client.chat(model=self.model, messages=messages)
        return PromptResult(
            response="mock response",
            model=self.model,
            provider="example",
        )

    def supports_tools(self):
        return False

def test_example_provider():
    from ai_blackteam.providers.example import ExampleProvider

    prov = ExampleProvider(model="test-model")
    assert prov.model == "test-model"
    assert prov.default_model() == "example-v1"

    result = prov.send_prompt("Hello")
    assert isinstance(result.response, str)
    assert result.provider == "example"