backend/src/agent/stream-handler.ts

import { streamText, convertToModelMessages } from "ai";
import { createOpenRouter } from "@openrouter/ai-sdk-provider";
import type { Request, Response } from "express";

export const DEFAULT_MODEL = "google/gemini-2.5-flash";

export function getProvider() {
  const apiKey = process.env.OPENROUTER_API_KEY;
  if (!apiKey) {
    throw new Error("OPENROUTER_API_KEY environment variable is required");
  }
  return createOpenRouter({ apiKey });
}

interface StreamChatOptions {
  systemPrompt: string;
  tools: Parameters<typeof streamText>[0]["tools"];
  logPrefix: string;
}

export async function streamChatResponse(
  req: Request,
  res: Response,
  { systemPrompt, tools, logPrefix }: StreamChatOptions,
) {
  try {
    const { messages, context, model } = req.body;

    if (!messages || !Array.isArray(messages)) {
      res.status(400).json({ error: "messages array is required" });
      return;
    }

    const provider = getProvider();
    const modelId = model || process.env.OPENROUTER_MODEL || DEFAULT_MODEL;
    const modelMessages = await convertToModelMessages(messages);

    const result = streamText({
      model: provider.chat(modelId),
      system: systemPrompt,
      messages: modelMessages,
      tools,
    });

    const webResponse = result.toUIMessageStreamResponse({
      onError: (error) => {
        console.error(`[${logPrefix}] stream error:`, error);
        return error instanceof Error ? error.message : "Stream error";
      },
    });

    res.writeHead(
      webResponse.status,
      Object.fromEntries(webResponse.headers.entries()),
    );
    const reader = webResponse.body!.getReader();
    const pump = async (): Promise<void> => {
      const { done, value } = await reader.read();
      if (done) {
        res.end();
        return;
      }
      res.write(value);
      return pump();
    };
    await pump();
  } catch (error: unknown) {
    const message =
      error instanceof Error ? error.message : "Internal server error";
    console.error(`[${logPrefix}] error:`, message);

    if (!res.headersSent) {
      res.status(500).json({ error: message });
    }
  }
}