@inferrlm/react-native-mlx 0.2.0-inferrlm.2 → 0.4.1

package/src/tool-utils.ts

@@ -0,0 +1,74 @@
+import type { AnyMap } from 'react-native-nitro-modules'
+import type { z } from 'zod'
+import type { ToolDefinition, ToolParameter, ToolParameterType } from './specs/LLM.nitro'
+
+type ZodObjectSchema = z.ZodObject<z.core.$ZodShape>
+type InferArgs<T extends ZodObjectSchema> = z.infer<T>
+
+export interface TypeSafeToolDefinition<T extends ZodObjectSchema> {
+  name: string
+  description: string
+  arguments: T
+  handler: (args: InferArgs<T>) => Promise<Record<string, unknown>>
+}
+
+function getZodTypeString(zodType: z.ZodType): ToolParameterType {
+  const typeName = zodType._zod.def.type
+  switch (typeName) {
+    case 'string':
+      return 'string'
+    case 'number':
+    case 'int':
+      return 'number'
+    case 'boolean':
+      return 'boolean'
+    case 'array':
+      return 'array'
+    case 'object':
+      return 'object'
+    case 'optional':
+      return getZodTypeString((zodType as z.ZodOptional<z.ZodType>)._zod.def.innerType)
+    case 'default':
+      return getZodTypeString((zodType as z.ZodDefault<z.ZodType>)._zod.def.innerType)
+    default:
+      return 'string'
+  }
+}
+
+function isZodOptional(zodType: z.ZodType): boolean {
+  const typeName = zodType._zod.def.type
+  return typeName === 'optional' || typeName === 'default'
+}
+
+function zodSchemaToParameters(schema: ZodObjectSchema): ToolParameter[] {
+  const shape = schema._zod.def.shape
+  const parameters: ToolParameter[] = []
+
+  for (const [key, zodType] of Object.entries(shape)) {
+    const zType = zodType as z.ZodType
+    parameters.push({
+      name: key,
+      type: getZodTypeString(zType),
+      description: zType.description ?? '',
+      required: !isZodOptional(zType),
+    })
+  }
+
+  return parameters
+}
+
+export function createTool<T extends ZodObjectSchema>(
+  definition: TypeSafeToolDefinition<T>,
+): ToolDefinition {
+  return {
+    name: definition.name,
+    description: definition.description,
+    parameters: zodSchemaToParameters(definition.arguments),
+    handler: async (args: AnyMap) => {
+      const argsObj = args as unknown as Record<string, unknown>
+      const parsedArgs = definition.arguments.parse(argsObj)
+      const result = await definition.handler(parsedArgs)
+      return result as unknown as AnyMap
+    },
+  }
+}

package/src/tts.ts

@@ -0,0 +1,60 @@
+import { NitroModules } from 'react-native-nitro-modules'
+import type {
+  TTS as TTSSpec,
+  TTSLoadOptions,
+  TTSGenerateOptions,
+} from './specs/TTS.nitro'
+
+let instance: TTSSpec | null = null
+
+function getInstance(): TTSSpec {
+  if (!instance) {
+    instance = NitroModules.createHybridObject<TTSSpec>('TTS')
+  }
+  return instance
+}
+
+export const TTS = {
+  load(modelId: string, options?: TTSLoadOptions): Promise<void> {
+    return getInstance().load(modelId, options)
+  },
+
+  generate(
+    text: string,
+    options?: TTSGenerateOptions
+  ): Promise<ArrayBuffer> {
+    return getInstance().generate(text, options)
+  },
+
+  stream(
+    text: string,
+    onAudioChunk: (audio: ArrayBuffer) => void,
+    options?: TTSGenerateOptions
+  ): Promise<void> {
+    return getInstance().stream(text, onAudioChunk, options)
+  },
+
+  stop(): void {
+    getInstance().stop()
+  },
+
+  unload(): void {
+    getInstance().unload()
+  },
+
+  get isLoaded(): boolean {
+    return getInstance().isLoaded
+  },
+
+  get isGenerating(): boolean {
+    return getInstance().isGenerating
+  },
+
+  get modelId(): string {
+    return getInstance().modelId
+  },
+
+  get sampleRate(): number {
+    return getInstance().sampleRate
+  },
+}