lihuigu commited on
Commit
0b6004e
·
1 Parent(s): 0b619bd

update get api help

Browse files
Files changed (1) hide show
  1. src/utils/llms_api.py +11 -7
src/utils/llms_api.py CHANGED
@@ -1,8 +1,9 @@
1
  from .api import HelperCompany
2
- import openai
3
  import re
 
4
  from .header import get_dir, Prompt, ConfigReader
5
  import traceback
 
6
 
7
  TAG_moti = "Motivations:"
8
  TAG_contr = "Details:"
@@ -32,14 +33,17 @@ class APIHelper(object):
32
  super(APIHelper, self).__init__()
33
  self.config = config
34
  self.__checkout_config__()
35
- self.summarizer = self.get_helper(config, config.used_llms_apis.summarization)
36
- self.generator = self.get_helper(config, config.used_llms_apis.generation)
37
  self.prompt = Prompt(get_dir(config.ARTICLE.summarizing_prompt))
38
 
39
- def get_helper(self, config, alias):
40
- api_config = config[alias]
41
- return HelperCompany.get()[api_config.type](
42
- api_config.api_key, api_config.model, api_config.base_url, timeout=None
 
 
 
43
  )
44
 
45
  def __checkout_config__(self):
 
1
  from .api import HelperCompany
 
2
  import re
3
+ import os
4
  from .header import get_dir, Prompt, ConfigReader
5
  import traceback
6
+ import openai
7
 
8
  TAG_moti = "Motivations:"
9
  TAG_contr = "Details:"
 
33
  super(APIHelper, self).__init__()
34
  self.config = config
35
  self.__checkout_config__()
36
+ self.summarizer = self.get_helper()
37
+ self.generator = self.get_helper()
38
  self.prompt = Prompt(get_dir(config.ARTICLE.summarizing_prompt))
39
 
40
+ def get_helper():
41
+ MODEL_TYPE = os.environ["MODEL_TYPE"]
42
+ MODEL_NAME = os.environ["MODEL_NAME"]
43
+ MODEL_API_KEY = os.environ["MODEL_API_KEY"]
44
+ BASE_URL = os.environ["BASE_URL"]
45
+ return HelperCompany.get()[MODEL_TYPE](
46
+ MODEL_API_KEY, MODEL_NAME, BASE_URL, timeout=None
47
  )
48
 
49
  def __checkout_config__(self):