Coverage for graphqler / compiler / resolvers / llm / llm_resolver.py: 11%

1"""LLMResolver — base class for LLM-backed dependency resolvers. 

2 

3Responsibilities: 

4 1. Build a compact, token-efficient schema context string from the objects dict. 

5 2. Convert raw parsed mutations/queries into a simplified JSON representation 

6 that is easier for an LLM to process. 

7 3. Call the LLM via litellm (supports OpenAI, Anthropic, Ollama, LiteLLM proxy, 

8 and any other provider supported by litellm out of the box). 

9 4. Parse and validate the structured JSON response. 

10 5. Fall back to the classic ID-based resolver on any failure when 

11 LLM_RESOLVER_FALLBACK_TO_ID is True. 

12""" 

13 

14import json 

15import logging 

16from graphqler import config 

17 

18 

19def _get_litellm(): 

20 """Import and return litellm only when LLM resolution is actually used.""" 

21 try: 

22 import litellm 

23 except ImportError as exc: 

24 raise ImportError("litellm is required for USE_LLM=True. Install it with: uv add litellm") from exc 

25 return litellm 

26 

27logger = logging.getLogger(__name__) 

28 

29 

30class LLMResolver: 

31 """Base class — subclasses implement `resolve()` and call the helpers here.""" 

32 

33 # ── Schema context ──────────────────────────────────────────────────────── 

34 

35 def build_schema_context(self, objects: dict) -> str: 

36 """Return a compact, human-readable description of all objects. 

37 

38 Example output line: 

39 User: id(ID!), email(String!), name(String), posts([Post]) 

40 

41 Args: 

42 objects (dict): Compiled objects dict from ObjectListParser. 

43 

44 Returns: 

45 str: Multi-line string, one object per line. 

46 """ 

47 lines = ["Available GraphQL objects in this API:"] 

48 for obj_name, obj_body in objects.items(): 

49 fields = obj_body.get("fields", []) 

50 field_strs = [] 

51 for f in fields: 

52 type_str = self._field_type_str(f) 

53 field_strs.append(f"{f['name']}({type_str})") 

54 lines.append(f" {obj_name}: {', '.join(field_strs) if field_strs else '(no fields)'}") 

55 return "\n".join(lines) 

56 

57 def _field_type_str(self, field: dict) -> str: 

58 """Convert a field's kind/ofType structure to a compact type string like 'String!', '[Post]'.""" 

59 kind = field.get("kind", "") 

60 name = field.get("type") or field.get("name") or "" 

61 oftype = field.get("ofType") 

62 

63 if kind == "NON_NULL": 

64 inner = self._oftype_str(oftype) if oftype else name 

65 return f"{inner}!" 

66 elif kind == "LIST": 

67 inner = self._oftype_str(oftype) if oftype else name 

68 return f"[{inner}]" 

69 else: 

70 return name or kind 

71 

72 def _oftype_str(self, oftype: dict) -> str: 

73 if oftype is None: 

74 return "" 

75 kind = oftype.get("kind", "") 

76 name = oftype.get("name") or oftype.get("type") or "" 

77 inner = oftype.get("ofType") 

78 if kind == "NON_NULL": 

79 return f"{self._oftype_str(inner)}!" if inner else f"{name}!" 

80 elif kind == "LIST": 

81 return f"[{self._oftype_str(inner)}]" if inner else f"[{name}]" 

82 else: 

83 return name or kind 

84 

85 # ── Simplified endpoint representations ────────────────────────────────── 

86 

87 def simplify_endpoints(self, endpoints: dict) -> dict: 

88 """Convert raw parsed endpoints to a compact dict suitable for LLM prompts. 

89 

90 The simplified form replaces the nested ofType structures with readable 

91 type strings like "String!", "ID!", "[Post]". 

92 

93 Args: 

94 endpoints (dict): Raw queries or mutations from the parser. 

95 

96 Returns: 

97 dict: Simplified representation keyed by endpoint name. 

98 """ 

99 simplified = {} 

100 for name, body in endpoints.items(): 

101 inputs = {} 

102 for input_name, input_body in (body.get("inputs") or {}).items(): 

103 inputs[input_name] = self._input_type_str(input_body) 

104 

105 output_type = self._oftype_str(body.get("output")) if body.get("output") else "" 

106 simplified[name] = { 

107 "description": body.get("description") or "", 

108 "inputs": inputs, 

109 "output": output_type, 

110 } 

111 return simplified 

112 

113 def _input_type_str(self, input_body: dict) -> str: 

114 """Return a readable type string for an input, e.g. 'ID!', 'String', '[String!]'.""" 

115 kind = input_body.get("kind", "") 

116 oftype = input_body.get("ofType") 

117 if kind == "NON_NULL": 

118 inner = self._oftype_str(oftype) if oftype else "" 

119 return f"{inner}!" 

120 elif kind == "LIST": 

121 inner = self._oftype_str(oftype) if oftype else "" 

122 return f"[{inner}]" 

123 else: 

124 type_name = input_body.get("type") or input_body.get("name") or kind 

125 return type_name 

126 

127 # ── LLM call ───────────────────────────────────────────────────────────── 

128 

129 def _supports_json_mode(self) -> bool: 

130 """Return True if the configured model supports format object. 

131 

132 Handles both prefixed (e.g. "openai/gpt-4o-mini") and unprefixed 

133 (e.g. "gpt-4o-mini") model strings. When there is no provider prefix 

134 the full string is used as the model name and litellm infers the 

135 provider automatically. 

136 """ 

137 llm_model = config.LLM_MODEL 

138 if "/" in llm_model: 

139 provider, model = llm_model.split("/", 1) 

140 else: 

141 # No prefix — treat as an OpenAI model and let litellm infer the provider. 

142 provider = "openai" 

143 model = llm_model 

144 

145 if not model: 

146 # Malformed string like "openai/" — fall back to the full string. 

147 model = llm_model 

148 provider = "openai" 

149 

150 litellm = _get_litellm() 

151 does_support_json = litellm.supports_response_schema(model=model, custom_llm_provider=provider) 

152 return does_support_json 

153 

154 def _extract_json_from_text(self, text: str) -> dict: 

155 """Extract a JSON object from text that may be wrapped in markdown fences or prose. 

156 

157 Tries in order: 

158 1. Strip a leading ``` / ```json fence and trailing ``` then parse. 

159 2. Direct json.loads on the stripped text. 

160 3. Slice from the first '{' to the last '}' and parse. 

161 

162 Raises: 

163 ValueError: If no valid JSON object can be extracted. 

164 """ 

165 text = text.strip() 

166 

167 # Strip markdown code fences (```json ... ``` or ``` ... ```) 

168 if text.startswith("```"): 

169 end_of_first_line = text.find("\n") 

170 if end_of_first_line != -1: 

171 text = text[end_of_first_line + 1 :] 

172 if text.endswith("```"): 

173 text = text[:-3].rstrip() 

174 text = text.strip() 

175 

176 # Direct parse 

177 try: 

178 return json.loads(text) 

179 except json.JSONDecodeError: 

180 pass 

181 

182 # Find the outermost JSON object by scanning for first '{' … last '}' 

183 start = text.find("{") 

184 end = text.rfind("}") 

185 if start != -1 and end != -1 and end > start: 

186 try: 

187 return json.loads(text[start : end + 1]) 

188 except json.JSONDecodeError: 

189 pass 

190 

191 raise ValueError(f"Could not extract valid JSON from LLM response: {text[:300]}") 

192 

193 def call_llm(self, system_prompt: str, user_prompt: str) -> dict: 

194 """Send a prompt to the configured LLM and return parsed JSON. 

195 

196 Uses litellm so any provider (OpenAI, Anthropic, Ollama, LiteLLM proxy) 

197 works transparently — just change LLM_MODEL / LLM_BASE_URL in config. 

198 

199 * response_format is only passed for models that support JSON mode. 

200 * If the response cannot be parsed as JSON, retries up to LLM_MAX_RETRIES 

201 times, appending a correction turn asking for JSON-only output. 

202 

203 Args: 

204 system_prompt (str): System-role message. 

205 user_prompt (str): User-role message. 

206 

207 Returns: 

208 dict: Parsed JSON from the LLM response. 

209 

210 Raises: 

211 ImportError: If litellm is not installed. 

212 ValueError: If the LLM returns malformed JSON after all retries. 

213 Exception: Any litellm / network error propagates to the caller. 

214 """ 

215 

216 litellm = _get_litellm() 

217 

218 base_kwargs: dict = { 

219 "model": config.LLM_MODEL, 

220 } 

221 if config.LLM_API_KEY: 

222 base_kwargs["api_key"] = config.LLM_API_KEY 

223 if config.LLM_BASE_URL: 

224 base_kwargs["base_url"] = config.LLM_BASE_URL 

225 if self._supports_json_mode(): 

226 base_kwargs["format"] = 'json' 

227 

228 messages = [ 

229 {"role": "system", "content": system_prompt}, 

230 {"role": "user", "content": user_prompt}, 

231 ] 

232 

233 last_error: Exception | None = None 

234 max_attempts = config.LLM_MAX_RETRIES + 1 

235 

236 for attempt in range(max_attempts): 

237 logger.info(f"Calling LLM ({config.LLM_MODEL}) for dependency resolution (attempt {attempt + 1}/{max_attempts}) …") 

238 response = litellm.completion(**{**base_kwargs, "messages": messages}) 

239 raw = response.choices[0].message.content or "" 

240 

241 try: 

242 return self._extract_json_from_text(raw) 

243 except ValueError as exc: 

244 last_error = exc 

245 logger.warning(f"LLM returned non-JSON on attempt {attempt + 1}: {raw[:200]}") 

246 if attempt < max_attempts - 1: 

247 # Append a correction turn so the model sees its own bad output 

248 messages = messages + [ 

249 {"role": "assistant", "content": raw}, 

250 { 

251 "role": "user", 

252 "content": ( 

253 "Your previous response was not valid JSON. " 

254 "Please respond with ONLY a valid JSON object — " 

255 "no markdown, no code fences, no explanation, just the raw JSON." 

256 ), 

257 }, 

258 ] 

259 

260 raise ValueError(f"LLM returned non-JSON after {max_attempts} attempt(s): {last_error}") from last_error 

261 

262 # ── Result merging ──────────────────────────────────────────────────────── 

263 

264 def merge_with_classic(self, llm_result: dict, classic_result: dict, endpoint_names: list[str]) -> dict: 

265 """Merge LLM result with classic result. 

266 

267 LLM result is authoritative for endpoints it returns. 

268 Classic result fills in any endpoints the LLM omitted. 

269 

270 Args: 

271 llm_result (dict): Structured output from the LLM (validated). 

272 classic_result (dict): Output from the classic resolver. 

273 endpoint_names (list[str]): All endpoint names to ensure complete coverage. 

274 

275 Returns: 

276 dict: Merged result with all endpoints populated. 

277 """ 

278 import copy 

279 merged = copy.deepcopy(classic_result) 

280 for name in endpoint_names: 

281 if name in llm_result: 

282 llm_entry = llm_result[name] 

283 merged[name]["hardDependsOn"] = llm_entry.get("hardDependsOn", {}) 

284 merged[name]["softDependsOn"] = llm_entry.get("softDependsOn", {}) 

285 if "mutationType" in llm_entry: 

286 merged[name]["mutationType"] = llm_entry["mutationType"] 

287 return merged 

288 

289 # ── Validation ──────────────────────────────────────────────────────────── 

290 

291 def validate_llm_mutation_result(self, raw: dict, endpoint_names: list[str], objects: dict) -> dict: 

292 """Validate and clean the LLM response for mutations. 

293 

294 - Removes entries for unknown endpoint names (hallucinations) 

295 - Removes dependency values that are not known object names 

296 - Falls back to "UNKNOWN" mutationType when missing 

297 

298 Args: 

299 raw (dict): Raw parsed JSON from the LLM. 

300 endpoint_names (list[str]): Valid mutation names. 

301 objects (dict): Valid object names. 

302 

303 Returns: 

304 dict: Cleaned result. 

305 """ 

306 valid = {} 

307 known_endpoints = set(endpoint_names) 

308 known_objects = set(objects.keys()) 

309 for name, entry in raw.items(): 

310 if name not in known_endpoints: 

311 logger.debug(f"LLM hallucinated unknown mutation '{name}', skipping") 

312 continue 

313 valid[name] = { 

314 "mutationType": entry.get("mutationType", "UNKNOWN") if entry.get("mutationType") in ("CREATE", "UPDATE", "DELETE", "UNKNOWN") else "UNKNOWN", 

315 "hardDependsOn": {k: v for k, v in entry.get("hardDependsOn", {}).items() if v in known_objects}, 

316 "softDependsOn": {k: v for k, v in entry.get("softDependsOn", {}).items() if v in known_objects}, 

317 } 

318 return valid 

319 

320 def validate_llm_query_result(self, raw: dict, endpoint_names: list[str], objects: dict) -> dict: 

321 """Validate and clean the LLM response for queries. 

322 

323 Args: 

324 raw (dict): Raw parsed JSON from the LLM. 

325 endpoint_names (list[str]): Valid query names. 

326 objects (dict): Valid object names. 

327 

328 Returns: 

329 dict: Cleaned result. 

330 """ 

331 valid = {} 

332 known_endpoints = set(endpoint_names) 

333 known_objects = set(objects.keys()) 

334 for name, entry in raw.items(): 

335 if name not in known_endpoints: 

336 logger.debug(f"LLM hallucinated unknown query '{name}', skipping") 

337 continue 

338 valid[name] = { 

339 "hardDependsOn": {k: v for k, v in entry.get("hardDependsOn", {}).items() if v in known_objects}, 

340 "softDependsOn": {k: v for k, v in entry.get("softDependsOn", {}).items() if v in known_objects}, 

341 } 

342 return valid