ai 7.0.113 → 7.0.114

package/src/telemetry/filter-included-context.ts

@@ -1,5 +1,13 @@
-import type { Context } from '@ai-sdk/provider-utils';
-import type { IncludedContext } from './telemetry-options';
+import type {
+  Context,
+  InferToolContext,
+  InferToolSetContext,
+  ToolSet,
+} from '@ai-sdk/provider-utils';
+import type {
+  IncludedContext,
+  IncludedToolsContext,
+} from './telemetry-options';
 
 /**
  * Returns a shallow copy of the runtime context with only top-level
@@ -22,3 +30,50 @@ export function filterIncludedContext<CONTEXT extends Context>({
     ),
   );
 }
+
+export function filterToolsContext<TOOLS extends ToolSet>({
+  toolsContext,
+  includeToolsContext,
+}: {
+  toolsContext: InferToolSetContext<TOOLS>;
+  includeToolsContext: IncludedToolsContext<TOOLS>;
+}): InferToolSetContext<TOOLS> {
+  if (includeToolsContext == null) {
+    return {} as InferToolSetContext<TOOLS>;
+  }
+
+  return Object.fromEntries(
+    Object.entries(toolsContext).map(([toolName, toolContext]) => [
+      toolName,
+      filterToolContext({
+        toolName,
+        toolContext,
+        includeToolsContext,
+      }),
+    ]),
+  ) as InferToolSetContext<TOOLS>;
+}
+
+export function filterToolContext<TOOLS extends ToolSet>({
+  toolName,
+  toolContext,
+  includeToolsContext,
+}: {
+  toolName: string;
+  toolContext: unknown;
+  includeToolsContext: IncludedToolsContext<TOOLS>;
+}) {
+  const includeToolContext = (
+    includeToolsContext as
+      | Record<
+          string,
+          IncludedContext<InferToolContext<TOOLS[typeof toolName]>>
+        >
+      | undefined
+  )?.[toolName];
+
+  return filterIncludedContext({
+    context: toolContext as InferToolContext<TOOLS[typeof toolName]>,
+    includeContext: includeToolContext,
+  });
+}