Coverage for scripts/invoke_model.py: 37%

1#!/usr/bin/env python3 

2"""Invoke an optional model backend without third-party Python dependencies.""" 

3 

4import argparse 

5import json 

6import os 

7import sys 

8from pathlib import Path 

9from urllib.error import HTTPError, URLError 

10from urllib.request import Request, urlopen 

11 

12 

13def post_json(url, payload, api_key=None): 

14 headers = {"Content-Type": "application/json"} 

15 if api_key: 15 ↛ 17line 15 didn't jump to line 17 because the condition on line 15 was always true

16 headers["Authorization"] = f"Bearer {api_key}" 

17 request = Request(url, json.dumps(payload).encode(), headers, method="POST") 

18 with urlopen(request, timeout=120) as response: 

19 return json.loads(response.read().decode()) 

20 

21 

22def main(): 

23 parser = argparse.ArgumentParser() 

24 parser.add_argument( 

25 "--provider", 

26 choices=("openai-compatible", "ollama", "prompt"), 

27 default="prompt", 

28 ) 

29 parser.add_argument("--prompt-file", type=Path, required=True) 

30 args = parser.parse_args() 

31 prompt = args.prompt_file.read_text(encoding="utf-8") 

32 

33 if args.provider == "prompt": 

34 print(prompt) 

35 return 

36 

37 model = os.environ.get("TEST_WRITER_MODEL") 

38 if not model: 

39 parser.error("TEST_WRITER_MODEL is required for model calls") 

40 

41 try: 

42 if args.provider == "openai-compatible": 

43 base = os.environ.get("TEST_WRITER_BASE_URL", "http://127.0.0.1:8000/v1") 

44 data = post_json( 

45 f"{base.rstrip('/')}/chat/completions", 

46 {"model": model, "messages": [{"role": "user", "content": prompt}]}, 

47 os.environ.get("TEST_WRITER_API_KEY"), 

48 ) 

49 print(data["choices"][0]["message"]["content"]) 

50 else: 

51 base = os.environ.get("TEST_WRITER_OLLAMA_URL", "http://127.0.0.1:11434") 

52 data = post_json( 

53 f"{base.rstrip('/')}/api/chat", 

54 { 

55 "model": model, 

56 "messages": [{"role": "user", "content": prompt}], 

57 "stream": False, 

58 }, 

59 ) 

60 print(data["message"]["content"]) 

61 except (HTTPError, URLError, KeyError, json.JSONDecodeError) as exc: 

62 print(f"Model invocation failed: {exc}", file=sys.stderr) 

63 raise SystemExit(1) from exc 

64 

65 

66if __name__ == "__main__": 66 ↛ 67line 66 didn't jump to line 67 because the condition on line 66 was never true

67 main()