diff --git a/package.json b/package.json
index e616f40..a849b81 100644
--- a/package.json
+++ b/package.json
@@ -1,6 +1,6 @@
 {
   "name": "use-ai-lib",
-  "version": "0.0.2",
+  "version": "0.0.3-alpha",
   "description": "A React hooks library for building AI-powered apps as simple as possible.",
   "type": "module",
   "source": "src/index.ts",
diff --git a/src/hooks/useAIModel.tsx b/src/hooks/useAIModel.tsx
index 47f9436..57c5722 100644
--- a/src/hooks/useAIModel.tsx
+++ b/src/hooks/useAIModel.tsx
@@ -1,7 +1,7 @@
-import type { DeepPartial, Schema } from "@ai-sdk/ui-utils";
+import type { DeepPartial } from "@ai-sdk/ui-utils";
 import type { LanguageModel } from "ai";
 import { useEffect, useMemo } from "react";
-import type { ZodTypeDef, Schema as zSchema } from "zod";
+import type { z } from "zod";
 import { useModelContext } from "../provider";
 import {
 	useGenerateObject,
@@ -15,7 +15,8 @@ interface Options<OBJECT> extends Prompt {
 	/**
 	 * The schema of the object that the model should generate. Use 'zod' to declare.
 	 */
-	schema?: zSchema<OBJECT, ZodTypeDef, OBJECT> | Schema<OBJECT>;
+	// biome-ignore lint/suspicious/noExplicitAny:
+	schema?: z.Schema<OBJECT, z.ZodTypeDef, any>;
 	/**
 	 * Streams the output or not.
 	 */
@@ -24,7 +25,9 @@ interface Options<OBJECT> extends Prompt {
 	 * Do something when AI data is generated.
 	 * Use this callback to get the data during streaming rather than through the final 'data'.
 	 */
-	onSuccess?: (data: string | DeepPartial<OBJECT> | OBJECT) => void;
+	onSuccess?: (
+		data: OBJECT extends string ? string : DeepPartial<OBJECT> | OBJECT,
+	) => void;
 }
 
 interface UseAIModel<D> {
@@ -94,6 +97,7 @@ export function useAIModel<D = string>(
 			...prompt,
 		},
 		{
+			onSuccess: onSuccess as (data: string) => void,
 			enabled: !!stream && !schema && !emptyInput,
 		},
 	);
@@ -104,7 +108,7 @@ export function useAIModel<D = string>(
 		isFetching: isObjectFetching,
 		isError: isObjectError,
 		error: objectError,
-	} = useGenerateObject<D>(
+	} = useGenerateObject<z.infer<typeof schema>>(
 		{
 			model,
 			// biome-ignore lint/style/noNonNullAssertion:
@@ -121,7 +125,7 @@ export function useAIModel<D = string>(
 		isFetching: isStreamObjectFetching,
 		isError: isStreamObjectError,
 		error: streamObjectError,
-	} = useStreamObject(
+	} = useStreamObject<z.infer<typeof schema>>(
 		{
 			model,
 			// biome-ignore lint/style/noNonNullAssertion:
@@ -129,7 +133,7 @@ export function useAIModel<D = string>(
 			...prompt,
 		},
 		{
-			onSuccess,
+			onSuccess: onSuccess as (data: DeepPartial<D> | D) => void,
 			enabled: !!stream && !!schema && !emptyInput,
 		},
 	);
@@ -140,7 +144,7 @@ export function useAIModel<D = string>(
 	// biome-ignore lint: onSuccess usually won't change
 	useEffect(() => {
 		if (!data) return;
-		onSuccess?.(data);
+		(onSuccess as (data: string | D) => void)?.(data);
 	}, [data]);
 
 	return {
diff --git a/src/queries/useStreamObject.tsx b/src/queries/useStreamObject.tsx
index 739205a..f130880 100644
--- a/src/queries/useStreamObject.tsx
+++ b/src/queries/useStreamObject.tsx
@@ -1,4 +1,4 @@
-import type { Schema } from "@ai-sdk/ui-utils";
+import type { DeepPartial } from "@ai-sdk/ui-utils";
 import { useQuery } from "@tanstack/react-query";
 import { streamObject } from "ai";
 import type { LanguageModel } from "ai";
@@ -6,10 +6,9 @@ import type { ZodTypeDef, Schema as zSchema } from "zod";
 import type { CallSettings, Prompt, TelemetrySettings } from "../types";
 
 // use some of options
-type Options = {
+type Options<OBJECT> = {
 	enabled?: boolean;
-	// biome-ignore lint/suspicious/noExplicitAny: TODO
-	onSuccess?: (data: any) => void;
+	onSuccess?: (data: DeepPartial<OBJECT> | OBJECT) => void;
 };
 
 // There are three overloads, use this one currently.
@@ -24,7 +23,7 @@ The language model to use.
 The schema of the object that the model should generate.
 */
 		// biome-ignore lint/suspicious/noExplicitAny: TODO
-		schema: zSchema<OBJECT, ZodTypeDef, any> | Schema<OBJECT>;
+		schema: zSchema<OBJECT, ZodTypeDef, any>;
 		/**
 Optional name of the output that should be generated.
 Used by some providers for additional LLM guidance, e.g.
@@ -64,7 +63,7 @@ Callback that is called when the LLM response and the final object validation ar
 
 export function useStreamObject<OBJECT>(
 	params: StreamObjectParams<OBJECT>,
-	options?: Options,
+	options?: Options<OBJECT>,
 ) {
 	const query = useQuery<OBJECT>({
 		queryKey: ["streamObject", JSON.stringify(params.messages)],