Skip to content
Open
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
56 changes: 50 additions & 6 deletions backend/main.py
Original file line number Diff line number Diff line change
Expand Up @@ -30,6 +30,9 @@

app = FastAPI(title="PDDL RLHF API", version="1.0.0")

# Global variable to track temporary credentials file for cleanup
_temp_credentials_file = None

# CORS configuration for Vercel frontend
app.add_middleware(
CORSMiddleware,
Expand All @@ -39,6 +42,17 @@
allow_headers=["*"],
)

def cleanup_temp_credentials():
"""Clean up temporary credentials file if it exists."""
global _temp_credentials_file
if _temp_credentials_file and os.path.exists(_temp_credentials_file):
try:
os.remove(_temp_credentials_file)
logger.info(f"✅ Cleaned up temporary credentials file: {_temp_credentials_file}")
_temp_credentials_file = None
except Exception as e:
logger.error(f"❌ Failed to clean up temporary credentials file: {str(e)}")

@app.on_event("startup")
async def startup_event():
"""Initialize the application on startup."""
Expand All @@ -53,6 +67,13 @@ async def startup_event():
logger.info("✅ Startup complete!")
logger.info("=" * 60)

@app.on_event("shutdown")
async def shutdown_event():
"""Clean up resources on application shutdown."""
logger.info("🔄 Shutting down PDDL RLHF API...")
cleanup_temp_credentials()
logger.info("✅ Shutdown complete!")

# Configuration
PROJECT_ID = os.getenv("GOOGLE_CLOUD_PROJECT", "deep-research-467303")
LOCATION = os.getenv("GOOGLE_CLOUD_LOCATION", "us-central1")
Expand All @@ -62,6 +83,7 @@ async def startup_event():
credentials_json = os.getenv("GOOGLE_APPLICATION_CREDENTIALS_JSON")
if credentials_json:
import tempfile
import atexit
try:
# Clean and extract JSON
credentials_json = credentials_json.strip()
Expand All @@ -70,17 +92,39 @@ async def startup_event():
if start_idx != -1 and end_idx != -1:
credentials_json = credentials_json[start_idx:end_idx+1]

# Parse and write to file
# Parse and write to file with restrictive permissions
credentials_dict = json.loads(credentials_json)
credentials_file = tempfile.NamedTemporaryFile(mode='w', delete=False, suffix='.json')
json.dump(credentials_dict, credentials_file, indent=2)
credentials_file.flush()
credentials_file.close()

os.environ["GOOGLE_APPLICATION_CREDENTIALS"] = credentials_file.name
logger.info(f"✅ Google Cloud credentials loaded from environment variable")
# Store file path globally for cleanup
_temp_credentials_file = credentials_file.name

try:
# Write credentials to file
json.dump(credentials_dict, credentials_file, indent=2)
credentials_file.flush()
credentials_file.close()

# Set restrictive file permissions (owner read/write only - 600)
os.chmod(_temp_credentials_file, 0o600)

os.environ["GOOGLE_APPLICATION_CREDENTIALS"] = _temp_credentials_file
logger.info(f"✅ Google Cloud credentials loaded from environment variable")

# Register cleanup handler for normal Python exit
atexit.register(cleanup_temp_credentials)

except Exception as e:
# Clean up on error
credentials_file.close()
if os.path.exists(_temp_credentials_file):
os.remove(_temp_credentials_file)
_temp_credentials_file = None
raise

except Exception as e:
logger.error(f"❌ Failed to setup credentials: {str(e)}")
_temp_credentials_file = None

# Initialize Google Vertex AI client (will be initialized on first use if credentials are available)
genai_client = None
Expand Down