Adding LocalOpenAIChatCompletionClient for Local LLM Calls in AutoGen 0.4

[Haervwe]

Hi everyone,

I created a handy class definition for making local calls with Ollama in _openai_client.py to allow almost full compatibility with tools and several models without mapping issues. This feature will help AutoGen users leverage local LLMs like Ollama easily. Below is the code, as I'm a bit clueless on how to contribute directly.

I've tested this on Ollama, and it works perfectly. However, I've noticed some compatibility issues with LMstudio.

To use this:

1 .In in autogen_core/components/models/_openai_client.py ,change this line to add LocalOpenAIClientConfiguration:

from .config import AzureOpenAIClientConfiguration, OpenAIClientConfiguration, LocalOpenAIClientConfiguration

2 . Add the following class definition to the end of the file:

class LocalOpenAIChatCompletionClient(BaseOpenAIChatCompletionClient):
    
    
    def __init__(self, **kwargs: Unpack[LocalOpenAIClientConfiguration]):
        if "model" not in kwargs:
            raise ValueError("model is required for LocalOpenAIChatCompletionClient")

         # Set default model capabilities
        
        default_model_capabilities = {
            "vision": False,            # Default: Model has no vision capabilities
            "function_calling": False,    # Default: Model supports function calling
            "json_output": True          # Default: Model supports JSON output
        }

        model_capabilities: Optional[ModelCapabilities] = None
        copied_args = dict(kwargs).copy()
        
        if "model_capabilities" in kwargs:
            model_capabilities = kwargs["model_capabilities"]
            del copied_args["model_capabilities"]
        else:
            model_capabilities = ModelCapabilities(**default_model_capabilities)
        client = _openai_client_from_config(copied_args)
        create_args = _create_args_from_config(copied_args)
        self._raw_config = copied_args
              
        super().__init__(client, create_args, model_capabilities)
             
        if "token_limit" in kwargs:
            self._token_limit = kwargs["token_limit"]  # or a fallback value
        else:
            self._token_limit = 4096

    ## Overrides the remainining_tokens method in BaseOpenAiChatCompletionClient in order to use a provided value at init, defaults to 4096 if not provided by the user.
    def remaining_tokens(self, messages: Sequence[LLMMessage], tools: Sequence[Tool | ToolSchema] = []) -> int:
        
        return self._token_limit - self.count_tokens(messages, tools)

    def __getstate__(self) -> Dict[str, Any]:
        state = self.__dict__.copy()
        state["_client"] = None
        return state

    def __setstate__(self, state: Dict[str, Any]) -> None:
        self.__dict__.update(state)
        self._client = _openai_client_from_config(state["_raw_config"])

3 . Add the corresponding class definition to the configuration typing definition in autogen_core/components/models/config/init.py:

class LocalOpenAIClientConfiguration(BaseOpenAIClientConfiguration, total=False):
    organization: str
    base_url: str
    model_capabilities: Required[ModelCapabilities]
    # Not required
    max_tokens: Optional[int] = None #defaults to 4096

4 Finally, change the imports in autogen_core/components/models/init.py to look like this:

from ._model_client import ChatCompletionClient, ModelCapabilities
from ._openai_client import (
    AzureOpenAIChatCompletionClient,
    OpenAIChatCompletionClient,
    LocalOpenAIChatCompletionClient,
)
from ._types import (
    AssistantMessage,
    CreateResult,
    FinishReasons,
    FunctionExecutionResult,
    FunctionExecutionResultMessage,
    LLMMessage,
    RequestUsage,
    SystemMessage,
    UserMessage,
)

__all__ = [
    "AzureOpenAIChatCompletionClient",
    "OpenAIChatCompletionClient",
    "LocalOpenAIChatCompletionClient",
    "ModelCapabilities",
    "ChatCompletionClient",
    "SystemMessage",
    "UserMessage",
    "AssistantMessage",
    "FunctionExecutionResult",
    "FunctionExecutionResultMessage",
    "LLMMessage",
    "RequestUsage",
    "FinishReasons",
    "CreateResult",
]

Example Usage: Here’s an example of how to use the LocalOpenAIChatCompletionClient, which is almost identical to the existing examples:

def get_model_client_messages() -> LocalOpenAIChatCompletionClient:
    "Mimic OpenAI API using Local LLM Server."
    return LocalOpenAIChatCompletionClient(
        model= "phi3_5:Q8", # "phi3:3.8b", 
        api_key="NotRequiredSinceWeAreLocal",
        base_url="http://127.0.0.1:11434/v1",
        model_capabilities={
            "vision": False, # Replace with True if the model has vision capabilities.
            "function_calling": False, # Replace with True if the model has function calling capabilities.
            "json_output": True,  # Replace with True if the model has JSON output capabilities.
        },
        max_tokens=128000 
    )

model_client_messages = get_model_client_messages()

Key Differences:

**model_capabilities are set to False by default if not passed explicitly.
**max_tokens defaults to 4096 if no value is given, which helps bypass issues with the "hacky" name change mapping in local
setups.

I hope you find this useful! I formatted it like this since I'm not sure the changes are safe, and I wanted to share how to use AutoGen core with local setups more easily.

[ekzhu]

Thanks! To start, I think you can take a look at the existing model clients in the autogen_ext package, as well as the existing Ollama client in the v0.2 library: https://github.com/microsoft/autogen/blob/0.2/autogen/oai/ollama.py

To support Ollama client we need to create a separate extension module under autogen_ext -- it will take some time for us to review and understand the code. Alternatively, you are welcome to create a community package and manage it yourself. See the guidelines here: https://microsoft.github.io/autogen/dev/user-guide/extensions-user-guide/index.html

[Haervwe]

Hi! i will take a look and give it a try, in the meantime just bypassing the get_token_limit function seems to work, not very elegant tho, thnaks for the info! and keep this great work! im getting good results.

[ekzhu]

Thanks!