Spaces:
Running
Running
Merge branch 'master'
Browse files- .gitignore +80 -0
- README.md +229 -0
- gemini_proxy.py +582 -0
- 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
|