Eduardo David Paredes Vara
530 lines
19 KiB
Python
530 lines
19 KiB
Python
"""
|
|
LiteLLM custom provider — OpenAI Codex via chatgpt.com/backend-api.
|
|
|
|
Auth tokens are read from CODEX_AUTH_FILE (default /root/.codex/auth.json)
|
|
and auto-refreshed using the OpenAI OAuth endpoint before expiry.
|
|
"""
|
|
import base64
|
|
import json
|
|
import os
|
|
import time
|
|
from typing import AsyncIterator, Callable, Iterator, Optional, Union
|
|
|
|
import httpx
|
|
import litellm
|
|
|
|
from litellm.llms.custom_llm import CustomLLM, CustomLLMError
|
|
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
|
|
from litellm.types.utils import GenericStreamingChunk, Usage
|
|
from litellm.utils import ModelResponse
|
|
|
|
AUTH_FILE = os.environ.get("CODEX_AUTH_FILE", "/root/.codex/auth.json")
|
|
|
|
# Self-register so the Router can validate this provider before custom_llm_setup() runs
|
|
if "codex" not in litellm.provider_list:
|
|
litellm.provider_list.append("codex")
|
|
if "codex" not in litellm._custom_providers:
|
|
litellm._custom_providers.append("codex")
|
|
TOKEN_URL = "https://auth.openai.com/oauth/token"
|
|
API_URL = "https://chatgpt.com/backend-api/codex/responses"
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Auth helpers
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def _load_auth() -> dict:
|
|
with open(AUTH_FILE) as f:
|
|
return json.load(f)
|
|
|
|
|
|
def _save_auth(data: dict) -> None:
|
|
with open(AUTH_FILE, "w") as f:
|
|
json.dump(data, f, indent=2)
|
|
|
|
|
|
def _jwt_exp(token: str) -> int:
|
|
try:
|
|
payload = token.split(".")[1]
|
|
payload += "=" * (4 - len(payload) % 4)
|
|
return json.loads(base64.urlsafe_b64decode(payload)).get("exp", 0)
|
|
except Exception:
|
|
return 0
|
|
|
|
|
|
def _client_id(auth: dict) -> str:
|
|
"""Resolve the OAuth client from config or the ID token audience."""
|
|
override = os.environ.get("CODEX_CLIENT_ID")
|
|
if override:
|
|
return override
|
|
|
|
try:
|
|
payload = auth["tokens"]["id_token"].split(".")[1]
|
|
payload += "=" * (-len(payload) % 4)
|
|
audience = json.loads(base64.urlsafe_b64decode(payload))["aud"]
|
|
if isinstance(audience, str) and audience:
|
|
return audience
|
|
if (
|
|
isinstance(audience, list)
|
|
and len(audience) == 1
|
|
and isinstance(audience[0], str)
|
|
and audience[0]
|
|
):
|
|
return audience[0]
|
|
except (KeyError, IndexError, TypeError, ValueError, json.JSONDecodeError):
|
|
pass
|
|
|
|
raise ValueError(
|
|
"Codex OAuth client ID is unavailable: set CODEX_CLIENT_ID or "
|
|
"provide an id_token with a single aud value"
|
|
)
|
|
|
|
|
|
def _refresh(auth: dict) -> dict:
|
|
resp = httpx.post(
|
|
TOKEN_URL,
|
|
json={
|
|
"client_id": _client_id(auth),
|
|
"grant_type": "refresh_token",
|
|
"refresh_token": auth["tokens"]["refresh_token"],
|
|
},
|
|
timeout=30,
|
|
)
|
|
resp.raise_for_status()
|
|
data = resp.json()
|
|
auth["tokens"]["access_token"] = data["access_token"]
|
|
if "refresh_token" in data:
|
|
auth["tokens"]["refresh_token"] = data["refresh_token"]
|
|
auth["last_refresh"] = time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime())
|
|
_save_auth(auth)
|
|
return auth
|
|
|
|
|
|
def _get_auth(force: bool = False) -> dict:
|
|
auth = _load_auth()
|
|
if force or time.time() > _jwt_exp(auth["tokens"]["access_token"]) - 60:
|
|
auth = _refresh(auth)
|
|
return auth
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Request / response helpers
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def _split_messages(messages: list) -> tuple[str, list]:
|
|
"""Convert OpenAI chat messages into Responses API input items."""
|
|
instructions_parts = []
|
|
input_msgs = []
|
|
for msg in messages:
|
|
role = msg.get("role", "user")
|
|
content = msg.get("content", "")
|
|
if role in ("system", "developer"):
|
|
instructions_parts.append(content)
|
|
elif role == "assistant" and msg.get("tool_calls"):
|
|
if content:
|
|
input_msgs.append({"role": role, "content": content})
|
|
for tool_call in msg["tool_calls"]:
|
|
if not isinstance(tool_call, dict):
|
|
tool_call = tool_call.model_dump()
|
|
function = tool_call.get("function", {})
|
|
input_msgs.append({
|
|
"type": "function_call",
|
|
"call_id": tool_call["id"],
|
|
"name": function["name"],
|
|
"arguments": function.get("arguments", ""),
|
|
})
|
|
elif role == "tool":
|
|
input_msgs.append({
|
|
"type": "function_call_output",
|
|
"call_id": msg["tool_call_id"],
|
|
"output": content,
|
|
})
|
|
else:
|
|
input_msgs.append({"role": role, "content": content})
|
|
return "\n\n".join(instructions_parts), input_msgs
|
|
|
|
|
|
def _resolve_model(model: str) -> str:
|
|
"""Translate internal cx-* names back to the real Codex API model names (gpt-*)."""
|
|
if model.startswith("cx-"):
|
|
return "gpt-" + model[3:]
|
|
return model
|
|
|
|
|
|
def _build_body(model: str, messages: list, optional_params: dict, stream: bool) -> dict:
|
|
instructions, input_msgs = _split_messages(messages)
|
|
body: dict = {
|
|
"model": _resolve_model(model),
|
|
"instructions": instructions or "You are a helpful assistant.",
|
|
"input": input_msgs,
|
|
"stream": True, # API only supports streaming; non-stream calls aggregate internally
|
|
"store": False,
|
|
"reasoning": {"effort": optional_params.get("reasoning_effort", "medium")},
|
|
}
|
|
|
|
tools = []
|
|
for tool in optional_params.get("tools") or []:
|
|
if not isinstance(tool, dict):
|
|
tool = tool.model_dump()
|
|
if tool.get("type") == "function" and "function" in tool:
|
|
function = tool["function"]
|
|
tools.append({
|
|
"type": "function",
|
|
"name": function["name"],
|
|
"description": function.get("description", ""),
|
|
"parameters": function.get("parameters", {"type": "object"}),
|
|
**({"strict": function["strict"]} if "strict" in function else {}),
|
|
})
|
|
else:
|
|
tools.append(tool)
|
|
if tools:
|
|
body["tools"] = tools
|
|
|
|
tool_choice = optional_params.get("tool_choice")
|
|
if isinstance(tool_choice, dict) and "function" in tool_choice:
|
|
function = tool_choice["function"]
|
|
body["tool_choice"] = {"type": "function", "name": function["name"]}
|
|
elif tool_choice is not None:
|
|
body["tool_choice"] = tool_choice
|
|
if "parallel_tool_calls" in optional_params:
|
|
body["parallel_tool_calls"] = optional_params["parallel_tool_calls"]
|
|
return body
|
|
|
|
|
|
def _auth_headers(auth: dict) -> dict:
|
|
return {
|
|
"Authorization": f"Bearer {auth['tokens']['access_token']}",
|
|
"ChatGPT-Account-Id": auth["tokens"]["account_id"],
|
|
"Content-Type": "application/json",
|
|
"User-Agent": "opencode/1.0",
|
|
"originator": "opencode",
|
|
}
|
|
|
|
|
|
def _accumulate_tool_chunk(tool_calls: dict, chunk: GenericStreamingChunk) -> None:
|
|
tool_use = chunk.get("tool_use")
|
|
if not tool_use:
|
|
return
|
|
index = tool_use.get("index", 0)
|
|
current = tool_calls.setdefault(index, {
|
|
"id": "",
|
|
"type": "function",
|
|
"function": {"name": "", "arguments": ""},
|
|
})
|
|
if tool_use.get("id"):
|
|
current["id"] = tool_use["id"]
|
|
function = tool_use.get("function") or {}
|
|
if function.get("name"):
|
|
current["function"]["name"] = function["name"]
|
|
current["function"]["arguments"] += function.get("arguments") or ""
|
|
|
|
|
|
def _event_chunk(event: dict, state: dict) -> Optional[GenericStreamingChunk]:
|
|
"""Translate a Responses API SSE event into LiteLLM's generic chunk."""
|
|
etype = event.get("type", "")
|
|
if etype == "response.output_text.delta":
|
|
return GenericStreamingChunk(
|
|
text=event.get("delta", ""),
|
|
is_finished=False,
|
|
finish_reason="",
|
|
usage=None,
|
|
)
|
|
if etype == "response.output_item.added":
|
|
item = event.get("item", {})
|
|
if item.get("type") == "function_call":
|
|
index = event.get("output_index", 0)
|
|
state["tool_calls"] = True
|
|
state.setdefault("tools", {})[index] = item
|
|
return GenericStreamingChunk(
|
|
text="",
|
|
tool_use={
|
|
"index": index,
|
|
"id": item.get("call_id"),
|
|
"type": "function",
|
|
"function": {
|
|
"name": item.get("name", ""),
|
|
"arguments": item.get("arguments", ""),
|
|
},
|
|
},
|
|
is_finished=False,
|
|
finish_reason="",
|
|
usage=None,
|
|
)
|
|
if etype == "response.function_call_arguments.delta":
|
|
index = event.get("output_index", 0)
|
|
return GenericStreamingChunk(
|
|
text="",
|
|
tool_use={
|
|
"index": index,
|
|
"id": None,
|
|
"type": "function",
|
|
"function": {"name": None, "arguments": event.get("delta", "")},
|
|
},
|
|
is_finished=False,
|
|
finish_reason="",
|
|
usage=None,
|
|
)
|
|
if etype in ("response.completed", "response.done"):
|
|
usage_data = event.get("response", {}).get("usage", {})
|
|
return GenericStreamingChunk(
|
|
text="",
|
|
is_finished=True,
|
|
finish_reason="tool_calls" if state.get("tool_calls") else "stop",
|
|
usage={
|
|
"prompt_tokens": usage_data.get("input_tokens", 0),
|
|
"completion_tokens": usage_data.get("output_tokens", 0),
|
|
"total_tokens": usage_data.get("input_tokens", 0)
|
|
+ usage_data.get("output_tokens", 0),
|
|
},
|
|
)
|
|
return None
|
|
|
|
|
|
def _iter_sse(response: httpx.Response) -> Iterator[GenericStreamingChunk]:
|
|
state: dict = {}
|
|
for line in response.iter_lines():
|
|
if not line or line.startswith(":") or not line.startswith("data: "):
|
|
continue
|
|
raw = line[6:]
|
|
if raw == "[DONE]":
|
|
break
|
|
try:
|
|
event = json.loads(raw)
|
|
except json.JSONDecodeError:
|
|
continue
|
|
chunk = _event_chunk(event, state)
|
|
if chunk is not None:
|
|
yield chunk
|
|
if chunk is not None and chunk["is_finished"]:
|
|
break
|
|
|
|
|
|
async def _aiter_sse(response: httpx.Response) -> AsyncIterator[GenericStreamingChunk]:
|
|
state: dict = {}
|
|
async for line in response.aiter_lines():
|
|
if not line or line.startswith(":") or not line.startswith("data: "):
|
|
continue
|
|
raw = line[6:]
|
|
if raw == "[DONE]":
|
|
break
|
|
try:
|
|
event = json.loads(raw)
|
|
except json.JSONDecodeError:
|
|
continue
|
|
chunk = _event_chunk(event, state)
|
|
if chunk is not None:
|
|
yield chunk
|
|
if chunk is not None and chunk["is_finished"]:
|
|
break
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# LiteLLM CustomLLM subclass
|
|
# ---------------------------------------------------------------------------
|
|
|
|
class CodexProvider(CustomLLM):
|
|
|
|
def completion(
|
|
self,
|
|
model: str,
|
|
messages: list,
|
|
api_base: str,
|
|
custom_prompt_dict: dict,
|
|
model_response: ModelResponse,
|
|
print_verbose: Callable,
|
|
encoding,
|
|
api_key,
|
|
logging_obj,
|
|
optional_params: dict,
|
|
acompletion=None,
|
|
litellm_params=None,
|
|
logger_fn=None,
|
|
headers={},
|
|
timeout: Optional[Union[float, httpx.Timeout]] = None,
|
|
client: Optional[HTTPHandler] = None,
|
|
) -> ModelResponse:
|
|
model_name = model.split("/", 1)[-1] if "/" in model else model
|
|
body = _build_body(model_name, messages, optional_params, stream=True)
|
|
text_parts = []
|
|
tool_calls: dict = {}
|
|
prompt_tokens = completion_tokens = 0
|
|
for attempt in range(2):
|
|
try:
|
|
auth = _get_auth(force=(attempt == 1))
|
|
except Exception as e:
|
|
raise CustomLLMError(status_code=401, message=f"Codex auth error: {e}")
|
|
try:
|
|
with httpx.Client(timeout=120) as c:
|
|
with c.stream("POST", API_URL, json=body, headers=_auth_headers(auth)) as resp:
|
|
resp.raise_for_status()
|
|
for chunk in _iter_sse(resp):
|
|
text_parts.append(chunk["text"])
|
|
_accumulate_tool_chunk(tool_calls, chunk)
|
|
if chunk.get("usage"):
|
|
prompt_tokens = chunk["usage"].get("prompt_tokens", 0)
|
|
completion_tokens = chunk["usage"].get("completion_tokens", 0)
|
|
break
|
|
except httpx.HTTPStatusError as e:
|
|
if attempt == 0 and e.response.status_code == 401:
|
|
text_parts.clear()
|
|
tool_calls.clear()
|
|
continue
|
|
raise CustomLLMError(status_code=e.response.status_code, message=str(e))
|
|
|
|
model_response.choices[0].message.content = "".join(text_parts) # type: ignore
|
|
if tool_calls:
|
|
model_response.choices[0].message.tool_calls = [ # type: ignore
|
|
tool_calls[index] for index in sorted(tool_calls)
|
|
]
|
|
model_response.choices[0].finish_reason = ( # type: ignore
|
|
"tool_calls" if tool_calls else "stop"
|
|
)
|
|
model_response.model = _resolve_model(model_name)
|
|
model_response.usage = Usage(
|
|
prompt_tokens=prompt_tokens,
|
|
completion_tokens=completion_tokens,
|
|
total_tokens=prompt_tokens + completion_tokens,
|
|
)
|
|
return model_response
|
|
|
|
def streaming(
|
|
self,
|
|
model: str,
|
|
messages: list,
|
|
api_base: str,
|
|
custom_prompt_dict: dict,
|
|
model_response: ModelResponse,
|
|
print_verbose: Callable,
|
|
encoding,
|
|
api_key,
|
|
logging_obj,
|
|
optional_params: dict,
|
|
acompletion=None,
|
|
litellm_params=None,
|
|
logger_fn=None,
|
|
headers={},
|
|
timeout: Optional[Union[float, httpx.Timeout]] = None,
|
|
client: Optional[HTTPHandler] = None,
|
|
) -> Iterator[GenericStreamingChunk]:
|
|
model_name = model.split("/", 1)[-1] if "/" in model else model
|
|
body = _build_body(model_name, messages, optional_params, stream=True)
|
|
for attempt in range(2):
|
|
try:
|
|
auth = _get_auth(force=(attempt == 1))
|
|
except Exception as e:
|
|
raise CustomLLMError(status_code=401, message=f"Codex auth error: {e}")
|
|
try:
|
|
with httpx.Client(timeout=120) as c:
|
|
with c.stream("POST", API_URL, json=body, headers=_auth_headers(auth)) as resp:
|
|
resp.raise_for_status()
|
|
yield from _iter_sse(resp)
|
|
return
|
|
except httpx.HTTPStatusError as e:
|
|
if attempt == 0 and e.response.status_code == 401:
|
|
continue
|
|
raise CustomLLMError(status_code=e.response.status_code, message=str(e))
|
|
|
|
async def acompletion(
|
|
self,
|
|
model: str,
|
|
messages: list,
|
|
api_base: str,
|
|
custom_prompt_dict: dict,
|
|
model_response: ModelResponse,
|
|
print_verbose: Callable,
|
|
encoding,
|
|
api_key,
|
|
logging_obj,
|
|
optional_params: dict,
|
|
acompletion=None,
|
|
litellm_params=None,
|
|
logger_fn=None,
|
|
headers={},
|
|
timeout: Optional[Union[float, httpx.Timeout]] = None,
|
|
client: Optional[AsyncHTTPHandler] = None,
|
|
) -> ModelResponse:
|
|
model_name = model.split("/", 1)[-1] if "/" in model else model
|
|
body = _build_body(model_name, messages, optional_params, stream=True)
|
|
text_parts = []
|
|
tool_calls: dict = {}
|
|
prompt_tokens = completion_tokens = 0
|
|
for attempt in range(2):
|
|
try:
|
|
auth = _get_auth(force=(attempt == 1))
|
|
except Exception as e:
|
|
raise CustomLLMError(status_code=401, message=f"Codex auth error: {e}")
|
|
try:
|
|
async with httpx.AsyncClient(timeout=120) as c:
|
|
async with c.stream("POST", API_URL, json=body, headers=_auth_headers(auth)) as resp:
|
|
resp.raise_for_status()
|
|
async for chunk in _aiter_sse(resp):
|
|
text_parts.append(chunk["text"])
|
|
_accumulate_tool_chunk(tool_calls, chunk)
|
|
if chunk.get("usage"):
|
|
prompt_tokens = chunk["usage"].get("prompt_tokens", 0)
|
|
completion_tokens = chunk["usage"].get("completion_tokens", 0)
|
|
break
|
|
except httpx.HTTPStatusError as e:
|
|
if attempt == 0 and e.response.status_code == 401:
|
|
text_parts.clear()
|
|
tool_calls.clear()
|
|
continue
|
|
raise CustomLLMError(status_code=e.response.status_code, message=str(e))
|
|
|
|
model_response.choices[0].message.content = "".join(text_parts) # type: ignore
|
|
if tool_calls:
|
|
model_response.choices[0].message.tool_calls = [ # type: ignore
|
|
tool_calls[index] for index in sorted(tool_calls)
|
|
]
|
|
model_response.choices[0].finish_reason = ( # type: ignore
|
|
"tool_calls" if tool_calls else "stop"
|
|
)
|
|
model_response.model = _resolve_model(model_name)
|
|
model_response.usage = Usage(
|
|
prompt_tokens=prompt_tokens,
|
|
completion_tokens=completion_tokens,
|
|
total_tokens=prompt_tokens + completion_tokens,
|
|
)
|
|
return model_response
|
|
|
|
async def astreaming(
|
|
self,
|
|
model: str,
|
|
messages: list,
|
|
api_base: str,
|
|
custom_prompt_dict: dict,
|
|
model_response: ModelResponse,
|
|
print_verbose: Callable,
|
|
encoding,
|
|
api_key,
|
|
logging_obj,
|
|
optional_params: dict,
|
|
acompletion=None,
|
|
litellm_params=None,
|
|
logger_fn=None,
|
|
headers={},
|
|
timeout: Optional[Union[float, httpx.Timeout]] = None,
|
|
client: Optional[AsyncHTTPHandler] = None,
|
|
) -> AsyncIterator[GenericStreamingChunk]:
|
|
model_name = model.split("/", 1)[-1] if "/" in model else model
|
|
body = _build_body(model_name, messages, optional_params, stream=True)
|
|
for attempt in range(2):
|
|
try:
|
|
auth = _get_auth(force=(attempt == 1))
|
|
except Exception as e:
|
|
raise CustomLLMError(status_code=401, message=f"Codex auth error: {e}")
|
|
try:
|
|
async with httpx.AsyncClient(timeout=120) as c:
|
|
async with c.stream("POST", API_URL, json=body, headers=_auth_headers(auth)) as resp:
|
|
resp.raise_for_status()
|
|
async for chunk in _aiter_sse(resp):
|
|
yield chunk
|
|
return
|
|
except httpx.HTTPStatusError as e:
|
|
if attempt == 0 and e.response.status_code == 401:
|
|
continue
|
|
raise CustomLLMError(status_code=e.response.status_code, message=str(e))
|
|
|
|
|
|
codex_provider = CodexProvider()
|