Coverage for src / quber / prompts / loader.py: 85%

1""" 

2Prompt loader utility for loading agent prompts from TOML files. 

3""" 

4 

5import sys 

6from pathlib import Path 

7from typing import Optional 

8 

9if sys.version_info >= (3, 11): 

10 import tomllib 

11else: 

12 try: 

13 import tomli as tomllib # type: ignore 

14 except ImportError as e: 

15 raise ImportError("tomli package required for Python < 3.11. Install with: pip install tomli") from e 

16 

17 

18def load_prompt(processor_name: str, prompt_key: str, prompts_dir: Optional[Path] = None) -> str: 

19 """ 

20 Load a specific prompt from a processor's TOML file. 

21 

22 Args: 

23 processor_name: Name of the processor (e.g., 'table_inference') 

24 prompt_key: Key of the prompt to load (e.g., 'system', 'user_analysis') 

25 prompts_dir: Optional directory containing prompts (defaults to this package directory) 

26 

27 Returns: 

28 The prompt content as a string 

29 

30 Raises: 

31 FileNotFoundError: If the TOML file does not exist 

32 KeyError: If the prompt key is not found in the TOML file 

33 ValueError: If the prompt content is empty 

34 

35 Example: 

36 >>> system_prompt = load_prompt("table_inference", "system") 

37 """ 

38 if prompts_dir is None: 

39 prompts_dir = Path(__file__).parent 

40 

41 toml_file = prompts_dir / f"{processor_name}.toml" 

42 

43 if not toml_file.exists(): 

44 raise FileNotFoundError(f"Prompt file not found: {toml_file}") 

45 

46 with open(toml_file, "rb") as f: 

47 data = tomllib.load(f) 

48 

49 try: 

50 # Access prompts.{prompt_key}.text 

51 prompt_data = data["prompts"][prompt_key] 

52 content = prompt_data["text"].strip() 

53 except KeyError as e: 

54 available_keys = list(data.get("prompts", {}).keys()) 

55 raise KeyError( 

56 f"Prompt key '{prompt_key}' not found in {toml_file}. Available keys: {available_keys}" 

57 ) from e 

58 

59 if not content: 

60 raise ValueError(f"Prompt content is empty for key '{prompt_key}' in {toml_file}") 

61 

62 return content