bibibi12345 commited on
Commit
1ccec3d
·
2 Parent(s): f622962 49b8075

Merge branch 'master'

Browse files
Files changed (4) hide show
  1. .gitignore +80 -0
  2. README.md +229 -0
  3. gemini_proxy.py +582 -0
  4. requirements.txt +7 -0
.gitignore ADDED
@@ -0,0 +1,80 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Credential files - should never be committed
2
+ oauth_creds.json
3
+
4
+ # Environment configuration
5
+ .env
6
+
7
+ # Python
8
+ __pycache__/
9
+ *.py[cod]
10
+ *$py.class
11
+ *.so
12
+ .Python
13
+ build/
14
+ develop-eggs/
15
+ dist/
16
+ downloads/
17
+ eggs/
18
+ .eggs/
19
+ lib/
20
+ lib64/
21
+ parts/
22
+ sdist/
23
+ var/
24
+ wheels/
25
+ pip-wheel-metadata/
26
+ share/python-wheels/
27
+ *.egg-info/
28
+ .installed.cfg
29
+ *.egg
30
+ MANIFEST
31
+
32
+ # PyInstaller
33
+ *.manifest
34
+ *.spec
35
+
36
+ # Installer logs
37
+ pip-log.txt
38
+ pip-delete-this-directory.txt
39
+
40
+ # Unit test / coverage reports
41
+ htmlcov/
42
+ .tox/
43
+ .nox/
44
+ .coverage
45
+ .coverage.*
46
+ .cache
47
+ nosetests.xml
48
+ coverage.xml
49
+ *.cover
50
+ *.py,cover
51
+ .hypothesis/
52
+ .pytest_cache/
53
+
54
+ # Virtual environments
55
+ .env
56
+ .venv
57
+ env/
58
+ venv/
59
+ ENV/
60
+ env.bak/
61
+ venv.bak/
62
+
63
+ # IDE
64
+ .vscode/
65
+ .idea/
66
+ *.swp
67
+ *.swo
68
+ *~
69
+
70
+ # OS
71
+ .DS_Store
72
+ .DS_Store?
73
+ ._*
74
+ .Spotlight-V100
75
+ .Trashes
76
+ ehthumbs.db
77
+ Thumbs.db
78
+
79
+ # Logs
80
+ *.log
README.md ADDED
@@ -0,0 +1,229 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Gemini CLI to API Proxy
2
+
3
+ A proxy server that converts Google's Gemini CLI authentication to standard API format, allowing you to use Gemini models with any OpenAI-compatible client.
4
+
5
+ ## Features
6
+
7
+ - OAuth 2.0 authentication with Google Cloud
8
+ - Automatic project ID detection and caching
9
+ - Support for both streaming and non-streaming requests
10
+ - Converts Google's internal API format to standard Gemini API format
11
+ - Credential caching for seamless restarts
12
+
13
+ ## Prerequisites
14
+
15
+ - Python 3.7 or higher
16
+ - Google Cloud account with Gemini API access
17
+ - Required Python packages (see `requirements.txt`)
18
+
19
+ ## Installation
20
+
21
+ 1. Clone or download this repository
22
+ 2. Install the required dependencies:
23
+ ```bash
24
+ pip install -r requirements.txt
25
+ ```
26
+
27
+ ## Setup and Usage
28
+
29
+ ### First Time Setup
30
+
31
+ 1. **(Optional) Configure your settings:**
32
+ If you know your Google Cloud project ID, you can set it in the `.env` file to skip automatic detection:
33
+ ```bash
34
+ # Edit the .env file
35
+ GEMINI_PROJECT_ID=your-project-id
36
+
37
+ # Optional: Change the default port
38
+ GEMINI_PORT=8888
39
+ ```
40
+
41
+ 2. **Start the proxy server:**
42
+ ```bash
43
+ python gemini_proxy.py
44
+ ```
45
+
46
+ 3. **Authenticate with Google:**
47
+ - On first run, the proxy will display an authentication URL
48
+ - Open the URL in your browser and sign in with your Google account
49
+ - Grant the necessary permissions
50
+ - The browser will show "Authentication successful!" when complete
51
+ - The proxy will automatically save your credentials for future use
52
+
53
+ 4. **Project ID Detection:**
54
+ - If `GEMINI_PROJECT_ID` is set in the `.env` file, it will be used
55
+ - Otherwise, the proxy will automatically detect and cache your Google Cloud project ID
56
+ - This only happens once - subsequent runs will use the cached project ID
57
+
58
+ ### Regular Usage
59
+
60
+ After initial setup, simply run:
61
+ ```bash
62
+ python gemini_proxy.py
63
+ ```
64
+
65
+ The proxy server will start on `http://localhost:8888` and display:
66
+ ```
67
+ Starting Gemini proxy server on http://localhost:8888
68
+ Send your Gemini API requests to this address.
69
+ ```
70
+
71
+ ### Using with API Clients
72
+
73
+ Configure your Gemini API client to use `http://localhost:8888` as the base URL. The proxy accepts standard Gemini API requests and handles the authentication automatically.
74
+
75
+ Example request:
76
+ ```bash
77
+ curl -X POST "http://localhost:8888/v1/models/gemini-pro:generateContent?key=123456" \
78
+ -H "Content-Type: application/json" \
79
+ -d '{
80
+ "contents": [{
81
+ "parts": [{"text": "Hello, how are you?"}]
82
+ }]
83
+ }'
84
+ ```
85
+
86
+ **Note:** The proxy supports multiple authentication methods. The `key` query parameter is the most compatible with standard Gemini clients.
87
+
88
+ ### Safety Settings
89
+
90
+ The proxy automatically sets default safety settings to `BLOCK_NONE` for all categories if no safety settings are specified in the request. This provides maximum flexibility for content generation. The default categories are:
91
+
92
+ - `HARM_CATEGORY_HARASSMENT`
93
+ - `HARM_CATEGORY_HATE_SPEECH`
94
+ - `HARM_CATEGORY_SEXUALLY_EXPLICIT`
95
+ - `HARM_CATEGORY_DANGEROUS_CONTENT`
96
+
97
+ You can override these defaults by including your own `safetySettings` in the request payload.
98
+
99
+ ## Authentication
100
+
101
+ The proxy supports multiple authentication methods for maximum compatibility:
102
+
103
+ - **Default Password:** `123456`
104
+ - **Configuration:** Set `GEMINI_AUTH_PASSWORD` in `.env` file to change the password
105
+
106
+ ### Authentication Methods
107
+
108
+ 1. **API Key (Query Parameter)** - Compatible with standard Gemini clients:
109
+ ```bash
110
+ curl -X POST "http://localhost:8888/v1/models/gemini-pro:generateContent?key=123456" \
111
+ -H "Content-Type: application/json" \
112
+ -d '{"contents": [{"parts": [{"text": "Hello!"}]}]}'
113
+ ```
114
+
115
+ 2. **Bearer Token** - Standard API token format:
116
+ ```bash
117
+ curl -X POST http://localhost:8888/v1/models/gemini-pro:generateContent \
118
+ -H "Authorization: Bearer 123456" \
119
+ -H "Content-Type: application/json" \
120
+ -d '{"contents": [{"parts": [{"text": "Hello!"}]}]}'
121
+ ```
122
+
123
+ 3. **HTTP Basic Authentication** - Traditional username/password:
124
+ ```bash
125
+ curl -u "user:123456" -X POST http://localhost:8888/v1/models/gemini-pro:generateContent \
126
+ -H "Content-Type: application/json" \
127
+ -d '{"contents": [{"parts": [{"text": "Hello!"}]}]}'
128
+ ```
129
+
130
+ **Python examples:**
131
+ ```python
132
+ import requests
133
+ from requests.auth import HTTPBasicAuth
134
+
135
+ # Method 1: Query parameter
136
+ response = requests.post(
137
+ "http://localhost:8888/v1/models/gemini-pro:generateContent?key=123456",
138
+ json={"contents": [{"parts": [{"text": "Hello!"}]}]}
139
+ )
140
+
141
+ # Method 2: Bearer token
142
+ response = requests.post(
143
+ "http://localhost:8888/v1/models/gemini-pro:generateContent",
144
+ headers={"Authorization": "Bearer 123456"},
145
+ json={"contents": [{"parts": [{"text": "Hello!"}]}]}
146
+ )
147
+
148
+ # Method 3: Basic auth
149
+ response = requests.post(
150
+ "http://localhost:8888/v1/models/gemini-pro:generateContent",
151
+ auth=HTTPBasicAuth("user", "123456"),
152
+ json={"contents": [{"parts": [{"text": "Hello!"}]}]}
153
+ )
154
+ ```
155
+
156
+ ## Configuration
157
+
158
+ The proxy uses the following configuration:
159
+ - **Port:** 8888 (default, configurable via `.env`)
160
+ - **Credential file:** `oauth_creds.json` (automatically created)
161
+ - **Configuration file:** `.env` (optional settings)
162
+ - **Scopes:** Cloud Platform, User Info (email/profile), OpenID
163
+
164
+ ### Configuration File (.env)
165
+
166
+ You can create a `.env` file in the same directory as the script to configure the proxy:
167
+
168
+ ```bash
169
+ # Set your Google Cloud Project ID to skip automatic detection
170
+ GEMINI_PROJECT_ID=my-gcp-project-123
171
+
172
+ # Set a custom port (default is 8888)
173
+ GEMINI_PORT=9000
174
+
175
+ # Set authentication password (default is 123456)
176
+ GEMINI_AUTH_PASSWORD=your-secure-password
177
+ ```
178
+
179
+ **Note:** The `.env` file is automatically excluded from version control via `.gitignore`.
180
+
181
+ ## File Structure
182
+
183
+ - `gemini_proxy.py` - Main proxy server
184
+ - `oauth_creds.json` - Cached OAuth credentials and project ID (auto-generated)
185
+ - `requirements.txt` - Python dependencies
186
+ - `.env` - Configuration file (optional, create as needed)
187
+ - `.gitignore` - Prevents credential and config files from being committed
188
+
189
+ ## Troubleshooting
190
+
191
+ ### Port Already in Use
192
+ If you see "error while attempting to bind on address", another instance is already running. Stop the existing process or use a different port.
193
+
194
+ ### Authentication Issues (Google OAuth)
195
+ - Delete `oauth_creds.json` and restart to re-authenticate
196
+ - Ensure your Google account has access to Google Cloud and Gemini API
197
+ - Check that the required scopes are granted during authentication
198
+
199
+ ### Authentication Issues (Proxy Access)
200
+ - If you get 401 Unauthorized errors, check your authentication method
201
+ - Default password is `123456` unless changed in `.env` file
202
+ - Try different authentication methods:
203
+ - Query parameter: `?key=123456`
204
+ - Bearer token: `Authorization: Bearer 123456`
205
+ - Basic auth: `Authorization: Basic base64(user:123456)`
206
+ - Most Gemini clients work best with the query parameter method (`?key=password`)
207
+
208
+ ### Project ID Issues
209
+ - The proxy automatically detects your project ID on first run
210
+ - If detection fails, you can manually set `GEMINI_PROJECT_ID` in the `.env` file
211
+ - Check your Google Cloud project permissions if auto-detection fails
212
+ - Delete `oauth_creds.json` to force re-detection
213
+ - Project ID in `.env` file takes priority over cached and auto-detected project IDs
214
+
215
+ ## Security Notes
216
+
217
+ - **Never commit `oauth_creds.json`** - it contains sensitive authentication tokens
218
+ - The `.gitignore` file is configured to prevent accidental commits
219
+ - Credentials are stored locally and refreshed automatically when expired
220
+ - The proxy runs on localhost only for security
221
+
222
+ ## API Compatibility
223
+
224
+ This proxy converts between:
225
+ - **Input:** Standard Gemini API format
226
+ - **Output:** Standard Gemini API responses
227
+ - **Internal:** Google's Cloud Code Assist API format
228
+
229
+ The conversion is transparent to API clients.
gemini_proxy.py ADDED
@@ -0,0 +1,582 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import os
2
+ import json
3
+ import requests
4
+ import re
5
+ import uvicorn
6
+ import base64
7
+ from datetime import datetime
8
+ from fastapi import FastAPI, Request, Response, HTTPException, Depends
9
+ from fastapi.responses import StreamingResponse
10
+ from fastapi.security import HTTPBasic, HTTPBasicCredentials
11
+ from http.server import BaseHTTPRequestHandler, HTTPServer
12
+ from urllib.parse import urlparse, parse_qs
13
+ import ijson
14
+ from dotenv import load_dotenv
15
+
16
+ from google.oauth2.credentials import Credentials
17
+ from google_auth_oauthlib.flow import Flow
18
+ from google.auth.transport.requests import Request as GoogleAuthRequest
19
+
20
+ # Load environment variables from .env file
21
+ load_dotenv()
22
+
23
+ # --- Configuration ---
24
+ CLIENT_ID = "681255809395-oo8ft2oprdrnp9e3aqf6av3hmdib135j.apps.googleusercontent.com"
25
+ CLIENT_SECRET = "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"
26
+ SCOPES = [
27
+ "https://www.googleapis.com/auth/cloud-platform",
28
+ "https://www.googleapis.com/auth/userinfo.email",
29
+ "https://www.googleapis.com/auth/userinfo.profile",
30
+ "openid",
31
+ ]
32
+ GEMINI_DIR = os.path.dirname(os.path.abspath(__file__)) # Same directory as the script
33
+ CREDENTIAL_FILE = os.path.join(GEMINI_DIR, "oauth_creds.json")
34
+ CODE_ASSIST_ENDPOINT = "https://cloudcode-pa.googleapis.com"
35
+ GEMINI_PORT = int(os.getenv("GEMINI_PORT", "8888")) # Default to 8888 if not set
36
+ GEMINI_AUTH_PASSWORD = os.getenv("GEMINI_AUTH_PASSWORD", "123456") # Default password
37
+
38
+ # --- Global State ---
39
+ credentials = None
40
+ user_project_id = None
41
+
42
+ app = FastAPI()
43
+ security = HTTPBasic()
44
+
45
+ def authenticate_user(request: Request):
46
+ """Authenticate the user with multiple methods."""
47
+ # Check for API key in query parameters first (for Gemini client compatibility)
48
+ api_key = request.query_params.get("key")
49
+ if api_key and api_key == GEMINI_AUTH_PASSWORD:
50
+ return "api_key_user"
51
+
52
+ # Check for API key in Authorization header (Bearer token format)
53
+ auth_header = request.headers.get("authorization", "")
54
+ if auth_header.startswith("Bearer ") and auth_header[7:] == GEMINI_AUTH_PASSWORD:
55
+ return "bearer_user"
56
+
57
+ # Check for HTTP Basic Authentication
58
+ if auth_header.startswith("Basic "):
59
+ try:
60
+ encoded_credentials = auth_header[6:]
61
+ decoded_credentials = base64.b64decode(encoded_credentials).decode('utf-8')
62
+ username, password = decoded_credentials.split(':', 1)
63
+ if password == GEMINI_AUTH_PASSWORD:
64
+ return username
65
+ except Exception:
66
+ pass
67
+
68
+ # If none of the authentication methods work
69
+ raise HTTPException(
70
+ status_code=401,
71
+ detail="Invalid authentication credentials. Use HTTP Basic Auth, Bearer token, or 'key' query parameter.",
72
+ headers={"WWW-Authenticate": "Basic"},
73
+ )
74
+
75
+ # Helper class to adapt a generator of bytes into a file-like object
76
+ # that ijson can read from.
77
+ class _GeneratorStream:
78
+ def __init__(self, generator):
79
+ self.generator = generator
80
+ self.buffer = b''
81
+
82
+ def read(self, size=-1):
83
+ # This read implementation is crucial for streaming.
84
+ # It must not block to read the entire stream if size is -1.
85
+ if size == -1:
86
+ # If asked to read all, return what's in the buffer and get one more chunk.
87
+ try:
88
+ self.buffer += next(self.generator)
89
+ except StopIteration:
90
+ pass
91
+ data = self.buffer
92
+ self.buffer = b''
93
+ return data
94
+
95
+ # Otherwise, read from the generator until we have enough bytes.
96
+ while len(self.buffer) < size:
97
+ try:
98
+ self.buffer += next(self.generator)
99
+ except StopIteration:
100
+ # Generator is exhausted.
101
+ break
102
+
103
+ data = self.buffer[:size]
104
+ self.buffer = self.buffer[size:]
105
+ return data
106
+
107
+ class _OAuthCallbackHandler(BaseHTTPRequestHandler):
108
+ auth_code = None
109
+ def do_GET(self):
110
+ query_components = parse_qs(urlparse(self.path).query)
111
+ code = query_components.get("code", [None])[0]
112
+ if code:
113
+ _OAuthCallbackHandler.auth_code = code
114
+ self.send_response(200)
115
+ self.send_header("Content-type", "text/html")
116
+ self.end_headers()
117
+ self.wfile.write(b"<h1>Authentication successful!</h1><p>You can close this window and restart the proxy.</p>")
118
+ else:
119
+ self.send_response(400)
120
+ self.send_header("Content-type", "text/html")
121
+ self.end_headers()
122
+ self.wfile.write(b"<h1>Authentication failed.</h1><p>Please try again.</p>")
123
+
124
+ def get_user_project_id(creds):
125
+ """Gets the user's project ID from environment variable, cache, or by probing the API."""
126
+ global user_project_id
127
+ if user_project_id:
128
+ return user_project_id
129
+
130
+ # First, check for environment variable override
131
+ env_project_id = os.getenv("GEMINI_PROJECT_ID")
132
+ if env_project_id:
133
+ user_project_id = env_project_id
134
+ print(f"Using project ID from environment variable: {user_project_id}")
135
+ # Save the environment project ID to cache for consistency
136
+ save_credentials(creds, user_project_id)
137
+ return user_project_id
138
+
139
+ # Second, try to load project ID from credential file
140
+ if os.path.exists(CREDENTIAL_FILE):
141
+ try:
142
+ with open(CREDENTIAL_FILE, "r") as f:
143
+ creds_data = json.load(f)
144
+ cached_project_id = creds_data.get("project_id")
145
+ if cached_project_id:
146
+ user_project_id = cached_project_id
147
+ print(f"Loaded project ID from cache: {user_project_id}")
148
+ return user_project_id
149
+ except Exception as e:
150
+ print(f"Could not load project ID from cache: {e}")
151
+
152
+ # If not found in environment or cache, probe for it
153
+ print("Project ID not found in environment or cache. Probing for user project ID...")
154
+ headers = {
155
+ "Authorization": f"Bearer {creds.token}",
156
+ "Content-Type": "application/json",
157
+ }
158
+
159
+ probe_payload = {
160
+ "cloudaicompanionProject": "gcp-project",
161
+ "metadata": {
162
+ "ideType": "VSCODE",
163
+ "pluginType": "GEMINI"
164
+ }
165
+ }
166
+
167
+ try:
168
+ resp = requests.post(
169
+ f"{CODE_ASSIST_ENDPOINT}/v1internal:loadCodeAssist",
170
+ data=json.dumps(probe_payload),
171
+ headers=headers,
172
+ )
173
+ resp.raise_for_status()
174
+ data = resp.json()
175
+ user_project_id = data.get("cloudaicompanionProject")
176
+ if not user_project_id:
177
+ raise ValueError("Could not find 'cloudaicompanionProject' in loadCodeAssist response.")
178
+ print(f"Successfully fetched user project ID: {user_project_id}")
179
+
180
+ # Save the project ID to the credential file for future use
181
+ save_credentials(creds, user_project_id)
182
+ print("Project ID saved to credential file for future use.")
183
+
184
+ return user_project_id
185
+ except requests.exceptions.HTTPError as e:
186
+ print(f"Error fetching project ID: {e.response.text}")
187
+ raise
188
+
189
+ def save_credentials(creds, project_id=None):
190
+ os.makedirs(GEMINI_DIR, exist_ok=True)
191
+ creds_data = {
192
+ "access_token": creds.token,
193
+ "refresh_token": creds.refresh_token,
194
+ "scope": " ".join(creds.scopes),
195
+ "token_type": "Bearer",
196
+ "expiry_date": creds.expiry.isoformat() if creds.expiry else None,
197
+ }
198
+
199
+ # If project_id is provided, save it; otherwise preserve existing project_id
200
+ if project_id:
201
+ creds_data["project_id"] = project_id
202
+ elif os.path.exists(CREDENTIAL_FILE):
203
+ try:
204
+ with open(CREDENTIAL_FILE, "r") as f:
205
+ existing_data = json.load(f)
206
+ if "project_id" in existing_data:
207
+ creds_data["project_id"] = existing_data["project_id"]
208
+ except Exception:
209
+ pass # If we can't read existing file, just continue without project_id
210
+
211
+ with open(CREDENTIAL_FILE, "w") as f:
212
+ json.dump(creds_data, f)
213
+
214
+ def get_credentials():
215
+ """Loads credentials from cache or initiates the OAuth 2.0 flow."""
216
+ global credentials
217
+
218
+ if credentials:
219
+ if credentials.valid:
220
+ return credentials
221
+ if credentials.expired and credentials.refresh_token:
222
+ print("Refreshing expired credentials...")
223
+ try:
224
+ credentials.refresh(GoogleAuthRequest())
225
+ save_credentials(credentials)
226
+ print("Credentials refreshed successfully.")
227
+ return credentials
228
+ except Exception as e:
229
+ print(f"Could not refresh token: {e}. Attempting to load from file.")
230
+
231
+ if os.path.exists(CREDENTIAL_FILE):
232
+ try:
233
+ with open(CREDENTIAL_FILE, "r") as f:
234
+ creds_data = json.load(f)
235
+
236
+ # Load project ID if available
237
+ global user_project_id
238
+ cached_project_id = creds_data.get("project_id")
239
+ if cached_project_id:
240
+ user_project_id = cached_project_id
241
+ print(f"Loaded project ID from credential file: {user_project_id}")
242
+
243
+ expiry = None
244
+ expiry_str = creds_data.get("expiry_date")
245
+ if expiry_str:
246
+ if not isinstance(expiry_str, str) or not expiry_str.strip():
247
+ expiry = None
248
+ elif expiry_str.endswith('Z'):
249
+ expiry_str = expiry_str[:-1] + '+00:00'
250
+ expiry = datetime.fromisoformat(expiry_str)
251
+ else:
252
+ expiry = datetime.fromisoformat(expiry_str)
253
+
254
+ credentials = Credentials(
255
+ token=creds_data.get("access_token"),
256
+ refresh_token=creds_data.get("refresh_token"),
257
+ token_uri="https://oauth2.googleapis.com/token",
258
+ client_id=CLIENT_ID,
259
+ client_secret=CLIENT_SECRET,
260
+ scopes=SCOPES,
261
+ expiry=expiry
262
+ )
263
+
264
+ if credentials.expired and credentials.refresh_token:
265
+ print("Loaded credentials from file are expired. Refreshing...")
266
+ credentials.refresh(GoogleAuthRequest())
267
+ save_credentials(credentials)
268
+
269
+ print("Successfully loaded credentials from cache.")
270
+ return credentials
271
+ except Exception as e:
272
+ print(f"Could not load cached credentials: {e}. Starting new login.")
273
+
274
+ client_config = {
275
+ "installed": {
276
+ "client_id": CLIENT_ID,
277
+ "client_secret": CLIENT_SECRET,
278
+ "auth_uri": "https://accounts.google.com/o/oauth2/auth",
279
+ "token_uri": "https://oauth2.googleapis.com/token",
280
+ }
281
+ }
282
+ flow = Flow.from_client_config(
283
+ client_config, scopes=SCOPES, redirect_uri="http://localhost:8080"
284
+ )
285
+ auth_url, _ = flow.authorization_url(access_type="offline", prompt="consent")
286
+ print(f"\nPlease open this URL in your browser to log in:\n{auth_url}\n")
287
+
288
+ server = HTTPServer(("", 8080), _OAuthCallbackHandler)
289
+ server.handle_request()
290
+
291
+ auth_code = _OAuthCallbackHandler.auth_code
292
+ if not auth_code:
293
+ print("Failed to retrieve authorization code.")
294
+ return None
295
+
296
+ flow.fetch_token(code=auth_code)
297
+ credentials = flow.credentials
298
+ save_credentials(credentials)
299
+ print("Authentication successful! Credentials saved.")
300
+ return credentials
301
+
302
+
303
+ @app.post("/{full_path:path}")
304
+ async def proxy_request(request: Request, full_path: str, username: str = Depends(authenticate_user)):
305
+ creds = get_credentials()
306
+ if not creds:
307
+ return Response(content="Authentication failed. Please restart the proxy to log in.", status_code=500)
308
+
309
+ # Only check if credentials are properly formed and not expired locally
310
+ if not creds.valid:
311
+ if creds.expired and creds.refresh_token:
312
+ print("Credentials expired locally. Refreshing...")
313
+ try:
314
+ creds.refresh(GoogleAuthRequest())
315
+ save_credentials(creds)
316
+ print("Credentials refreshed successfully.")
317
+ except Exception as e:
318
+ print(f"Could not refresh token during request: {e}")
319
+ return Response(content="Token refresh failed. Please restart the proxy to re-authenticate.", status_code=500)
320
+ else:
321
+ print("Credentials are invalid locally and cannot be refreshed.")
322
+ return Response(content="Invalid credentials. Please restart the proxy to re-authenticate.", status_code=500)
323
+
324
+ proj_id = get_user_project_id(creds)
325
+ if not proj_id:
326
+ return Response(content="Failed to get user project ID.", status_code=500)
327
+
328
+ post_data = await request.body()
329
+ path = f"/{full_path}"
330
+ model_name_from_url = None
331
+ action = None
332
+
333
+ model_match = re.match(r"/(v\d+(?:beta)?)/models/([^:]+):(\w+)", path)
334
+
335
+ is_streaming = False
336
+ if model_match:
337
+ model_name_from_url = model_match.group(2)
338
+ action = model_match.group(3)
339
+ target_url = f"{CODE_ASSIST_ENDPOINT}/v1internal:{action}"
340
+ if "stream" in action.lower():
341
+ is_streaming = True
342
+ else:
343
+ target_url = f"{CODE_ASSIST_ENDPOINT}{path}"
344
+
345
+ # Remove authentication query parameters before forwarding to Google API
346
+ query_params = dict(request.query_params)
347
+ # Remove our authentication parameters
348
+ query_params.pop("key", None)
349
+
350
+ # Add remaining query parameters to target URL if any
351
+ if query_params:
352
+ from urllib.parse import urlencode
353
+ target_url += "?" + urlencode(query_params)
354
+
355
+ try:
356
+ incoming_json = json.loads(post_data)
357
+ final_model = model_name_from_url if model_match else incoming_json.get("model")
358
+
359
+ # Set default safety settings to BLOCK_NONE if not specified by user
360
+ safety_settings = incoming_json.get("safetySettings")
361
+ if not safety_settings:
362
+ safety_settings = [
363
+ {"category": "HARM_CATEGORY_HARASSMENT", "threshold": "BLOCK_NONE"},
364
+ {"category": "HARM_CATEGORY_HATE_SPEECH", "threshold": "BLOCK_NONE"},
365
+ {"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT", "threshold": "BLOCK_NONE"},
366
+ {"category": "HARM_CATEGORY_DANGEROUS_CONTENT", "threshold": "BLOCK_NONE"},
367
+ {"category": "HARM_CATEGORY_CIVIC_INTEGRITY", "threshold": "BLOCK_NONE"}
368
+ ]
369
+
370
+ structured_payload = {
371
+ "model": final_model,
372
+ "project": proj_id,
373
+ "request": {
374
+ "contents": incoming_json.get("contents"),
375
+ "systemInstruction": incoming_json.get("systemInstruction"),
376
+ "cachedContent": incoming_json.get("cachedContent"),
377
+ "tools": incoming_json.get("tools"),
378
+ "toolConfig": incoming_json.get("toolConfig"),
379
+ "safetySettings": safety_settings,
380
+ "generationConfig": incoming_json.get("generationConfig"),
381
+ },
382
+ }
383
+ structured_payload["request"] = {
384
+ k: v
385
+ for k, v in structured_payload["request"].items()
386
+ if v is not None
387
+ }
388
+ final_post_data = json.dumps(structured_payload)
389
+ except (json.JSONDecodeError, AttributeError):
390
+ final_post_data = post_data
391
+
392
+ headers = {
393
+ "Authorization": f"Bearer {creds.token}",
394
+ "Content-Type": "application/json",
395
+ # We remove 'Accept-Encoding' to allow the server to send gzip,
396
+ # which it seems to stream correctly. We will decompress on the fly.
397
+ }
398
+
399
+ if is_streaming:
400
+ async def stream_generator():
401
+ try:
402
+ print(f"[STREAM] Starting streaming request to: {target_url}")
403
+ print(f"[STREAM] Request payload size: {len(final_post_data)} bytes")
404
+
405
+ # Make the initial streaming request
406
+ resp = requests.post(target_url, data=final_post_data, headers=headers, stream=True)
407
+ print(f"[STREAM] Response status: {resp.status_code}")
408
+ print(f"[STREAM] Response headers: {dict(resp.headers)}")
409
+
410
+ # If we get a 401, try refreshing the token and retry once
411
+ if resp.status_code == 401 and creds.refresh_token:
412
+ print("[STREAM] Received 401 from Google API. Attempting to refresh token and retry...")
413
+ resp.close() # Close the failed response
414
+ try:
415
+ creds.refresh(GoogleAuthRequest())
416
+ save_credentials(creds)
417
+ print("[STREAM] Token refreshed successfully. Retrying streaming request...")
418
+
419
+ # Update headers with new token and retry
420
+ headers["Authorization"] = f"Bearer {creds.token}"
421
+ resp = requests.post(target_url, data=final_post_data, headers=headers, stream=True)
422
+ print(f"[STREAM] Retry request status: {resp.status_code}")
423
+ except Exception as e:
424
+ print(f"[STREAM] Could not refresh token after 401 error: {e}")
425
+ error_message = json.dumps({"error": {"message": "Token refresh failed after 401 error. Please restart the proxy to re-authenticate."}})
426
+ yield f"data: {error_message}\n\n"
427
+ return
428
+
429
+ with resp:
430
+ resp.raise_for_status()
431
+
432
+ buffer = ""
433
+ brace_count = 0
434
+ in_array = False
435
+ chunk_count = 0
436
+ total_bytes = 0
437
+ objects_yielded = 0
438
+
439
+ print(f"[STREAM] Starting to process chunks...")
440
+
441
+ for chunk in resp.iter_content(chunk_size=1024):
442
+ if isinstance(chunk, bytes):
443
+ chunk = chunk.decode('utf-8', errors='replace')
444
+ chunk_count += 1
445
+ chunk_size = len(chunk) if chunk else 0
446
+ total_bytes += chunk_size
447
+
448
+ buffer += chunk
449
+
450
+ # Process complete JSON objects from the buffer
451
+ processing_iterations = 0
452
+ while buffer:
453
+ processing_iterations += 1
454
+ if processing_iterations > 100: # Prevent infinite loops
455
+ break
456
+
457
+ buffer = buffer.lstrip()
458
+
459
+ if not buffer:
460
+ break
461
+
462
+ # Handle array start
463
+ if buffer.startswith('[') and not in_array:
464
+ buffer = buffer[1:].lstrip()
465
+ in_array = True
466
+ continue
467
+
468
+ # Handle array end
469
+ if buffer.startswith(']'):
470
+ break
471
+
472
+ # Skip commas between objects
473
+ if buffer.startswith(','):
474
+ buffer = buffer[1:].lstrip()
475
+ continue
476
+
477
+ # Look for complete JSON objects
478
+ if buffer.startswith('{'):
479
+ brace_count = 0
480
+ in_string = False
481
+ escape_next = False
482
+ end_pos = -1
483
+
484
+ for i, char in enumerate(buffer):
485
+ if escape_next:
486
+ escape_next = False
487
+ continue
488
+ if char == '\\':
489
+ escape_next = True
490
+ continue
491
+ if char == '"' and not escape_next:
492
+ in_string = not in_string
493
+ continue
494
+ if not in_string:
495
+ if char == '{':
496
+ brace_count += 1
497
+ elif char == '}':
498
+ brace_count -= 1
499
+ if brace_count == 0:
500
+ end_pos = i + 1
501
+ break
502
+
503
+ if end_pos > 0:
504
+ # Found complete JSON object
505
+ json_str = buffer[:end_pos]
506
+ buffer = buffer[end_pos:].lstrip()
507
+
508
+
509
+ try:
510
+ obj = json.loads(json_str)
511
+
512
+ if "response" in obj:
513
+ response_chunk = obj["response"]
514
+ objects_yielded += 1
515
+ response_json = json.dumps(response_chunk)
516
+ yield f"data: {response_json}\n\n"
517
+ except json.JSONDecodeError as e:
518
+ continue
519
+ else:
520
+ # Incomplete object, wait for more data
521
+ break
522
+ else:
523
+ # Skip unexpected characters
524
+ buffer = buffer[1:]
525
+
526
+ except requests.exceptions.RequestException as e:
527
+ print(f"Error during streaming request: {e}")
528
+ error_message = json.dumps({"error": {"message": f"Upstream request failed: {e}"}})
529
+ yield f"data: {error_message}\n\n"
530
+ except Exception as e:
531
+ print(f"An unexpected error occurred during streaming: {e}")
532
+ error_message = json.dumps({"error": {"message": f"An unexpected error occurred: {e}"}})
533
+ yield f"data: {error_message}\n\n"
534
+
535
+ return StreamingResponse(stream_generator(), media_type="text/event-stream")
536
+ else:
537
+ # Make the request
538
+ resp = requests.post(target_url, data=final_post_data, headers=headers)
539
+
540
+ # If we get a 401, try refreshing the token and retry once
541
+ if resp.status_code == 401 and creds.refresh_token:
542
+ print("Received 401 from Google API. Attempting to refresh token and retry...")
543
+ try:
544
+ creds.refresh(GoogleAuthRequest())
545
+ save_credentials(creds)
546
+ print("Token refreshed successfully. Retrying request...")
547
+
548
+ # Update headers with new token and retry
549
+ headers["Authorization"] = f"Bearer {creds.token}"
550
+ resp = requests.post(target_url, data=final_post_data, headers=headers)
551
+ print(f"Retry request status: {resp.status_code}")
552
+ except Exception as e:
553
+ print(f"Could not refresh token after 401 error: {e}")
554
+ return Response(content="Token refresh failed after 401 error. Please restart the proxy to re-authenticate.", status_code=500)
555
+
556
+ if resp.status_code == 200:
557
+ try:
558
+ google_api_response = resp.json()
559
+ # The actual response is nested under the "response" key
560
+ standard_gemini_response = google_api_response.get("response")
561
+ # The standard client expects a list containing the response object
562
+ return Response(content=json.dumps([standard_gemini_response]), status_code=200, media_type="application/json")
563
+ except (json.JSONDecodeError, AttributeError) as e:
564
+ print(f"Error converting to standard Gemini format: {e}")
565
+ # Fallback to sending the original content if conversion fails
566
+ return Response(content=resp.content, status_code=resp.status_code, media_type=resp.headers.get("Content-Type"))
567
+ else:
568
+ return Response(content=resp.content, status_code=resp.status_code, media_type=resp.headers.get("Content-Type"))
569
+
570
+
571
+ if __name__ == "__main__":
572
+ print("Initializing credentials...")
573
+ creds = get_credentials()
574
+ if creds:
575
+ get_user_project_id(creds)
576
+ print(f"\nStarting Gemini proxy server on http://localhost:{GEMINI_PORT}")
577
+ print("Send your Gemini API requests to this address.")
578
+ print(f"Authentication required - Password: {GEMINI_AUTH_PASSWORD}")
579
+ print("Use HTTP Basic Authentication with any username and the password above.")
580
+ uvicorn.run(app, host="0.0.0.0", port=GEMINI_PORT)
581
+ else:
582
+ print("\nCould not obtain credentials. Please authenticate and restart the server.")
requirements.txt ADDED
@@ -0,0 +1,7 @@
 
 
 
 
 
 
 
 
1
+ fastapi
2
+ uvicorn
3
+ google-auth
4
+ google-auth-oauthlib
5
+ requests
6
+ ijson
7
+ python-dotenv