Fix Response.Part, AllParts, StreamPart not inferring Schema properly if toolkit is WithHandler (#5685)

Author: janglad

.changeset/thick-steaks-see.md

@@ -0,0 +1,5 @@
+---
+"@effect/ai": patch
+---
+
+Fix Response.Part, AllParts, StreamPart not inferring Schema properly if toolkit is WithHandler

packages/ai/ai/src/Response.ts

@@ -200,7 +200,7 @@ export type AllPartsEncoded =
  */
 export const AllParts = <T extends Toolkit.Any | Toolkit.WithHandler<any>>(
   toolkit: T
-): Schema.Schema<AllParts<Toolkit.Tools<T>>, AllPartsEncoded> => {
+): Schema.Schema<AllParts<T extends Toolkit.Any ? Toolkit.Tools<T> : Toolkit.WithHandlerTools<T>>, AllPartsEncoded> => {
   const toolCalls: Array<Schema.Schema<ToolCallPart<string, any>, ToolCallPartEncoded>> = []
   const toolCallResults: Array<Schema.Schema<ToolResultPart<string, any, any>, ToolResultPartEncoded>> = []
   for (const tool of Object.values(toolkit.tools as Record<string, Tool.Any>)) {
@@ -281,7 +281,7 @@ export type PartEncoded =
  */
 export const Part = <T extends Toolkit.Any | Toolkit.WithHandler<any>>(
   toolkit: T
-): Schema.Schema<Part<Toolkit.Tools<T>>, PartEncoded> => {
+): Schema.Schema<Part<T extends Toolkit.Any ? Toolkit.Tools<T> : Toolkit.WithHandlerTools<T>>, PartEncoded> => {
   const toolCalls: Array<Schema.Schema<ToolCallPart<string, any>, ToolCallPartEncoded>> = []
   const toolCallResults: Array<Schema.Schema<ToolResultPart<string, any, any>, ToolResultPartEncoded>> = []
   for (const tool of Object.values(toolkit.tools as Record<string, Tool.Any>)) {
@@ -365,7 +365,12 @@ export type StreamPartEncoded =
  */
 export const StreamPart = <T extends Toolkit.Any | Toolkit.WithHandler<any>>(
   toolkit: T
-): Schema.Schema<StreamPart<Toolkit.Tools<T>>, StreamPartEncoded> => {
+): Schema.Schema<
+  StreamPart<
+    T extends Toolkit.Any ? Toolkit.Tools<T> : Toolkit.WithHandlerTools<T>
+  >,
+  StreamPartEncoded
+> => {
   const toolCalls: Array<Schema.Schema<ToolCallPart<string, any>, ToolCallPartEncoded>> = []
   const toolCallResults: Array<Schema.Schema<ToolResultPart<string, any, any>, ToolResultPartEncoded>> = []
   for (const tool of Object.values(toolkit.tools as Record<string, Tool.Any>)) {

packages/ai/ai/src/Toolkit.ts

@@ -199,6 +199,15 @@ export type HandlersFrom<Tools extends Record<string, Tool.Any>> = {
   >
 }
 
+/**
+ * A utility type which can be used to extract the tools from a toolkit with handlers.
+ *
+ * @since 1.0.0
+ * @category Utility Types
+ */
+
+export type WithHandlerTools<T> = T extends WithHandler<infer Tools> ? Tools : never
+
 /**
  * A toolkit instance with registered handlers ready for tool execution.
  *