feat: add useOllamaBridge hook - uses extension to proxy Ollama (no CORS config needed)

src/hooks/use-ollama-bridge.ts

@@ -0,0 +1,316 @@
+"use client"
+
+import { useState, useCallback, useRef, useEffect } from "react"
+
+interface UseOllamaBridgeOptions {
+    ollamaUrl?: string
+    onError?: (error: Error) => void
+    onFinish?: (output: string) => void
+}
+
+interface UseOllamaBridgeReturn {
+    /** Whether the extension bridge is available */
+    bridgeAvailable: boolean
+    /** Whether Ollama is running (detected via extension) */
+    ollamaRunning: boolean | null
+    /** Available models from Ollama */
+    models: string[]
+    /** Accumulated completion text */
+    completion: string
+    /** Whether a stream is currently in progress */
+    isLoading: boolean
+    /** Check if extension + Ollama are available */
+    checkConnection: () => Promise<boolean>
+    /** Fetch available models */
+    refreshModels: () => Promise<string[]>
+    /** Stream a generate request through the extension bridge */
+    generate: (prompt: string, model: string, options?: Record<string, unknown>) => Promise<string>
+    /** Stream a chat request through the extension bridge */
+    chat: (messages: Array<{ role: string; content: string }>, model: string, options?: Record<string, unknown>) => Promise<string>
+    /** Abort the current stream */
+    stop: () => void
+    /** Reset completion to empty */
+    reset: () => void
+}
+
+/**
+ * Hook to use Ollama through the OpenPrompt browser extension.
+ * 
+ * The extension acts as a CORS-free bridge between the website and the user's
+ * local Ollama instance. No need for OLLAMA_ORIGINS configuration!
+ * 
+ * Falls back to direct connection if extension is not available.
+ */
+export function useOllamaBridge(options: UseOllamaBridgeOptions = {}): UseOllamaBridgeReturn {
+    const { ollamaUrl = "http://localhost:11434", onError, onFinish } = options
+
+    const [bridgeAvailable, setBridgeAvailable] = useState(false)
+    const [ollamaRunning, setOllamaRunning] = useState<boolean | null>(null)
+    const [models, setModels] = useState<string[]>([])
+    const [completion, setCompletion] = useState("")
+    const [isLoading, setIsLoading] = useState(false)
+    const requestIdRef = useRef<string | null>(null)
+    const accumulatedRef = useRef("")
+
+    // Generate unique request ID
+    const generateId = () => `req_${Date.now()}_${Math.random().toString(36).slice(2, 8)}`
+
+    // Detect if the bridge is available
+    useEffect(() => {
+        const handleBridgeReady = (event: MessageEvent) => {
+            if (event.data?.type === 'OLLAMA_BRIDGE_READY') {
+                setBridgeAvailable(true)
+            }
+        }
+
+        // Check for attribute (already loaded)
+        if (document.documentElement.getAttribute('data-ollama-bridge') === 'true') {
+            setBridgeAvailable(true)
+        }
+
+        window.addEventListener('message', handleBridgeReady)
+        return () => window.removeEventListener('message', handleBridgeReady)
+    }, [])
+
+    // Listen for stream responses
+    useEffect(() => {
+        const handleMessage = (event: MessageEvent) => {
+            if (event.origin !== window.location.origin) return
+            const msg = event.data
+            if (!msg || !msg.id || msg.id !== requestIdRef.current) return
+
+            if (msg.type === 'OLLAMA_STREAM_CHUNK' && msg.chunk) {
+                accumulatedRef.current += msg.chunk
+                setCompletion(accumulatedRef.current)
+            } else if (msg.type === 'OLLAMA_STREAM_END') {
+                setIsLoading(false)
+                onFinish?.(accumulatedRef.current)
+            } else if (msg.type === 'OLLAMA_ERROR') {
+                setIsLoading(false)
+                onError?.(new Error(msg.error || 'Ollama bridge error'))
+            }
+        }
+
+        window.addEventListener('message', handleMessage)
+        return () => window.removeEventListener('message', handleMessage)
+    }, [onError, onFinish])
+
+    const checkConnection = useCallback(async (): Promise<boolean> => {
+        if (!bridgeAvailable) {
+            // Try direct connection as fallback
+            try {
+                const response = await fetch(`${ollamaUrl}/api/tags`, { signal: AbortSignal.timeout(3000) })
+                const isRunning = response.ok
+                setOllamaRunning(isRunning)
+                return isRunning
+            } catch {
+                setOllamaRunning(false)
+                return false
+            }
+        }
+
+        return new Promise((resolve) => {
+            const id = generateId()
+            const timeout = setTimeout(() => {
+                window.removeEventListener('message', handler)
+                setOllamaRunning(false)
+                resolve(false)
+            }, 5000)
+
+            const handler = (event: MessageEvent) => {
+                if (event.data?.id === id) {
+                    clearTimeout(timeout)
+                    window.removeEventListener('message', handler)
+                    const running = event.data.data?.running ?? false
+                    setOllamaRunning(running)
+                    resolve(running)
+                }
+            }
+
+            window.addEventListener('message', handler)
+            window.postMessage({ type: 'OLLAMA_PING', id, ollamaUrl }, '*')
+        })
+    }, [bridgeAvailable, ollamaUrl])
+
+    const refreshModels = useCallback(async (): Promise<string[]> => {
+        if (!bridgeAvailable) {
+            // Direct connection fallback
+            try {
+                const response = await fetch(`${ollamaUrl}/api/tags`)
+                const data = await response.json()
+                const modelNames = data.models?.map((m: { name: string }) => m.name) || []
+                setModels(modelNames)
+                return modelNames
+            } catch {
+                return []
+            }
+        }
+
+        return new Promise((resolve) => {
+            const id = generateId()
+            const timeout = setTimeout(() => {
+                window.removeEventListener('message', handler)
+                resolve([])
+            }, 10000)
+
+            const handler = (event: MessageEvent) => {
+                if (event.data?.id === id) {
+                    clearTimeout(timeout)
+                    window.removeEventListener('message', handler)
+
+                    if (event.data.type === 'OLLAMA_ERROR') {
+                        onError?.(new Error(event.data.error))
+                        resolve([])
+                    } else {
+                        const modelNames = event.data.data?.models?.map((m: { name: string }) => m.name) || []
+                        setModels(modelNames)
+                        resolve(modelNames)
+                    }
+                }
+            }
+
+            window.addEventListener('message', handler)
+            window.postMessage({ type: 'OLLAMA_LIST_MODELS', id, ollamaUrl }, '*')
+        })
+    }, [bridgeAvailable, ollamaUrl, onError])
+
+    const generate = useCallback(async (
+        prompt: string,
+        model: string,
+        genOptions?: Record<string, unknown>
+    ): Promise<string> => {
+        const id = generateId()
+        requestIdRef.current = id
+        accumulatedRef.current = ""
+        setCompletion("")
+        setIsLoading(true)
+
+        if (!bridgeAvailable) {
+            // Direct fallback (requires CORS)
+            try {
+                const response = await fetch(`${ollamaUrl}/api/generate`, {
+                    method: 'POST',
+                    headers: { 'Content-Type': 'application/json' },
+                    body: JSON.stringify({ model, prompt, stream: true, ...(genOptions && { options: genOptions }) }),
+                })
+
+                const reader = response.body?.getReader()
+                if (!reader) throw new Error('No stream')
+
+                const decoder = new TextDecoder()
+                while (true) {
+                    const { done, value } = await reader.read()
+                    if (done) break
+                    const lines = decoder.decode(value, { stream: true }).split('\n').filter(l => l.trim())
+                    for (const line of lines) {
+                        try {
+                            const json = JSON.parse(line)
+                            if (json.response) {
+                                accumulatedRef.current += json.response
+                                setCompletion(accumulatedRef.current)
+                            }
+                        } catch { /* skip */ }
+                    }
+                }
+
+                setIsLoading(false)
+                onFinish?.(accumulatedRef.current)
+                return accumulatedRef.current
+            } catch (err) {
+                setIsLoading(false)
+                const error = err instanceof Error ? err : new Error('Failed to connect to Ollama')
+                onError?.(error)
+                throw error
+            }
+        }
+
+        // Use extension bridge
+        window.postMessage({
+            type: 'OLLAMA_GENERATE',
+            id,
+            ollamaUrl,
+            model,
+            prompt,
+            options: genOptions,
+        }, '*')
+
+        // Return a promise that resolves when streaming is done
+        return new Promise((resolve, reject) => {
+            const handler = (event: MessageEvent) => {
+                if (event.data?.id !== id) return
+
+                if (event.data.type === 'OLLAMA_STREAM_END') {
+                    window.removeEventListener('message', handler)
+                    resolve(accumulatedRef.current)
+                } else if (event.data.type === 'OLLAMA_ERROR') {
+                    window.removeEventListener('message', handler)
+                    reject(new Error(event.data.error))
+                }
+            }
+            window.addEventListener('message', handler)
+        })
+    }, [bridgeAvailable, ollamaUrl, onError, onFinish])
+
+    const chat = useCallback(async (
+        messages: Array<{ role: string; content: string }>,
+        model: string,
+        chatOptions?: Record<string, unknown>
+    ): Promise<string> => {
+        const id = generateId()
+        requestIdRef.current = id
+        accumulatedRef.current = ""
+        setCompletion("")
+        setIsLoading(true)
+
+        window.postMessage({
+            type: 'OLLAMA_CHAT',
+            id,
+            ollamaUrl,
+            model,
+            messages,
+            options: chatOptions,
+        }, '*')
+
+        return new Promise((resolve, reject) => {
+            const handler = (event: MessageEvent) => {
+                if (event.data?.id !== id) return
+
+                if (event.data.type === 'OLLAMA_STREAM_END') {
+                    window.removeEventListener('message', handler)
+                    resolve(accumulatedRef.current)
+                } else if (event.data.type === 'OLLAMA_ERROR') {
+                    window.removeEventListener('message', handler)
+                    reject(new Error(event.data.error))
+                }
+            }
+            window.addEventListener('message', handler)
+        })
+    }, [ollamaUrl])
+
+    const stop = useCallback(() => {
+        if (requestIdRef.current) {
+            window.postMessage({ type: 'OLLAMA_ABORT', id: requestIdRef.current }, '*')
+            requestIdRef.current = null
+            setIsLoading(false)
+        }
+    }, [])
+
+    const reset = useCallback(() => {
+        setCompletion("")
+        accumulatedRef.current = ""
+    }, [])
+
+    return {
+        bridgeAvailable,
+        ollamaRunning,
+        models,
+        completion,
+        isLoading,
+        checkConnection,
+        refreshModels,
+        generate,
+        chat,
+        stop,
+        reset,
+    }
+}