@mupt-ai/dari-gateway 0.1.0 → 0.2.1

package/src/providers/pi.ts

@@ -1,3 +1,22 @@
+import {
+  getSupportedThinkingLevels,
+  normalizeContext,
+  type Api,
+  type AssistantMessageEventStream,
+  type Context,
+  type Model,
+  type ProviderStreams,
+} from "@mupt-ai/pi-ai";
+import { anthropicMessagesApi } from "@mupt-ai/pi-ai/api/anthropic-messages.lazy";
+import { azureOpenAIResponsesApi } from "@mupt-ai/pi-ai/api/azure-openai-responses.lazy";
+import { bedrockConverseStreamApi } from "@mupt-ai/pi-ai/api/bedrock-converse-stream.lazy";
+import { googleGenerativeAIApi } from "@mupt-ai/pi-ai/api/google-generative-ai.lazy";
+import { googleVertexApi } from "@mupt-ai/pi-ai/api/google-vertex.lazy";
+import { mistralConversationsApi } from "@mupt-ai/pi-ai/api/mistral-conversations.lazy";
+import { openAICodexResponsesApi } from "@mupt-ai/pi-ai/api/openai-codex-responses.lazy";
+import { openAICompletionsApi } from "@mupt-ai/pi-ai/api/openai-completions.lazy";
+import { openAIResponsesApi } from "@mupt-ai/pi-ai/api/openai-responses.lazy";
+import { getBuiltinModel } from "@mupt-ai/pi-ai/providers/all";
 import {
   REASONING_EFFORTS,
   isRecord,
@@ -7,29 +26,35 @@ import {
   type ReasoningEffort,
   type RouterModelPrice,
 } from "@mupt-ai/dari-router/policy-engine";
-import { HttpError } from "../http/errors";
-import { effectiveContextWindow } from "../models/context_limits";
-import { defaultBaseUrl } from "../models/provider_defaults";
-import { contextWindowOverride, inputModalitiesOverride, runtimeCompat } from "../models/runtime_overrides";
-import { BEDROCK_API } from "./model_capabilities";
-import { buildContext } from "./pi_context";
-import { providerIdentityFromModel } from "../interaction/provider_continuation_state";
-import { buildOptions, supportedThinkingLevels } from "./pi_options";
-import type { ProviderFetch } from "./response_capture";
+import { HttpError } from "../http/errors.js";
+import { effectiveContextWindow } from "../models/context_limits.js";
+import { defaultBaseUrl } from "../models/provider_defaults.js";
+import { contextWindowOverride, inputModalitiesOverride, runtimeCompat } from "../models/runtime_overrides.js";
+import { BEDROCK_API } from "./model_capabilities.js";
+import { buildContext } from "./pi_context.js";
+import { providerIdentityFromModel } from "../interaction/provider_continuation_state.js";
+import { buildOptions, supportedThinkingLevels } from "./pi_options.js";
+import type { ProviderFetch } from "./response_capture.js";
 import {
   eventsWithSafeguardResults,
   forwardsAnthropicSafeguards,
   messageWithSafeguardResults,
   safeguardResultsTap,
-} from "./anthropic_safeguards";
+} from "./anthropic_safeguards.js";
+import {
+  eventsWithProviderReportedCost,
+  messageWithProviderReportedCost,
+  providerReportedCostTap,
+  reportsProviderCost,
+} from "./provider_reported_cost.js";
 import {
   eventsWithSubscriptionCacheWrites,
   messageWithSubscriptionCacheWrites,
   modelPricingFromPiModel,
   type SubscriptionUsageContext,
-} from "./subscription_usage";
-import type { InteractionRequest, ProviderIdentity } from "../interaction/types";
-import type { ResolvedProviderAuth, ResolvedRouterModelConfig } from "../types";
+} from "./subscription_usage.js";
+import type { InteractionRequest, ProviderIdentity } from "../interaction/types.js";
+import type { ResolvedProviderAuth, ResolvedRouterModelConfig } from "../types.js";
 
 type PiAi = {
   completeSimple: (model: Record<string, unknown>, context: Record<string, unknown>, options: Record<string, unknown>) => Promise<Record<string, unknown>>;
@@ -43,6 +68,12 @@ export type ModelRuntimeMetadata = {
   supportedThinkingLevels: string[];
 };
 
+// A platform-supplied endpoint for one call (see ProviderCredential).
+export type ProviderEndpoint = {
+  baseUrl?: string;
+  headers?: Record<string, string>;
+};
+
 export type ProviderCompletion = {
   message: Record<string, unknown>;
   source: ProviderIdentity;
@@ -55,14 +86,40 @@ export type ProviderEventStream = {
 
 export type ModelMetadataLookup = (modelId: string, pricing?: ModelPricing | null) => ModelRuntimeMetadata;
 
-let piAiPromise: Promise<PiAi> | null = null;
+// Pi's built-in providers stream through these API implementations unchanged.
+// Dari always passes resolved credentials, so it dispatches on the API rather
+// than through a Models collection with its own auth resolution.
+const PI_APIS: Record<string, ProviderStreams> = {
+  "anthropic-messages": anthropicMessagesApi(),
+  "azure-openai-responses": azureOpenAIResponsesApi(),
+  "bedrock-converse-stream": bedrockConverseStreamApi(),
+  "google-generative-ai": googleGenerativeAIApi(),
+  "google-vertex": googleVertexApi(),
+  "mistral-conversations": mistralConversationsApi(),
+  "openai-codex-responses": openAICodexResponsesApi(),
+  "openai-completions": openAICompletionsApi(),
+  "openai-responses": openAIResponsesApi(),
+};
 
-async function piAi(): Promise<PiAi> {
-  // Pi 0.80 keeps the legacy global dispatch and catalog surface under /compat.
-  piAiPromise ??= import("@mupt-ai/pi-ai/compat") as unknown as Promise<PiAi>;
-  return piAiPromise;
+function piStreamSimple(
+  model: Record<string, unknown>,
+  context: Record<string, unknown>,
+  options: Record<string, unknown>,
+): AssistantMessageEventStream {
+  const api = PI_APIS[String(model.api)];
+  if (!api) throw new Error(`No Pi API implementation for ${String(model.api)}.`);
+  return api.streamSimple(model as unknown as Model<Api>, normalizeContext(context as unknown as Context), options);
 }
 
+export const pi: PiAi = {
+  completeSimple: async (model, context, options) =>
+    await piStreamSimple(model, context, options).result() as unknown as Record<string, unknown>,
+  streamSimple: (model, context, options) =>
+    piStreamSimple(model, context, options) as unknown as AsyncIterable<Record<string, unknown>>,
+  getModel: (provider, id) => getBuiltinModel(provider as never, id as never) as unknown as Record<string, unknown> | undefined,
+  getSupportedThinkingLevels: (model) => getSupportedThinkingLevels(model as unknown as Model<Api>),
+};
+
 export async function completeInteractionModel(args: {
   model: string;
   request: InteractionRequest;
@@ -70,18 +127,19 @@ export async function completeInteractionModel(args: {
   auth?: ResolvedProviderAuth;
   pricing?: ModelPricing | null;
   config?: ResolvedRouterModelConfig;
+  endpoint?: ProviderEndpoint;
   sessionId?: string;
   selectedReasoningEffort?: ReasoningEffort;
   timeoutMs?: number;
   signal?: AbortSignal;
   fetch?: ProviderFetch;
 }): Promise<ProviderCompletion> {
-  const pi = await piAi();
-  const model = buildModel(pi, args.model, args.pricing, args.config);
+  const model = modelAtEndpoint(buildModel(pi, args.model, args.pricing, args.config), args.endpoint);
   const source = providerIdentityFromModel(model);
   const context = buildContext(args.request, source);
+  const costTap = reportsProviderCost(model) ? providerReportedCostTap(args.fetch) : undefined;
   const safeguards = forwardsAnthropicSafeguards(args.request, model)
-    ? safeguardResultsTap(args.fetch)
+    ? safeguardResultsTap(costTap?.fetch ?? args.fetch)
     : undefined;
   const options = buildOptions(
     args.request,
@@ -93,16 +151,19 @@ export async function completeInteractionModel(args: {
     args.selectedReasoningEffort,
     args.timeoutMs,
     args.signal,
-    safeguards?.fetch ?? args.fetch,
+    safeguards?.fetch ?? costTap?.fetch ?? args.fetch,
   );
   const usageContext = subscriptionUsageContext(args.model, model, source);
   const message = options.emitProviderEvents === true
     ? await completeProviderEventStream(pi.streamSimple(model, context, options))
     : await pi.completeSimple(model, context, options);
   return {
-    message: messageWithSubscriptionCacheWrites(
-      messageWithSafeguardResults(message, safeguards?.results()),
-      usageContext
+    message: messageWithProviderReportedCost(
+      messageWithSubscriptionCacheWrites(
+        messageWithSafeguardResults(message, safeguards?.results()),
+        usageContext
+      ),
+      costTap?.cost(),
     ),
     source,
   };
@@ -167,16 +228,17 @@ export async function streamInteractionModel(args: {
   auth?: ResolvedProviderAuth;
   pricing?: ModelPricing | null;
   config?: ResolvedRouterModelConfig;
+  endpoint?: ProviderEndpoint;
   sessionId?: string;
   selectedReasoningEffort?: ReasoningEffort;
   signal?: AbortSignal;
   fetch?: ProviderFetch;
 }): Promise<ProviderEventStream> {
-  const pi = await piAi();
-  const model = buildModel(pi, args.model, args.pricing, args.config);
+  const model = modelAtEndpoint(buildModel(pi, args.model, args.pricing, args.config), args.endpoint);
   const source = providerIdentityFromModel(model);
+  const costTap = reportsProviderCost(model) ? providerReportedCostTap(args.fetch) : undefined;
   const safeguards = forwardsAnthropicSafeguards(args.request, model)
-    ? safeguardResultsTap(args.fetch)
+    ? safeguardResultsTap(costTap?.fetch ?? args.fetch)
     : undefined;
   const events = pi.streamSimple(
     model,
@@ -191,20 +253,34 @@ export async function streamInteractionModel(args: {
       args.selectedReasoningEffort,
       undefined,
       args.signal,
-      safeguards?.fetch ?? args.fetch,
+      safeguards?.fetch ?? costTap?.fetch ?? args.fetch,
     )
   );
+  const withCacheWrites = eventsWithSubscriptionCacheWrites(
+    safeguards ? eventsWithSafeguardResults(events, safeguards) : events,
+    subscriptionUsageContext(args.model, model, source)
+  );
   return {
-    events: eventsWithSubscriptionCacheWrites(
-      safeguards ? eventsWithSafeguardResults(events, safeguards) : events,
-      subscriptionUsageContext(args.model, model, source)
-    ),
+    events: costTap ? eventsWithProviderReportedCost(withCacheWrites, costTap) : withCacheWrites,
     source,
   };
 }
 
-export async function modelContextWindow(modelId: string): Promise<number | null> {
-  const pi = await piAi();
+function modelAtEndpoint(
+  model: Record<string, unknown>,
+  endpoint: ProviderEndpoint | undefined,
+): Record<string, unknown> {
+  if (endpoint === undefined) return model;
+  return {
+    ...model,
+    ...(endpoint.baseUrl ? { baseUrl: endpoint.baseUrl } : {}),
+    ...(endpoint.headers
+      ? { headers: { ...(isRecord(model.headers) ? model.headers : {}), ...endpoint.headers } }
+      : {}),
+  };
+}
+
+export function modelContextWindow(modelId: string): number | null {
   const provider = providerForModel(modelId);
   const nativeId = nativeModelId(modelId);
   const registered = resolveModelMetadata(pi, provider, modelId, nativeId);
@@ -215,11 +291,10 @@ export async function modelContextWindow(modelId: string): Promise<number | null
  * Pi's catalog price for a canonical model served by `provider`, in USD per
  * million tokens. Null when Pi does not list the model.
  */
-export async function registeredModelPrice(
+export function registeredModelPrice(
   modelId: string,
   provider: string = providerForModel(modelId),
-): Promise<RouterModelPrice | null> {
-  const pi = await piAi();
+): RouterModelPrice | null {
   const registered = resolveModelMetadata(pi, provider, modelId, nativeModelId(modelId));
   const cost = registered?.cost;
   if (!isRecord(cost) || typeof cost.input !== "number" || typeof cost.output !== "number") {
@@ -233,10 +308,9 @@ export async function registeredModelPrice(
   };
 }
 
-export async function modelMetadataLookup(
+export function modelMetadataLookup(
   configs?: Record<string, ResolvedRouterModelConfig>
-): Promise<ModelMetadataLookup> {
-  const pi = await piAi();
+): ModelMetadataLookup {
   return (modelId: string, pricing?: ModelPricing | null) => {
     const model = buildModel(pi, modelId, pricing, configs?.[modelId]);
     return {

package/src/providers/pi_context.ts

@@ -1,7 +1,7 @@
 import { encodeOpenAIReasoningSignature } from "@mupt-ai/dari-router/protocols";
 import { isRecord } from "@mupt-ai/dari-router/policy-engine";
 
-import { HttpError } from "../http/errors";
+import { HttpError } from "../http/errors.js";
 import type {
   InteractionContent,
   InteractionItem,
@@ -9,14 +9,14 @@ import type {
   InteractionTool,
   ProviderContinuationState,
   ProviderIdentity,
-} from "../interaction/types";
-import { parseDataImageUrl } from "../interaction/image";
-import { providerToolCallId, toolWireName } from "../interaction/tool_identity";
+} from "../interaction/types.js";
+import { parseDataImageUrl } from "../interaction/image.js";
+import { providerToolCallId, toolWireName } from "../interaction/tool_identity.js";
 import {
   compatibleProviderContinuation,
   sameProviderIdentity,
-} from "../interaction/provider_continuation_state";
-import { isOpenAIResponsesApi } from "./model_capabilities";
+} from "../interaction/provider_continuation_state.js";
+import { isOpenAIResponsesApi } from "./model_capabilities.js";
 
 const UNKNOWN_SOURCE: ProviderIdentity = {
   provider: "dari",