update endpoint helper files

infer_cloud.py

@@ -4,6 +4,7 @@ import os
 import subprocess
 import sys
 import time
+from pathlib import Path
 
 import requests
 from huggingface_hub import InferenceClient, get_token
@@ -11,7 +12,26 @@ from huggingface_hub import InferenceClient, get_token
 from infer_local import build_instruction_prompt, build_structured_result
 
 
+REQUIRED_OUTPUT_KEYS = {
+    "code",
+    "explanation",
+    "confidence",
+    "important_tokens",
+    "relevancy_score",
+    "hallucination",
+    "hallucination_check_reason",
+    "latency_ms",
+}
+
+
+def is_structured_result(payload):
+    return isinstance(payload, dict) and REQUIRED_OUTPUT_KEYS.issubset(payload.keys())
+
+
 def normalize_hf_response(response):
+    if is_structured_result(response):
+        return json.dumps(response, ensure_ascii=False)
+
     if isinstance(response, str):
         return response
 
@@ -26,6 +46,8 @@ def normalize_hf_response(response):
         return str(first)
 
     if isinstance(response, dict):
+        if "code" in response and "explanation" in response:
+            return json.dumps(response, ensure_ascii=False)
         return str(response.get("generated_text", response.get("text", "")))
 
     return str(response)
@@ -59,13 +81,13 @@ def call_direct_inference_api(repo_id, token, prompt_text, generation_kwargs):
     return body
 
 
-def call_endpoint_url(endpoint_url, token, prompt_text, generation_kwargs):
+def call_endpoint_url(endpoint_url, token, user_prompt, generation_kwargs):
     headers = {"Content-Type": "application/json"}
     if token:
         headers["Authorization"] = f"Bearer {token}"
 
     payload = {
-        "inputs": prompt_text,
+        "inputs": user_prompt,
         "parameters": generation_kwargs,
         "options": {"wait_for_model": True},
     }
@@ -97,9 +119,10 @@ def run_local_fallback(args, reason):
         ),
         file=sys.stderr,
     )
+    script_path = Path(__file__).resolve().with_name("infer_local.py")
     cmd = [
         sys.executable,
-        "infer_local.py",
+        str(script_path),
         "--model-path",
         args.fallback_model_path,
         "--prompt",
@@ -117,6 +140,8 @@ def run_local_fallback(args, reason):
                 str(args.top_p),
             ]
         )
+    if args.allow_downloads:
+        cmd.append("--allow-downloads")
     completed = subprocess.run(cmd, check=True, text=True, capture_output=True)
     if completed.stderr:
         print(completed.stderr, file=sys.stderr, end="")
@@ -149,6 +174,11 @@ def main():
     parser.add_argument("--temperature", type=float, default=0.25)
     parser.add_argument("--top-p", type=float, default=0.9)
     parser.add_argument("--do-sample", action="store_true")
+    parser.add_argument(
+        "--allow-downloads",
+        action="store_true",
+        help="Allow local fallback inference to download missing model files.",
+    )
     args = parser.parse_args()
     if args.no_local_fallback:
         args.fallback_model_path = ""
@@ -171,7 +201,7 @@ def main():
     start_time = time.perf_counter()
     if args.endpoint_url:
         try:
-            response = call_endpoint_url(args.endpoint_url, token, prompt_text, generation_kwargs)
+            response = call_endpoint_url(args.endpoint_url, token, args.prompt, generation_kwargs)
         except Exception as exc:
             run_local_fallback(args, str(exc))
             return
@@ -199,6 +229,10 @@ def main():
                 return
     latency_ms = int((time.perf_counter() - start_time) * 1000)
 
+    if is_structured_result(response):
+        print(json.dumps(response, indent=2, ensure_ascii=False))
+        return
+
     generated_text = normalize_hf_response(response).strip()
     if generated_text.startswith(prompt_text):
         generated_text = generated_text[len(prompt_text) :].strip()