Spaces:
Sleeping
Sleeping
File size: 5,409 Bytes
f7d4608 935efff f7d4608 22481bd |
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 |
import os
import json
from typing import Union, BinaryIO, Optional
from openai import OpenAI
from google import genai
from google.genai import types
from application.utils import logger
from application.schemas.response_schema import RESPONSE_FORMAT,GEMINI_RESPONSE_FORMAT
logger = logger.get_logger()
client = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))
# --- Constants ---
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."
)
# --- OpenAI Helpers ---
def get_files() -> list:
"""Retrieve all files from OpenAI client."""
try:
files = client.files.list()
logger.info(f"Retrieved {len(files.data)} files.")
return files.data
except Exception as e:
logger.error(f"Failed to retrieve files: {e}")
raise
def get_or_create_file(file_input: BinaryIO, client) -> object:
"""
Retrieve a file from OpenAI by name or upload it if not present.
Args:
file_input: File-like object with `.name` attribute.
client: OpenAI client instance.
Returns:
File object.
"""
file_name = getattr(file_input, 'name', None)
if not file_name:
raise ValueError("File input must have a 'name' attribute.")
try:
for file in get_files():
if file.filename == file_name:
logger.info(f"File '{file_name}' already exists with ID: {file.id}")
return client.files.retrieve(file.id)
logger.info(f"Uploading new file '{file_name}'...")
new_file = client.files.create(file=(file_name, file_input), purpose="assistants")
logger.info(f"File uploaded successfully with ID: {new_file.id}")
return new_file
except Exception as e:
logger.error(f"Error during get_or_create_file: {e}")
raise
def delete_file_by_size(size: int, client):
"""
Deletes files from OpenAI that match a given byte size.
Args:
size: File size in bytes to match for deletion.
client: OpenAI client instance.
"""
try:
files = get_files()
for file in files:
if file.bytes == size:
client.files.delete(file.id)
logger.info(f"File {file.filename} deleted (size matched: {size} bytes).")
else:
logger.info(f"File {file.filename} skipped (size mismatch).")
except Exception as e:
logger.error(f"Failed to delete files: {e}")
raise
# --- Main Function ---
def extract_emissions_data_as_json(
api: str,
model: str,
file_input: Union[BinaryIO, bytes]
) -> Optional[dict]:
"""
Extract ESG data from PDF using OpenAI or Gemini APIs.
Args:
api: 'openai' or 'gemini'
model: Model name (e.g. gpt-4o, gemini-pro)
file_input: File-like object or bytes of the PDF.
Returns:
Parsed ESG data as dict or None if failed.
"""
try:
if api.lower() == "openai":
file = get_or_create_file(file_input, client)
logger.info("[OpenAI] Sending content for generation...")
response = client.chat.completions.create(
model=model,
messages=[{
"role": "user",
"content": [
{"type": "file", "file": {"file_id": file.id}},
{"type": "text", "text": PROMPT}
]
}],
response_format=RESPONSE_FORMAT
)
result = response.choices[0].message.content
logger.info("ESG data extraction successful.")
return result
elif api.lower() == "gemini":
client = genai.Client(api_key=os.getenv("gemini_api_key"))
file_bytes = file_input.read()
logger.info("[Gemini] Sending content for generation...")
response = client.models.generate_content(
model=model,
contents=[
types.Part.from_bytes(data=file_bytes, mime_type="application/pdf"),
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}
else:
logger.error(f"Unsupported API: {api}")
return None
except Exception as e:
logger.exception("Error during ESG data extraction.")
return None
def list_all_files():
"""Lists all files currently uploaded to OpenAI."""
try:
files = get_files()
for file in files:
logger.info(f"File ID: {file.id}, Name: {file.filename}, Size: {file.bytes} bytes")
except Exception as e:
logger.error(f"Failed to list files: {e}") |