Skip to content

Commit bcf9cbe

Browse files
committed
Add support for Open Router API in Coder class
This commit introduces the ability to use the Open Router API in the Coder class if an `openrouter_api_key` is provided in the configuration. This allows for access to models with larger token limits. If the Open Router API key is not provided, the code falls back to using the OpenAI API key. The commit also refactors the `get_llm_model_name` method to support model selection based on the provided API key. 🚀🔑
1 parent 5d187e8 commit bcf9cbe

File tree

1 file changed

+49
-28
lines changed

1 file changed

+49
-28
lines changed

aicodebot/coder.py

Lines changed: 49 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -63,50 +63,71 @@ def get_llm(
6363
callbacks=None,
6464
):
6565
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
6677

6778
return ChatOpenAI(
68-
openai_api_key=config["openai_api_key"],
79+
openai_api_key=openai.api_key,
6980
model=model_name,
7081
max_tokens=response_token_size,
7182
verbose=verbose,
7283
temperature=temperature,
7384
streaming=streaming,
7485
callbacks=callbacks,
86+
tiktoken_model_name=tiktoken_model_name,
87+
model_kwargs={"headers": headers},
7588
)
7689

7790
@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
8597

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()
87117

88118
# For some unknown reason, tiktoken often underestimates the token size by ~10%, so let's buffer
89119
token_size = int(token_size * 1.1)
90120

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
110131

111132
@staticmethod
112133
def get_token_length(text, model="gpt-3.5-turbo"):

0 commit comments

Comments
 (0)