50 lines
1.8 KiB
Python
50 lines
1.8 KiB
Python
"""
|
|
Creates different AI instances, based on the given configuration.
|
|
"""
|
|
|
|
import argparse
|
|
from typing import cast, Optional
|
|
from .configuration import Config, AIConfig, OpenAIConfig
|
|
from .ai import AI, AIError
|
|
from .ais.openai import OpenAI
|
|
|
|
|
|
def create_ai(args: argparse.Namespace, config: Config, # noqa: 11
|
|
def_ai: Optional[str] = None,
|
|
def_model: Optional[str] = None) -> AI:
|
|
"""
|
|
Creates an AI subclass instance from the given arguments and configuration file.
|
|
If AI has not been set in the arguments, it searches for the ID 'default'. If
|
|
that is not found, it uses the first AI in the list. It's also possible to
|
|
specify a default AI and model using 'def_ai' and 'def_model'.
|
|
"""
|
|
ai_conf: AIConfig
|
|
if hasattr(args, 'AI') and args.AI:
|
|
try:
|
|
ai_conf = config.ais[args.AI]
|
|
except KeyError:
|
|
raise AIError(f"AI ID '{args.AI}' does not exist in this configuration")
|
|
elif def_ai:
|
|
ai_conf = config.ais[def_ai]
|
|
elif 'default' in config.ais:
|
|
ai_conf = config.ais['default']
|
|
else:
|
|
try:
|
|
ai_conf = next(iter(config.ais.values()))
|
|
except StopIteration:
|
|
raise AIError("No AI found in this configuration")
|
|
|
|
if ai_conf.name == 'openai':
|
|
ai = OpenAI(cast(OpenAIConfig, ai_conf))
|
|
if hasattr(args, 'model') and args.model:
|
|
ai.config.model = args.model
|
|
elif def_model:
|
|
ai.config.model = def_model
|
|
if hasattr(args, 'max_tokens') and args.max_tokens:
|
|
ai.config.max_tokens = args.max_tokens
|
|
if hasattr(args, 'temperature') and args.temperature:
|
|
ai.config.temperature = args.temperature
|
|
return ai
|
|
else:
|
|
raise AIError(f"AI '{args.AI}' is not supported")
|