Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

More fixes to proxy #1237

Merged
merged 2 commits into from
Dec 12, 2024
Merged

More fixes to proxy #1237

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
d72586f
More fixes to proxy
aaronvg Dec 12, 2024
e38f259
fix lint
aaronvg Dec 12, 2024
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
22 changes: 11 additions & 11 deletions engine/baml-runtime/src/internal/llm_client/traits/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -635,9 +635,10 @@ async fn to_base64_with_inferred_mime_type(
Ok((base64, mime_type))
} else {
Err(anyhow::anyhow!(
"Failed to fetch media: {}, {}",
"Failed to fetch media: {} {}, {}",
response.status(),
response.text().await.unwrap_or_default()
media_url.url,
response.text().await.unwrap_or_default(),
))
}
}
Expand All @@ -664,15 +665,14 @@ async fn fetch_with_proxy(
let client = reqwest::Client::new();

let request = if let Some(proxy) = proxy_url {
client
.get(format!(
"{}{}",
proxy,
url.parse::<url::Url>()
.map_err(|e| anyhow::anyhow!("Failed to parse URL: {}", e))?
.path()
))
.header("baml-original-url", url)
let new_proxy_url = format!(
"{}{}",
proxy,
url.parse::<url::Url>()
.map_err(|e| anyhow::anyhow!("Failed to parse URL: {}", e))?
.path()
);
client.get(new_proxy_url).header("baml-original-url", url)
} else {
client.get(url)
};
Expand Down
60 changes: 30 additions & 30 deletions typescript/vscode-ext/packages/vscode/src/extension.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,3 +1,4 @@
/* eslint-disable @typescript-eslint/no-misused-promises */
import * as vscode from 'vscode'
import axios from 'axios'
import glooLens from './LanguageToBamlCodeLensProvider'
Expand Down Expand Up @@ -193,11 +194,21 @@ export function activate(context: vscode.ExtensionContext) {
createProxyMiddleware({
changeOrigin: true,
pathRewrite: (path, req) => {
// Ensure the URL does not end with a slash
console.log('reqmethod', req.method)

// Remove the path in the case of images. Since we request things differently for image GET requests, where we add the url to localhost:4500/actual-url.png
// to prevent caching issues with Rust reqwest.
// But for normal completion POST requests, we always call localhost:4500.
// The original url is always in baml-original-url header.

// Check for file extensions and set path to empty string.
if (/\.[a-zA-Z0-9]+$/.test(path) && req.method === 'GET') {
return ''
}
// Remove trailing slash
if (path.endsWith('/')) {
return path.slice(0, -1)
}
console.log('pathRewrite', path, req)
return path
},
router: (req) => {
Expand All @@ -209,11 +220,14 @@ export function activate(context: vscode.ExtensionContext) {
delete req.headers['origin']

// Ensure the URL does not end with a slash
console.log('originalUrl1', originalUrl)
if (originalUrl.endsWith('/')) {
originalUrl = originalUrl.slice(0, -1)
}
console.log('returning original url', originalUrl)
return new URL(originalUrl).origin
// return new URL(originalUrl).toString()

return originalUrl
} else {
console.log('baml-original-url header is missing or invalid')
throw new Error('baml-original-url header is missing or invalid')
Expand All @@ -222,33 +236,19 @@ export function activate(context: vscode.ExtensionContext) {
logger: console,
on: {
proxyReq: (proxyReq, req, res) => {
console.log('proxying request')

try {
const bamlOriginalUrl = req.headers['baml-original-url']
if (bamlOriginalUrl === undefined) {
return
}
const targetUrl = new URL(bamlOriginalUrl)
// proxyReq.path = targetUrl.pathname
// proxyReq.p
// It is very important that we ONLY resolve against API_KEY_INJECTION_ALLOWED
// by using the URL origin! (i.e. NOT using str.startsWith - the latter can still
// leak API keys to malicious subdomains e.g. https://api.openai.com.evil.com)
// const headers = API_KEY_INJECTION_ALLOWED[proxyOrigin]
// if (headers === undefined) {
// return
// }
// for (const [header, value] of Object.entries(headers)) {
// proxyReq.setHeader(header, value)
// }
// proxyReq.removeHeader('origin')
// proxyReq.setHeader('Origin', targetUrl.origin)
console.info('Proxying an LLM request (to bypass CORS)', { proxyReq, req, res })
} catch (err) {
// This is not console.warn because it's not important
console.log('baml-original-url is not parsable', err)
}
// const bamlOriginalUrl = req.headers['baml-original-url']
// if (typeof bamlOriginalUrl === 'string') {
// const targetUrl = new URL(bamlOriginalUrl)
// // Copy all original headers except those we want to modify/remove
// Object.entries(req.headers).forEach(([key, value]) => {
// if (key !== 'host' && key !== 'origin' && key !== 'baml-original-url') {
// proxyReq.setHeader(key, value)
// }
// })
// // Set the correct origin and host headers
// proxyReq.setHeader('origin', targetUrl.origin)
// proxyReq.setHeader('host', targetUrl.host)
// }
},
proxyRes: (proxyRes, req, res) => {
proxyRes.headers['Access-Control-Allow-Origin'] = '*'
Expand Down
Loading