Spaces:
Sleeping
Sleeping
| """ | |
| Google Gemini API client wrapper for structured extraction and text generation. | |
| """ | |
| import os | |
| import logging | |
| import json | |
| import time | |
| from typing import List, Dict, Any, Optional, Union | |
| from tenacity import retry, stop_after_attempt, wait_exponential, retry_if_exception_type | |
| import google.generativeai as genai | |
| from google.api_core import exceptions | |
| # Setup Logging | |
| logger = logging.getLogger("ingestion.gemini") | |
| class GeminiClient: | |
| """ | |
| Wrapper for Gemini Pro/Flash models providing structured JSON extraction, | |
| translation, and embedding capabilities. | |
| """ | |
| def __init__(self, api_key: str): | |
| """Initializes the Gemini module with the provided API key.""" | |
| if not api_key: | |
| logger.warning("GEMINI_API_KEY is empty. API calls will fail.") | |
| genai.configure(api_key=api_key) | |
| self.api_key = api_key | |
| def generate_structured( | |
| self, | |
| prompt: str, | |
| response_schema: dict, | |
| model_name: str = "gemini-2.5-flash", | |
| temperature: float = 0.2 | |
| ) -> dict: | |
| """ | |
| Calls Gemini to generate a structured JSON response matching the provided schema. | |
| """ | |
| logger.debug(f"Generating structured data. Prompt snippet: {prompt[:200]}...") | |
| model = genai.GenerativeModel(model_name) | |
| # Configure model for JSON response | |
| generation_config = { | |
| "response_mime_type": "application/json", | |
| "response_schema": response_schema, | |
| "temperature": temperature | |
| } | |
| response = model.generate_content(prompt, generation_config=generation_config) | |
| try: | |
| return json.loads(response.text) | |
| except json.JSONDecodeError as e: | |
| logger.error(f"Failed to parse Gemini JSON response: {response.text}") | |
| raise e | |
| def translate_zh_to_en(self, mandarin_text: str) -> str: | |
| """Translates Mandarin text to English using a simple prompt.""" | |
| prompt = f"Translate the following Chinese text to English:\n\n{mandarin_text}" | |
| return self.generate_text(prompt) | |
| def embed(self, texts: Union[str, List[str]], model: str = "models/text-embedding-004") -> List[List[float]]: | |
| """ | |
| Generates vector embeddings for the provided text(s). | |
| Returns a list of vectors. | |
| """ | |
| if isinstance(texts, str): | |
| texts = [texts] | |
| result = genai.embed_content( | |
| model=model, | |
| content=texts, | |
| task_type="retrieval_document" | |
| ) | |
| # Result is {'embedding': [[...], [...]]} | |
| return result["embedding"] | |
| def generate_text(self, prompt: str, model_name: str = "gemini-2.5-flash") -> str: | |
| """Generates plain text from a prompt.""" | |
| model = genai.GenerativeModel(model_name) | |
| response = model.generate_content(prompt) | |
| return response.text | |
| # Global instance | |
| gemini = GeminiClient(api_key=os.getenv("GEMINI_API_KEY", "")) | |