File size: 2,575 Bytes
aeb21b0 2203a4c aeb21b0 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 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 |
from dataclasses import dataclass
import json
import os
from pathlib import Path
from dotenv import load_dotenv
@dataclass
class InputConfig:
model_name: str = 'openai/gpt-3.5-turbo'
model_api_key: str = 'sk-xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx'
max_tokens: int = 1024
max_steps: int = 3
temperature: float = 0.2
timeout: float = 30.0
sleeptime: float = 0.0
force_max_steps: bool = True
@classmethod
def load(cls, env_file, config_file):
# Load env variables
load_dotenv(env_file)
env_dict = {
'model_name': os.getenv('MODEL_NAME', 'not set'),
'model_api_key': os.getenv('MODEL_API_KEY', 'not set')
}
# Load config JSON
with open(config_file, 'r') as f:
config_dict = json.load(f)
# Combine both
return cls(
model_name=env_dict.get('model_name', cls.model_name),
model_api_key=env_dict.get('model_api_key', cls.model_api_key),
max_tokens=config_dict.get('max_tokens', cls.max_tokens),
max_steps=config_dict.get('max_steps', cls.max_steps),
temperature=config_dict.get('temperature', cls.temperature),
timeout=config_dict.get('timeout', cls.timeout),
sleeptime=config_dict.get('sleeptime', cls.sleeptime),
force_max_steps=config_dict.get('force_max_steps', cls.force_max_steps)
)
def save(self, env_file, config_file):
# Read existing env content if it exists
env_vars = {}
if os.path.exists(env_file):
with open(env_file, 'r') as f:
for line in f:
if line.strip(): # Ignore empty lines
key, value = line.strip().split('=', 1)
env_vars[key] = value
# Update the necessary keys
env_vars['MODEL_API_KEY'] = self.model_api_key
env_vars['MODEL_NAME'] = self.model_name
# Write back to the .env file
with open(env_file, 'w') as f:
for key, value in env_vars.items():
f.write(f'{key}={value}\n')
# Save other parameters to input_config.json
config_dict = {
'max_tokens': self.max_tokens,
'max_steps': self.max_steps,
'temperature': self.temperature,
'timeout': self.timeout,
'sleeptime': self.sleeptime,
'force_max_steps': self.force_max_steps
}
with open(config_file, 'w') as f:
json.dump(config_dict, f, indent=4)
|