diff --git a/readme.md b/readme.md
index 26646db..7d8cb89 100644
--- a/readme.md
+++ b/readme.md
@@ -15,6 +15,10 @@
 
 TODO
 
+## TODO
+
+https://github.com/colinhacks/zod#writing-generic-functions
+
 ## Use Cases
 
 https://platform.openai.com/examples
diff --git a/src/llm.ts b/src/llm.ts
index 24528c7..f43ceb8 100644
--- a/src/llm.ts
+++ b/src/llm.ts
@@ -1,4 +1,4 @@
-import type { ZodType } from 'zod'
+import { ZodRawShape, ZodTypeAny, z } from 'zod'
 
 import * as types from './types'
 
@@ -66,21 +66,33 @@ export class Agentic {
   }
 }
 
-export abstract class BaseLLMCallBuilder<TInput, TOutput, TModelParams> {
+export abstract class BaseLLMCallBuilder<
+  TInput extends ZodRawShape | ZodTypeAny = ZodTypeAny,
+  TOutput extends ZodRawShape | ZodTypeAny = ZodTypeAny,
+  TModelParams extends Record<string, any> = Record<string, any>
+> {
   _options: types.BaseLLMOptions<TInput, TOutput, TModelParams>
 
   constructor(options: types.BaseLLMOptions<TInput, TOutput, TModelParams>) {
     this._options = options
   }
 
-  input(inputSchema: ZodType<TInput>) {
-    this._options.input = inputSchema
-    return this
+  input<U extends ZodRawShape | ZodTypeAny = TInput>(
+    inputSchema: U
+  ): BaseLLMCallBuilder<U, TOutput, TModelParams> {
+    ;(
+      this as unknown as BaseLLMCallBuilder<U, TOutput, TModelParams>
+    )._options.input = inputSchema
+    return this as unknown as BaseLLMCallBuilder<U, TOutput, TModelParams>
   }
 
-  output(outputSchema: ZodType<TOutput>) {
-    this._options.output = outputSchema
-    return this
+  output<U extends ZodRawShape | ZodTypeAny = TOutput>(
+    outputSchema: U
+  ): BaseLLMCallBuilder<TInput, U, TModelParams> {
+    ;(
+      this as unknown as BaseLLMCallBuilder<TInput, U, TModelParams>
+    )._options.output = outputSchema
+    return this as unknown as BaseLLMCallBuilder<TInput, U, TModelParams>
   }
 
   examples(examples: types.LLMExample[]) {
@@ -93,21 +105,21 @@ export abstract class BaseLLMCallBuilder<TInput, TOutput, TModelParams> {
     return this
   }
 
-  abstract call(input?: TInput): Promise<TOutput>
+  abstract call(
+    input?: types.ParsedData<TInput>
+  ): Promise<types.ParsedData<TOutput>>
 
   // TODO
-  // abstract stream(
+  // abstract stream({
   //   input: TInput,
   //   onProgress: types.ProgressFunction
-  // ): Promise<TOutput>
-
-  // abstract stream(onProgress: types.ProgressFunction): Promise<TOutput>
+  // }): Promise<TOutput>
 }
 
 export abstract class ChatModelBuilder<
-  TInput,
-  TOutput,
-  TModelParams
+  TInput extends ZodRawShape | ZodTypeAny = ZodTypeAny,
+  TOutput extends ZodRawShape | ZodTypeAny = ZodTypeAny,
+  TModelParams extends Record<string, any> = Record<string, any>
 > extends BaseLLMCallBuilder<TInput, TOutput, TModelParams> {
   _messages: types.ChatMessage[]
 
@@ -118,7 +130,10 @@ export abstract class ChatModelBuilder<
   }
 }
 
-export class OpenAIChatModelBuilder<TInput, TOutput> extends ChatModelBuilder<
+export class OpenAIChatModelBuilder<
+  TInput extends ZodRawShape | ZodTypeAny = ZodTypeAny,
+  TOutput extends ZodRawShape | ZodTypeAny = ZodTypeAny
+> extends ChatModelBuilder<
   TInput,
   TOutput,
   Omit<types.openai.ChatCompletionParams, 'messages'>
@@ -141,7 +156,11 @@ export class OpenAIChatModelBuilder<TInput, TOutput> extends ChatModelBuilder<
     this._client = client
   }
 
-  override async call(input?: TInput): Promise<TOutput> {
+  override async call(
+    input?: types.ParsedData<TInput>
+  ): Promise<types.ParsedData<TOutput>> {
+    // this._options.output?.describe
     // TODO
+    return true as types.ParsedData<TOutput>
   }
 }
diff --git a/src/temp.ts b/src/temp.ts
new file mode 100644
index 0000000..c735131
--- /dev/null
+++ b/src/temp.ts
@@ -0,0 +1,20 @@
+import dotenv from 'dotenv-safe'
+import { OpenAIClient } from 'openai-fetch'
+import { z } from 'zod'
+
+import { Agentic } from './llm'
+
+dotenv.config()
+async function main() {
+  const openai = new OpenAIClient({ apiKey: process.env.OPENAI_API_KEY! })
+  const $ = new Agentic(openai)
+
+  const ex0 = await $.gpt4(`give me a single boolean value`)
+    .output(z.boolean())
+    // .retry({ attempts: 3 })
+    .call()
+
+  console.log(ex0)
+}
+
+main()
diff --git a/src/temp2.ts b/src/temp2.ts
new file mode 100644
index 0000000..9bae58e
--- /dev/null
+++ b/src/temp2.ts
@@ -0,0 +1,28 @@
+import { ZodRawShape, ZodType, ZodTypeAny, z } from 'zod'
+
+import * as types from './types'
+
+class Test<T extends ZodRawShape | ZodTypeAny = ZodTypeAny> {
+  _schema: T
+
+  schema<U extends ZodRawShape | ZodTypeAny>(schema: U): Test<U> {
+    ;(this as unknown as Test<U>)._schema = schema
+    return this as unknown as Test<U>
+  }
+
+  call(value: types.ParsedData<T>): types.ParsedData<T> {
+    const finalSchema =
+      this._schema instanceof ZodType ? this._schema : z.object(this._schema)
+    return finalSchema.parse(value)
+  }
+}
+
+async function main() {
+  const t = new Test()
+  const t2 = t.schema(z.string())
+  const t3 = t2.call('foo')
+
+  console.log(t3)
+}
+
+main()
diff --git a/src/types.ts b/src/types.ts
index 830e13a..d0c1720 100644
--- a/src/types.ts
+++ b/src/types.ts
@@ -1,43 +1,67 @@
 import * as openai from 'openai-fetch'
-import type { ZodType } from 'zod'
+import {
+  SafeParseReturnType,
+  ZodObject,
+  ZodRawShape,
+  ZodSchema,
+  ZodTypeAny,
+  output,
+  z
+} from 'zod'
 
 export { openai }
 
+export type ParsedData<T extends ZodRawShape | ZodTypeAny> =
+  T extends ZodTypeAny
+    ? output<T>
+    : T extends ZodRawShape
+    ? output<ZodObject<T>>
+    : never
+
+export type SafeParsedData<T extends ZodRawShape | ZodTypeAny> =
+  T extends ZodTypeAny
+    ? SafeParseReturnType<z.infer<T>, ParsedData<T>>
+    : T extends ZodRawShape
+    ? SafeParseReturnType<ZodObject<T>, ParsedData<T>>
+    : never
+
 export interface BaseLLMOptions<
-  TInput = any,
-  TOutput = any,
-  TModelParams = Record<string, any>
+  TInput extends ZodRawShape | ZodTypeAny = ZodTypeAny,
+  TOutput extends ZodRawShape | ZodTypeAny = ZodTypeAny,
+  TModelParams extends Record<string, any> = Record<string, any>
 > {
   provider?: string
   model?: string
   modelParams?: TModelParams
   timeoutMs?: number
 
-  input?: ZodType<TInput>
-  output?: ZodType<TOutput>
+  input?: TInput
+  output?: TOutput
   examples?: LLMExample[]
   retryConfig?: LLMRetryConfig
 }
 
 export interface LLMOptions<
-  TInput = any,
-  TOutput = any,
-  TModelParams = Record<string, any>
+  TInput extends ZodRawShape | ZodTypeAny = ZodTypeAny,
+  TOutput extends ZodRawShape | ZodTypeAny = ZodTypeAny,
+  TModelParams extends Record<string, any> = Record<string, any>
 > extends BaseLLMOptions<TInput, TOutput, TModelParams> {
   promptTemplate?: string
   promptPrefix?: string
   promptSuffix?: string
 }
 
+export type ChatMessageRole = 'user' | 'system' | 'assistant' | 'tool'
+
 export interface ChatMessage {
-  role: 'user' | 'system' | 'assistant' | 'tool'
+  role: ChatMessageRole
   content: string
 }
 
 export interface ChatModelOptions<
-  TInput = any,
-  TOutput = any,
-  TModelParams = Record<string, any>
+  TInput extends ZodRawShape | ZodTypeAny = ZodTypeAny,
+  TOutput extends ZodRawShape | ZodTypeAny = ZodTypeAny,
+  TModelParams extends Record<string, any> = Record<string, any>
 > extends BaseLLMOptions<TInput, TOutput, TModelParams> {
   messages: ChatMessage[]
 }
@@ -51,3 +75,5 @@ export interface LLMRetryConfig {
   attempts: number
   strategy: string
 }
+
+// export type ProgressFunction = (partialResponse: ChatMessage) => void