✨ Explicitely type clients

lib/chats/index.ts

@@ -202,7 +202,12 @@ export class Chat {
     }
 }
 
-export default function client(clientOptions: InputClientOptions = {}) {
+export type ChatClient = {
+    createChat: (systemPrompt?: string) => Promise<string>;
+    Chat: typeof Chat;
+};
+
+export default function client(clientOptions: InputClientOptions = {}): ChatClient {
     return {
         createChat: (systemPrompt?: string) => createChat(systemPrompt, clientOptions),
         Chat: class C extends Chat {

lib/generate.ts

@@ -200,7 +200,18 @@ export function generateStream(
     });
 }
 
-export default function client(clientOptions: InputClientOptions = {}) {
+export type GenerationClient = {
+    generateWithTokenUsage: (
+        task: string,
+        options?: GenerationOptions,
+    ) => Promise<GenerationResult>;
+    generate: (task: string, options?: GenerationOptions) => Promise<string>;
+    generateWithInfo: (task: string, options?: GenerationOptions) => Promise<GenerationResult>;
+    generateStream: (task: string, options?: GenerationOptions) => Readable;
+    generateStreamWithInfos: (task: string, options?: GenerationOptions) => Readable;
+};
+
+export default function client(clientOptions: InputClientOptions = {}): GenerationClient {
     return {
         generateWithTokenUsage: (task: string, options: GenerationOptions = {}) =>
             generateWithTokenUsage(task, options, clientOptions),

lib/index.ts

@@ -9,17 +9,25 @@ import generateClient, {
     generate,
     generateWithTokenUsage,
     GenerationOptions,
+    GenerationClient,
 } from "./generate";
 import generateWithTypeClient, {
     generateWithType,
     generateWithTypeWithTokenUsage,
+    GenerationWithTypeClient,
 } from "./probabilistic_helpers/generateWithType";
-import transcribeClient, { transcribe } from "./transcribe";
-import chatClient, { Chat } from "./chats";
-import memoryClient, { Memory, createMemory, updateMemory, getAllMemories } from "./memory";
+import transcribeClient, { transcribe, TranscribeClient } from "./transcribe";
+import chatClient, { Chat, ChatClient } from "./chats";
+import memoryClient, {
+    Memory,
+    createMemory,
+    updateMemory,
+    getAllMemories,
+    MemoryClient,
+} from "./memory";
 import { splitString, tokenCount } from "./split";
 import { InputClientOptions } from "./clientOpts";
-import kvClient, { get as KVGet, set as KVSet } from "./kv";
+import kvClient, { get as KVGet, set as KVSet, KVClient } from "./kv";
 
 // Export types and models
 export type { TokenUsage, Ressource, GenerationResult } from "./generate";
@@ -52,7 +60,13 @@ export {
     kv,
 };
 
-function client(co: InputClientOptions) {
+type Client = GenerationClient &
+    GenerationWithTypeClient &
+    TranscribeClient &
+    MemoryClient &
+    ChatClient & { kv: KVClient };
+
+function client(co: InputClientOptions): Client {
     return {
         ...generateClient(co),
         ...generateWithTypeClient(co),
@@ -227,8 +241,8 @@ export default Polyfact;
 const reactMutex = new Mutex();
 
 export function usePolyfact({ project, endpoint }: { project: string; endpoint?: string }): {
-    polyfact: ReturnType<typeof client>;
-    login: (input: { provider: "github" }) => Promise<void>;
+    polyfact: Client | undefined;
+    login: ((input: { provider: "github" }) => Promise<void>) | undefined;
     loading: boolean;
 } {
     if (typeof window === "undefined") {

lib/kv.ts

@@ -56,7 +56,12 @@ export async function get(key: string, clientOptions: InputClientOptions = {}):
     }
 }
 
-export default function client(clientOptions: InputClientOptions = {}) {
+export type KVClient = {
+    get: (key: string) => Promise<string>;
+    set: (key: string, value: string) => Promise<void>;
+};
+
+export default function client(clientOptions: InputClientOptions = {}): KVClient {
     return {
         get: (key: string) => get(key, clientOptions),
         set: (key: string, value: string) => set(key, value, clientOptions),

lib/memory.ts

@@ -101,7 +101,14 @@ class Memory {
 
 export { createMemory, updateMemory, getAllMemories, Memory };
 
-export default function client(clientOptions: InputClientOptions = {}) {
+export type MemoryClient = {
+    createMemory: () => Promise<{ id: string }>;
+    updateMemory: (id: string, input: string, maxToken?: number) => Promise<{ success: boolean }>;
+    getAllMemories: () => Promise<{ ids: string[] }>;
+    Memory: () => Memory;
+};
+
+export default function client(clientOptions: InputClientOptions = {}): MemoryClient {
     return {
         createMemory: () => createMemory(clientOptions),
         updateMemory: (id: string, input: string, maxToken?: number) =>

lib/probabilistic_helpers/generateWithType.ts

@@ -133,7 +133,20 @@ export async function generateWithType<T extends t.Props>(
     return res.result;
 }
 
-export default function client(clientOptions: InputClientOptions = {}) {
+export type GenerationWithTypeClient = {
+    generateWithTypeWithTokenUsage: <T extends t.Props>(
+        task: string,
+        type: t.TypeC<T>,
+        options?: GenerationOptions,
+    ) => Promise<{ result: t.TypeOf<t.TypeC<T>>; tokenUsage: { input: number; output: number } }>;
+    generateWithType: <T extends t.Props>(
+        task: string,
+        type: t.TypeC<T>,
+        options?: GenerationOptions,
+    ) => Promise<t.TypeOf<t.TypeC<T>>>;
+};
+
+export default function client(clientOptions: InputClientOptions = {}): GenerationWithTypeClient {
     return {
         generateWithTypeWithTokenUsage: <T extends t.Props>(
             task: string,

lib/transcribe.ts

@@ -43,7 +43,12 @@ export async function transcribe(
         throw e;
     }
 }
-export default function client(clientOptions: InputClientOptions = {}) {
+
+export type TranscribeClient = {
+    transcribe: (file: Buffer | Readable) => Promise<string>;
+};
+
+export default function client(clientOptions: InputClientOptions = {}): TranscribeClient {
     return {
         transcribe: (file: Buffer | Readable) => transcribe(file, clientOptions),
     };