RAG Gemini et Google Drive

gemini.ts

import { GoogleGenAI } from '@google/genai'
import { mkdirSync, readFileSync, writeFileSync, existsSync } from 'node:fs'
import { join } from 'node:path'

const DATA_DIR = '.gemini-rag'
const STATE_FILE = join(DATA_DIR, 'state.json')

type State = {
  storeName: string
  files: Record<string, { name: string; geminiFile: string; mimeType: string }>
}

function loadState(): State | null {
  if (!existsSync(STATE_FILE)) return null
  return JSON.parse(readFileSync(STATE_FILE, 'utf8'))
}

function saveState(state: State) {
  mkdirSync(DATA_DIR, { recursive: true })
  writeFileSync(STATE_FILE, JSON.stringify(state, null, 2))
}

function client() {
  const apiKey = process.env.GEMINI_API_KEY
  if (!apiKey) throw new Error('GEMINI_API_KEY is not set')
  return new GoogleGenAI({ apiKey })
}

async function getOrCreateStore(): Promise<State> {
  const existing = loadState()
  if (existing) return existing
  const ai = client()
  const store = await ai.fileSearchStores.create({
    config: { displayName: 'gemini-rag' },
  })
  const state: State = { storeName: store.name!, files: {} }
  saveState(state)
  return state
}

export async function getIndexedFiles() {
  const state = loadState()
  if (!state) return []
  return Object.entries(state.files).map(([driveId, f]) => ({
    driveId,
    name: f.name,
    mimeType: f.mimeType,
  }))
}

/**
 * Uploads a Buffer to the FileSearchStore and waits for indexing to finish.
 * Skips files that are already indexed (by Drive file id).
 */
export async function indexFile(args: {
  driveId: string
  name: string
  data: Buffer
  mimeType: string
}) {
  const state = await getOrCreateStore()
  if (state.files[args.driveId]) return state.files[args.driveId]

  const ai = client()
  const blob = new Blob([new Uint8Array(args.data)], { type: args.mimeType })

  let op = await ai.fileSearchStores.uploadToFileSearchStore({
    file: blob,
    fileSearchStoreName: state.storeName,
    config: { displayName: args.name, mimeType: args.mimeType },
  })

  while (!op.done) {
    await new Promise((r) => setTimeout(r, 2000))
    op = await ai.operations.get({ operation: op })
  }

  const geminiFile = (op.response as any)?.name ?? args.name
  state.files[args.driveId] = {
    name: args.name,
    geminiFile,
    mimeType: args.mimeType,
  }
  saveState(state)
  return state.files[args.driveId]
}

/** Asks Gemini with the FileSearchStore as a tool. Returns the model's text. */
export async function ask(question: string) {
  const state = loadState()
  if (!state || Object.keys(state.files).length === 0) {
    throw new Error('No files indexed yet. Index some Drive files first.')
  }
  const ai = client()
  const res = await ai.models.generateContent({
    model: 'gemini-2.5-flash',
    contents: question,
    config: {
      tools: [{ fileSearch: { fileSearchStoreNames: [state.storeName] } }],
    },
  })
  return res.text ?? ''
}

google.ts

import { google } from 'googleapis'
import { getSession, setSession, type Session } from './session'

const SCOPES = [
  'https://www.googleapis.com/auth/drive.readonly',
  'https://www.googleapis.com/auth/userinfo.email',
]

function oauthClient() {
  return new google.auth.OAuth2(
    process.env.GOOGLE_CLIENT_ID,
    process.env.GOOGLE_CLIENT_SECRET,
    process.env.GOOGLE_REDIRECT_URI,
  )
}

export function authUrl() {
  return oauthClient().generateAuthUrl({
    access_type: 'offline',
    prompt: 'consent',
    scope: SCOPES,
  })
}

export async function exchangeCode(code: string): Promise<Session> {
  const { tokens } = await oauthClient().getToken(code)
  return {
    access_token: tokens.access_token!,
    refresh_token: tokens.refresh_token ?? undefined,
    expiry_date: tokens.expiry_date ?? undefined,
  }
}

/**
 * Returns an authed Drive client for the current session. Refreshes the
 * access token when needed and persists it back to the cookie.
 */
export function driveClient() {
  const session = getSession()
  if (!session) throw new Error('Not signed in to Google')

  const auth = oauthClient()
  auth.setCredentials(session)
  auth.on('tokens', (tokens) => {
    setSession({
      access_token: tokens.access_token ?? session.access_token,
      refresh_token: tokens.refresh_token ?? session.refresh_token,
      expiry_date: tokens.expiry_date ?? session.expiry_date,
    })
  })
  return google.drive({ version: 'v3', auth })
}

/** Mime types Gemini File Search can index well. */
const SUPPORTED_PREFIXES = ['application/pdf', 'text/', 'application/json']
const GOOGLE_DOC_EXPORTS: Record<string, { mimeType: string; ext: string }> = {
  'application/vnd.google-apps.document': { mimeType: 'application/pdf', ext: 'pdf' },
  'application/vnd.google-apps.presentation': { mimeType: 'application/pdf', ext: 'pdf' },
  'application/vnd.google-apps.spreadsheet': { mimeType: 'text/csv', ext: 'csv' },
}

export type DriveFile = {
  id: string
  name: string
  mimeType: string
  modifiedTime?: string
  supported: boolean
}

export async function listFiles(query?: string): Promise<DriveFile[]> {
  const drive = driveClient()
  const q = [
    "trashed = false",
    query ? `name contains '${query.replace(/'/g, "\\'")}'` : null,
  ]
    .filter(Boolean)
    .join(' and ')

  const res = await drive.files.list({
    q,
    pageSize: 50,
    fields: 'files(id, name, mimeType, modifiedTime)',
    orderBy: 'modifiedTime desc',
  })

  return (res.data.files ?? []).map((f) => ({
    id: f.id!,
    name: f.name ?? '(untitled)',
    mimeType: f.mimeType ?? 'application/octet-stream',
    modifiedTime: f.modifiedTime ?? undefined,
    supported:
      f.mimeType! in GOOGLE_DOC_EXPORTS ||
      SUPPORTED_PREFIXES.some((p) => f.mimeType!.startsWith(p)),
  }))
}

/** Downloads a Drive file as a Buffer, exporting Google Docs to PDF/CSV. */
export async function downloadFile(
  fileId: string,
): Promise<{ data: Buffer; mimeType: string; suggestedExt: string }> {
  const drive = driveClient()
  const meta = await drive.files.get({ fileId, fields: 'mimeType,name' })
  const mimeType = meta.data.mimeType!

  const exportAs = GOOGLE_DOC_EXPORTS[mimeType]
  if (exportAs) {
    const res = await drive.files.export(
      { fileId, mimeType: exportAs.mimeType },
      { responseType: 'arraybuffer' },
    )
    return {
      data: Buffer.from(res.data as ArrayBuffer),
      mimeType: exportAs.mimeType,
      suggestedExt: exportAs.ext,
    }
  }

  const res = await drive.files.get(
    { fileId, alt: 'media' },
    { responseType: 'arraybuffer' },
  )
  const ext = mimeType.startsWith('text/') ? 'txt' : mimeType.split('/').pop() ?? 'bin'
  return {
    data: Buffer.from(res.data as ArrayBuffer),
    mimeType,
    suggestedExt: ext,
  }
}

index.tsx

import { useState } from 'react'
import { createFileRoute, useRouter } from '@tanstack/react-router'
import { createServerFn } from '@tanstack/react-start'
import { authUrl, downloadFile, listFiles, type DriveFile } from '~/lib/google'
import { ask, getIndexedFiles, indexFile } from '~/lib/gemini'
import { clearSession, getSession } from '~/lib/session'

// --- Server functions ----------------------------------------------------

const getStatus = createServerFn({ method: 'GET' }).handler(async () => {
  return {
    signedIn: !!getSession(),
    indexed: await getIndexedFiles(),
  }
})

const getLoginUrl = createServerFn({ method: 'GET' }).handler(() => authUrl())

const signOut = createServerFn({ method: 'POST' }).handler(() => {
  clearSession()
})

const browseDrive = createServerFn({ method: 'GET' })
  .validator((d: { q?: string }) => d)
  .handler(async ({ data }) => listFiles(data.q))

const indexSelected = createServerFn({ method: 'POST' })
  .validator((d: { ids: string[] }) => d)
  .handler(async ({ data }) => {
    const drive = await listFiles()
    const byId = new Map(drive.map((f) => [f.id, f]))
    for (const id of data.ids) {
      const meta = byId.get(id)
      if (!meta) continue
      const file = await downloadFile(id)
      await indexFile({
        driveId: id,
        name: meta.name,
        data: file.data,
        mimeType: file.mimeType,
      })
    }
    return getIndexedFiles()
  })

const chat = createServerFn({ method: 'POST' })
  .validator((d: { question: string }) => d)
  .handler(async ({ data }) => ask(data.question))

// --- Route ---------------------------------------------------------------

export const Route = createFileRoute('/')({
  component: Home,
  loader: () => getStatus(),
})

type Status = Awaited<ReturnType<typeof getStatus>>

function Home() {
  const router = useRouter()
  const status = Route.useLoaderData() as Status

  if (!status.signedIn) return <SignedOut />
  return <SignedIn status={status} onChange={() => router.invalidate()} />
}

function SignedOut() {
  async function login() {
    window.location.href = await getLoginUrl()
  }
  return (
    <div className="container">
      <h1>Gemini RAG</h1>
      <div className="card">
        <p>Sign in with Google to pick files from your Drive and chat with them.</p>
        <button className="btn" onClick={login}>Connect Google Drive</button>
      </div>
    </div>
  )
}

function SignedIn({ status, onChange }: { status: Status; onChange: () => void }) {
  return (
    <div className="container">
      <div className="row" style={{ justifyContent: 'space-between' }}>
        <h1>Gemini RAG</h1>
        <button
          className="btn secondary"
          onClick={async () => {
            await signOut()
            onChange()
          }}
        >
          Sign out
        </button>
      </div>
      <Indexer indexed={status.indexed} onChange={onChange} />
      <Chat disabled={status.indexed.length === 0} />
    </div>
  )
}

// --- Indexer (browse Drive + index selected) -----------------------------

function Indexer({
  indexed,
  onChange,
}: {
  indexed: Status['indexed']
  onChange: () => void
}) {
  const [files, setFiles] = useState<DriveFile[] | null>(null)
  const [selected, setSelected] = useState<Set<string>>(new Set())
  const [query, setQuery] = useState('')
  const [loading, setLoading] = useState(false)
  const [error, setError] = useState<string | null>(null)

  const indexedIds = new Set(indexed.map((f) => f.driveId))

  async function browse() {
    setLoading(true)
    setError(null)
    try {
      setFiles(await browseDrive({ data: { q: query || undefined } }))
    } catch (e: any) {
      setError(e.message ?? 'Failed to list files')
    } finally {
      setLoading(false)
    }
  }

  async function runIndex() {
    setLoading(true)
    setError(null)
    try {
      await indexSelected({ data: { ids: Array.from(selected) } })
      setSelected(new Set())
      onChange()
    } catch (e: any) {
      setError(e.message ?? 'Failed to index')
    } finally {
      setLoading(false)
    }
  }

  function toggle(id: string) {
    const next = new Set(selected)
    next.has(id) ? next.delete(id) : next.add(id)
    setSelected(next)
  }

  return (
    <div className="card">
      <h2>1. Pick Drive files to index</h2>
      <div className="row" style={{ marginBottom: 8 }}>
        <input
          placeholder="Search by name (optional)"
          value={query}
          onChange={(e) => setQuery(e.target.value)}
          style={{ flex: 1, padding: 8, border: '1px solid #d1d5db', borderRadius: 6 }}
          onKeyDown={(e) => e.key === 'Enter' && browse()}
        />
        <button className="btn secondary" onClick={browse} disabled={loading}>
          Browse
        </button>
      </div>

      {error && <p className="error">{error}</p>}

      {files && (
        <>
          <div style={{ maxHeight: 240, overflowY: 'auto', marginBottom: 8 }}>
            {files.length === 0 && <p className="muted">No files found.</p>}
            {files.map((f) => {
              const already = indexedIds.has(f.id)
              return (
                <div key={f.id} className="file">
                  <input
                    type="checkbox"
                    id={f.id}
                    checked={selected.has(f.id)}
                    disabled={!f.supported || already}
                    onChange={() => toggle(f.id)}
                  />
                  <label htmlFor={f.id}>
                    {f.name}
                    {!f.supported && <span className="tag">unsupported</span>}
                    {already && <span className="tag">indexed</span>}
                  </label>
                </div>
              )
            })}
          </div>
          <button
            className="btn"
            onClick={runIndex}
            disabled={selected.size === 0 || loading}
          >
            {loading ? 'Indexing…' : `Index ${selected.size} file(s)`}
          </button>
        </>
      )}

      {indexed.length > 0 && (
        <>
          <h2>Indexed</h2>
          {indexed.map((f) => (
            <div key={f.driveId} className="muted">• {f.name}</div>
          ))}
        </>
      )}
    </div>
  )
}

// --- Chat -----------------------------------------------------------------

type Msg = { role: 'user' | 'assistant'; text: string }

function Chat({ disabled }: { disabled: boolean }) {
  const [messages, setMessages] = useState<Msg[]>([])
  const [input, setInput] = useState('')
  const [loading, setLoading] = useState(false)

  async function send() {
    if (!input.trim() || loading) return
    const question = input.trim()
    setInput('')
    setMessages((m) => [...m, { role: 'user', text: question }])
    setLoading(true)
    try {
      const answer = await chat({ data: { question } })
      setMessages((m) => [...m, { role: 'assistant', text: answer }])
    } catch (e: any) {
      setMessages((m) => [
        ...m,
        { role: 'assistant', text: `Error: ${e.message ?? e}` },
      ])
    } finally {
      setLoading(false)
    }
  }

  return (
    <div className="card">
      <h2>2. Ask a question</h2>
      {disabled && <p className="muted">Index at least one file to start chatting.</p>}
      <div className="chat">
        {messages.map((m, i) => (
          <div key={i} className={`msg ${m.role}`}>{m.text}</div>
        ))}
        {loading && <div className="msg assistant">Thinking…</div>}
      </div>
      <textarea
        placeholder="Ask something about your files…"
        value={input}
        onChange={(e) => setInput(e.target.value)}
        onKeyDown={(e) => {
          if (e.key === 'Enter' && !e.shiftKey) {
            e.preventDefault()
            send()
          }
        }}
        disabled={disabled}
        style={{ marginTop: 12 }}
      />
      <div className="row" style={{ marginTop: 8, justifyContent: 'flex-end' }}>
        <button className="btn" onClick={send} disabled={disabled || loading}>
          Send
        </button>
      </div>
    </div>
  )
}