@@ -63,50 +63,71 @@ def get_llm(
63
63
callbacks = None ,
64
64
):
65
65
config = read_config ()
66
+ if "openrouter_api_key" in config :
67
+ # If the openrouter_api_key is set, use the Open Router API
68
+ # OpenRouter allows for access to many models that have larger token limits
69
+ openai .api_key = config ["openrouter_api_key" ]
70
+ openai .api_base = "https://openrouter.ai/api/v1"
71
+ headers = {"HTTP-Referer" : "https://aicodebot.dev" , "X-Title" : "AICodeBot" }
72
+ tiktoken_model_name = model_name .replace ("openai/" , "" )
73
+ else :
74
+ openai .api_key = config ["openai_api_key" ]
75
+ headers = None
76
+ tiktoken_model_name = model_name
66
77
67
78
return ChatOpenAI (
68
- openai_api_key = config [ "openai_api_key" ] ,
79
+ openai_api_key = openai . api_key ,
69
80
model = model_name ,
70
81
max_tokens = response_token_size ,
71
82
verbose = verbose ,
72
83
temperature = temperature ,
73
84
streaming = streaming ,
74
85
callbacks = callbacks ,
86
+ tiktoken_model_name = tiktoken_model_name ,
87
+ model_kwargs = {"headers" : headers },
75
88
)
76
89
77
90
@staticmethod
78
- def get_llm_model_name (token_size = 0 ):
79
- model_options = {
80
- "gpt-4" : 8192 ,
81
- "gpt-4-32k" : 32768 ,
82
- "gpt-3.5-turbo" : 4096 ,
83
- "gpt-3.5-turbo-16k" : 16384 ,
84
- }
91
+ def get_llm_headers ():
92
+ config = read_config ()
93
+ if "openrouter_api_key" in config :
94
+ return {"HTTP-Referer" : "https://aicodebot.dev" , "X-Title" : "AICodeBot" }
95
+ else :
96
+ return None
85
97
86
- engines = Coder .get_openai_supported_engines ()
98
+ @staticmethod
99
+ def get_llm_model_name (token_size = 0 ):
100
+ config = read_config ()
101
+ if "openrouter_api_key" in config :
102
+ model_options = {
103
+ "openai/gpt-4" : 8192 ,
104
+ "openai/gpt-4-32k" : 32768 ,
105
+ # Not working yet "anthropic/claude-2": 100_000,
106
+ }
107
+ supported_engines = model_options .keys ()
108
+ else :
109
+ model_options = {
110
+ "gpt-4" : 8192 ,
111
+ "gpt-4-32k" : 32768 ,
112
+ "gpt-3.5-turbo" : 4096 ,
113
+ "gpt-3.5-turbo-16k" : 16384 ,
114
+ }
115
+ # Pull the list of supported engines from the OpenAI API for this key
116
+ supported_engines = Coder .get_openai_supported_engines ()
87
117
88
118
# For some unknown reason, tiktoken often underestimates the token size by ~10%, so let's buffer
89
119
token_size = int (token_size * 1.1 )
90
120
91
- # Try to use GPT-4 if it is supported and the token size is small enough
92
- if "gpt-4" in engines and token_size <= model_options ["gpt-4" ]:
93
- logger .info (f"Using GPT-4 for token size { token_size } " )
94
- return "gpt-4"
95
- elif "gpt-4-32k" in engines and token_size <= model_options ["gpt-4-32k" ]:
96
- logger .info (f"Using GPT-4-32k for token size { token_size } " )
97
- return "gpt-4-32k"
98
- elif token_size <= model_options ["gpt-3.5-turbo" ]:
99
- logger .info (f"Using GPT-3.5-turbo for token size { token_size } " )
100
- return "gpt-3.5-turbo"
101
- elif token_size <= model_options ["gpt-3.5-turbo-16k" ]:
102
- logger .info (f"Using GPT-3.5-turbo-16k for token size { token_size } " )
103
- return "gpt-3.5-turbo-16k"
104
- else :
105
- logger .critical (
106
- f"🛑 The context is too large ({ token_size } )"
107
- "for the any of the models supported by your Open AI API key. 😞"
108
- )
109
- return None
121
+ for model , max_tokens in model_options .items ():
122
+ if model in supported_engines and token_size <= max_tokens :
123
+ logger .info (f"Using { model } for token size { token_size } " )
124
+ return model
125
+
126
+ logger .critical (f"The context is too large ({ token_size } ) for any of the models supported by your API key. 😞" )
127
+ if "openrouter_api_key" not in config :
128
+ logger .critical ("If you provide an Open Router API key, you can access larger models, up to 32k tokens" )
129
+
130
+ return None
110
131
111
132
@staticmethod
112
133
def get_token_length (text , model = "gpt-3.5-turbo" ):
0 commit comments