Skip to content

Commit f09b284

Browse files
committed
fix for headers and add test
1 parent 3ef4812 commit f09b284

File tree

3 files changed

+20
-3
lines changed

3 files changed

+20
-3
lines changed

tale/llm/llm_io.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,8 @@ def __init__(self, config: dict = None, backend_config: dict = None):
2121
self.headers = headers
2222
self.io_adapter = LlamaCppAdapter(self.url, backend_config['STREAM_ENDPOINT'], config.get('USER_START', ''), config.get('USER_END', ''), config.get('SYSTEM_START', ''), config.get('PROMPT_END', ''))
2323
else:
24-
headers['Authorization'] = f"Bearer {backend_config['API_PASSWORD']}"
24+
if 'API_PASSWORD' in backend_config and backend_config['API_PASSWORD']:
25+
headers['Authorization'] = f"Bearer {backend_config['API_PASSWORD']}"
2526
self.headers = headers
2627
self.io_adapter = KoboldCppAdapter(self.url, backend_config['STREAM_ENDPOINT'], backend_config['DATA_ENDPOINT'], config.get('USER_START', ''), config.get('USER_END', ''), config.get('SYSTEM_START', ''), config.get('PROMPT_END', ''))
2728

tests/test_json_story.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@
1212
from tale.mob_spawner import MobSpawner
1313

1414
class TestJsonStory():
15+
wearable.wearbles_story = []
1516
driver = IFDriver(screen_delay=99, gui=False, web=True, wizard_override=True)
1617
driver.game_clock = util.GameDateTime(datetime.datetime(year=2023, month=1, day=1), 1)
1718
story = JsonStory('tests/files/world_story/', parse_utils.load_story_config(parse_utils.load_json('tests/files/world_story/story_config.json')))
@@ -74,8 +75,7 @@ def test_load_story(self):
7475

7576
assert self.story.day_cycle
7677
assert self.story.random_events
77-
78-
assert len(wearable.wearbles_story) == 1
78+
7979
assert wearable.wearbles_story[0]['name'] == 'fur jacket'
8080

8181

tests/test_llm_io.py

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -76,6 +76,22 @@ def test_set_prompt_llama_cpp(self):
7676
assert('<context>context</context>' in result['messages'][1]['content'])
7777
assert(result['messages'][0]['content'] != '<context>context</context>')
7878

79+
80+
def test_password_in_header(self):
81+
config_file = self._load_config()
82+
config_file['BACKEND'] = 'kobold_cpp'
83+
84+
backend_config = self._load_backend_config('kobold_cpp')
85+
io_util = IoUtil(config=config_file, backend_config=backend_config)
86+
87+
assert not io_util.headers
88+
89+
backend_config = self._load_backend_config('kobold_cpp')
90+
backend_config['API_PASSWORD'] = 'test_password'
91+
io_util = IoUtil(config=config_file, backend_config=backend_config)
92+
93+
assert io_util.headers['Authorization'] == f"Bearer {backend_config['API_PASSWORD']}"
94+
7995
@responses.activate
8096
def test_error_response(self):
8197
config_file = self._load_config()

0 commit comments

Comments
 (0)