route.ts

import { openai } from '@ai-sdk/openai'
import { convertToModelMessages, stepCountIs, streamText, tool, type UIMessage } from 'ai'
import { sql } from 'drizzle-orm'
import { z } from 'zod'
import { db, docsEmbeddings } from '@/lib/db'
import { generateSearchEmbedding } from '@/lib/embeddings'

export const runtime = 'nodejs'
export const maxDuration = 30

/** Model used for the Ask AI chat. Override with OPENAI_CHAT_MODEL in the environment. */
const CHAT_MODEL = process.env.OPENAI_CHAT_MODEL || 'gpt-5.4-mini'

/** Max documentation chunks returned per search to ground an answer. */
const SEARCH_LIMIT = 6

/**
 * Abuse guards. This endpoint proxies a paid LLM, so an unauthenticated public
 * route is a target for scripted "free inference". These bounds cap the cost of
 * any single request; durable per-IP rate limiting, a provider spend cap, and
 * edge bot protection are provisioned separately (see the PR checklist).
 *
 * The size cap counts only user-authored text — NOT the conversation history,
 * assistant turns, or retrieved doc chunks we add via the searchDocs tool, which
 * legitimately grow large over a multi-turn chat.
 */
const MAX_MESSAGES = 200
const MAX_USER_INPUT_CHARS = 400_000
const MAX_OUTPUT_TOKENS = 4000
const MAX_STEPS = 6

/** Total length of user-authored text across the conversation. */
function userInputChars(messages: UIMessage[]): number {
  let total = 0
  for (const message of messages) {
    if (message.role !== 'user') continue
    for (const part of message.parts) {
      if (part.type === 'text') total += part.text.length
    }
  }
  return total
}

/**
 * Reject obvious cross-origin calls. Same-origin browser requests send an
 * `Origin` header matching the host; we allow those, plus any host in
 * DOCS_ALLOWED_ORIGINS (comma-separated). Requests with no Origin (e.g. curl)
 * are allowed through to the cost caps rather than blocked, since Origin is
 * trivially spoofable and is a filter, not a security boundary.
 */
function isAllowedOrigin(req: Request): boolean {
  const origin = req.headers.get('origin')
  if (!origin) return true

  let originHost: string
  try {
    originHost = new URL(origin).host
  } catch {
    return false
  }

  const requestHost = req.headers.get('x-forwarded-host') ?? req.headers.get('host')
  if (originHost === requestHost) return true

  const allowlist = (process.env.DOCS_ALLOWED_ORIGINS ?? '')
    .split(',')
    .map((value) => value.trim())
    .filter(Boolean)
  return allowlist.includes(originHost)
}

const SYSTEM_PROMPT = `You are the documentation assistant for Sim — the open-source AI workspace where teams build, deploy, and manage AI agents.

Answer questions about Sim using the documentation. Always call the searchDocs tool before answering anything specific about Sim's features, configuration, or usage — do not answer from memory. Base your answer only on the returned documentation; if the docs do not cover the question, say so plainly rather than guessing.

Guidelines:
- Be direct and concrete. Lead with the answer, then the detail.
- Reference the relevant pages by their titles so the user knows where to read more.
- When you show configuration or code, keep it minimal and correct.
- The agent is called "Sim" and the chat surface is "Chat" — never say "Mothership" or "copilot".
- If a question is unrelated to Sim, briefly say it's outside the docs' scope.`

/**
 * Vector search over the docs embeddings, returning the most relevant chunks
 * with their source links so the model can ground and cite its answer.
 */
async function searchDocs(query: string) {
  const embedding = await generateSearchEmbedding(query)
  const vectorLiteral = JSON.stringify(embedding)

  const rows = await db
    .select({
      title: docsEmbeddings.headerText,
      url: docsEmbeddings.sourceLink,
      content: docsEmbeddings.chunkText,
      similarity: sql<number>`1 - (${docsEmbeddings.embedding} <=> ${vectorLiteral}::vector)`,
    })
    .from(docsEmbeddings)
    .orderBy(sql`${docsEmbeddings.embedding} <=> ${vectorLiteral}::vector`)
    .limit(SEARCH_LIMIT)

  return rows.map((row) => ({
    title: row.title,
    url: row.url,
    content: row.content,
  }))
}

export async function POST(req: Request) {
  if (!isAllowedOrigin(req)) {
    return new Response('Forbidden', { status: 403 })
  }

  const { messages }: { messages: UIMessage[] } = await req.json()

  if (!Array.isArray(messages) || messages.length === 0 || messages.length > MAX_MESSAGES) {
    return new Response('Invalid request', { status: 400 })
  }
  if (userInputChars(messages) > MAX_USER_INPUT_CHARS) {
    return new Response('Request too large', { status: 413 })
  }

  const result = streamText({
    model: openai(CHAT_MODEL),
    system: SYSTEM_PROMPT,
    messages: convertToModelMessages(messages),
    stopWhen: stepCountIs(MAX_STEPS),
    maxOutputTokens: MAX_OUTPUT_TOKENS,
    tools: {
      searchDocs: tool({
        description:
          'Search the Sim documentation for relevant content. Use this before answering any question about Sim.',
        inputSchema: z.object({
          query: z.string().describe('A focused natural-language search query.'),
        }),
        execute: async ({ query }) => searchDocs(query),
      }),
    },
  })

  return result.toUIMessageStreamResponse()
}