Create genai_client.py

genai_client.py

@@ -0,0 +1,250 @@
+"""
+Compatibility wrapper for google.generativeai surfaces.
+
+- Reads API key from env var GOOGLE_API_KEY
+- Detects older global-style surfaces (genai.configure / GenerativeModel / genai.responses)
+  and newer client-style surfaces (NewClient / genai.Client / client.responses)
+- Exposes:
+    - create_response(**kwargs) -> response object (SDK response)
+    - upload_file(path, mime_type=None) -> file handle/object for use with responses/files
+    - set_model_default(model_name)
+- Minimal surface: adapt your existing calls to import these helpers and call them.
+"""
+
+import os
+import logging
+from pathlib import Path
+from typing import Optional, Any, Dict
+
+logger = logging.getLogger(__name__)
+logger.setLevel(logging.INFO)
+
+# Load api key from env
+API_KEY = os.getenv("GOOGLE_API_KEY") or os.getenv("GEMINI_API_KEY") or os.getenv("GOOGLE_GENAI_API_KEY")
+if not API_KEY:
+    logger.warning("No GOOGLE_API_KEY (or GEMINI_API_KEY) found in environment. Set it before calling the API.")
+
+# Import SDK and detect available surfaces
+try:
+    import google.generativeai as genai
+except Exception as e:
+    raise RuntimeError("google.generativeai import failed: " + str(e))
+
+def _list_attrs():
+    return [k for k in dir(genai) if not k.startswith("_")]
+
+_HAS_RESPONSES = hasattr(genai, "responses") or hasattr(genai, "Responses")
+_HAS_NEWCLIENT = hasattr(genai, "NewClient") or hasattr(genai, "Client") or hasattr(genai, "client")
+_HAS_GENERATIVEMODEL = hasattr(genai, "GenerativeModel") or hasattr(genai, "GenerativeModelV1")
+
+logger.info("genai attributes detected: %s", _list_attrs())
+logger.info("has responses: %s, has NewClient/Client: %s, has GenerativeModel: %s",
+            _HAS_RESPONSES, _HAS_NEWCLIENT, _HAS_GENERATIVEMODEL)
+
+# Configure API key for older-style SDKs if possible
+try:
+    # many 0.x versions use genai.configure or genai.api_key
+    if hasattr(genai, "configure"):
+        try:
+            genai.configure(api_key=API_KEY)
+        except Exception:
+            # some versions expect env var only
+            pass
+    elif hasattr(genai, "api_key"):
+        try:
+            genai.api_key = API_KEY
+        except Exception:
+            pass
+except Exception:
+    pass
+
+# Create a client when available (newer unified client)
+_client = None
+if _HAS_NEWCLIENT:
+    try:
+        if hasattr(genai, "NewClient"):
+            _client = genai.NewClient(api_key=API_KEY) if API_KEY else genai.NewClient()
+        elif hasattr(genai, "Client"):
+            # some variants expose Client
+            _client = genai.Client(api_key=API_KEY) if API_KEY else genai.Client()
+        elif hasattr(genai, "client"):
+            # wrapper namespace
+            _client = genai.client.NewClient(api_key=API_KEY) if API_KEY else genai.client.NewClient()
+        logger.info("Initialized genai client: %s", type(_client))
+    except Exception as e:
+        logger.warning("Failed to initialize genai client: %s", e)
+        _client = None
+
+_DEFAULT_MODEL = os.getenv("DEFAULT_MODEL", "gemini-2.5-flash-lite")
+
+def set_model_default(model_name: str):
+    global _DEFAULT_MODEL
+    _DEFAULT_MODEL = model_name
+
+def _read_bytes(path: str) -> bytes:
+    return Path(path).read_bytes()
+
+def upload_file(path: str, mime_type: Optional[str] = None) -> Any:
+    """
+    Upload a file to the SDK's file service if supported, or return a dict
+    that can be passed to create_response as raw bytes depending on the SDK.
+    Returns the uploaded file handle/object expected by the SDK in create_response.
+    """
+    path = str(path)
+    data = _read_bytes(path)
+    filename = Path(path).name
+
+    # New client: client.files.upload or client.files.create depending on version
+    if _client is not None:
+        # try common upload surfaces
+        for attr in ("files", "Files", "file_manager", "FileManager"):
+            files_container = getattr(_client, attr, None)
+            if files_container:
+                # try upload method names
+                for m in ("upload", "upload_file", "create", "create_file", "upload_bytes"):
+                    if hasattr(files_container, m):
+                        func = getattr(files_container, m)
+                        try:
+                            logger.info("Uploading via client.%s.%s", attr, m)
+                            return func(file_bytes=data, filename=filename, mime_type=mime_type)
+                        except TypeError:
+                            # try alternative signature
+                            try:
+                                return func(path)
+                            except Exception:
+                                pass
+                        except Exception as e:
+                            logger.warning("client file upload failed: %s", e)
+                # fallback: some SDKs expect an object literal
+                try:
+                    return {"name": filename, "bytes": data, "mime_type": mime_type}
+                except Exception:
+                    pass
+
+    # Older/global-style surfaces: genai.Files or genai.upload_file or genai.responses.upload
+    for candidate in ("Files", "files", "upload_file", "upload", "responses"):
+        svc = getattr(genai, candidate, None)
+        if svc:
+            # try common methods
+            for m in ("upload", "upload_file", "create"):
+                if hasattr(svc, m):
+                    try:
+                        func = getattr(svc, m)
+                        logger.info("Uploading via genai.%s.%s", candidate, m)
+                        return func(path)  # many older wrappers accept path
+                    except Exception:
+                        pass
+            # fallback return dict
+            try:
+                return {"name": filename, "bytes": data, "mime_type": mime_type}
+            except Exception:
+                pass
+
+    # Last resort: return bytes dict for manual usage in create_response
+    return {"name": filename, "bytes": data, "mime_type": mime_type}
+
+def create_response(*, model: Optional[str] = None, input: Optional[Any] = None, **kwargs) -> Any:
+    """
+    Unified call to generate a response.
+
+    Common kwargs you'll want to pass: multimodal contents (contents), files, temperature, max_output_tokens
+    Example:
+        create_response(model="gemini-2.5-flash-lite", input="Hello")
+        create_response(contents=["Describe this image", {"name":"img.png","bytes":b"..."}])
+    """
+    model = model or _DEFAULT_MODEL
+    # Prefer client.responses.create / client.models.generate_content depending on SDK
+    if _client is not None:
+        # newer unified client surfaces
+        # many SDKs expose client.responses.create or client.models.generate_content
+        if hasattr(_client, "responses") and hasattr(_client.responses, "create"):
+            logger.info("Calling client.responses.create with model=%s", model)
+            payload = kwargs.copy()
+            payload.setdefault("model", model)
+            if input is not None:
+                # some SDKs accept 'input' or 'contents'
+                payload.setdefault("input", input)
+            return _client.responses.create(**payload)
+        if hasattr(_client, "models") and hasattr(_client.models, "generate_content"):
+            logger.info("Calling client.models.generate_content with model=%s", model)
+            payload = kwargs.copy()
+            payload.setdefault("model", model)
+            if input is not None:
+                payload.setdefault("contents", input)
+            return _client.models.generate_content(**payload)
+
+    # Global/genai module-style surfaces
+    if hasattr(genai, "responses") and hasattr(genai.responses, "create"):
+        logger.info("Calling genai.responses.create with model=%s", model)
+        payload = kwargs.copy()
+        payload.setdefault("model", model)
+        if input is not None:
+            payload.setdefault("input", input)
+        return genai.responses.create(**payload)
+
+    # Older GenerativeModel style
+    if _HAS_GENERATIVEMODEL:
+        try:
+            # instantiate model wrapper and call generate_content / start_chat style
+            ModelClass = getattr(genai, "GenerativeModel", None) or getattr(genai, "GenerativeModelV1", None)
+            if ModelClass:
+                m = ModelClass(model)
+                logger.info("Calling GenerativeModel.generate_content with model=%s", model)
+                if input is not None:
+                    return m.generate_content(input if not isinstance(input, (list, tuple)) else input)
+                # pass kwargs as best-effort
+                if "contents" in kwargs:
+                    return m.generate_content(kwargs["contents"])
+                return m.generate_content(kwargs)
+        except Exception as e:
+            logger.warning("GenerativeModel path failed: %s", e)
+
+    raise RuntimeError("No compatible google.generativeai surface found. See logs for details.")
+
+# Small helper that returns a plain text result if present (convenience)
+def extract_text(resp: Any) -> str:
+    """
+    Attempt to extract a concise text result from many SDK response shapes.
+    """
+    if resp is None:
+        return ""
+    # New unified: resp.output[0].content or resp.output_text
+    try:
+        if hasattr(resp, "output_text"):
+            return getattr(resp, "output_text") or ""
+        if hasattr(resp, "output"):
+            out = resp.output
+            # output may be list of dicts with 'content'/'text'
+            if isinstance(out, (list, tuple)) and len(out) > 0:
+                first = out[0]
+                if isinstance(first, dict):
+                    return first.get("content") or first.get("text") or ""
+                if hasattr(first, "content"):
+                    return getattr(first, "content") or ""
+        # older: resp.text
+        if hasattr(resp, "text"):
+            return getattr(resp, "text") or ""
+        # dict-like
+        if isinstance(resp, dict):
+            for k in ("output_text", "text", "result", "content", "output"):
+                if k in resp:
+                    val = resp[k]
+                    if isinstance(val, str):
+                        return val
+                    if isinstance(val, list) and len(val) > 0 and isinstance(val[0], str):
+                        return val[0]
+    except Exception:
+        pass
+    # fallback: string representation
+    try:
+        return str(resp)
+    except Exception:
+        return ""
+
+# Exported API
+__all__ = [
+    "create_response",
+    "upload_file",
+    "set_model_default",
+    "extract_text",
+]