Coverage for astrocyte/harness/localconfig.py: 99%

1"""The ``astrocyte.yaml`` that ``astrocyte setup`` writes for a local install. 

2 

3Provider choice is detected, never assumed, and every choice is explained in 

4the file itself so a developer can see *why* their memory uses what it uses: 

5 

6* **Completions** — the ``claude`` CLI when present (subscription auth, no API 

7 key; most people running ``astrocyte setup`` for a coding agent have it), 

8 otherwise OpenAI when ``OPENAI_API_KEY`` is set. 

9* **Embeddings** — a local model (``local_embeddings``, fastembed preferred) 

10 at its native 384-dim width; OpenAI's only if no local backend is installed. 

11 

12Setup never overwrites an existing config: it is the user's file. 

13""" 

14 

15from __future__ import annotations 

16 

17import importlib.util 

18import json 

19import os 

20import shutil 

21import sys 

22from dataclasses import dataclass, field 

23from datetime import UTC, datetime 

24from pathlib import Path 

25from typing import Any 

26 

27 

28class SetupError(Exception): 

29 """Setup cannot proceed; the message says what to do.""" 

30 

31 

32@dataclass(frozen=True) 

33class ProviderChoice: 

34 llm_provider: str 

35 llm_provider_config: dict[str, Any] 

36 llm_why: str 

37 embedding_provider: str | None = None 

38 embedding_provider_config: dict[str, Any] = field(default_factory=dict) 

39 embedding_why: str = "" 

40 

41 

42def _module_available(name: str) -> bool: 

43 if name in sys.modules: 

44 return sys.modules[name] is not None 

45 try: 

46 return importlib.util.find_spec(name) is not None 

47 except (ImportError, ValueError): 

48 return False 

49 

50 

51def local_embedding_backend() -> str | None: 

52 for module, label in (("fastembed", "fastembed"), ("sentence_transformers", "sentence-transformers")): 

53 if _module_available(module): 

54 return label 

55 return None 

56 

57 

58def choose_providers(env: dict[str, str] | None = None) -> ProviderChoice: 

59 env = dict(os.environ if env is None else env) 

60 backend = local_embedding_backend() 

61 

62 if shutil.which("claude"): 

63 llm, llm_cfg = "claude_cli", {"model": "haiku"} 

64 llm_why = "the Claude Code CLI on your PATH (subscription auth, no API key)" 

65 elif env.get("OPENAI_API_KEY"): 

66 if not _module_available("openai"): 

67 raise SetupError( 

68 "OPENAI_API_KEY is set but the OpenAI SDK isn't installed here. Install it with: " 

69 "uv tool install --force 'astrocyte[local]' (or pip install 'astrocyte[openai]')." 

70 ) 

71 llm, llm_cfg = "openai", {"model": "gpt-4o-mini"} 

72 llm_why = "OpenAI, because OPENAI_API_KEY is set" 

73 else: 

74 raise SetupError( 

75 "No language model available for fact extraction. Either install Claude Code " 

76 "(https://claude.com/claude-code) so `claude` is on your PATH, or set OPENAI_API_KEY." 

77 ) 

78 

79 if backend: 

80 return ProviderChoice( 

81 llm, llm_cfg, llm_why, 

82 embedding_provider="local_embeddings", 

83 # 0 = native width (384 for bge-small): SQLite has no fixed column, 

84 # and padding to 1536 would cost ~4x in search time and disk. 

85 embedding_provider_config={"pad_to": 0}, 

86 embedding_why=f"a local bge-small model via {backend} (no API key; nothing leaves this machine)", 

87 ) 

88 if llm == "openai": 

89 return ProviderChoice(llm, llm_cfg, llm_why, embedding_why="OpenAI (no local embedding backend installed)") 

90 raise SetupError( 

91 "No embedding model available. Install the local one with: pip install 'astrocyte[local]' " 

92 "(or set OPENAI_API_KEY)." 

93 ) 

94 

95 

96def _yaml_scalar(value: Any) -> str: 

97 # JSON scalars are valid YAML and quote paths/specials safely. 

98 return json.dumps(value) 

99 

100 

101def _yaml_block(key: str, cfg: dict[str, Any]) -> str: 

102 if not cfg: 

103 return "" 

104 lines = [f"{key}:"] + [f" {k}: {_yaml_scalar(v)}" for k, v in cfg.items()] 

105 return "\n".join(lines) + "\n" 

106 

107 

108def render_config(choice: ProviderChoice, db_path: Path) -> str: 

109 stamp = datetime.now(UTC).strftime("%Y-%m-%d") 

110 out = [ 

111 f"# Written by `astrocyte setup` on {stamp}. Yours to edit — setup never overwrites it.", 

112 "# Check it any time with: astrocyte doctor", 

113 "provider_tier: storage", 

114 "", 

115 "# Every memory lives in one local SQLite file. No server to run.", 

116 "vector_store: sqlite", 

117 _yaml_block("vector_store_config", {"path": str(db_path)}).rstrip(), 

118 "", 

119 f"# Fact extraction uses {choice.llm_why}.", 

120 f"llm_provider: {choice.llm_provider}", 

121 _yaml_block("llm_provider_config", choice.llm_provider_config).rstrip(), 

122 ] 

123 if choice.embedding_provider: 

124 out += [ 

125 "", 

126 f"# Embeddings use {choice.embedding_why}.", 

127 f"embedding_provider: {choice.embedding_provider}", 

128 _yaml_block("embedding_provider_config", choice.embedding_provider_config).rstrip(), 

129 ] 

130 elif choice.embedding_why: 

131 out += ["", f"# Embeddings use {choice.embedding_why}."] 

132 return "\n".join(line for line in out if line is not None) + "\n" 

133 

134 

135def write_config(path: Path, text: str) -> None: 

136 path.parent.mkdir(parents=True, exist_ok=True) 

137 tmp = path.with_name(path.name + ".tmp") 

138 tmp.write_text(text, encoding="utf-8") 

139 os.replace(tmp, path)