Skills › Scientific Agent Skills › latex-posters › generate_schematic_ai.py
generate_schematic_ai.py
857 lines · skills/latex-posters/scripts/generate_schematic_ai.py at 5f2bc03
1#!/usr/bin/env python3
2"""
3AI-powered scientific schematic generation using Nano Banana 2.
5This script uses a smart iterative refinement approach:
61. Generate initial image with Nano Banana 2
72. AI quality review using Gemini 3.7 Flash for scientific critique
83. Only regenerate if quality is below threshold for document type
94. Repeat until quality meets standards (max iterations)
11Requirements:
12 - OPENROUTER_API_KEY environment variable
13 - requests library
15Usage:
16 python generate_schematic_ai.py "Create a flowchart showing CONSORT participant flow" -o flowchart.png
17 python generate_schematic_ai.py "Neural network architecture diagram" -o architecture.png --iterations 2
18 python generate_schematic_ai.py "Simple block diagram" -o diagram.png --doc-type poster
19"""
21import argparse
22import base64
23import json
24import os
25import re
26import sys
27import time
28from pathlib import Path
29from typing import Optional, Dict, Any, List, NamedTuple
31try:
32 import requests
33except ImportError:
34 print("Error: requests library not found. Install with: uv pip install requests")
35 sys.exit(1)
38class ReviewResult(NamedTuple):
39 """The outcome of one quality review.
41 A named tuple rather than a bare tuple because the review has several
42 failure modes -- the model returns no choices, the request raises, the
43 answer arrives in an unexpected shape -- and each of them used to be a
44 chance to return the wrong number of values to the caller.
46 `score` is None whenever no usable number was parsed, and `reviewed` is
47 False whenever the review never produced an answer at all. Neither case
48 invents a score: a diagram whose quality was not measured is reported as
49 unmeasured, not as passing.
50 """
52 critique: str
53 score: Optional[float]
54 needs_improvement: bool
55 reviewed: bool
56 error: Optional[str] = None
59# Markdown decoration the reviewer routinely wraps its answer in. Gemini writes
60# "**SCORE:** 8.5" far more often than the bare "SCORE: 8.5" the rubric asks
61# for, and a pattern that does not allow for it silently scores every diagram
62# at whatever the default happens to be.
63_MARKUP = re.compile(r"[*_`#]")
65_SCORE_PATTERN = re.compile(r"\bSCORE\s*:?\s*(\d+(?:\.\d+)?)", re.IGNORECASE)
66_LOOSE_SCORE_PATTERN = re.compile(
67 r"(?:score|rating|quality)\s*[:\s]\s*(\d+(?:\.\d+)?)\s*(?:/\s*10)?", re.IGNORECASE
68)
69# The underscore is optional because normalization strips it out of
70# NEEDS_IMPROVEMENT along with the surrounding markdown.
71_VERDICT_PATTERN = re.compile(
72 r"^\s*VERDICT\s*:?\s*(ACCEPTABLE|NEEDS[_ ]?IMPROVEMENT)", re.IGNORECASE | re.MULTILINE
73)
76def _normalize_review_text(text: str) -> str:
77 """Strip markdown decoration so the rubric's fields can be found."""
78 return _MARKUP.sub("", text or "")
81def _parse_score(text: str) -> Optional[float]:
82 """Return the reviewer's 0-10 score, or None when there isn't a usable one.
84 None rather than a default: a fabricated score is indistinguishable from a
85 real one once it reaches the review log, and a default near the middle of
86 the range quietly passes lenient document types while failing strict ones.
87 """
88 normalized = _normalize_review_text(text)
89 match = _SCORE_PATTERN.search(normalized) or _LOOSE_SCORE_PATTERN.search(normalized)
90 if not match:
91 return None
92 score = float(match.group(1))
93 # A number outside the rubric's range means something other than the score
94 # was matched -- a figure count, a percentage, an iteration number.
95 return score if 0.0 <= score <= 10.0 else None
98def _parse_verdict(text: str) -> Optional[str]:
99 """Return ACCEPTABLE / NEEDS_IMPROVEMENT from the verdict line, or None.
101 Anchored to the start of a line. The rubric prompt spells out both tokens
102 when it explains the response format, so a reviewer that echoes those
103 instructions back would fail a perfectly good diagram under a plain
104 substring search.
105 """
106 match = _VERDICT_PATTERN.search(_normalize_review_text(text))
107 if not match:
108 return None
109 verdict = match.group(1).upper().replace(" ", "").replace("_", "")
110 return "NEEDS_IMPROVEMENT" if verdict == "NEEDSIMPROVEMENT" else "ACCEPTABLE"
113def _resolve_api_key(explicit: Optional[str] = None) -> Optional[str]:
114 """Resolve the OpenRouter key from --api-key, the environment, then any .env file.
116 The .env scan walks up from the working directory and finally checks the
117 script's own directory, so running from anywhere inside a project picks up
118 the key at its root. Only the standard library is used: python-dotenv is a
119 common omission, and a missing optional dependency should not read as a
120 missing credential.
121 """
122 if explicit:
123 return explicit
125 from_env = os.environ.get("OPENROUTER_API_KEY", "").strip()
126 if from_env:
127 return from_env
129 cwd = Path.cwd()
130 for directory in [cwd, *cwd.parents, Path(__file__).resolve().parent]:
131 env_file = directory / ".env"
132 if not env_file.is_file():
133 continue
134 try:
135 content = env_file.read_text(encoding="utf-8", errors="replace")
136 except OSError:
137 continue
138 for raw in content.splitlines():
139 line = raw.strip()
140 if line.startswith("#") or "=" not in line:
141 continue
142 name, _, value = line.partition("=")
143 if name.strip() == "OPENROUTER_API_KEY":
144 value = value.strip().strip('"').strip("'")
145 if value:
146 return value
148 return None
151class ScientificSchematicGenerator:
152 """Generate scientific schematics using AI with smart iterative refinement.
154 Uses Gemini 3.7 Flash for quality review to determine if regeneration is needed.
155 Multiple passes only occur if the generated schematic doesn't meet the
156 quality threshold for the target document type.
157 """
159 # Quality thresholds by document type (score out of 10)
160 # Higher thresholds for more formal publications
161 QUALITY_THRESHOLDS = {
162 "journal": 8.5, # Nature, Science, etc. - highest standards
163 "conference": 8.0, # Conference papers - high standards
164 "poster": 7.0, # Academic posters - good quality
165 "presentation": 6.5, # Slides/talks - clear but less formal
166 "report": 7.5, # Technical reports - professional
167 "grant": 8.0, # Grant proposals - must be compelling
168 "thesis": 8.0, # Dissertations - formal academic
169 "preprint": 7.5, # arXiv, etc. - good quality
170 "default": 7.5, # Default threshold
171 }
173 # Scientific diagram best practices prompt template
174 SCIENTIFIC_DIAGRAM_GUIDELINES = """
175Create a high-quality scientific diagram with these requirements:
177VISUAL QUALITY:
178- Clean white or light background (no textures or gradients)
179- High contrast for readability and printing
180- Professional, publication-ready appearance
181- Sharp, clear lines and text
182- Adequate spacing between elements to prevent crowding
184TYPOGRAPHY:
185- Clear, readable sans-serif fonts (Arial, Helvetica style)
186- Minimum 10pt font size for all labels
187- Consistent font sizes throughout
188- All text horizontal or clearly readable
189- No overlapping text
191SCIENTIFIC STANDARDS:
192- Accurate representation of concepts
193- Clear labels for all components
194- Include scale bars, legends, or axes where appropriate
195- Use standard scientific notation and symbols
196- Include units where applicable
198ACCESSIBILITY:
199- Colorblind-friendly color palette (use Okabe-Ito colors if using color)
200- High contrast between elements
201- Redundant encoding (shapes + colors, not just colors)
202- Works well in grayscale
204LAYOUT:
205- Logical flow (left-to-right or top-to-bottom)
206- Clear visual hierarchy
207- Balanced composition
208- Appropriate use of whitespace
209- No clutter or unnecessary decorative elements
211IMPORTANT - NO FIGURE NUMBERS:
212- Do NOT include "Figure 1:", "Fig. 1", or any figure numbering in the image
213- Do NOT add captions or titles like "Figure: ..." at the top or bottom
214- Figure numbers and captions are added separately in the document/LaTeX
215- The diagram should contain only the visual content itself
216"""
218 def __init__(self, api_key: Optional[str] = None, verbose: bool = False):
219 """
220 Initialize the generator.
222 Args:
223 api_key: OpenRouter API key (or use OPENROUTER_API_KEY env var)
224 verbose: Print detailed progress information
225 """
226 # Priority: 1) explicit api_key param, 2) environment variable, 3) .env file
227 self.api_key = _resolve_api_key(api_key)
229 if not self.api_key:
230 raise ValueError(
231 "OPENROUTER_API_KEY not found. Please either:\n"
232 " 1. Set the OPENROUTER_API_KEY environment variable\n"
233 " 2. Add OPENROUTER_API_KEY to your .env file\n"
234 " 3. Pass api_key parameter to the constructor\n"
235 "Get your API key from: https://openrouter.ai/keys"
236 )
238 self.verbose = verbose
239 self._last_error = None # Track last error for better reporting
240 self.base_url = "https://openrouter.ai/api/v1"
241 # Nano Banana 2 - Google's advanced image generation model. The slug must
242 # be an image-output model; a text-only chat model is rejected with
243 # "No endpoints found that support the requested output modalities".
244 # https://openrouter.ai/google/gemini-3.1-flash-image
245 self.image_model = "google/gemini-3.1-flash-image"
246 # Gemini 3.7 Flash for quality review - excellent vision and reasoning
247 self.review_model = "google/gemini-3.7-flash"
249 def _log(self, message: str):
250 """Log message if verbose mode is enabled."""
251 if self.verbose:
252 print(f"[{time.strftime('%H:%M:%S')}] {message}")
254 def _post_request(self, endpoint: str, payload: Dict[str, Any]) -> Dict[str, Any]:
255 """Send one non-streaming request; never automatically replay a paid call."""
256 headers = {
257 "Authorization": f"Bearer {self.api_key}",
258 "Content-Type": "application/json",
259 "X-Title": "Scientific Schematic Generator",
260 }
261 self._log(f"POST {endpoint} using {payload['model']}")
262 try:
263 response = requests.post(
264 f"{self.base_url}/{endpoint}", headers=headers,
265 json=payload, timeout=120,
266 )
267 try:
268 response_json = response.json()
269 except ValueError:
270 raise RuntimeError(
271 f"API returned non-JSON data (HTTP {response.status_code})"
272 ) from None
273 if not isinstance(response_json, dict):
274 raise RuntimeError("API response must be a JSON object")
275 if response.status_code != 200 or "error" in response_json:
276 error = response_json.get("error", {})
277 detail = error.get("message", "request rejected") if isinstance(error, dict) else str(error)
278 # Providers occasionally echo input; do not dump whole bodies or credentials.
279 detail = str(detail).replace(self.api_key, "[redacted]")[:500]
280 raise RuntimeError(f"API request failed (HTTP {response.status_code}): {detail}")
281 return response_json
282 except requests.exceptions.Timeout:
283 raise RuntimeError("API request timed out after 120 seconds") from None
284 except requests.exceptions.RequestException as exc:
285 detail = str(exc).replace(self.api_key, "[redacted]")
286 raise RuntimeError(f"API request failed: {detail}") from None
288 def _make_request(self, model: str, messages: List[Dict[str, Any]],
289 modalities: Optional[List[str]] = None) -> Dict[str, Any]:
290 """Use chat completions for the text-and-image vision review."""
291 payload = {"model": model, "messages": messages}
292 if modalities:
293 payload["modalities"] = modalities
294 return self._post_request("chat/completions", payload)
296 def _extract_image_from_response(self, response: Dict[str, Any]) -> bytes:
297 """Decode the dedicated Image API response and enforce this helper's PNG contract.
299 The provider may omit media_type; the PNG signature still has to match.
300 This is a format check, not a full image-decoder or scientific-content check.
301 """
302 data = response.get("data")
303 if not isinstance(data, list) or not data or not isinstance(data[0], dict):
304 raise RuntimeError("Image API returned no data[0] image")
305 item = data[0]
306 if item.get("media_type") not in (None, "image/png"):
307 raise RuntimeError(f"Expected PNG output; received {item['media_type']}")
308 encoded = item.get("b64_json")
309 if not isinstance(encoded, str) or not encoded:
310 raise RuntimeError("Image API returned no b64_json image")
311 try:
312 image_data = base64.b64decode(encoded, validate=True)
313 except ValueError:
314 raise RuntimeError("Image API returned invalid base64 image data") from None
315 if not image_data.startswith(b"\x89PNG\r\n\x1a\n"):
316 raise RuntimeError("Image API output does not have a PNG signature")
317 return image_data
319 def _image_to_base64(self, image_path: str) -> str:
320 """
321 Convert image file to base64 data URL.
323 Args:
324 image_path: Path to image file
326 Returns:
327 Base64 data URL string
328 """
329 with open(image_path, "rb") as f:
330 image_data = f.read()
332 # Determine image type from extension
333 ext = Path(image_path).suffix.lower()
334 mime_type = {
335 ".png": "image/png",
336 ".jpg": "image/jpeg",
337 ".jpeg": "image/jpeg",
338 ".gif": "image/gif",
339 ".webp": "image/webp"
340 }.get(ext, "image/png")
342 base64_data = base64.b64encode(image_data).decode("utf-8")
343 return f"data:{mime_type};base64,{base64_data}"
345 def generate_image(self, prompt: str) -> Optional[bytes]:
346 """
347 Generate an image using Nano Banana 2.
349 Args:
350 prompt: Description of the diagram to generate
352 Returns:
353 Image bytes or None if generation failed
354 """
355 self._last_error = None # Reset error
357 try:
358 # Image generation has its own body and response, distinct from chat.
359 # n=1 is supported by both current Gemini image-provider endpoints.
360 # output_format is absent from their capability maps: use the default
361 # and validate the returned MIME/signature rather than send that knob.
362 response = self._post_request("images", {
363 "model": self.image_model, "prompt": prompt, "n": 1,
364 })
365 image_data = self._extract_image_from_response(response)
366 self._log(f"[OK] Generated PNG ({len(image_data)} bytes)")
367 return image_data
368 except RuntimeError as exc:
369 self._last_error = str(exc)
370 self._log(f"[FAIL] Generation failed: {self._last_error}")
371 return None
373 def review_image(self, image_path: str, original_prompt: str,
374 iteration: int, doc_type: str = "default",
375 max_iterations: int = 2) -> ReviewResult:
376 """
377 Review the generated image for quality, using the vision review model.
379 Args:
380 image_path: Path to the generated image
381 original_prompt: Original user prompt
382 iteration: Current iteration number
383 doc_type: Document type (journal, poster, presentation, etc.)
384 max_iterations: Maximum iterations allowed
386 Returns:
387 A ReviewResult. When the review cannot be completed -- the request
388 fails, or the model answers with nothing usable -- `reviewed` is
389 False and `score` is None rather than a stand-in number, and
390 `needs_improvement` is False: the fault is in the reviewer, not in
391 the diagram, so regenerating would not help.
392 """
393 image_data_url = self._image_to_base64(image_path)
395 # Get quality threshold for this document type
396 threshold = self.QUALITY_THRESHOLDS.get(doc_type.lower(),
397 self.QUALITY_THRESHOLDS["default"])
399 review_prompt = f"""You are an expert reviewer evaluating a scientific diagram for publication quality.
401ORIGINAL REQUEST: {original_prompt}
403DOCUMENT TYPE: {doc_type} (quality threshold: {threshold}/10)
404ITERATION: {iteration}/{max_iterations}
406Carefully evaluate this diagram on these criteria:
4081. **Scientific Accuracy** (0-2 points)
409 - Correct representation of concepts
410 - Proper notation and symbols
411 - Accurate relationships shown
4132. **Clarity and Readability** (0-2 points)
414 - Easy to understand at a glance
415 - Clear visual hierarchy
416 - No ambiguous elements
4183. **Label Quality** (0-2 points)
419 - All important elements labeled
420 - Labels are readable (appropriate font size)
421 - Consistent labeling style
4234. **Layout and Composition** (0-2 points)
424 - Logical flow (top-to-bottom or left-to-right)
425 - Balanced use of space
426 - No overlapping elements
4285. **Professional Appearance** (0-2 points)
429 - Publication-ready quality
430 - Clean, crisp lines and shapes
431 - Appropriate colors/contrast
433RESPOND IN THIS EXACT FORMAT:
434SCORE: [total score 0-10]
436STRENGTHS:
437- [strength 1]
438- [strength 2]
440ISSUES:
441- [issue 1 if any]
442- [issue 2 if any]
444VERDICT: [ACCEPTABLE or NEEDS_IMPROVEMENT]
446If score >= {threshold}, the diagram is ACCEPTABLE for {doc_type} publication.
447If score < {threshold}, mark as NEEDS_IMPROVEMENT with specific suggestions."""
449 messages = [
450 {
451 "role": "user",
452 "content": [
453 {
454 "type": "text",
455 "text": review_prompt
456 },
457 {
458 "type": "image_url",
459 "image_url": {
460 "url": image_data_url
461 }
462 }
463 ]
464 }
465 ]
467 try:
468 response = self._make_request(
469 model=self.review_model,
470 messages=messages
471 )
473 # Extract text response
474 choices = response.get("choices", [])
475 if not choices:
476 # A normal enough outcome -- a rate limit, a content filter, a
477 # provider hiccup. The diagram itself is fine; only its review
478 # is missing, and saying so beats inventing a score.
479 reason = "the review model returned no choices"
480 self._log(f"[WARN] Review unavailable: {reason}")
481 return ReviewResult(
482 f"Review unavailable: {reason}.",
483 None, False, reviewed=False, error=reason,
484 )
486 message = choices[0].get("message", {})
487 content = message.get("content", "")
489 # Some models put their analysis in a separate reasoning field
490 reasoning = message.get("reasoning", "")
491 if reasoning and not content:
492 content = reasoning
494 if isinstance(content, list):
495 # Extract text from content blocks
496 text_parts = []
497 for block in content:
498 if isinstance(block, dict) and block.get("type") == "text":
499 text_parts.append(block.get("text", ""))
500 content = "\n".join(text_parts)
502 score = _parse_score(content)
503 verdict = _parse_verdict(content)
505 if score is None and verdict is None:
506 reason = "no score or verdict found in the review"
507 self._log(f"[WARN] Review unusable: {reason}")
508 return ReviewResult(
509 content if content else f"Review unusable: {reason}.",
510 None, False, reviewed=True, error=reason,
511 )
513 needs_improvement = verdict == "NEEDS_IMPROVEMENT" or (
514 score is not None and score < threshold
515 )
517 shown = f"{score}/10" if score is not None else "not stated"
518 self._log(f"[OK] Review complete (Score: {shown}, Threshold: {threshold}/10)")
519 self._log(f" Verdict: {'Needs improvement' if needs_improvement else 'Acceptable'}")
521 return ReviewResult(
522 content if content else "Review returned no text",
523 score,
524 needs_improvement,
525 reviewed=True,
526 error=None if score is not None else "no score found in the review",
527 )
528 except Exception as e:
529 # A failed review must not fail the run -- the image is already
530 # generated and saved -- but it must not read as a pass either.
531 self._log(f"[WARN] Review failed: {str(e)}")
532 return ReviewResult(
533 f"Review failed: {str(e)}",
534 None, False, reviewed=False, error=str(e),
535 )
537 def improve_prompt(self, original_prompt: str, critique: str,
538 iteration: int) -> str:
539 """
540 Improve the generation prompt based on critique.
542 Args:
543 original_prompt: Original user prompt
544 critique: Review critique from previous iteration
545 iteration: Current iteration number
547 Returns:
548 Improved prompt for next generation
549 """
550 improved_prompt = f"""{self.SCIENTIFIC_DIAGRAM_GUIDELINES}
552USER REQUEST: {original_prompt}
554ITERATION {iteration}: Based on previous feedback, address these specific improvements:
555{critique}
557Generate an improved version that addresses all the critique points while maintaining scientific accuracy and professional quality."""
559 return improved_prompt
561 def generate_iterative(self, user_prompt: str, output_path: str,
562 iterations: int = 2,
563 doc_type: str = "default") -> Dict[str, Any]:
564 """
565 Generate scientific schematic with smart iterative refinement.
567 Only regenerates if the quality score is below the threshold for the
568 specified document type. This saves API calls and time when the first
569 generation is already good enough.
571 Args:
572 user_prompt: User's description of desired diagram
573 output_path: Path to save final image
574 iterations: Maximum refinement iterations (default: 2, max: 2)
575 doc_type: Document type for quality threshold (journal, poster, etc.)
577 Returns:
578 Dictionary with generation results and metadata
579 """
580 if isinstance(iterations, bool) or not isinstance(iterations, int) or not 1 <= iterations <= 2:
581 raise ValueError("iterations must be 1 or 2")
582 if doc_type.lower() not in self.QUALITY_THRESHOLDS:
583 raise ValueError(f"Unknown document type: {doc_type}")
584 output_path = Path(output_path)
585 if output_path.suffix.lower() != ".png":
586 raise ValueError("Output path must end in .png; format conversion is not supported")
587 output_dir = output_path.parent
588 output_dir.mkdir(parents=True, exist_ok=True)
590 base_name = output_path.stem
591 extension = output_path.suffix or ".png"
593 # Get quality threshold for this document type
594 threshold = self.QUALITY_THRESHOLDS.get(doc_type.lower(),
595 self.QUALITY_THRESHOLDS["default"])
597 results = {
598 "user_prompt": user_prompt,
599 "doc_type": doc_type,
600 "quality_threshold": threshold,
601 "image_model": self.image_model,
602 "review_model": self.review_model,
603 "quality_met": False,
604 "termination_reason": None,
605 "iterations": [],
606 "final_image": None,
607 # None, not 0.0 -- a run whose review never completed has no score,
608 # and a number here would be read as one the reviewer gave.
609 "final_score": None,
610 "final_reviewed": False,
611 "success": False,
612 "early_stop": False,
613 "early_stop_reason": None
614 }
616 current_prompt = f"""{self.SCIENTIFIC_DIAGRAM_GUIDELINES}
618USER REQUEST: {user_prompt}
620Generate a publication-quality scientific diagram that meets all the guidelines above."""
622 print(f"\n{'='*60}")
623 print(f"Generating Scientific Schematic")
624 print(f"{'='*60}")
625 print(f"Description: {user_prompt}")
626 print(f"Document Type: {doc_type}")
627 print(f"Quality Threshold: {threshold}/10")
628 print(f"Max Iterations: {iterations}")
629 print(f"Output: {output_path}")
630 print(f"{'='*60}\n")
632 for i in range(1, iterations + 1):
633 print(f"\n[Iteration {i}/{iterations}]")
634 print("-" * 40)
636 # Generate image
637 print(f"Generating image...")
638 image_data = self.generate_image(current_prompt)
640 if not image_data:
641 error_msg = getattr(self, '_last_error', 'Image generation failed - no image data returned')
642 print(f"[FAIL] Generation failed: {error_msg}")
643 results["iterations"].append({
644 "iteration": i,
645 "success": False,
646 "error": error_msg
647 })
648 # A failed refinement must not discard the earlier saved image.
649 # Stop here rather than replay a request after an ambiguous timeout.
650 results["termination_reason"] = "generation_failed"
651 break
653 # Save iteration image
654 iter_path = output_dir / f"{base_name}_v{i}{extension}"
655 with open(iter_path, "wb") as f:
656 f.write(image_data)
657 print(f"[OK] Saved: {iter_path}")
659 # Review the image with the vision review model
660 print(f"Reviewing image with {self.review_model}...")
661 review = self.review_image(
662 str(iter_path), user_prompt, i, doc_type, iterations
663 )
664 if review.score is not None:
665 print(f"[OK] Score: {review.score}/10 (threshold: {threshold}/10)")
666 else:
667 print(f"[WARN] Review unavailable - image kept, quality not verified")
668 print(f" Reason: {review.error}")
670 # Save iteration results
671 iteration_result = {
672 "iteration": i,
673 "image_path": str(iter_path),
674 "prompt": current_prompt,
675 "critique": review.critique,
676 "score": review.score,
677 "reviewed": review.reviewed,
678 "review_error": review.error,
679 "needs_improvement": review.needs_improvement,
680 "success": True
681 }
682 results["iterations"].append(iteration_result)
683 # Preserve the latest successful image even if a later API call fails.
684 results["final_image"] = str(iter_path)
685 results["final_score"] = review.score
686 results["final_reviewed"] = review.reviewed and review.score is not None
687 results["success"] = True
688 results["quality_met"] = (
689 results["final_reviewed"] and review.score >= threshold
690 and not review.needs_improvement
691 )
693 # Check if quality is acceptable - STOP EARLY if so
694 if not review.needs_improvement:
695 if review.score is not None:
696 print(f"\n[OK] Quality meets {doc_type} threshold ({review.score} >= {threshold})")
697 print(f" No further iterations needed!")
698 reason = (f"Quality score {review.score} meets threshold "
699 f"{threshold} for {doc_type}")
700 else:
701 # Regenerating cannot fix a reviewer that did not answer.
702 print(f"\n[WARN] Stopping without a verified score - review the image yourself")
703 reason = f"Review did not produce a score: {review.error}"
704 results["final_image"] = str(iter_path)
705 results["final_score"] = review.score
706 results["final_reviewed"] = review.reviewed and review.score is not None
707 results["success"] = True
708 results["early_stop"] = True
709 results["early_stop_reason"] = reason
710 results["termination_reason"] = (
711 "quality_threshold_met" if results["quality_met"] else "review_unavailable"
712 )
713 break
715 # If this is the last iteration, we're done regardless
716 if i == iterations:
717 print(f"\n[WARN] Maximum iterations reached")
718 results["termination_reason"] = "max_iterations"
719 results["final_image"] = str(iter_path)
720 results["final_score"] = review.score
721 results["final_reviewed"] = review.reviewed and review.score is not None
722 results["success"] = True
723 break
725 # Quality below threshold - improve prompt for next iteration
726 print(f"\n[WARN] Review requests improvement (score: {review.score}; threshold: {threshold})")
727 print(f"Improving prompt based on feedback...")
728 current_prompt = self.improve_prompt(user_prompt, review.critique, i + 1)
730 # Copy final version to output path
731 if results["success"] and results["final_image"]:
732 final_iter_path = Path(results["final_image"])
733 if final_iter_path != output_path:
734 import shutil
735 shutil.copy(final_iter_path, output_path)
736 print(f"\n[OK] Final image: {output_path}")
738 # Save review log
739 log_path = output_dir / f"{base_name}_review_log.json"
740 with open(log_path, "w", encoding="utf-8") as f:
741 json.dump(results, f, indent=2)
742 print(f"[OK] Review log: {log_path}")
744 print(f"\n{'='*60}")
745 print(f"Generation Complete!")
746 if results["final_score"] is not None:
747 print(f"Final Score: {results['final_score']}/10")
748 else:
749 print(f"Final Score: unavailable (the review did not produce one)")
750 if results["early_stop"]:
751 print(f"Iterations Used: {len([r for r in results['iterations'] if r.get('success')])}/{iterations} (early stop)")
752 print(f"{'='*60}\n")
754 return results
757def main():
758 """Command-line interface."""
759 parser = argparse.ArgumentParser(
760 description="Generate scientific schematics using AI with smart iterative refinement",
761 formatter_class=argparse.RawDescriptionHelpFormatter,
762 epilog="""
763Examples:
764 # Generate a flowchart for a journal paper
765 python generate_schematic_ai.py "CONSORT participant flow diagram" -o flowchart.png --doc-type journal
767 # Generate neural network architecture for presentation (lower threshold)
768 python generate_schematic_ai.py "Transformer encoder-decoder architecture" -o transformer.png --doc-type presentation
770 # Generate with custom max iterations for poster
771 python generate_schematic_ai.py "Biological signaling pathway" -o pathway.png --iterations 2 --doc-type poster
773 # Verbose output
774 python generate_schematic_ai.py "Circuit diagram" -o circuit.png -v
776Document Types (quality thresholds):
777 journal 8.5/10 - Nature, Science, peer-reviewed journals
778 conference 8.0/10 - Conference papers
779 thesis 8.0/10 - Dissertations, theses
780 grant 8.0/10 - Grant proposals
781 preprint 7.5/10 - arXiv, bioRxiv, etc.
782 report 7.5/10 - Technical reports
783 poster 7.0/10 - Academic posters
784 presentation 6.5/10 - Slides, talks
785 default 7.5/10 - General purpose
787Note: Multiple iterations only occur if quality is BELOW the threshold.
788 If the first generation meets the threshold, no extra API calls are made.
790Environment:
791 OPENROUTER_API_KEY OpenRouter API key (required)
792 """
793 )
795 parser.add_argument("prompt", help="Description of the diagram to generate")
796 parser.add_argument("-o", "--output", required=True,
797 help="Output image path (e.g., diagram.png)")
798 parser.add_argument("--iterations", type=int, default=2,
799 help="Maximum refinement iterations (default: 2, max: 2)")
800 parser.add_argument("--doc-type", default="default",
801 choices=["journal", "conference", "poster", "presentation",
802 "report", "grant", "thesis", "preprint", "default"],
803 help="Document type for quality threshold (default: default)")
804 parser.add_argument("--api-key", help="OpenRouter API key (or set OPENROUTER_API_KEY)")
805 parser.add_argument("-v", "--verbose", action="store_true",
806 help="Verbose output")
808 args = parser.parse_args()
810 # Check for API key — resolves --api-key, the environment, then any .env file
811 api_key = _resolve_api_key(args.api_key)
812 if not api_key:
813 print("Error: OPENROUTER_API_KEY not found")
814 print("\nSet it with:")
815 print(" export OPENROUTER_API_KEY='your_api_key'")
816 print("\nOr add OPENROUTER_API_KEY=your_api_key to a .env file")
817 print("Or provide via --api-key flag")
818 sys.exit(1)
820 # Validate iterations - enforce max of 2
821 if args.iterations < 1 or args.iterations > 2:
822 print("Error: Iterations must be between 1 and 2")
823 sys.exit(1)
825 try:
826 generator = ScientificSchematicGenerator(api_key=api_key, verbose=args.verbose)
827 results = generator.generate_iterative(
828 user_prompt=args.prompt,
829 output_path=args.output,
830 iterations=args.iterations,
831 doc_type=args.doc_type
832 )
834 if results["success"]:
835 print(f"\n[OK] Success! Image saved to: {args.output}")
836 used = len([r for r in results['iterations'] if r.get('success')])
837 if results.get("final_reviewed"):
838 if results.get("quality_met"):
839 print(f" (Completed in {used} iteration(s) - quality threshold met)")
840 else:
841 print(" (Image saved; review did not meet the quality criteria.)")
842 else:
843 # The image is real; the quality claim is not. Say which.
844 print(f" (Completed in {used} iteration(s) - quality NOT verified,"
845 f" the review produced no score. Check the image yourself.)")
846 sys.exit(0)
847 else:
848 print(f"\n[FAIL] Generation failed. Check review log for details.")
849 sys.exit(1)
850 except Exception as e:
851 print(f"\n[FAIL] Error: {str(e)}")
852 sys.exit(1)
855if __name__ == "__main__":
856 main()