confluence_web/agent/model-provider.ts

import type {
  Context,
  Model,
  AssistantMessage,
  SimpleStreamOptions,
  TextContent,
  ToolCall
} from "@earendil-works/pi-ai";
import { createAssistantMessageEventStream } from "@earendil-works/pi-ai";
import type { StreamFn } from "@earendil-works/pi-agent-core";
import type {
  ModelDescriptor,
  ModelRequest,
  ModelResponse,
  MessageContract,
  ContentBlock,
  ToolDeclaration
} from "./types.js";

/**
 * Converts pi-ai Context into CONTRACTS ModelRequest format.
 * Invariants:
 * - System instruction is NOT duplicated into messages.
 * - Tool calls use type: "tool_call".
 * - Tool results use role: "tool", tool_call_id, is_error.
 * - Provider state is round-tripped on assistant messages.
 */
export function contextToModelRequest(context: Context): ModelRequest {
  const messages: MessageContract[] = [];

  for (const m of context.messages) {
    if (m.role === "user") {
      let textContent: string = "";
      if (typeof m.content === "string") {
        textContent = m.content;
      } else if (Array.isArray(m.content)) {
        textContent = m.content
          .filter((c: any) => c.type === "text")
          .map((c: any) => c.text)
          .join("\n");
      }
      messages.push({
        role: "user",
        content: [{ type: "text", text: textContent }]
      });
    } else if (m.role === "assistant") {
      const contentBlocks: ContentBlock[] = [];
      for (const block of m.content) {
        if (block.type === "text") {
          contentBlocks.push({
            type: "text",
            text: block.text
          });
        } else if (block.type === "toolCall") {
          contentBlocks.push({
            type: "tool_call",
            id: block.id,
            name: block.name,
            arguments: block.arguments || {}
          });
        }
        // Thinking blocks are not part of shared model boundary
      }

      const assistantMsg: MessageContract = {
        role: "assistant",
        content: contentBlocks,
        ...((m as any).provider_state !== undefined ? { provider_state: (m as any).provider_state } : {})
      };
      messages.push(assistantMsg);
    } else if (m.role === "toolResult") {
      let text = "";
      if (typeof m.content === "string") {
        text = m.content;
      } else if (Array.isArray(m.content)) {
        text = m.content
          .filter((c: any) => c.type === "text")
          .map((c: any) => c.text)
          .join("\n");
      }

      messages.push({
        role: "tool",
        tool_call_id: m.toolCallId,
        name: m.toolName,
        content: text,
        is_error: Boolean(m.isError)
      });
    }
  }

  const tools: ToolDeclaration[] = (context.tools || []).map((t) => ({
    name: t.name,
    description: t.description || "",
    input_schema: (t.parameters as any) || {}
  }));

  return {
    messages,
    tools
  };
}

/**
 * Converts CONTRACTS ModelResponse to pi-ai AssistantMessage.
 */
export function modelResponseToAssistantMessage(
  response: ModelResponse,
  model: Model<any>
): AssistantMessage {
  const content: (TextContent | ToolCall)[] = [];

  for (const block of response.content) {
    if (block.type === "text") {
      content.push({
        type: "text",
        text: block.text
      });
    } else if (block.type === "tool_call") {
      content.push({
        type: "toolCall",
        id: block.id,
        name: block.name,
        arguments: block.arguments || {}
      });
    }
  }

  let stopReason: AssistantMessage["stopReason"] = "stop";
  if (response.stop_reason === "tool_calls") {
    stopReason = "toolUse";
  } else if (response.stop_reason === "length") {
    stopReason = "length";
  } else if (response.stop_reason === "stop") {
    stopReason = "stop";
  }

  const assistantMessage: AssistantMessage = {
    role: "assistant",
    content,
    api: model.api,
    provider: model.provider,
    model: model.id,
    usage: {
      input: response.usage?.input_tokens ?? 0,
      output: response.usage?.output_tokens ?? 0,
      cacheRead: 0,
      cacheWrite: 0,
      totalTokens: (response.usage?.input_tokens ?? 0) + (response.usage?.output_tokens ?? 0),
      cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }
    },
    stopReason,
    timestamp: Date.now()
  };

  if (response.provider_state !== undefined) {
    (assistantMessage as any).provider_state = response.provider_state;
  }

  return assistantMessage;
}

export type SendModelRequestFn = (request: ModelRequest) => Promise<ModelResponse>;

/**
 * Creates a Model definition and StreamFn that routes model calls across the bridge.
 */
export function createModelProvider(
  descriptor: ModelDescriptor,
  sendModelRequest: SendModelRequestFn
): { model: Model<any>; streamFn: StreamFn } {
  const model: Model<any> = {
    id: descriptor.id,
    name: descriptor.id,
    api: "custom" as any,
    provider: "backend",
    baseUrl: "",
    reasoning: false,
    input: ["text"],
    cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },
    contextWindow: descriptor.context_window_tokens,
    maxTokens: descriptor.max_output_tokens
  };

  const streamFn: StreamFn = async (_model, context, options) => {
    const stream = createAssistantMessageEventStream();

    (async () => {
      try {
        if (options?.signal?.aborted) {
          const abortedMsg: AssistantMessage = {
            role: "assistant",
            content: [],
            api: model.api,
            provider: model.provider,
            model: model.id,
            usage: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, totalTokens: 0, cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 } },
            stopReason: "aborted",
            errorMessage: "Request aborted",
            timestamp: Date.now()
          };
          stream.push({ type: "error", reason: "aborted", error: abortedMsg });
          stream.end(abortedMsg);
          return;
        }

        const modelRequest = contextToModelRequest(context);
        const response = await sendModelRequest(modelRequest);
        const assistantMessage = modelResponseToAssistantMessage(response, model);

        if (options?.signal?.aborted) {
          const abortedMsg: AssistantMessage = {
            role: "assistant",
            content: [],
            api: model.api,
            provider: model.provider,
            model: model.id,
            usage: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, totalTokens: 0, cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 } },
            stopReason: "aborted",
            errorMessage: "Request aborted",
            timestamp: Date.now()
          };
          stream.push({ type: "error", reason: "aborted", error: abortedMsg });
          stream.end(abortedMsg);
          return;
        }

        // Emit streaming events
        const partial: AssistantMessage = {
          ...assistantMessage,
          content: [],
          stopReason: "pending" as any
        };
        stream.push({ type: "start", partial: { ...partial } });

        for (let i = 0; i < assistantMessage.content.length; i++) {
          const block = assistantMessage.content[i];
          if (block.type === "text") {
            partial.content = [...partial.content, { type: "text", text: "" }];
            stream.push({ type: "text_start", contentIndex: i, partial: { ...partial } });
            partial.content[i] = { type: "text", text: block.text };
            stream.push({ type: "text_delta", contentIndex: i, delta: block.text, partial: { ...partial } });
            stream.push({ type: "text_end", contentIndex: i, content: block.text, partial: { ...partial } });
          } else if (block.type === "toolCall") {
            partial.content = [...partial.content, { type: "toolCall", id: block.id, name: block.name, arguments: {} }];
            stream.push({ type: "toolcall_start", contentIndex: i, partial: { ...partial } });
            partial.content[i] = block;
            stream.push({ type: "toolcall_delta", contentIndex: i, delta: JSON.stringify(block.arguments), partial: { ...partial } });
            stream.push({ type: "toolcall_end", contentIndex: i, toolCall: block, partial: { ...partial } });
          }
        }

        const validDoneReason = assistantMessage.stopReason as "stop" | "toolUse" | "length";
        stream.push({ type: "done", reason: validDoneReason, message: assistantMessage });
        stream.end(assistantMessage);
      } catch (err: any) {
        const errorMsg: AssistantMessage = {
          role: "assistant",
          content: [],
          api: model.api,
          provider: model.provider,
          model: model.id,
          usage: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, totalTokens: 0, cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 } },
          stopReason: "error",
          errorMessage: "Model request failed.",
          timestamp: Date.now()
        };
        stream.push({ type: "error", reason: "error", error: errorMsg });
        stream.end(errorMsg);
      }
    })();

    return stream;
  };

  return { model, streamFn };
}