@dbx-tools/appkit-web-search 0.9.4 → 0.9.6

package/src/search.ts

@@ -28,9 +28,9 @@ import {
   getExecutionContext,
   ValidationError,
 } from "@databricks/appkit";
-import { invoke, resolve, serving } from "@dbx-tools/model";
-import { log, object, string } from "@dbx-tools/shared-core";
-import { openaiChat, openaiResponses } from "@dbx-tools/shared-model";
+import { invoke, resolve, modelCatalog } from "@dbx-tools/model";
+import { log, object, stringUtils } from "@dbx-tools/shared-core";
+import { openaiChat } from "@dbx-tools/shared-model";
 import { MODEL_ENV, SERVING_ENDPOINT_ENV, type ResolvedWebSearchConfig } from "./config.ts";
 import { toCallSettings, webSearchExecuteDefaults } from "./defaults.ts";
 import { detectWebSearchProvider, supportsWebSearch, webSearchToolSpec } from "./provider.ts";
@@ -43,10 +43,10 @@ import {
 import type { WebSearchCitation, WebSearchRequest, WebSearchResult } from "./schema.ts";
 import { runScrapeSearch } from "./scrape.ts";
 
-type WorkspaceClientLike = serving.WorkspaceClientLike & invoke.AuthenticatingClientLike;
+type WorkspaceClientLike = modelCatalog.WorkspaceClientLike & invoke.AuthenticatingClientLike;
 const logger = log.logger("web-search/search");
 const { resolveModel } = resolve;
-const { listServingEndpoints } = serving;
+const { listServingEndpoints } = modelCatalog;
 
 /**
  * How deep the grounding-metadata walk descends. Gemini nests its sources a
@@ -66,6 +66,8 @@ const RATE_LIMITED_STATUS = 429;
 export interface WebSearchContext {
   client: WorkspaceClientLike;
   host: string;
+  /** Trusted identity of the credential carried by `client`. */
+  cacheIdentity?: string;
 }
 
 /**
@@ -77,7 +79,8 @@ export interface WebSearchContext {
 export async function resolveWebSearchContext(): Promise<WebSearchContext> {
   const ctx = getExecutionContext();
   const host = (await ctx.client.config.getHost()).toString();
-  return { client: ctx.client, host };
+  const cacheIdentity = "userId" in ctx ? ctx.userId : ctx.serviceUserId;
+  return { client: ctx.client, host, cacheIdentity };
 }
 
 /**
@@ -104,7 +107,9 @@ async function resolveWebSearchModel(
   config: ResolvedWebSearchConfig,
   requested: string | undefined,
 ): Promise<string | null> {
-  const endpoints = await listServingEndpoints(ctx.client, ctx.host);
+  const endpoints = await listServingEndpoints(ctx.client, ctx.host, {
+    ...(ctx.cacheIdentity ? { cacheIdentity: ctx.cacheIdentity } : {}),
+  });
   // Only deployed, web-search-capable endpoints are candidates.
   const capable = endpoints.filter((e) => supportsWebSearch(e.name));
   const pinned = requested ?? config.model;
@@ -213,15 +218,37 @@ async function postServing(
 /* --------------------------- response extraction --------------------------- */
 
 /**
- * Extract answer text and citations from an OpenAI Responses API payload through
- * the shared reader in `@dbx-tools/shared-model`.
+ * Extract answer text and citations from an OpenAI Responses API payload.
  */
 function fromResponsesPayload(payload: Record<string, unknown>): {
   answer: string;
   citations: WebSearchCitation[];
 } {
-  const { text, citations } = openaiResponses.readResponsesOutput(payload);
-  return { answer: text, citations };
+  const output = Array.isArray(payload.output) ? payload.output : [];
+  const texts: string[] = [];
+  const citations: WebSearchCitation[] = [];
+  const seen = new Set<string>();
+  for (const item of output) {
+    if (!object.isRecord(item)) continue;
+    const parts = openaiChat.chatContentParts(item.content);
+    if (!parts) continue;
+    for (const part of parts) {
+      if (!object.isRecord(part)) continue;
+      const text = stringUtils.trimToEmpty(part.text);
+      if (text) texts.push(text);
+      const annotations = Array.isArray(part.annotations) ? part.annotations : [];
+      for (const annotation of annotations) {
+        if (!object.isRecord(annotation)) continue;
+        const url = stringUtils.trimToEmpty(annotation.url);
+        if (!url || seen.has(url)) continue;
+        seen.add(url);
+        const title = stringUtils.trimToEmpty(annotation.title);
+        citations.push({ url, ...(title ? { title } : {}) });
+      }
+    }
+  }
+  const answer = stringUtils.trimToEmpty(payload.output_text) || texts.join("\n").trim();
+  return { answer, citations };
 }
 
 /**
@@ -243,10 +270,10 @@ function fromChatPayload(payload: Record<string, unknown>): {
   const seen = new Set<string>();
   const visit = (v: unknown, depth: number): void => {
     if (depth > MAX_GROUNDING_WALK_DEPTH || !object.isRecord(v)) return;
-    const url = string.trimToEmpty(v.url) || string.trimToEmpty(v.uri);
+    const url = stringUtils.trimToEmpty(v.url) || stringUtils.trimToEmpty(v.uri);
     if (url && !seen.has(url)) {
       seen.add(url);
-      const title = string.trimToEmpty(v.title);
+      const title = stringUtils.trimToEmpty(v.title);
       citations.push({ url, ...(title ? { title } : {}) });
     }
     for (const val of Object.values(v)) {