make it work for direct providers

Author: tarekziade

.env.example

@@ -18,6 +18,12 @@ LLM_API_KEY=sk-...
 # Optional. If omitted, the first model returned by {LLM_API_BASE}/v1/models
 # (or {LLM_API_BASE}/models when the base already ends in /v1) is used.
 # LLM_MODEL=gpt-4o
+# Optional. Org slug sent as X-HF-Bill-To. Required when using a Hugging Face
+# token whose Inference Providers permission is scoped to an org, not the user.
+# LLM_BILL_TO=huggingface
+# Optional. Max completion tokens requested from the LLM. Bump for reasoning
+# models (e.g. Kimi-K2) that spend tokens on reasoning before emitting JSON.
+# LLM_MAX_TOKENS=4096
 
 # --- Review behaviour ---
 # Phrase that must appear in a PR/issue comment to trigger a review.

.github/workflows/ai-review.yml

@@ -37,3 +37,6 @@ jobs:
           llm_api_key: ${{ secrets.LLM_API_KEY }}
           llm_api_base: ${{ secrets.LLM_API_BASE }}
           llm_model: ${{ secrets.LLM_MODEL }}
+          llm_bill_to: ${{ secrets.LLM_BILL_TO }}
+          llm_max_tokens: ${{ secrets.LLM_MAX_TOKENS }}
+          max_diff_chars: ${{ secrets.MAX_DIFF_CHARS }}

action.yml

@@ -16,6 +16,13 @@ inputs:
   llm_model:
     description: 'Model identifier to pass in the chat-completion request. If omitted, the action auto-discovers the first model from {llm_api_base}/models.'
     required: false
+  llm_bill_to:
+    description: 'Optional org slug for routing inference billing (sent as the X-HF-Bill-To header). Required when using a Hugging Face token whose Inference Providers permission is scoped to an organization rather than the user.'
+    required: false
+  llm_max_tokens:
+    description: 'Maximum completion tokens requested from the LLM. Bump for reasoning models (e.g. Kimi-K2) that spend tokens on reasoning before emitting JSON.'
+    required: false
+    default: '4096'
   mention_trigger:
     description: 'Phrase that must appear in a comment to trigger a review.'
     required: false
@@ -64,6 +71,8 @@ runs:
         LLM_API_BASE: ${{ inputs.llm_api_base }}
         LLM_API_KEY: ${{ inputs.llm_api_key }}
         LLM_MODEL: ${{ inputs.llm_model }}
+        LLM_BILL_TO: ${{ inputs.llm_bill_to }}
+        LLM_MAX_TOKENS: ${{ inputs.llm_max_tokens }}
         MENTION_TRIGGER: ${{ inputs.mention_trigger }}
         REVIEW_EVENT: ${{ inputs.review_event }}
         MAX_DIFF_CHARS: ${{ inputs.max_diff_chars }}

config.py

@@ -3,6 +3,14 @@
 from typing import Optional
 
 
+def _int_env(name: str, default: int) -> int:
+    """Like int(os.environ[name]) with a default, but also treats an empty
+    string as "use default" so unset GitHub Action secrets (which forward as
+    "") don't blow up int parsing."""
+    raw = (os.environ.get(name) or "").strip()
+    return int(raw) if raw else default
+
+
 def _load_private_key() -> Optional[str]:
     inline = os.environ.get("GITHUB_PRIVATE_KEY")
     if inline:
@@ -25,6 +33,8 @@ class Config:
     llm_api_base: str
     llm_api_key: str
     llm_model: Optional[str]
+    llm_bill_to: Optional[str]
+    llm_max_tokens: int
 
     mention_trigger: str
     review_event: str
@@ -62,9 +72,11 @@ def from_env(cls, *, require_app: bool = True) -> "Config":
             llm_api_base=os.environ.get("LLM_API_BASE", "https://api.openai.com/v1").rstrip("/"),
             llm_api_key=os.environ["LLM_API_KEY"],
             llm_model=os.environ.get("LLM_MODEL") or None,
+            llm_bill_to=os.environ.get("LLM_BILL_TO") or None,
+            llm_max_tokens=_int_env("LLM_MAX_TOKENS", 4096),
             mention_trigger=os.environ.get("MENTION_TRIGGER", "@serge"),
             review_event=os.environ.get("REVIEW_EVENT", "COMMENT"),
-            max_diff_chars=int(os.environ.get("MAX_DIFF_CHARS", "200000")),
+            max_diff_chars=_int_env("MAX_DIFF_CHARS", 200000),
             review_rules_path=os.environ.get("REVIEW_RULES_PATH", ".ai/review-rules.md"),
             default_review_rules=os.environ.get(
                 "DEFAULT_REVIEW_RULES",

llm_client.py

@@ -15,21 +15,32 @@ class ChatCompletionClient:
     LM Studio, llama.cpp server, etc.
     """
 
-    def __init__(self, api_base: str, api_key: str, model: Optional[str] = None):
+    def __init__(
+        self,
+        api_base: str,
+        api_key: str,
+        model: Optional[str] = None,
+        bill_to: Optional[str] = None,
+    ):
         self.api_base = api_base.rstrip("/")
         self.api_key = api_key
         self.model = model
+        self.bill_to = bill_to or None
 
     def _api_base_v1(self) -> str:
         if self.api_base.endswith("/v1"):
             return self.api_base
         return f"{self.api_base}/v1"
 
     def _headers(self) -> dict[str, str]:
-        return {
+        headers = {
             "Authorization": f"Bearer {self.api_key}",
             "Content-Type": "application/json",
         }
+        if self.bill_to:
+            # HF Router: route inference billing to an org the token has access to.
+            headers["X-HF-Bill-To"] = self.bill_to
+        return headers
 
     def _discover_model(self) -> str:
         models_url = f"{self._api_base_v1()}/models"

reviewer.py

@@ -150,7 +150,9 @@ def run_review(cfg: Config, gh: GitHubClient, req: ReviewRequest) -> None:
 
     review_rules = _load_review_rules(gh, req.owner, req.repo, pr, cfg)
 
-    llm = ChatCompletionClient(cfg.llm_api_base, cfg.llm_api_key, cfg.llm_model)
+    llm = ChatCompletionClient(
+        cfg.llm_api_base, cfg.llm_api_key, cfg.llm_model, bill_to=cfg.llm_bill_to
+    )
     system_prompt = build_system_prompt(review_rules)
     user_prompt = build_user_prompt(
         repo_full_name=f"{req.owner}/{req.repo}",
@@ -169,6 +171,7 @@ def run_review(cfg: Config, gh: GitHubClient, req: ReviewRequest) -> None:
             {"role": "user", "content": user_prompt},
         ],
         response_format={"type": "json_object"},
+        max_tokens=cfg.llm_max_tokens,
     )
 
     try:

tests/test_llm_client.py

@@ -96,6 +96,39 @@ def test_complete_adds_v1_when_base_is_unversioned(self) -> None:
             "https://example.com/v1/chat/completions",
         )
 
+    def test_complete_sends_bill_to_header_when_configured(self) -> None:
+        with patch("llm_client.requests.get") as mock_get, patch(
+            "llm_client.requests.post"
+        ) as mock_post:
+            mock_post.return_value = Mock(
+                json=Mock(return_value={"choices": [{"message": {"content": "ok"}}]}),
+                raise_for_status=Mock(),
+            )
+
+            client = ChatCompletionClient(
+                "https://example.com/v1", "token", "fixed-model", bill_to="my-org"
+            )
+            client.complete([{"role": "user", "content": "hi"}])
+
+        mock_get.assert_not_called()
+        self.assertEqual(
+            mock_post.call_args.kwargs["headers"]["X-HF-Bill-To"], "my-org"
+        )
+
+    def test_complete_omits_bill_to_header_when_not_configured(self) -> None:
+        with patch("llm_client.requests.get"), patch(
+            "llm_client.requests.post"
+        ) as mock_post:
+            mock_post.return_value = Mock(
+                json=Mock(return_value={"choices": [{"message": {"content": "ok"}}]}),
+                raise_for_status=Mock(),
+            )
+
+            client = ChatCompletionClient("https://example.com/v1", "token", "fixed-model")
+            client.complete([{"role": "user", "content": "hi"}])
+
+        self.assertNotIn("X-HF-Bill-To", mock_post.call_args.kwargs["headers"])
+
     def test_complete_raises_when_discovery_returns_no_models(self) -> None:
         with patch("llm_client.requests.get") as mock_get, patch(
             "llm_client.requests.post"