|
32 | 32 | env_settings = toml.load(f"{app_base_path}/sidekick/configs/env.toml") |
33 | 33 | db_dialect = env_settings["DB-DIALECT"]["DB_TYPE"] |
34 | 34 | model_name = env_settings["MODEL_INFO"]["MODEL_NAME"] |
35 | | -h2o_remote_url = env_settings["MODEL_INFO"]["RECOMMENDATION_MODEL_REMOTE_URL"] |
36 | | -h2o_key = env_settings["MODEL_INFO"]["RECOMMENDATION_MODEL_API_KEY"] |
| 35 | +h2o_remote_url = env_settings["MODEL_INFO"]["H2OGPTE_URL"] |
| 36 | +h2o_key = env_settings["MODEL_INFO"]["H2OGPTE_API_TOKEN"] |
37 | 37 | # h2ogpt base model urls |
38 | | -h2ogpt_base_model_url = env_settings["MODEL_INFO"]["H2O_BASE_MODEL_URL"] |
39 | | -h2ogpt_base_model_key = env_settings["MODEL_INFO"]["H2O_BASE_MODEL_API_KEY"] |
| 38 | +h2ogpt_base_model_url = env_settings["MODEL_INFO"]["H2OGPT_URL"] |
| 39 | +h2ogpt_base_model_key = env_settings["MODEL_INFO"]["H2OGPT_API_TOKEN"] |
| 40 | + |
40 | 41 | self_correction_model = env_settings["MODEL_INFO"]["SELF_CORRECTION_MODEL"] |
| 42 | +recommendation_model = env_settings["MODEL_INFO"]['RECOMMENDATION_MODEL'] |
41 | 43 |
|
42 | 44 | os.environ["TOKENIZERS_PARALLELISM"] = "False" |
43 | 45 | # Env variables |
44 | | -if not os.getenv("H2O_BASE_MODEL_URL"): |
45 | | - os.environ["H2O_BASE_MODEL_URL"] = h2ogpt_base_model_url |
46 | | -if not os.getenv("H2O_BASE_MODEL_API_KEY"): |
47 | | - os.environ["H2O_BASE_MODEL_API_KEY"] = h2ogpt_base_model_key |
48 | | -if not os.getenv("RECOMMENDATION_MODEL_REMOTE_URL"): |
49 | | - os.environ["RECOMMENDATION_MODEL_REMOTE_URL"] = h2o_remote_url |
50 | | -if not os.getenv("RECOMMENDATION_MODEL_API_KEY"): |
51 | | - os.environ["RECOMMENDATION_MODEL_API_KEY"] = h2o_key |
| 46 | +if not os.getenv("H2OGPT_URL"): |
| 47 | + os.environ["H2OGPT_URL"] = h2ogpt_base_model_url |
| 48 | +if not os.getenv("H2OGPT_API_TOKEN"): |
| 49 | + os.environ["H2OGPT_API_TOKEN"] = h2ogpt_base_model_key |
| 50 | +if not os.getenv("H2OGPTE_URL"): |
| 51 | + os.environ["H2OGPTE_URL"] = h2o_remote_url |
| 52 | +if not os.getenv("H2OGPTE_API_TOKEN"): |
| 53 | + os.environ["H2OGPTE_API_TOKEN"] = h2o_key |
52 | 54 | if not os.getenv("SELF_CORRECTION_MODEL"): |
53 | 55 | os.environ["SELF_CORRECTION_MODEL"] = self_correction_model |
| 56 | +if not os.getenv("RECOMMENDATION_MODEL"): |
| 57 | + os.environ["RECOMMENDATION_MODEL"] = recommendation_model |
54 | 58 |
|
55 | 59 | def color(fore="", back="", text=None): |
56 | 60 | return f"{fore}{back}{text}{Style.RESET_ALL}" |
@@ -159,17 +163,17 @@ def recommend_suggestions(cache_path: str, table_name: str, n_qs: int=10): |
159 | 163 | r_url = _key = None |
160 | 164 | # First check for keys in env variables |
161 | 165 | logger.debug(f"Checking environment settings ...") |
162 | | - env_url = os.environ["RECOMMENDATION_MODEL_REMOTE_URL"] |
163 | | - env_key = os.environ["RECOMMENDATION_MODEL_API_KEY"] |
| 166 | + env_url = os.environ["H2OGPTE_URL"] |
| 167 | + env_key = os.environ["H2OGPTE_API_TOKEN"] |
164 | 168 | if env_url and env_key: |
165 | 169 | r_url = env_url |
166 | 170 | _key = env_key |
167 | 171 | elif Path(f"{app_base_path}/sidekick/configs/env.toml").exists(): |
168 | 172 | # Reload .env info |
169 | 173 | logger.debug(f"Checking configuration file ...") |
170 | 174 | env_settings = toml.load(f"{app_base_path}/sidekick/configs/env.toml") |
171 | | - r_url = env_settings["MODEL_INFO"]["RECOMMENDATION_MODEL_REMOTE_URL"] |
172 | | - _key = env_settings["MODEL_INFO"]["RECOMMENDATION_MODEL_API_KEY"] |
| 175 | + r_url = env_settings["MODEL_INFO"]["H2OGPTE_URL"] |
| 176 | + _key = env_settings["MODEL_INFO"]["H2OGPTE_API_TOKEN"] |
173 | 177 | else: |
174 | 178 | raise Exception("Model url or key is missing.") |
175 | 179 |
|
@@ -666,8 +670,8 @@ def ask( |
666 | 670 | logger.debug(f"Attempt: {attempt+1}") |
667 | 671 | _tmp = err.split("\n") |
668 | 672 | _err = _tmp[0].split("Error occurred:")[1] if len(_tmp) > 0 else None |
669 | | - env_url = os.environ["RECOMMENDATION_MODEL_REMOTE_URL"] |
670 | | - env_key = os.environ["RECOMMENDATION_MODEL_API_KEY"] |
| 673 | + env_url = os.environ["H2OGPTE_URL"] |
| 674 | + env_key = os.environ["H2OGPTE_API_TOKEN"] |
671 | 675 | corr_sql = sql_g.self_correction(input_query=_val, error_msg=_err, remote_url=env_url, client_key=env_key) |
672 | 676 | q_res, err = DBConfig.execute_query(query=corr_sql) |
673 | 677 | if not 'Error occurred'.lower() in str(err).lower(): |
|
0 commit comments