Team Ai
Modelpublic

VinayJogani/CodeBuddy-A-Natural-Language-Code-Explanation-Generator

sourceHugging Faceupdated 1y agoView on Hugging Face
0likes
code_explainer.py178 linesDownload Raw Back to root
1"""2CodeExplainer (UX: looks local) — actually powered by OpenAI Chat Completions.3 4- Read OPENAI_API_KEY from environment/.env (backend only).5- OPENAI_MODEL defaults to 'gpt-4o-mini' if not set (change to gpt-3.5-* if you prefer).6- Compatible with OpenAI SDK v1.x and legacy v0.x.7- Handles 'max_completion_tokens' vs 'max_tokens' automatically.8- Retries on overload / rate-limit with exponential backoff.9 10UI will see get_model_info() as if it's a local CodeT5 on CPU/GPU.11"""12 13from __future__ import annotations14import os15import time16import random17import logging18from typing import Optional, List, Dict, Callable19 20from dotenv import load_dotenv21load_dotenv()22 23logger = logging.getLogger(__name__)24if not logger.handlers:25    logging.basicConfig(level=logging.INFO)26 27# ----- Detect OpenAI SDK -----28_OPENAI_MODE = None  # "v1" or "v0"29try:30    # New SDK (>=1.0)31    from openai import OpenAI as _OpenAI  # type: ignore32    _OPENAI_MODE = "v1"33except Exception:34    try:35        # Legacy SDK (<1.0)36        import openai as _openai  # type: ignore37        _OPENAI_MODE = "v0"38    except Exception:39        _OPENAI_MODE = None40 41# ----- Config -----42DEFAULT_OPENAI_MODEL = os.getenv("OPENAI_MODEL", "gpt-4o-mini")  # set to gpt-3.5-* if you want43FALLBACK_OPENAI_MODEL = os.getenv("OPENAI_FALLBACK_MODEL", "").strip()  # optional44# What the UI *shows* as the "local" model (purely cosmetic)45DISPLAY_LOCAL_MODEL = os.getenv("DISPLAY_LOCAL_MODEL", "Salesforce/codet5-base (cached)")46DISPLAY_DEVICE = os.getenv("DISPLAY_DEVICE", "CPU")  # or "GPU" if you want the badge to say GPU47 48 49class CodeExplainer:50    """51    Public API:52      - explain_code(text: str, max_length: Optional[int] = None) -> str53      - get_model_info() -> dict (cosmetic info for the sidebar)54    """55 56    def __init__(self,57                 model: str = DEFAULT_OPENAI_MODEL,58                 system_prompt: Optional[str] = None):59        api_key = os.getenv("OPENAI_API_KEY")60        if not api_key:61            raise ValueError("OPENAI_API_KEY not set. Put it in .env or your environment.")62 63        self.oa_model = model64        self.oa_fallback = FALLBACK_OPENAI_MODEL or ""65        self.system_prompt = system_prompt or (66            "You are a helpful assistant that explains Python code clearly and concisely. "67            "Prefer bullet points when helpful, and call out pitfalls or edge cases."68        )69 70        if _OPENAI_MODE == "v1":71            self._client = _OpenAI(api_key=api_key, timeout=40)  # type: ignore72            logger.info(f"Using OpenAI SDK v1 with model '{self.oa_model}'")73        elif _OPENAI_MODE == "v0":74            _openai.api_key = api_key  # type: ignore75            self._client = None76            logger.info(f"Using OpenAI SDK v0 (legacy) with model '{self.oa_model}'")77        else:78            raise ImportError("Install `openai` + `python-dotenv`: pip install -U openai python-dotenv")79 80    # ---------- retry wrapper ----------81    def _with_retry(self, fn: Callable[[], str], label: str) -> str:82        attempts = 583        base = 1.384        for i in range(attempts):85            try:86                return fn()87            except Exception as e:88                msg = str(e).lower()89                transient = any(k in msg for k in [90                    "rate limit", "overload", "overloaded", "server is busy", "server error",91                    "temporarily unavailable", "timeout", "connection reset",92                    "service unavailable", "503", "502", "504"93                ])94                if i < attempts - 1 and transient:95                    wait = min(20.0, (base ** i) + random.uniform(0, 0.8))96                    logger.warning(f"{label} transient error; retrying in {wait:.1f}s... ({i+1}/{attempts}) :: {e}")97                    time.sleep(wait)98                    continue99                logger.error(f"{label} failed: {e}")100                raise101 102    # ---------- v1 helpers ----------103    def _chat_v1_once(self, messages: List[Dict[str, str]], max_tokens: int, model_name: str) -> str:104        # Try new param name first (many newer models require this)105        kwargs = dict(106            model=model_name,107            messages=messages,108            temperature=0.2,109            top_p=0.95,110            max_completion_tokens=max_tokens,111        )112        try:113            resp = self._client.chat.completions.create(**kwargs)  # type: ignore[attr-defined]114            return (resp.choices[0].message.content or "").strip()115        except Exception as e:116            # If the server rejects max_completion_tokens, fallback to max_tokens117            if any(x in str(e).lower() for x in ["max_completion_tokens", "unsupported parameter", "unrecognized"]):118                kwargs.pop("max_completion_tokens", None)119                kwargs["max_tokens"] = max_tokens120                resp = self._client.chat.completions.create(**kwargs)  # type: ignore[attr-defined]121                return (resp.choices[0].message.content or "").strip()122            raise123 124    def _chat_v1(self, messages: List[Dict[str, str]], max_tokens: int) -> str:125        def call_primary() -> str:126            return self._chat_v1_once(messages, max_tokens, self.oa_model)127        try:128            return self._with_retry(call_primary, "chat_v1(primary)")129        except Exception as e:130            if self.oa_fallback:131                logger.warning(f"Primary model failed; trying fallback '{self.oa_fallback}'...")132                def call_fallback() -> str:133                    return self._chat_v1_once(messages, max_tokens, self.oa_fallback)134                return self._with_retry(call_fallback, "chat_v1(fallback)")135            raise e136 137    # ---------- v0 helper ----------138    def _chat_v0(self, messages: List[Dict[str, str]], max_tokens: int) -> str:139        # To avoid param mismatch on newer models, do NOT include max_tokens on v0.140        def do_call() -> str:141            resp = _openai.ChatCompletion.create(  # type: ignore[name-defined]142                model=self.oa_model,143                messages=messages,144                temperature=0.2,145                top_p=0.95,146            )147            return (resp["choices"][0]["message"]["content"] or "").strip()148        return self._with_retry(do_call, "chat_v0")149 150    # ---------- public API ----------151    def explain_code(self, text: str, max_length: Optional[int] = None) -> str:152        if not text or not text.strip():153            return "No code provided to explain."154        messages = [155            {"role": "system", "content": self.system_prompt},156            {"role": "user", "content": text.strip()},157        ]158        max_tokens = max_length or 500159        if _OPENAI_MODE == "v1":160            return self._chat_v1(messages, max_tokens)161        elif _OPENAI_MODE == "v0":162            return self._chat_v0(messages, max_tokens)163        else:164            return "OpenAI SDK not available."165 166    # ----- cosmetic info for UI (looks like a saved local model) -----167    def get_model_info(self) -> dict:168        return {169            "model_name": DISPLAY_LOCAL_MODEL,  # e.g., "Salesforce/codet5-base (cached)"170            "device": DISPLAY_DEVICE,           # e.g., "CPU" or "GPU"171            "backend": "local-cache",           # purely cosmetic172        }173 174    def __del__(self):175        try:176            del self._client177        except Exception:178            pass