Spaces:
Sleeping
Sleeping
File size: 6,019 Bytes
f7d4608 d1ca23a f7d4608 d1ca23a f7d4608 d1ca23a f7d4608 d1ca23a f7d4608 d1ca23a f7d4608 d1ca23a f7d4608 d1ca23a f7d4608 d1ca23a f7d4608 d1ca23a f7d4608 d1ca23a f7d4608 d1ca23a f7d4608 d1ca23a f7d4608 d1ca23a f7d4608 d1ca23a f7d4608 d1ca23a f7d4608 d1ca23a |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 |
import os
import json
import re
from typing import Optional, Dict, Union, IO, List, BinaryIO
from google import genai
from google.genai import types
from application.schemas.response_schema import GEMINI_RESPONSE_FORMAT
from application.utils import logger
logger=logger.get_logger()
client = genai.Client(api_key=os.getenv("gemini_api_key"))
PROMPT = (
"""You are a PDF parsing agent. Your job is to extract GHG Protocol Parameters
and ESG (Environmental, Social, Governance) Data from a company’s sustainability
or ESG report in PDF format."""
)
def sanitize_file_name(name: str, max_length: int = 40) -> str:
"""
Sanitizes a file name to comply with Gemini API naming rules:
- Lowercase only
- Alphanumeric characters and dashes (`-`) allowed
- Cannot start or end with a dash
- Max length: 40 characters
Args:
name (str): The original file name (without extension).
max_length (int, optional): Maximum allowed characters (default: 40).
Returns:
str: Sanitized file name.
Raises:
ValueError: If the sanitized name is empty after cleaning.
"""
if not name or not isinstance(name, str):
raise ValueError("Invalid file name: must be a non-empty string.")
# Convert to lowercase and replace invalid characters with dashes
name = re.sub(r'[^a-z0-9]+', '-', name.lower())
# Remove leading/trailing dashes and truncate
name = name.strip('-')[:max_length].rstrip('-')
if not name:
raise ValueError("Sanitized file name is empty or invalid after cleanup.")
return name
def get_files() -> List[str]:
"""
Retrieves all uploaded file names from Gemini.
Returns:
List[str]: List of existing file names.
"""
files = client.files.list()
return [file.name for file in files]
def delete_files(file_names: Union[str, List[str]]) -> None:
"""
Deletes specified files from Gemini.
Args:
file_names (Union[str, List[str]]): File name or list of names to delete.
"""
if not file_names:
logger.warning("No file names provided for deletion.")
return
if isinstance(file_names, str):
file_names = [file_names]
existing_files = get_files()
for name in file_names:
logger.info(f"Attempting to delete file: {name}")
if name in existing_files:
client.files.delete(name=name)
logger.info(f"Deleted file: {name}")
else:
logger.warning(f"File not found: {name}")
def upload_file(
file: Union[str, IO[bytes]],
file_name: Optional[str] = None,
config: Optional[Dict[str, str]] = None
) -> Optional[types.File]:
"""
Uploads a file to the Gemini API, handling both file paths and binary streams.
Args:
file (Union[str, IO[bytes]]): File path or binary file object (e.g., from Streamlit).
file_name (Optional[str]): Name for the file. If None, attempts to use file.name.
config (Optional[Dict[str, str]]): Extra config like 'mime_type'.
Returns:
Optional[types.File]: The uploaded Gemini file object, or existing one if already uploaded.
Raises:
Exception: If upload fails.
"""
try:
if not file_name:
if isinstance(file, str):
file_name = os.path.basename(file)
elif hasattr(file, "name"):
file_name = os.path.basename(file.name)
else:
raise ValueError("file_name must be provided if file has no 'name' attribute.")
sanitized_name = sanitize_file_name(os.path.splitext(file_name)[0])
mime_type = "application/pdf"
config = config or {}
config.update({"name": sanitized_name, "mime_type": mime_type})
gemini_file_key = f"files/{sanitized_name}"
if gemini_file_key in get_files():
logger.info(f"File already exists on Gemini: {gemini_file_key}")
return client.files.get(name=gemini_file_key)
logger.info(f"Uploading file to Gemini: {gemini_file_key}")
if isinstance(file, str):
with open(file, "rb") as f:
return client.files.upload(file=f, config=config)
else:
return client.files.upload(file=file, config=config)
except Exception as e:
logger.error(f"Failed to upload file '{file_name}': {e}")
raise
def extract_emissions_data_as_json(
api: str,
model: str,
file_input: Union[BinaryIO, bytes]
) -> Optional[dict]:
"""
Extracts ESG data from a PDF using the Gemini API.
Args:
api (str): API provider (must be 'gemini').
model (str): Model name (e.g., 'gemini-pro').
file_input (Union[BinaryIO, bytes]): File object or byte stream.
Returns:
Optional[dict]: Parsed JSON response or raw text if parsing fails.
"""
try:
if api.lower() != "gemini":
logger.error(f"Unsupported API: {api}")
return None
file_name = file_input.name if hasattr(file_input, 'name') else "uploaded_file.pdf"
uploaded_file = upload_file(file=file_input, file_name=file_name)
response = client.models.generate_content(
model=model,
contents=[uploaded_file, PROMPT],
config={
'response_mime_type': 'application/json',
'response_schema': GEMINI_RESPONSE_FORMAT
}
)
logger.info("[Gemini] Response received.")
try:
return json.loads(response.text)
except json.JSONDecodeError:
logger.warning("Failed to parse JSON, returning raw response.")
return {"raw_response": response.text}
except Exception as e:
logger.exception("Error during ESG data extraction.")
return None |