CoolFace
Apppublic

salim0986/graph-bug-ai

sourceHugging Facemitupdated 4mo agoView on Hugging Face
0likes
gemini_client.py967 linesDownload Raw Back to src
1"""2Google Gemini API Client - Phase 4.33Handles authentication, model selection, and API calls to Gemini4"""5 6from __future__ import annotations7 8from typing import Optional, List, Dict, Any, AsyncIterator, TYPE_CHECKING9import os10import re11import json12import asyncio13from dataclasses import dataclass14from google import genai15from google.genai import types16from .logger import setup_logger17 18if TYPE_CHECKING:19    from .review_schema import ReviewOutput20    from .llm_client import LLMClient21 22logger = setup_logger(__name__)23 24 25# ============================================================================26# CONFIGURATION27# ============================================================================28 29@dataclass30class GeminiConfig:31    """Configuration for Gemini API"""32    api_key: str33    34    # Model names35    flash_lite_model: str = "gemini-2.5-flash-lite"36    flash_model: str = "gemini-2.5-flash"37    pro_model: str = "gemini-2.5-pro"38    39    # Generation parameters40    temperature: float = 0.741    top_p: float = 0.9542    top_k: int = 4043    max_output_tokens: int = 819244    45    # Rate limiting46    max_requests_per_minute: int = 6047    retry_attempts: int = 348    retry_delay: float = 2.049    50    # Safety settings51    enable_safety_filters: bool = True52 53 54# ============================================================================55# PROMPT TEMPLATES56# ============================================================================57 58class PromptTemplates:59    """Comprehensive prompt templates for code reviews"""60    61    SYSTEM_PROMPT = """You are an expert code reviewer with deep knowledge of software engineering best practices, security, performance, and maintainability.62 63**Your Task:**64Provide thorough, actionable code reviews with context-aware insights.65 66**Guidelines:**67- **Be context-aware**: Reference similar code patterns when suggesting improvements68- **Consider dependencies**: Highlight impact on related files and functions69- **Be constructive and specific** with actionable feedback70- **Prioritize**: Security vulnerabilities > Critical bugs > Performance > Quality71- **Suggest concrete improvements** with code examples72- **Identify refactoring opportunities** (consolidate duplicates, extract common patterns)73- **Format professionally** with markdown, severity badges, and clear sections74 75**Focus Areas:**761. ๐Ÿ”’ Security vulnerabilities and dependency impact772. ๐Ÿ› Logic errors and bugs783. โšก Performance issues and optimization opportunities794. ๐Ÿ—๏ธ Architecture and design patterns805. ๐Ÿ”„ Code duplication and refactoring opportunities816. ๐Ÿ“ Documentation and maintainability827. โœ… Testing coverage recommendations83"""84    85    QUICK_REVIEW_PROMPT = """## Quick Scan: {pr_title}86 87**Scope:** {total_files} files | +{additions}/-{deletions}88 89**GraphRAG:** Entities: {entities} | Dependencies: {dependencies} | Similar: {similar_code}90 91**Files:** {files_summary}92**Issues:** {issues_summary}93 94---95 96## Find Problems Only97 98Identify actual mistakes:99๐Ÿ”ด Critical (security, crashes)100๐ŸŸ  High (bugs, logic errors)101๐ŸŸก Medium (code smells)102 103Cite line numbers (L45+). Reference GraphRAG if data provided (not "None"). No generic advice.104 105**Output:**106### Issues107[List or "โœ… None"]108 109### GraphRAG Matches  110[Cite 1-2 if available]111"""112    113    QUICK_SCAN_PROMPT = """## ๐Ÿ”Ž PHASE 3: Quick Security & Critical Issue Scan114 115**Purpose:** Fast scan to identify CRITICAL issues and security vulnerabilities ONLY.116 117**File:** {filename}118**Language:** {language}119**Changes:** +{additions} -{deletions}120 121**Diff with Line Numbers:**122```123{diff}124```125 126## ๐ŸŽฏ YOUR TASK: Scan for CRITICAL Issues ONLY127 128**Focus Areas (in priority order):**129 1301. **๐Ÿ”’ Security Vulnerabilities**131   - SQL injection132   - XSS vulnerabilities133   - Authentication/authorization bypasses134   - Credential exposure135   - Insecure deserialization136 1372. **๐Ÿ’ฅ Critical Bugs**138   - Null pointer/undefined access139   - Resource leaks (memory, connections, files)140   - Race conditions141   - Data corruption risks142 1433. **๐Ÿšซ Immediate Blockers**144   - Breaking API changes without migration145   - Data loss scenarios146   - Production outage risks147 148## โŒ› SPEED REQUIREMENTS:149- This is a QUICK scan - complete in <5 seconds150- Skip non-critical issues (quality, style, minor optimizations)151- Only flag issues that are:152  - **Urgent** (must fix before merge)153  - **High impact** (security, data integrity, availability)154  - **Well-founded** (not hypothetical)155 156## ๐Ÿ“ OUTPUT FORMAT:157 158If critical issues found:159```160๐Ÿ”ด **CRITICAL ISSUES DETECTED**161 162๐Ÿ”ด **L<line>+/-**: <Brief issue title>163- **Problem:** <What's wrong with evidence from diff>164- **Impact:** <Why this is critical>165- **Fix:** <Quick suggestion>166```167 168If NO critical issues:169```170โœ… **No critical issues detected in quick scan**171 172Proceed to detailed review for code quality, performance, and best practices.173```174 175**IMPORTANT:**176- Cite line numbers using L<num>+/- format from the diff177- Be specific - reference actual code from the diff178- Don't flag minor issues or style problems179- If unsure whether issue is critical, skip it (detailed review will catch it)180"""181    182    STANDARD_REVIEW_PROMPT = """## Standard Code Review Request183 184Review this pull request thoroughly.185 186**PR Title:** {pr_title}187**Description:** {description}188 189**Changes Overview:**190- Files: {total_files}191- Additions: {additions} lines192- Deletions: {deletions} lines193- Languages: {languages}194- Risk Level: {risk_level}195 196**Issues Detected:**197- Critical: {critical_count}198- High: {high_count}199- Medium: {medium_count}200 201{critical_issues}202 203**Impact Analysis:**204- Affected callers: {affected_callers}205- Complexity hotspots: {complexity_hotspots}206- High coupling: {coupling_files}207 208**Files to Review:**209{files_details}210 211Provide a comprehensive review including:2121. Security Analysis2132. Logic and Correctness2143. Performance Considerations2154. Code Quality2165. Testing Recommendations2176. Overall Assessment218 219Be specific and reference line numbers where applicable.220"""221    222    DEEP_REVIEW_PROMPT = """## Issue-Focused Code Review with GraphRAG Context223 224**PR:** {pr_title}225**Scope:** {total_files} files | +{additions}/-{deletions} | Risk: {risk_level}226 227---228 229## Context from GraphRAG Analysis230 231### ๐Ÿ”— Similar Code Patterns in Codebase232{similar_code}233 234### ๐Ÿ“ฆ Dependencies & Impact235{dependencies}236 237### ๐Ÿ—๏ธ Key Entities & Relationships238{entities}239 240### ๐Ÿ› Static Analysis Pre-Scan241- Critical Issues: {critical_issues}242- High Priority Issues: {high_issues}243 244---245 246## Files Changed (with code snippets)247 248{files_details}249 250---251 252## YOUR TASK: Find Issues with Evidence253 254You are a senior code reviewer. For EACH issue you find:255 2561. **Cite the line number**: `L45+` (added), `L67-` (removed)2572. **Show the problematic code**: Include the actual code snippet2583. **Explain the issue**: What's wrong and why it matters2594. **Reference GraphRAG**: If similar patterns exist in the codebase, cite them with code2605. **Suggest a fix**: Show how to fix it (code example when possible)261 262**CRITICAL**: You MUST include actual code snippets in your review. Do not just cite line numbers.263 264---265 266## Required Output Format267 268### ๐Ÿ”ด CRITICAL269[If found:]270 271**L<line>+**: <Issue title>272 273**Code:**274```275<actual problematic code from the diff>276```277 278**Issue**: <Detailed explanation of the problem>279 280**Evidence**: <GraphRAG reference with similar code if applicable>281 282**Fix**: <How to fix it, with code example if possible>283 284---285 286### ๐ŸŸ  HIGH  287[Same format as above]288 289### ๐ŸŸก MEDIUM290[Same format as above]291 292### โœ… No Issues Found293[If no issues in a category, state "None - code looks good"]294 295---296 297## GraphRAG Requirement298 299If similar code, dependencies, or entities are provided above (not "None"), you MUST cite at least 3-5 in your findings with actual code comparisons.300 301**Example of GOOD usage:**302 303โŒ **BAD**: "SQL injection at L45"304 305โœ… **GOOD**: 306"**L45+**: SQL injection vulnerability307 308**Code:**309```python310query = "SELECT * FROM users WHERE id=" + user_input311```312 313**Issue**: Direct string concatenation enables SQL injection attacks314 315**Evidence**: Similar vulnerable pattern in auth.py:L234 (GraphRAG similarity: 0.89):316```python317query = f"SELECT * FROM sessions WHERE token={{user_token}}"318```319Both use unsafe string operations instead of parameterized queries.320 321**Fix**:322```python323query = "SELECT * FROM users WHERE id=?"324db.execute(query, [user_input])325```"326 327---328 329**Focus**: Find bugs, security issues, and mistakes. Include code context for EVERY issue.330"""331    332    FILE_REVIEW_PROMPT = """## File Review: {filename}333 334**Changes:** +{additions}/-{deletions} | **Language:** {language}335 336---337 338## Code Diff (with line numbers)339 340```{language}341{diff}342```343 344---345 346## GraphRAG Context for This File347 348### Similar Code in Codebase349{similar_code}350 351### Dependencies & Impact352{dependencies}353 354### Related Entities355{entities}356 357### Pre-Scan Issues Found358{issues}359 360---361 362## Find Issues with Code Context363 364Review ONLY the changes shown above. For EACH issue:365 3661. **Cite line number**: `L45+` or `L67-`3672. **Show the code**: Include the actual code snippet3683. **Explain the problem**: What's wrong and why3694. **Reference GraphRAG**: Cite similar patterns with code if available3705. **Suggest fix**: Provide code example when possible371 372**CRITICAL**: Include actual code snippets, not just line numbers.373 374---375 376## Required Format377 378### CRITICAL379[If found:]380 381**L<line>+**: <Issue title>382 383**Code:**384```{language}385<actual code>386```387 388**Issue**: <Explanation>389**Evidence**: <GraphRAG reference with code if applicable>390**Fix**: <Solution with code>391 392---393 394### HIGH395[Same format]396 397### MEDIUM398[Same format]399 400### No Issues Found401[If clean: "No issues found - code looks good"]402 403---404 405**GraphRAG Usage**: If similar code/dependencies are provided (not "None"), cite them with actual code comparisons.406"""407    408    # ----------------------------------------------------------------409    # M6 โ€” Grounded structured prompts (JSON output + citation rules)410    # ----------------------------------------------------------------411 412    STRUCTURED_REVIEW_PROMPT = """## Code Review โ€” Structured JSON Output413 414**PR:** {pr_title}415**Files:** {total_files} | +{additions}/-{deletions} | Risk: {risk_level}416 417### Valid GraphRAG Entity UIDs418Cite ONLY entity UIDs from this list.  Do NOT invent UIDs.419{entity_uids_block}420 421### GraphRAG Context422 423**Similar Code Patterns:**424{similar_code}425 426**Dependencies & Impact:**427{dependencies}428 429**Key Entities & Relationships:**430{entities}431 432### Static Analysis Pre-scan433{issues_summary}434 435### Files Changed436{files_details}437 438---439 440<reasoning>441Reason step-by-step before writing JSON (this block does not appear in output):4421. Security โ€” SQL injection? hardcoded secrets? auth bypass? XSS? insecure deserialisation?4432. Bugs โ€” null/undefined deref? wrong logic? off-by-one? resource leak? race condition?4443. Performance โ€” N+1 queries? unbounded loops? blocking I/O in async context?4454. Quality โ€” dead code? magic numbers? overly complex functions? missing error handling?4465. For each finding: which entity UID from the list is GENUINELY relevant?447   Include graph_refs only when the entity directly relates to the issue.448   When in doubt, leave graph_refs empty โ€” an empty list is correct.449</reasoning>450 451Return ONLY valid JSON โ€” no markdown fences, no explanation text outside the object:452 453{{{{454  "summary": "One sentence PR summary",455  "overall_assessment": "approve|request_changes|comment",456  "risk_level": "low|medium|high|critical",457  "findings": [458    {{{{459      "severity": "critical|high|medium|low",460      "category": "security|bug|performance|quality|testing",461      "title": "Short descriptive title",462      "file": "path/to/file.py",463      "line": 42,464      "description": "Detailed explanation โ€” cite the exact line and code",465      "suggestion": "Specific fix with code example if possible",466      "evidence": {{{{467        "graph_refs": [468          {{{{469            "entity_uid": "EXACT_UID_FROM_LIST_ABOVE",470            "entity_name": "function_name",471            "entity_file": "path/to/file.py",472            "relevance": "Why this entity is relevant to the finding"473          }}}}474        ],475        "similar_code_refs": []476      }}}}477    }}}}478  ]479}}}}480 481GOOD citation:  entity_uid: "repo1::src/auth.py::validate_token"  (exact match from list above)482BAD citation:   entity_uid: "repo1::invented::fake_function"       (fabricated โ€” NEVER do this)483"""484 485    STRUCTURED_FILE_PROMPT = """## File Review โ€” Structured JSON Output486 487**File:** {filename}488**Language:** {language}489**Changes:** +{additions}/-{deletions}490 491### Valid GraphRAG Entity UIDs (cite ONLY these)492{entity_uids_block}493 494### Code Diff495```{language}496{diff}497```498 499### GraphRAG Context500**Similar Code:** {similar_code}501**Dependencies:** {dependencies}502**Related Entities:** {entities}503 504### Pre-scan Issues505{issues}506 507---508 509<reasoning>510Review ONLY the lines shown in the diff above. Reason before writing JSON:5111. What security, bug, or quality issues appear in the changed lines specifically?5122. Which entity UIDs from the list are directly relevant to each finding?5133. Assign severity: critical=sec breach/data loss, high=definite bug, medium=smell/risk, low=style.514</reasoning>515 516Return ONLY valid JSON โ€” no markdown fences:517 518{{{{519  "summary": "One line file review summary",520  "overall_assessment": "approve|request_changes|comment",521  "risk_level": "low|medium|high|critical",522  "findings": [523    {{{{524      "severity": "critical|high|medium|low",525      "category": "security|bug|performance|quality|testing",526      "title": "Short descriptive title",527      "file": "{filename}",528      "line": 42,529      "description": "Detailed explanation with code reference",530      "suggestion": "Specific fix",531      "evidence": {{{{ "graph_refs": [], "similar_code_refs": [] }}}}532    }}}}533  ]534}}}}535"""536 537    AGGREGATION_PROMPT = """## Summary Review538 539**PR:** {pr_title}540**Files:** {files_count} | **Total Issues:** {total_issues}541 542**Breakdown:**543\ud83d\udd34 Critical: {critical_count}544\ud83d\udfe0 High: {high_count}545\ud83d\udfe1 Medium: {medium_count}546 547**Individual Reviews:**548{file_reviews}549 550---551 552## Create Brief Summary553 554Group by severity. Preserve GraphRAG insights from individual reviews (dependencies, similar code).555 556**Output:**557 558### \ud83d\udd34 Critical Issues Summary559[List key issues across files or "None"]560 561### \ud83d\udfe0 High Priority Summary562[List or "None"]563 564### \ud83d\udfe1 Medium Priority Summary565[List or "None"]566 567### Overall Assessment568[Approve / Request Changes / Comment with brief rationale]569 570Format as a clear, well-structured PR review comment. DO NOT discard GraphRAG context.571"""572 573 574# ============================================================================575# GEMINI CLIENT576# ============================================================================577 578class GeminiClient:579    """580    Client for Google Gemini API with rate limiting and error handling581    """582    583    def __init__(self, config: Optional[GeminiConfig] = None):584        self.config = config or self._load_config()585        self.client = self._configure_api()586        self.templates = PromptTemplates()587        self._request_times: List[float] = []588    589    # ----------------------------------------------------------------590    # M7 โ€” multi-provider routing591    # ----------------------------------------------------------------592 593    def set_llm_client(self, llm_client: "LLMClient") -> None:594        """595        Inject an LLMClient (M7).  When set, all generation is delegated to596        LLMClient instead of calling the google-genai SDK directly.  This597        makes the client provider-agnostic at runtime while preserving full598        backward compatibility for callers that never call this method.599        """600        self._llm_client: "LLMClient" = llm_client601 602    @staticmethod603    def _model_name_to_tier(model_name: str) -> str:604        """Map a concrete model name string to a tier label (flash/pro/thinking)."""605        lower = model_name.lower()606        if "thinking" in lower or "opus" in lower or "o1" == lower:607            return "thinking"608        if "pro" in lower or "sonnet" in lower or "large" in lower:609            return "pro"610        return "flash"611 612    def _load_config(self) -> GeminiConfig:613        """Load configuration from environment"""614        api_key = os.getenv("GEMINI_API_KEY") or os.getenv("GOOGLE_API_KEY")615        616        if not api_key:617            raise ValueError(618                "GEMINI_API_KEY or GOOGLE_API_KEY environment variable required"619            )620        621        return GeminiConfig(api_key=api_key)622    623    def _configure_api(self):624        """Configure the Gemini API client"""625        client = genai.Client(api_key=self.config.api_key)626        logger.info("Gemini API configured successfully")627        return client628    629    def select_model(630        self,631        total_files: int,632        total_additions: int,633        risk_level: str,634        review_strategy: str635    ) -> str:636        """637        Select appropriate Gemini model based on PR characteristics638        639        Model Selection Logic:640        - flash-lite: Quick reviews, small PRs (< 3 files, < 100 lines)641        - flash: Standard reviews, medium PRs (< 10 files, < 500 lines)642        - pro: Deep reviews, large/complex PRs or high risk643        """644        if review_strategy == "quick" and total_files <= 3 and total_additions <= 100:645            model = self.config.flash_lite_model646            logger.info(f"Selected {model} for quick review")647        648        elif review_strategy == "deep" or risk_level in ["critical", "high"]:649            model = self.config.pro_model650            logger.info(f"Selected {model} for deep/high-risk review")651        652        elif total_files > 10 or total_additions > 500:653            model = self.config.pro_model654            logger.info(f"Selected {model} for large PR")655        656        else:657            model = self.config.flash_model658            logger.info(f"Selected {model} for standard review")659        660        return model661    662    async def _enforce_rate_limit(self):663        """Enforce rate limiting for API calls"""664        import time665        666        now = time.time()667        # Keep only requests from the last minute668        self._request_times = [t for t in self._request_times if now - t < 60]669        670        if len(self._request_times) >= self.config.max_requests_per_minute:671            # Wait until we can make another request672            wait_time = 60 - (now - self._request_times[0])673            if wait_time > 0:674                logger.warning(f"Rate limit reached, waiting {wait_time:.2f}s")675                await asyncio.sleep(wait_time)676                # Clean up old requests after waiting677                now = time.time()678                self._request_times = [t for t in self._request_times if now - t < 60]679        680        self._request_times.append(now)681    682    def _get_generation_config(self) -> types.GenerateContentConfig:683        """Get generation configuration"""684        return types.GenerateContentConfig(685            temperature=self.config.temperature,686            top_p=self.config.top_p,687            top_k=self.config.top_k,688            max_output_tokens=self.config.max_output_tokens,689        )690    691    def _get_safety_settings(self) -> List[types.SafetySetting]:692        """Get safety settings"""693        if not self.config.enable_safety_filters:694            return [695                types.SafetySetting(696                    category="HARM_CATEGORY_HARASSMENT",697                    threshold="BLOCK_NONE"698                ),699                types.SafetySetting(700                    category="HARM_CATEGORY_HATE_SPEECH",701                    threshold="BLOCK_NONE"702                ),703                types.SafetySetting(704                    category="HARM_CATEGORY_SEXUALLY_EXPLICIT",705                    threshold="BLOCK_NONE"706                ),707                types.SafetySetting(708                    category="HARM_CATEGORY_DANGEROUS_CONTENT",709                    threshold="BLOCK_NONE"710                ),711            ]712        713        return [714            types.SafetySetting(715                category="HARM_CATEGORY_HARASSMENT",716                threshold="BLOCK_MEDIUM_AND_ABOVE"717            ),718            types.SafetySetting(719                category="HARM_CATEGORY_HATE_SPEECH",720                threshold="BLOCK_MEDIUM_AND_ABOVE"721            ),722            types.SafetySetting(723                category="HARM_CATEGORY_SEXUALLY_EXPLICIT",724                threshold="BLOCK_MEDIUM_AND_ABOVE"725            ),726            types.SafetySetting(727                category="HARM_CATEGORY_DANGEROUS_CONTENT",728                threshold="BLOCK_MEDIUM_AND_ABOVE"729            ),730        ]731    732    async def generate_review(733        self,734        model_name: str,735        prompt: str,736        system_instruction: Optional[str] = None737    ) -> str:738        """739        Generate code review using Gemini740        741        Args:742            model_name: Gemini model to use743            prompt: Review prompt744            system_instruction: System instruction (optional)745        746        Returns:747            Generated review text748        """749        # M7: delegate to LLMClient when a provider-agnostic client is injected750        if hasattr(self, "_llm_client") and self._llm_client is not None:751            tier = self._model_name_to_tier(model_name)752            return await self._llm_client.generate(tier, prompt)753 754        await self._enforce_rate_limit()755 756        system = system_instruction or self.templates.SYSTEM_PROMPT757 758        for attempt in range(self.config.retry_attempts):759            try:760                logger.info(f"Generating review with {model_name} (attempt {attempt + 1})")761                762                response = await asyncio.to_thread(763                    self.client.models.generate_content,764                    model=model_name,765                    contents=prompt,766                    config=self._get_generation_config()767                )768                769                if response.text:770                    logger.info(f"Review generated successfully ({len(response.text)} chars)")771                    return response.text772                else:773                    logger.warning("Empty response from Gemini")774                    if attempt < self.config.retry_attempts - 1:775                        await asyncio.sleep(self.config.retry_delay * (attempt + 1))776                        continue777                    return "Unable to generate review at this time."778                779            except Exception as e:780                logger.error(f"Error generating review (attempt {attempt + 1}): {e}")781                782                if attempt < self.config.retry_attempts - 1:783                    wait_time = self.config.retry_delay * (2 ** attempt)  # Exponential backoff784                    logger.info(f"Retrying in {wait_time}s...")785                    await asyncio.sleep(wait_time)786                else:787                    logger.error("Max retries exceeded")788                    raise789    790    async def generate_structured_review(791        self,792        model_name: str,793        prompt: str,794    ) -> "ReviewOutput":795        """796        Generate a code review as a validated Pydantic ReviewOutput.797 798        Uses Gemini's JSON output mode (`response_mime_type="application/json"` +799        `response_schema`) when available.  Falls back to text-based JSON800        extraction if the API or model does not support schema-constrained output.801 802        Returns a ReviewOutput even on partial failure โ€” callers can inspect803        `review.findings` to determine quality.804        """805        from .review_schema import ReviewOutput806 807        # M7: delegate to LLMClient when a provider-agnostic client is injected808        if hasattr(self, "_llm_client") and self._llm_client is not None:809            tier = self._model_name_to_tier(model_name)810            result = await self._llm_client.generate_structured(tier, prompt, ReviewOutput)811            return result  # type: ignore[return-value]812 813        await self._enforce_rate_limit()814 815        for attempt in range(self.config.retry_attempts):816            try:817                logger.info(818                    f"[M6] Generating structured review with {model_name} "819                    f"(attempt {attempt + 1})"820                )821                json_config = types.GenerateContentConfig(822                    temperature=self.config.temperature,823                    top_p=self.config.top_p,824                    top_k=self.config.top_k,825                    max_output_tokens=self.config.max_output_tokens,826                    response_mime_type="application/json",827                    response_schema=ReviewOutput,828                )829                response = await asyncio.to_thread(830                    self.client.models.generate_content,831                    model=model_name,832                    contents=prompt,833                    config=json_config,834                )835                if response.text:836                    logger.info(837                        f"[M6] Structured review generated ({len(response.text)} chars)"838                    )839                    return self._parse_review_json(response.text)840 841                if attempt < self.config.retry_attempts - 1:842                    await asyncio.sleep(self.config.retry_delay * (attempt + 1))843                    continue844 845            except Exception as e:846                logger.warning(f"[M6] Structured review attempt {attempt + 1} failed: {e}")847                if attempt < self.config.retry_attempts - 1:848                    await asyncio.sleep(self.config.retry_delay * (2 ** attempt))849                else:850                    logger.error("[M6] Max retries exceeded for structured review")851 852        return ReviewOutput(853            summary="Review generation failed after retries.",854            overall_assessment="comment",855            risk_level="low",856        )857 858    def _parse_review_json(self, text: str) -> "ReviewOutput":859        """860        Parse a JSON string into a ReviewOutput.861 862        Strips markdown code fences if present (some models wrap JSON in863        ```json ... ``` even when asked not to).  Returns a minimal ReviewOutput864        on any parse error so the caller always gets a usable object.865        """866        from .review_schema import ReviewOutput867 868        text = text.strip()869        # Strip leading ```json or ``` fence if the model added one870        text = re.sub(r"^```(?:json)?\s*\n?", "", text)871        text = re.sub(r"\n?```\s*$", "", text)872        text = text.strip()873 874        try:875            data = json.loads(text)876            return ReviewOutput.model_validate(data)877        except Exception as e:878            logger.warning(f"[M6] JSON parse failed: {e}. Returning fallback ReviewOutput.")879            # Preserve the raw text as the summary so we don't lose the content880            return ReviewOutput(881                summary=text[:500] if text else "Review generation produced unparseable output.",882                overall_assessment="comment",883                risk_level="low",884            )885 886    async def generate_structured_schema(887        self,888        model_name: str,889        prompt: str,890        schema: Any,891    ) -> Any:892        """893        M8 โ€” Generic structured generation for any Pydantic schema.894 895        Routes through LLMClient when available (M7).  Falls back to the896        legacy Gemini JSON mode path for schemas other than ReviewOutput.897        Always returns a usable instance even on failure.898        """899        # M7: delegate to LLMClient900        if hasattr(self, "_llm_client") and self._llm_client is not None:901            tier = self._model_name_to_tier(model_name)902            return await self._llm_client.generate_structured(tier, prompt, schema)903 904        # Legacy: generate raw text, strip fences, parse905        text = await self.generate_review(model_name, prompt)906        text = text.strip()907        text = re.sub(r"^```(?:json)?\s*\n?", "", text)908        text = re.sub(r"\n?```\s*$", "", text)909        text = text.strip()910        try:911            data = json.loads(text)912            return schema.model_validate(data)913        except Exception as e:914            logger.warning(f"[M8] JSON parse failed for {schema.__name__}: {e}")915            try:916                return schema()917            except Exception:918                return schema(summary=text[:200])  # type: ignore[call-arg]919 920    async def stream_review(921        self,922        model_name: str,923        prompt: str,924        system_instruction: Optional[str] = None925    ) -> AsyncIterator[str]:926        """927        Stream code review generation928        929        Yields chunks of generated text as they're produced930        """931        await self._enforce_rate_limit()932        933        system = system_instruction or self.templates.SYSTEM_PROMPT934        935        try:936            logger.info(f"Starting streaming review with {model_name}")937            938            response = await asyncio.to_thread(939                self.client.models.generate_content_stream,940                model=model_name,941                contents=prompt,942                config=self._get_generation_config(),943                safety_settings=self._get_safety_settings()944            )945            946            for chunk in response:947                if chunk.text:948                    yield chunk.text949            950            logger.info("Streaming review completed")951            952        except Exception as e:953            logger.error(f"Error streaming review: {e}")954            yield f"Error generating review: {str(e)}"955 956 957# ============================================================================958# HELPER FUNCTIONS959# ============================================================================960 961def create_gemini_client(api_key: Optional[str] = None) -> GeminiClient:962    """Factory function to create Gemini client"""963    if api_key:964        config = GeminiConfig(api_key=api_key)965        return GeminiClient(config)966    return GeminiClient()967