polyfact - npm Package Compare versions

Comparing version 0.1.25 to 0.1.26

examples/memory.ts

@@ -1,2 +0,3 @@
-import { createMemory, updateMemory, getAllMemories } from "../lib/index";
+import { generateWithInfo } from "../lib/generate";
+import { createMemory, updateMemory, getAllMemories, generate } from "../lib/index";
 
@@ -7,3 +8,3 @@ (async () => {
 
-    const result = await updateMemory(memory.id, "<input-data>");
+    const result = await updateMemory(memory.id, "<input-data> is here", 0);
     console.log("Updated memory:", result);
@@ -13,2 +14,5 @@
     console.log("All memories:", memories);
+
+    const result2 = await generateWithInfo("<input-data>", { memoryId: memory.id, infos: true });
+    console.log("Generated:", result2);
 })();

examples/repl.ts

@@ -14,2 +14,3 @@ import * as readline from "node:readline/promises";
         const stream = chat.sendMessageStream(userInput);
+
         stream.pipe(output);
@@ -16,0 +17,0 @@ await new Promise((res) => stream.on("end", res));

lib/chats/index.ts

 import axios from "axios";
 import * as t from "polyfact-io-ts";
 import { Readable, PassThrough } from "stream";
-import { generateStream, generateWithTokenUsage, GenerationOptions } from "../generate";
+import {
+    generateStream,
+    generateStreamWithInfos,
+    generateWithTokenUsage,
+    GenerationOptions,
+    GenerationResult,
+} from "../generate";
 import { ClientOptions, defaultOptions } from "../clientOpts";
@@ -62,3 +68,3 @@ import { Memory } from "../memory";
         options: GenerationOptions = {},
-    ): Promise<{ result: string; tokenUsage: { input: number; output: number } }> {
+    ): Promise<GenerationResult> {
         const chatId = await this.chatId;
@@ -90,2 +96,43 @@
 
+    sendMessageStreamWithInfos(message: string, options: GenerationOptions = {}): Readable {
+        const resultStream = new Readable({
+            read() {},
+            objectMode: true,
+        });
+        const bufs: Buffer[] = [];
+
+        (async () => {
+            const chatId = await this.chatId;
+
+            if (this.autoMemory && !options.memory && !options.memoryId) {
+                options.memory = this.autoMemory;
+            }
+
+            const result = generateStreamWithInfos(
+                message,
+                { ...options, chatId },
+                this.clientOptions,
+            );
+
+            result.on("infos", (data) => {
+                resultStream.emit("infos", data);
+            });
+
+            result.on("data", (d) => {
+                bufs.push(d);
+                resultStream.push(d);
+            });
+            result.on("end", () => {
+                resultStream.push(null);
+                if (this.autoMemory) {
+                    const totalResult = Buffer.concat(bufs).toString("utf8");
+                    this.autoMemory.add(`Human: ${message}`);
+                    this.autoMemory.add(`AI: ${totalResult}`);
+                }
+            });
+        })();
+
+        return resultStream;
+    }
+
     sendMessageStream(message: string, options: GenerationOptions = {}): Readable {
@@ -92,0 +139,0 @@ const resultStream = new PassThrough();

lib/generate.ts

@@ -27,3 +27,7 @@ import axios from "axios";
 
-const ResultType = t.type({
+const PartialResultType = t.partial({
+    ressources: t.array(t.type({ id: t.string, content: t.string, similarity: t.number })),
+});
+
+const Required = t.type({
     result: t.string,
@@ -36,2 +40,4 @@ token_usage: t.type({
 
+const GenerationAPIResponse = t.intersection([Required, PartialResultType]);
+
 export type GenerationOptions = {
@@ -45,7 +51,29 @@ provider?: "openai" | "cohere";
 
+export type GenerationOptionsWithInfos = GenerationOptions & { infos?: boolean };
+
+export type TokenUsage = {
+    input: number;
+    output: number;
+};
+
+export type Ressource = {
+    similarity: number;
+    id: string;
+    content: string;
+};
+
+export type GenerationResult = {
+    result: string;
+    tokenUsage: {
+        input: number;
+        output: number;
+    };
+    ressources?: Ressource[];
+};
+
 export async function generateWithTokenUsage(
     task: string,
-    options: GenerationOptions = {},
+    options: GenerationOptionsWithInfos = {},
     clientOptions: Partial<ClientOptions> = {},
-): Promise<{ result: string; tokenUsage: { input: number; output: number } }> {
+): Promise<GenerationResult> {
     const { token, endpoint } = defaultOptions(clientOptions);
@@ -60,2 +88,3 @@ const requestBody: {
         stop: GenerationOptions["stop"];
+        infos: boolean;
     } = {
@@ -67,2 +96,3 @@ task,
         stop: options?.stop || [],
+        infos: options?.infos || false,
     };
@@ -78,9 +108,11 @@
 
-        const responseData = res.data;
-
-        if (!ResultType.is(responseData)) {
+        if (!GenerationAPIResponse.is(res)) {
             throw new GenerationError();
         }
 
-        return { result: responseData.result, tokenUsage: responseData.token_usage };
+        return {
+            result: res.result,
+            tokenUsage: res.token_usage,
+            ressources: res.ressources,
+        };
     } catch (e) {
@@ -104,6 +136,18 @@ if (e instanceof Error) {
 
-export function generateStream(
+export async function generateWithInfo(
     task: string,
-    options: GenerationOptions = {},
+    options: GenerationOptionsWithInfos = {},
     clientOptions: Partial<ClientOptions> = {},
+): Promise<GenerationResult> {
+    options.infos = true;
+    const res = await generateWithTokenUsage(task, options, clientOptions);
+
+    return res;
+}
+
+function stream(
+    task: string,
+    options: GenerationOptionsWithInfos = {},
+    clientOptions: Partial<ClientOptions> = {},
+    onMessage: (data: any, resultStream: Readable) => void,
 ): Readable {
@@ -122,2 +166,3 @@ const resultStream = new Readable({
             stop: GenerationOptions["stop"];
+            infos?: boolean;
         } = {
@@ -129,24 +174,54 @@ task,
             stop: options?.stop || [],
+            infos: options?.infos || false,
         };
 
         const { token, endpoint } = defaultOptions(clientOptions);
-        const ws = new WebSocket(`${endpoint.replace("http", "ws")}/stream`, {
-            headers: {
-                "X-Access-Token": token,
-            },
-        });
+        const ws = new WebSocket(`${endpoint.replace("http", "ws")}/stream?token=${token}`);
 
         ws.onopen = () => ws.send(JSON.stringify(requestBody));
+        ws.onmessage = (data: any) => onMessage(data, resultStream);
+    })();
+    return resultStream;
+}
 
-        ws.onmessage = (data: any) => {
+export function generateStreamWithInfos(
+    task: string,
+    options: GenerationOptions = {},
+    clientOptions: Partial<ClientOptions> = {},
+): Readable {
+    return stream(
+        task,
+        { ...options, infos: true },
+        clientOptions,
+        (data: any, resultStream: Readable) => {
             if (data.data === "") {
                 resultStream.push(null);
+            } else if (data.data.startsWith("[INFOS]:")) {
+                try {
+                    const potentialRessources = JSON.parse(data.data.replace("[INFOS]:", ""));
+                    resultStream.emit("infos", potentialRessources);
+                } catch (e) {
+                    resultStream.push("");
+                }
             } else {
                 resultStream.push(data.data);
             }
-        };
-    })();
-    return resultStream;
+        },
+    );
 }
 
+export function generateStream(
+    task: string,
+    options: GenerationOptions = {},
+    clientOptions: Partial<ClientOptions> = {},
+): Readable {
+    return stream(task, options, clientOptions, (data: any, resultStream: Readable) => {
+        if (data.data === "") {
+            resultStream.push(null);
+        } else {
+            resultStream.push(data.data);
+        }
+    });
+}
+
 export default function client(clientOptions: Partial<ClientOptions> = {}) {
@@ -153,0 +228,0 @@ return {

lib/index.ts

 import * as t from "polyfact-io-ts";
-import generateClient, { generate, generateWithTokenUsage, GenerationOptions } from "./generate";
+
+import generateClient, {
+    generateWithInfo,
+    generateStreamWithInfos,
+    generate,
+    generateWithTokenUsage,
+    GenerationOptions,
+} from "./generate";
+
 import generateWithTypeClient, {
@@ -14,2 +22,7 @@ generateWithType,
 
+// Export types and models
+export type { TokenUsage, Ressource, GenerationResult } from "./generate";
+export * from "./helpers/models";
+
+// KV operations
 const kv = {
@@ -20,2 +33,3 @@ get: KVGet,
 
+// Export methods
 export {
@@ -26,2 +40,4 @@ generate,
     generateWithTypeWithTokenUsage,
+    generateWithInfo,
+    generateStreamWithInfos,
     splitString,
@@ -28,0 +44,0 @@ tokenCount,

lib/memory.ts

@@ -35,4 +35,4 @@ import axios from "axios";
     } catch (e) {
+        console.log(e);
         if (e instanceof Error) {
-            console.log("\n\n\n", e, "\n\n\n");
             throw new MemoryError(e.name);
@@ -39,0 +39,0 @@ }

package.json

 {
     "name": "polyfact",
-    "version": "0.1.25",
+    "version": "0.1.26",
     "main": "dist/index.js",
@@ -5,0 +5,0 @@ "types": "dist/index.d.ts",

No alert changes

Worsened metrics

Total package byte prevSize

319683

decreased by-66.06%

Number of package files

37

decreased by-32.73%

Lines of code

1604

decreased by-29.62%

No dependency changes