diff --git a/gptme/llm.py b/gptme/llm.py index 2c9a5bb6..b3b94cbd 100644 --- a/gptme/llm.py +++ b/gptme/llm.py @@ -54,7 +54,7 @@ def reply(messages: list[Message], model: str, stream: bool = False) -> Message: def _chat_complete(messages: list[Message], model: str) -> str: provider = _client_to_provider() - if provider == "openai": + if provider in ["openai", "azure", "openrouter"]: return chat_openai(messages, model) elif provider == "anthropic": return chat_anthropic(messages, model) @@ -64,7 +64,7 @@ def _chat_complete(messages: list[Message], model: str) -> str: def _stream(messages: list[Message], model: str) -> Iterator[str]: provider = _client_to_provider() - if provider == "openai": + if provider in ["openai", "azure", "openrouter"]: return stream_openai(messages, model) elif provider == "anthropic": return stream_anthropic(messages, model) diff --git a/gptme/llm_openai.py b/gptme/llm_openai.py index a9a0ab0e..e0b56499 100644 --- a/gptme/llm_openai.py +++ b/gptme/llm_openai.py @@ -10,6 +10,13 @@ logger = logging.getLogger(__name__) +# Shows in rankings on openrouter.ai +openrouter_headers = { + "HTTP-Referer": "https://github.com/ErikBjare/gptme", + "X-Title": "gptme", +} + + def init(llm: str, config): global openai @@ -50,6 +57,9 @@ def chat(messages: list[Message], model: str) -> str: messages=msgs2dicts(messages, openai=True), # type: ignore temperature=TEMPERATURE, top_p=TOP_P, + extra_headers=( + openrouter_headers if "openrouter.ai" in str(openai.base_url) else {} + ), ) content = response.choices[0].message.content assert content @@ -68,6 +78,9 @@ def stream(messages: list[Message], model: str) -> Generator[str, None, None]: # the llama-cpp-python server needs this explicitly set, otherwise unreliable results # TODO: make this better max_tokens=1000 if not model.startswith("gpt-") else 4096, + extra_headers=( + openrouter_headers if "openrouter.ai" in str(openai.base_url) else {} + ), ): if not chunk.choices: # type: ignore # Got a chunk with no choices, Azure always sends one of these at the start