fix for headers and add test

neph1 · Nov 4, 2024 · f09b284 · f09b284
1 parent 3ef4812
commit f09b284
Show file tree

Hide file tree

Showing 3 changed files with 20 additions and 3 deletions.
diff --git a/tale/llm/llm_io.py b/tale/llm/llm_io.py
@@ -21,7 +21,8 @@ def __init__(self, config: dict = None, backend_config: dict = None):
             self.headers = headers
             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', ''))
         else:
-            headers['Authorization'] = f"Bearer {backend_config['API_PASSWORD']}"
+            if 'API_PASSWORD' in backend_config and backend_config['API_PASSWORD']:
+                headers['Authorization'] = f"Bearer {backend_config['API_PASSWORD']}"
             self.headers = headers
             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', ''))
 

diff --git a/tests/test_json_story.py b/tests/test_json_story.py
@@ -12,6 +12,7 @@
 from tale.mob_spawner import MobSpawner
 
 class TestJsonStory():
+    wearable.wearbles_story = []
     driver = IFDriver(screen_delay=99, gui=False, web=True, wizard_override=True)
     driver.game_clock = util.GameDateTime(datetime.datetime(year=2023, month=1, day=1), 1)
     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):
 
         assert self.story.day_cycle
         assert self.story.random_events
-
-        assert len(wearable.wearbles_story) == 1
+
         assert wearable.wearbles_story[0]['name'] == 'fur jacket'
 
 

diff --git a/tests/test_llm_io.py b/tests/test_llm_io.py
@@ -76,6 +76,22 @@ def test_set_prompt_llama_cpp(self):
         assert('<context>context</context>' in result['messages'][1]['content'])
         assert(result['messages'][0]['content'] != '<context>context</context>')
 
+
+    def test_password_in_header(self):
+        config_file = self._load_config()
+        config_file['BACKEND'] = 'kobold_cpp'
+
+        backend_config = self._load_backend_config('kobold_cpp')
+        io_util = IoUtil(config=config_file, backend_config=backend_config)
+
+        assert not io_util.headers
+
+        backend_config = self._load_backend_config('kobold_cpp')
+        backend_config['API_PASSWORD'] = 'test_password'
+        io_util = IoUtil(config=config_file, backend_config=backend_config)
+
+        assert io_util.headers['Authorization'] == f"Bearer {backend_config['API_PASSWORD']}"
+
     @responses.activate
     def test_error_response(self):
         config_file = self._load_config()