Skip to content

feat: Integrate ONNX model via API endpoint #311

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

Open
wants to merge 6 commits into
base: main
Choose a base branch
from
Open

feat: Integrate ONNX model via API endpoint #311

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
384b2b3
feat: Integrate ONNX model via API endpoint
google-labs-jules[bot] Oct 2, 2025
ebc1982
feat: Add toggle to select ONNX model for chat
google-labs-jules[bot] Oct 3, 2025
30ac4aa
feat: Add settings toggle for ONNX model
google-labs-jules[bot] Oct 3, 2025
6d849f6
fix: Correct syntax error in model selection form
google-labs-jules[bot] Oct 3, 2025
0d86125
feat: Add settings toggle for ONNX model
google-labs-jules[bot] Oct 3, 2025
fe63b3f
feat: Integrate ONNX model with settings toggle and improved API
google-labs-jules[bot] Oct 3, 2025
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
90 changes: 90 additions & 0 deletions app/api/onnx/route.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,90 @@
import { NextResponse } from 'next/server';

// Simulate sending an error to a monitoring service (e.g., Sentry, DataDog)
const reportErrorToMonitoringService = (error: Error, context: object) => {
// In a real application, this would send the error to a service like Sentry.
console.log("Reporting error to monitoring service:", {
errorMessage: error.message,
stack: error.stack,
context,
});
};
Comment on lines +3 to +11
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🧹 Nitpick | 🔵 Trivial

Strengthen type safety for the monitoring helper.

The reportErrorToMonitoringService function accepts context: object, which is too permissive and loses type information. Consider defining a specific interface for the context to ensure consistent structure and improve maintainability.

Apply this diff to add type safety:

+interface ErrorContext {
+  query?: string | null;
+  fileName?: string;
+  fileSize?: number;
+  [key: string]: any;
+}
+
-const reportErrorToMonitoringService = (error: Error, context: object) => {
+const reportErrorToMonitoringService = (error: Error, context: ErrorContext) => {
   // In a real application, this would send the error to a service like Sentry.
   console.log("Reporting error to monitoring service:", {
     errorMessage: error.message,
     stack: error.stack,
     context,
   });
 };
📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
// Simulate sending an error to a monitoring service (e.g., Sentry, DataDog)
const reportErrorToMonitoringService = (error: Error, context: object) => {
// In a real application, this would send the error to a service like Sentry.
console.log("Reporting error to monitoring service:", {
errorMessage: error.message,
stack: error.stack,
context,
});
};
// Define a structured context type for error reporting
interface ErrorContext {
query?: string | null;
fileName?: string;
fileSize?: number;
[key: string]: any;
}
// Simulate sending an error to a monitoring service (e.g., Sentry, DataDog)
const reportErrorToMonitoringService = (error: Error, context: ErrorContext) => {
// In a real application, this would send the error to a service like Sentry.
console.log("Reporting error to monitoring service:", {
errorMessage: error.message,
stack: error.stack,
context,
});
};
🤖 Prompt for AI Agents 
In app/api/onnx/route.ts around lines 3 to 11, the monitoring helper currently
types context as a generic object; define a specific interface (e.g.,
MonitoringContext) with expected fields such as route?: string, requestId?:
string, userId?: string, and metadata?: Record<string, unknown> (adjust fields
to match real usage), update the function signature to
reportErrorToMonitoringService(error: Error, context: MonitoringContext),
replace loose object usages with the typed interface, and export the interface
if it will be reused elsewhere to enforce consistent structure and improve
maintainability.


export async function POST(request: Request) {
// 1. Input Validation for FormData
let formData;
try {
formData = await request.formData();
} catch (error) {
return NextResponse.json({ error: 'Invalid FormData body' }, { status: 400 });
}

const query = formData.get('query') as string | null;
const file = formData.get('file') as File | null;

if (!query && !file) {
return NextResponse.json({ error: 'Request must contain a non-empty "query" string or a "file"' }, { status: 400 });
}

if (query && (typeof query !== 'string' || query.trim() === '')) {
return NextResponse.json({ error: 'Field "query" must be a non-empty string if provided' }, { status: 400 });
}

// 2. Use Environment Variable for Endpoint and Provide Mock Fallback
const azureVmEndpoint = process.env.AZURE_ONNX_ENDPOINT;

if (!azureVmEndpoint) {
console.warn('AZURE_ONNX_ENDPOINT is not configured. Falling back to mock response.');
const mockResponse = {
prediction: "This is a simulated response because the Azure endpoint is not configured.",
confidence: 0.98,
input: {
query: query,
fileName: file?.name,
fileSize: file?.size,
}
};
return NextResponse.json(mockResponse);
}

// 3. Forward Request and Handle Errors
try {
// Pass the FormData directly to the fetch call.
// The browser will automatically set the 'Content-Type' to 'multipart/form-data' with the correct boundary.
const response = await fetch(azureVmEndpoint, {
method: 'POST',
body: formData,
});
Comment on lines +50 to +57
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🧹 Nitpick | 🔵 Trivial

Correct the misleading comment about Content-Type.

The comment states "The browser will automatically set the 'Content-Type' to 'multipart/form-data' with the correct boundary," but this code runs server-side in a Next.js API route, not in a browser. The Node.js fetch (via undici) will handle the boundary, but the phrasing is incorrect.

Apply this diff to clarify:

   try {
-    // Pass the FormData directly to the fetch call.
-    // The browser will automatically set the 'Content-Type' to 'multipart/form-data' with the correct boundary.
+    // Pass the FormData directly to fetch.
+    // The fetch implementation will automatically set 'Content-Type: multipart/form-data' with the appropriate boundary.
     const response = await fetch(azureVmEndpoint, {
       method: 'POST',
       body: formData,
     });
📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
// 3. Forward Request and Handle Errors
try {
// Pass the FormData directly to the fetch call.
// The browser will automatically set the 'Content-Type' to 'multipart/form-data' with the correct boundary.
const response = await fetch(azureVmEndpoint, {
method: 'POST',
body: formData,
});
// 3. Forward Request and Handle Errors
try {
// Pass the FormData directly to fetch.
// The fetch implementation will automatically set 'Content-Type: multipart/form-data' with the appropriate boundary.
const response = await fetch(azureVmEndpoint, {
method: 'POST',
body: formData,
});
🤖 Prompt for AI Agents 
In app/api/onnx/route.ts around lines 50 to 57, the inline comment incorrectly
says "The browser will automatically set the 'Content-Type'..." even though this
runs server-side; update the comment to correctly state that the Node.js fetch
implementation (undici) will handle setting the multipart/form-data Content-Type
and boundary when passing FormData, and remove any browser-specific wording so
the comment accurately reflects server-side behavior.


if (!response.ok) {
// Create an error with status to be caught by the catch block
const error = new Error(`Azure endpoint returned an error: ${response.statusText}`);
(error as any).status = response.status;
throw error;
}

const data = await response.json();
return NextResponse.json(data);

} catch (error: unknown) {
let errorMessage = 'An unknown error occurred';
let errorStatus = 500;
let errorContext: { [key: string]: any } = { query, fileName: file?.name };

if (error instanceof Error) {
errorMessage = error.message;
// Use status from the error if it exists, otherwise default to 500
errorStatus = (error as any).status || 500;
}

// In a production environment, send the error to a monitoring service
if (process.env.NODE_ENV === 'production') {
reportErrorToMonitoringService(error as Error, errorContext);
} else {
// Log detailed error in development
console.error('Error in ONNX proxy API:', error);
}

return NextResponse.json({ error: 'Internal Server Error', details: errorMessage }, { status: errorStatus });
}
}
46 changes: 26 additions & 20 deletions app/layout.tsx
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,8 @@ import { SpeedInsights } from "@vercel/speed-insights/next"
import { Toaster } from '@/components/ui/sonner'
import { MapToggleProvider } from '@/components/map-toggle-context'
import { ProfileToggleProvider } from '@/components/profile-toggle-context'
import { GeospatialModelProvider } from '@/lib/geospatial-model-context'
import ErrorBoundary from '@/components/error-boundary'
import { MapLoadingProvider } from '@/components/map-loading-context';
import ConditionalLottie from '@/components/conditional-lottie';

Expand Down Expand Up @@ -54,26 +56,30 @@ export default function RootLayout({
return (
<html lang="en" suppressHydrationWarning>
<body className={cn('font-sans antialiased', fontSans.variable)}>
<MapToggleProvider>
<ProfileToggleProvider>
<ThemeProvider
attribute="class"
defaultTheme="earth"
enableSystem
disableTransitionOnChange
themes={['light', 'dark', 'earth']}
>
<MapLoadingProvider>
<Header />
<ConditionalLottie />
{children}
<Sidebar />
<Footer />
<Toaster />
</MapLoadingProvider>
</ThemeProvider>
</ProfileToggleProvider>
</MapToggleProvider>
<ErrorBoundary>
<MapToggleProvider>
<ProfileToggleProvider>
<GeospatialModelProvider>
<ThemeProvider
attribute="class"
defaultTheme="earth"
enableSystem
disableTransitionOnChange
themes={['light', 'dark', 'earth']}
>
<MapLoadingProvider>
<Header />
<ConditionalLottie />
{children}
<Sidebar />
<Footer />
<Toaster />
</MapLoadingProvider>
</ThemeProvider>
</GeospatialModelProvider>
</ProfileToggleProvider>
</MapToggleProvider>
</ErrorBoundary>
<Analytics />
<SpeedInsights />
</body>
Expand Down
132 changes: 101 additions & 31 deletions components/chat-panel.tsx
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,7 @@
import { useEffect, useState, useRef, ChangeEvent, forwardRef, useImperativeHandle } from 'react'
import type { AI, UIState } from '@/app/actions'
import { useUIState, useActions } from 'ai/rsc'
// Removed import of useGeospatialToolMcp as it's no longer used/available
import { useGeospatialModel } from '@/lib/geospatial-model-context'
import { cn } from '@/lib/utils'
import { UserMessage } from './user-message'
import { Button } from './ui/button'
Expand All @@ -24,7 +24,7 @@ export interface ChatPanelRef {
export const ChatPanel = forwardRef<ChatPanelRef, ChatPanelProps>(({ messages, input, setInput }, ref) => {
const [, setMessages] = useUIState<typeof AI>()
const { submit, clearChat } = useActions()
// Removed mcp instance as it's no longer passed to submit
const { isGeospatialModelEnabled } = useGeospatialModel();
const [isMobile, setIsMobile] = useState(false)
const [selectedFile, setSelectedFile] = useState<File | null>(null)
const inputRef = useRef<HTMLTextAreaElement>(null)
Expand All @@ -37,7 +37,6 @@ export const ChatPanel = forwardRef<ChatPanelRef, ChatPanelProps>(({ messages, i
}
}));

// Detect mobile layout
useEffect(() => {
const checkMobile = () => {
setIsMobile(window.innerWidth <= 1024)
Expand Down Expand Up @@ -71,39 +70,111 @@ export const ChatPanel = forwardRef<ChatPanelRef, ChatPanelProps>(({ messages, i

const handleSubmit = async (e: React.FormEvent<HTMLFormElement>) => {
e.preventDefault()
if (!input && !selectedFile) {
return
}

const content: ({ type: 'text'; text: string } | { type: 'image'; image: string })[] = []
if (input) {
content.push({ type: 'text', text: input })
}
if (selectedFile && selectedFile.type.startsWith('image/')) {
content.push({
type: 'image',
image: URL.createObjectURL(selectedFile)
})
const trimmedInput = input.trim();
if (!trimmedInput && !selectedFile) {
return; // Prevent submission if both text and file are empty
}

setMessages(currentMessages => [
...currentMessages,
{
id: nanoid(),
component: <UserMessage content={content} />
if (isGeospatialModelEnabled) {
// Logic for Geospatial Model (ONNX)
const userMessageContent: ({ type: 'text'; text: string } | { type: 'image'; image: string })[] = [];
if (trimmedInput) {
userMessageContent.push({ type: 'text', text: `[ONNX Request]: ${trimmedInput}` });
}
if (selectedFile) {
userMessageContent.push({ type: 'image', image: URL.createObjectURL(selectedFile) });
}
])

const formData = new FormData(e.currentTarget)
if (selectedFile) {
formData.append('file', selectedFile)
}
setMessages(currentMessages => [
...currentMessages,
{ id: nanoid(), component: <UserMessage content={userMessageContent} /> }
]);

setInput('')
clearAttachment()
const onnxFormData = new FormData();
if (trimmedInput) {
onnxFormData.append('query', trimmedInput);
}
if (selectedFile) {
onnxFormData.append('file', selectedFile);
}

setInput('');
clearAttachment();

try {
const response = await fetch('/api/onnx', {
method: 'POST',
body: onnxFormData,
});

if (!response.ok) {
throw new Error(`ONNX API request failed with status ${response.status}`);
}

const data = await response.json();

// Validate the response shape
if (typeof data !== 'object' || data === null || !data.prediction) {
console.error("Unexpected response format from ONNX API:", data);
throw new Error('Unexpected response format from the ONNX model.');
}

const responseMessage = await submit(formData)
setMessages(currentMessages => [...currentMessages, responseMessage as any])
const predictionText = typeof data.prediction === 'string' ? data.prediction : JSON.stringify(data.prediction, null, 2);

setMessages(currentMessages => [
...currentMessages,
{
id: nanoid(),
component: (
<div className="p-4 my-2 border rounded-lg bg-muted">
<p><strong>ONNX Model Response:</strong></p>
<pre className="whitespace-pre-wrap">{predictionText}</pre>
</div>
)
}
]);

} catch (error) {
console.error('ONNX API call failed:', error);
setMessages(currentMessages => [
...currentMessages,
{
id: nanoid(),
component: (
<div className="p-4 my-2 border rounded-lg bg-destructive/20 text-destructive">
<p><strong>Error:</strong> {error instanceof Error ? error.message : 'Failed to get ONNX model response.'}</p>
</div>
)
}
]);
}
} else {
// Default chat submission logic
const defaultContent: ({ type: 'text'; text: string } | { type: 'image'; image: string })[] = [];
if (trimmedInput) {
defaultContent.push({ type: 'text', text: trimmedInput });
}
if (selectedFile) {
defaultContent.push({ type: 'image', image: URL.createObjectURL(selectedFile) });
}

setMessages(currentMessages => [
...currentMessages,
{ id: nanoid(), component: <UserMessage content={defaultContent} /> }
]);

const defaultFormData = new FormData(e.currentTarget);
if (selectedFile) {
defaultFormData.append('file', selectedFile);
}

setInput('');
clearAttachment();

const responseMessage = await submit(defaultFormData);
setMessages(currentMessages => [...currentMessages, responseMessage as any]);
}
}

const handleClear = async () => {
Expand All @@ -116,7 +187,6 @@ export const ChatPanel = forwardRef<ChatPanelRef, ChatPanelProps>(({ messages, i
inputRef.current?.focus()
}, [])

// New chat button (appears when there are messages)
if (messages.length > 0 && !isMobile) {
return (
<div
Expand Down Expand Up @@ -254,4 +324,4 @@ export const ChatPanel = forwardRef<ChatPanelRef, ChatPanelProps>(({ messages, i
</div>
)
})
ChatPanel.displayName = 'ChatPanel'
ChatPanel.displayName = 'ChatPanel'
51 changes: 51 additions & 0 deletions components/error-boundary.tsx
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,51 @@
'use client';

import React, { Component, ErrorInfo, ReactNode } from 'react';

interface Props {
children: ReactNode;
}

interface State {
hasError: boolean;
}

class ErrorBoundary extends Component<Props, State> {
public state: State = {
hasError: false,
};

public static getDerivedStateFromError(_: Error): State {
// Update state so the next render will show the fallback UI.
return { hasError: true };
}

public componentDidCatch(error: Error, errorInfo: ErrorInfo) {
// In a real application, you would log this to an error reporting service.
console.error("ErrorBoundary caught an error:", error, errorInfo);
}

public render() {
if (this.state.hasError) {
// You can render any custom fallback UI.
return (
<div className="flex flex-col items-center justify-center h-screen bg-background text-foreground p-4 text-center">
<h1 className="text-2xl font-bold mb-4">Something went wrong.</h1>
<p className="text-muted-foreground">
An unexpected error occurred. We&apos;ve logged the issue and will look into it.
</p>
<button
onClick={() => this.setState({ hasError: false })}
className="mt-6 px-4 py-2 bg-primary text-primary-foreground rounded-md hover:bg-primary/90"
>
Try again
</button>
Comment on lines +37 to +42
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟡 Minor

Add explicit type="button" to prevent unintended form submission.

The button lacks an explicit type attribute. If this error boundary is rendered inside a form context, the button's default type="submit" could trigger form submission.

Apply this diff:

           <button
+            type="button"
             onClick={() => this.setState({ hasError: false })}
             className="mt-6 px-4 py-2 bg-primary text-primary-foreground rounded-md hover:bg-primary/90"
           >
             Try again
           </button>
📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
<button
onClick={() => this.setState({ hasError: false })}
className="mt-6 px-4 py-2 bg-primary text-primary-foreground rounded-md hover:bg-primary/90"
>
Try again
</button>
<button
type="button"
onClick={() => this.setState({ hasError: false })}
className="mt-6 px-4 py-2 bg-primary text-primary-foreground rounded-md hover:bg-primary/90"
>
Try again
</button>
🧰 Tools
🪛 Biome (2.1.2)

[error] 37-40: Provide an explicit type prop for the button element.

The default type of a button is submit, which causes the submission of a form when placed inside a form element. This is likely not the behaviour that you want inside a React application.
Allowed button types are: submit, button or reset

(lint/a11y/useButtonType)

🤖 Prompt for AI Agents 
In components/error-boundary.tsx around lines 37 to 42, the "Try again" button
is missing an explicit type which can default to "submit" inside a form; update
the button to include type="button" so clicking it does not trigger form
submission, leaving the existing onClick and classes unchanged.

</div>
);
}

return this.props.children;
}
}

export default ErrorBoundary;
Loading