Skip to content

Teal Agents

Chat Completion Factories

Background

If you hadn't noticed, Teal Agents Framework is built on top of the open source Semantic Kernel library. Semantic Kernel leverages the concept of the Kernel to which certain capabilities are added, which can then be used to interact with LLMs for things like agents.

One of the core components you add to a Kernel is the ChatCompletionClient. This object provides connection details to the specific LLM you want to enable for a given Kernel (and in this context, an Agent). ChatCompletionClient's come in many flavors with the ability to connect to things like standard OpenAI APIs, Azure OpenAI and Anthropic APIs, Anthropic APIs, Ollama APIs, etc. Due to the early nature of this product, there is limited, in-built support for the different model hosting options (indeed only standard OpenAI at the time of writing).

ChatCompletion Factories

In the ska_types module, there is defined a ChatCompletionFactory abstract class. Within the chat_completion package, there is a default implementation of this class called DefaultChatCompletionFactory. Standard, widely supported ChatCompletions should be included here, and we welcome contributions to build out the available standard models.

Custom Factories

There is also the ability for users to add a supplementary ChatCompletionFactory of their own writing. To do so, you'll need to create a module with a class that extends the ChatCompletionFactory abstract class and implement three methods:

  • get_configs - Static method you should implement if your chat completions require additional configuration (via environment variables).
  • get_chat_completion_for_model_name - Method that returns an instance of ChatCompletionClientBase for connecting to a model with a given name. When implemented, the model name can then be used for an agent's model property in the agent config YAML.
  • get_model_type_for_name - Method that returns the appropriate ModelType for a given model name. This is currently only used to determine the correct way to calculate the token usage for a model response and if someone can think of a better way to do it, I'd like this method to go away.

Once you've created your custom factory, enable it by setting the two following environment variables:

  • TA_CUSTOM_CHAT_COMPLETION_FACTORY_MODULE - The relative path to the python file containing your factory class
  • TA_CUSTOM_CHAT_COMPLETION_FACTORY_CLASS_NAME - The class name of your custom factory class

Once enabled, you have the ability to use any model your factory can handle with a given agent by specifying the model name in the model field of the agent's YAML configuration.

Included is an example of a custom chat completion factory (non-working unless you have Azure OpenAI endpoints):

sk_agents.chat_completion.custom.example_custom_chat_completion_factory.ExampleCustomChatCompletionFactory

Bases: ChatCompletionFactory

Source code in src/sk_agents/chat_completion/custom/example_custom_chat_completion_factory.py 
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
class ExampleCustomChatCompletionFactory(ChatCompletionFactory):
    _OPENAI_MODELS: list[str] = [
        "gpt-35-turbo-1106",
        "gpt-35-turbo-0125",
        "gpt-4o-2024-05-13",
        "gpt-4o-2024-08-06",
        "gpt-4o-mini-2024-07-18",
        "gpt-4-turbo-2024-04-09",
    ]
    _ANTHROPIC_MODELS: list[str] = [
        "claude-3-5-sonnet-20240620",
        "claude-3-haiku-20240307",
    ]

    _GOOGLE_MODELS: list[str] = [
        "gemini-2-5-pro-preview-03-25",
        "gemini-2-5-flash-preview-04-17",
        "gemini-2-0-flash",
        "gemini-2-0-flash-lite",
    ]

    TA_BASE_URL = UtilConfig(
        env_name="TA_BASE_URL",
        is_required=False,
        default_value="https://<Your AI Service Endpoint>",
    )
    TA_API_VERSION = UtilConfig(
        env_name="TA_API_VERSION", is_required=False, default_value="2024-10-21"
    )

    _CONFIGS: list[UtilConfig] = [TA_BASE_URL, TA_API_VERSION]

    def __init__(self, app_config: AppConfig):
        super().__init__(app_config)
        self.api_key = app_config.get(TA_API_KEY.env_name)
        self.url_base = app_config.get(ExampleCustomChatCompletionFactory.TA_BASE_URL.env_name)
        self.api_version = app_config.get(
            ExampleCustomChatCompletionFactory.TA_API_VERSION.env_name
        )

    @staticmethod
    def get_configs() -> list[UtilConfig]:
        return ExampleCustomChatCompletionFactory._CONFIGS

    def get_chat_completion_for_model_name(
        self, model_name: str, service_id: str
    ) -> ChatCompletionClientBase:
        if model_name in ExampleCustomChatCompletionFactory._OPENAI_MODELS:
            return AzureChatCompletion(
                service_id=service_id,
                deployment_name=model_name,
                api_key=self.api_key,
                base_url=f"{self.url_base}/openai",
                api_version=self.api_version,
            )
        elif model_name in ExampleCustomChatCompletionFactory._ANTHROPIC_MODELS:
            return AnthropicChatCompletion(
                service_id=service_id,
                api_key="unused",
                ai_model_id=model_name,
                async_client=AsyncAnthropic(
                    api_key="unused",
                    base_url=f"{self.url_base}/anthropic/{model_name}-v1",
                    default_headers={"X-Custom-Header": self.api_key},
                ),
            )
        elif model_name in ExampleCustomChatCompletionFactory._GOOGLE_MODELS:
            return GoogleAIChatCompletion(
                service_id=service_id,
                deployment_name=model_name,
                api_key=self.api_key,
            )
        else:
            raise ValueError("Model type not supported")

    def get_model_type_for_name(self, model_name: str) -> ModelType:
        if model_name in ExampleCustomChatCompletionFactory._OPENAI_MODELS:
            return ModelType.OPENAI
        elif model_name in ExampleCustomChatCompletionFactory._ANTHROPIC_MODELS:
            return ModelType.ANTHROPIC
        elif model_name in ExampleCustomChatCompletionFactory._GOOGLE_MODELS:
            return ModelType.GOOGLE
        else:
            raise ValueError(f"Unknown model name {model_name}")

    def model_supports_structured_output(self, model_name: str) -> bool:
        if model_name in ExampleCustomChatCompletionFactory._OPENAI_MODELS:
            return True
        elif model_name in ExampleCustomChatCompletionFactory._ANTHROPIC_MODELS:
            return False
        elif model_name in ExampleCustomChatCompletionFactory._GOOGLE_MODELS:
            return True
        else:
            raise ValueError(f"Unknown model name {model_name}")

Notes

  1. A custom factory takes precedence over the default factory, so if your factory provides a chat completion with a model name matching one in the default factory, yours will be selected.
  2. While not exactly on-topic, there is an additional configuration property TA_STRUCTURED_OUTPUT_TRANSFORMER_MODEL which specifies which model should be used to convert agent results in to structured output whenever an output_type is defined. This defaults to openai-gpt-4o but if you've implemented a custom factory, you can override this property to be one of your models.