From 2ae5b3359225ede6fe8a1889641098c63e1854c1 Mon Sep 17 00:00:00 2001 From: Ian Arawjo Date: Sat, 22 Feb 2025 23:16:29 -0500 Subject: [PATCH] wip --- chainforge/react-server/src/JoinNode.tsx | 6 ++--- chainforge/react-server/src/PromptNode.tsx | 10 +++---- chainforge/react-server/src/SplitNode.tsx | 4 +-- chainforge/react-server/src/VisNode.tsx | 4 +-- .../src/backend/__test__/utils.test.ts | 3 --- .../react-server/src/backend/backend.ts | 1 + chainforge/react-server/src/backend/cache.ts | 26 ++++++++++++++++++- chainforge/react-server/src/backend/query.ts | 11 +++++--- chainforge/react-server/src/backend/typing.ts | 9 ++++--- chainforge/react-server/src/backend/utils.ts | 16 +++++------- 10 files changed, 57 insertions(+), 33 deletions(-) diff --git a/chainforge/react-server/src/JoinNode.tsx b/chainforge/react-server/src/JoinNode.tsx index 30d84a4..092d35f 100644 --- a/chainforge/react-server/src/JoinNode.tsx +++ b/chainforge/react-server/src/JoinNode.tsx @@ -78,8 +78,8 @@ const displayJoinedTexts = ( return textInfos.map((info, idx) => { const llm_name = typeof info !== "string" - ? typeof info.llm === "string" - ? info.llm + ? typeof info.llm === "string" || typeof info.llm === "number" + ? StringLookup.get(info.llm) : info.llm?.name : ""; const ps = ( @@ -326,7 +326,7 @@ const JoinNode: React.FC = ({ data, id }) => { let joined_texts: (TemplateVarInfo | string)[] = []; const [groupedRespsByLLM, nonLLMRespGroup] = groupResponsesBy( resp_objs, - (r) => (typeof r.llm === "string" ? r.llm : r.llm?.key), + (r) => (typeof r.llm === "string" || typeof r.llm === "number" ? StringLookup.get(r.llm) : r.llm?.key), ); // eslint-disable-next-line Object.entries(groupedRespsByLLM).forEach(([llm_key, resp_objs]) => { diff --git a/chainforge/react-server/src/PromptNode.tsx b/chainforge/react-server/src/PromptNode.tsx index 62a7ea1..9d28560 100644 --- a/chainforge/react-server/src/PromptNode.tsx +++ b/chainforge/react-server/src/PromptNode.tsx @@ -487,7 +487,7 @@ const PromptNode: React.FC = ({ (info: TemplateVarInfo) => { // Add to unique LLMs list, if necessary if ( - typeof info?.llm !== "string" && + (typeof info?.llm !== "string" && typeof info?.llm !== "number") && info?.llm?.name !== undefined && !llm_names.has(info.llm.name) ) { @@ -508,7 +508,7 @@ const PromptNode: React.FC = ({ // Append any present system message retroactively as the first message in the chat history: if ( - typeof info?.llm !== "string" && + (typeof info?.llm !== "string" && typeof info?.llm !== "number") && typeof info?.llm?.settings?.system_msg === "string" && updated_chat_hist[0].role !== "system" ) @@ -521,7 +521,7 @@ const PromptNode: React.FC = ({ messages: updated_chat_hist, fill_history: info.fill_history ?? {}, metavars: info.metavars ?? {}, - llm: typeof info?.llm === "string" ? info.llm : info?.llm?.name, + llm: (typeof info?.llm === "string" || typeof info?.llm === "number") ? (StringLookup.get(info.llm) ?? "(LLM lookup failed)") : info?.llm?.name, uid: uuid(), }; }, @@ -969,8 +969,8 @@ Soft failing by replacing undefined with empty strings.`, // Add a meta var to keep track of which LLM produced this response o.metavars[llm_metavar_key] = - typeof resp_obj.llm === "string" - ? resp_obj.llm + (typeof resp_obj.llm === "string" || typeof resp_obj.llm === "number") + ? (StringLookup.get(resp_obj.llm) ?? "(LLM lookup failed)") : resp_obj.llm.name; return o; }), diff --git a/chainforge/react-server/src/SplitNode.tsx b/chainforge/react-server/src/SplitNode.tsx index dae9f73..1d2c62e 100644 --- a/chainforge/react-server/src/SplitNode.tsx +++ b/chainforge/react-server/src/SplitNode.tsx @@ -28,7 +28,7 @@ import { } from "./backend/utils"; import { fromMarkdown } from "mdast-util-from-markdown"; -import StorageCache from "./backend/cache"; +import StorageCache, { StringLookup } from "./backend/cache"; import { ResponseBox } from "./ResponseBoxes"; import { Root, RootContent } from "mdast"; import { Dict, TemplateVarInfo } from "./backend/typing"; @@ -289,7 +289,7 @@ const SplitNode: React.FC = ({ data, id }) => { .map((resp_obj: TemplateVarInfo | string) => { if (typeof resp_obj === "string") return splitText(resp_obj, formatting, true); - const texts = splitText(resp_obj?.text ?? "", formatting, true); + const texts = splitText(StringLookup.get(resp_obj?.text) ?? "", formatting, true); if (texts !== undefined && texts.length >= 1) return texts.map( (t: string) => diff --git a/chainforge/react-server/src/VisNode.tsx b/chainforge/react-server/src/VisNode.tsx index 7ee6da5..66d1c69 100644 --- a/chainforge/react-server/src/VisNode.tsx +++ b/chainforge/react-server/src/VisNode.tsx @@ -225,8 +225,8 @@ const VisNode: React.FC = ({ data, id }) => { const get_llm = (resp_obj: LLMResponse) => { if (selectedLLMGroup === "LLM") - return typeof resp_obj.llm === "string" - ? resp_obj.llm + return typeof resp_obj.llm === "string" || typeof resp_obj.llm === "number" + ? (StringLookup.get(resp_obj.llm) ?? "(LLM lookup failed)") : resp_obj.llm?.name; else return resp_obj.metavars[selectedLLMGroup] as string; }; diff --git a/chainforge/react-server/src/backend/__test__/utils.test.ts b/chainforge/react-server/src/backend/__test__/utils.test.ts index e6e1dbe..1bd9233 100644 --- a/chainforge/react-server/src/backend/__test__/utils.test.ts +++ b/chainforge/react-server/src/backend/__test__/utils.test.ts @@ -17,7 +17,6 @@ test("merge response objects", () => { // Merging two response objects const A: RawLLMResponseObject = { responses: ["x", "y", "z"], - raw_response: ["x", "y", "z"], prompt: "this is a test", query: {}, llm: NativeLLM.OpenAI_ChatGPT, @@ -27,7 +26,6 @@ test("merge response objects", () => { }; const B: RawLLMResponseObject = { responses: ["a", "b", "c"], - raw_response: { B: "B" }, prompt: "this is a test 2", query: {}, llm: NativeLLM.OpenAI_ChatGPT, @@ -40,7 +38,6 @@ test("merge response objects", () => { expect(JSON.stringify(C.responses)).toBe( JSON.stringify(["x", "y", "z", "a", "b", "c"]), ); - expect(C.raw_response).toHaveLength(4); expect(Object.keys(C.vars)).toHaveLength(2); expect(Object.keys(C.vars)).toContain("varB1"); expect(Object.keys(C.metavars)).toHaveLength(1); diff --git a/chainforge/react-server/src/backend/backend.ts b/chainforge/react-server/src/backend/backend.ts index 520bcf2..fc3a5d1 100644 --- a/chainforge/react-server/src/backend/backend.ts +++ b/chainforge/react-server/src/backend/backend.ts @@ -260,6 +260,7 @@ function filterVarsByLLM(vars: PromptVarsDict, llm_key: string): Dict { typeof v === "number" || v?.llm === undefined || typeof v.llm === "string" || + typeof v.llm === "number" || v.llm.key === llm_key, ); }); diff --git a/chainforge/react-server/src/backend/cache.ts b/chainforge/react-server/src/backend/cache.ts index 7939d86..3cc72c4 100644 --- a/chainforge/react-server/src/backend/cache.ts +++ b/chainforge/react-server/src/backend/cache.ts @@ -164,7 +164,7 @@ export class StringLookup { return s.stringToIndex.get(str)!; // Return existing index } - // Add new string + // Add new string to the table const index = s.indexToString.length; s.indexToString.push(str); s.stringToIndex.set(str, index); @@ -179,6 +179,30 @@ export class StringLookup { return s.indexToString[index]; // O(1) lookup } + /** + * Transforms a Dict by interning all strings encountered, up to 1 level of depth, + * and returning the modified Dict with the strings as hash indexes instead. + */ + public static internDict(d: Dict, inplace?: boolean): Dict { + const newDict = inplace ? d : {} as Dict; + const entries = Object.entries(d); + + for (const [key, value] of entries) { + if (typeof value === "string") { + newDict[key] = StringLookup.intern(value); + } else if (Array.isArray(value) && value.every(v => typeof v === "string")) { + newDict[key] = value.map(v => StringLookup.intern(v)); + } else if (typeof value === "object" && value !== null) { + newDict[key] = StringLookup.internDict(value as Dict, inplace); + } else { + if (!inplace) + newDict[key] = value; + } + } + + return newDict as Map; + } + /** Serializes interned strings and their mappings */ public static toJSON() { const s = StringLookup.getInstance(); diff --git a/chainforge/react-server/src/backend/query.ts b/chainforge/react-server/src/backend/query.ts index c781a5e..ce28962 100644 --- a/chainforge/react-server/src/backend/query.ts +++ b/chainforge/react-server/src/backend/query.ts @@ -21,7 +21,7 @@ import { repairCachedResponses, compressBase64Image, } from "./utils"; -import StorageCache from "./cache"; +import StorageCache, { StringLookup } from "./cache"; import { UserForcedPrematureExit } from "./errors"; import { typecastSettingsDict } from "../ModelSettingSchemas"; @@ -131,7 +131,6 @@ export class PromptPipeline { query: query ?? {}, uid: uuid(), responses: extracted_resps, - raw_response: contains_imgs ? {} : response ?? {}, // don't double-store images llm, vars: mergeDicts(info, chat_history?.fill_history) ?? {}, metavars: mergeDicts(metavars, chat_history?.metavars) ?? {}, @@ -147,6 +146,9 @@ export class PromptPipeline { resp_obj, past_resp_obj, ) as RawLLMResponseObject; + + // Hash strings present in the response object, to improve performance + StringLookup.internDict(resp_obj, true); // Save the current state of cache'd responses to a JSON file // NOTE: We do this to save money --in case something breaks between calls, can ensure we got the data! @@ -291,7 +293,6 @@ export class PromptPipeline { query: cached_resp.query, uid: cached_resp.uid ?? uuid(), responses: extracted_resps.slice(0, n), - raw_response: cached_resp.raw_response, llm: cached_resp.llm || NativeLLM.OpenAI_ChatGPT, // We want to use the new info, since 'vars' could have changed even though // the prompt text is the same (e.g., "this is a tool -> this is a {x} where x='tool'") @@ -300,6 +301,10 @@ export class PromptPipeline { }; if (chat_history !== undefined) resp.chat_history = chat_history.messages; + + // Hash any strings in the cache'd response object, if any + StringLookup.internDict(resp, true); + yield resp; continue; } diff --git a/chainforge/react-server/src/backend/typing.ts b/chainforge/react-server/src/backend/typing.ts index 0da8daf..5b4fd48 100644 --- a/chainforge/react-server/src/backend/typing.ts +++ b/chainforge/react-server/src/backend/typing.ts @@ -185,7 +185,7 @@ export interface BaseLLMResponseObject { /** Any associated metavariables. */ metavars: Dict; /** The LLM to query (usually a dict of settings) */ - llm: string | LLMSpec; + llm: StringOrHash | LLMSpec; /** Optional: The chat history to pass the LLM */ chat_history?: ChatHistory; } @@ -195,7 +195,8 @@ export interface RawLLMResponseObject extends BaseLLMResponseObject { // A snapshot of the exact query (payload) sent to the LLM's API query: Dict; // The raw JSON response from the LLM - raw_response: Dict; + // NOTE: This is now deprecated since it wastes precious storage space. + // raw_response: Dict; // Extracted responses (1 or more) from raw_response responses: LLMResponseData[]; // Token lengths (if given) @@ -246,10 +247,10 @@ export interface TemplateVarInfo { image?: StringOrHash; // base-64 encoding fill_history?: Dict; metavars?: Dict; - associate_id?: string; + associate_id?: StringOrHash; prompt?: StringOrHash; uid?: ResponseUID; - llm?: string | LLMSpec; + llm?: StringOrHash | LLMSpec; chat_history?: ChatHistory; } diff --git a/chainforge/react-server/src/backend/utils.ts b/chainforge/react-server/src/backend/utils.ts index d68ad25..16f92da 100644 --- a/chainforge/react-server/src/backend/utils.ts +++ b/chainforge/react-server/src/backend/utils.ts @@ -48,7 +48,7 @@ import { fromModelId, ChatMessage as BedrockChatMessage, } from "@mirai73/bedrock-fm"; -import StorageCache from "./cache"; +import StorageCache, { StringLookup } from "./cache"; import Compressor from "compressorjs"; import { Models } from "@mirai73/bedrock-fm/lib/bedrock"; @@ -1904,13 +1904,8 @@ export function merge_response_objs( else if (!resp_obj_A && resp_obj_B) return resp_obj_B; resp_obj_A = resp_obj_A as RawLLMResponseObject; // required by typescript resp_obj_B = resp_obj_B as RawLLMResponseObject; - let raw_resp_A = resp_obj_A.raw_response; - let raw_resp_B = resp_obj_B.raw_response; - if (!Array.isArray(raw_resp_A)) raw_resp_A = [raw_resp_A]; - if (!Array.isArray(raw_resp_B)) raw_resp_B = [raw_resp_B]; const res: RawLLMResponseObject = { responses: resp_obj_A.responses.concat(resp_obj_B.responses), - raw_response: raw_resp_A.concat(raw_resp_B), prompt: resp_obj_B.prompt, query: resp_obj_B.query, llm: resp_obj_B.llm, @@ -2065,7 +2060,7 @@ export const stripLLMDetailsFromResponses = ( ): LLMResponse[] => resps.map((r) => ({ ...r, - llm: typeof r?.llm === "string" ? r?.llm : r?.llm?.name ?? "undefined", + llm: (typeof r?.llm === "string" || typeof r?.llm === "number" ? StringLookup.get(r?.llm) : r?.llm?.name) ?? "undefined", })); // NOTE: The typing is purposefully general since we are trying to cast to an expected format. @@ -2114,6 +2109,7 @@ export const tagMetadataWithLLM = (input_data: LLMResponsesByVarDict) => { typeof r === "number" || !r?.llm || typeof r.llm === "string" || + typeof r.llm === "number" || !r.llm.key ) return r; @@ -2130,14 +2126,14 @@ export const extractLLMLookup = ( (StringOrHash | TemplateVarInfo | BaseLLMResponseObject | LLMResponse)[] >, ) => { - const llm_lookup: Dict = {}; + const llm_lookup: Dict = {}; Object.values(input_data).forEach((resp_objs) => { resp_objs.forEach((r) => { const llm_name = typeof r === "string" || typeof r === "number" ? undefined - : !r.llm || typeof r.llm === "string" - ? r.llm + : !r.llm || typeof r.llm === "string" || typeof r.llm === "number" + ? StringLookup.get(r.llm) : r.llm.key; if ( typeof r === "string" || typeof r === "number" ||