add format

receptron · Jun 20, 2024 · e365601 · e365601
1 parent 8f690ec
commit e365601
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 4 deletions.
diff --git a/src/utils/agentlist.ts b/src/utils/agentlist.ts
@@ -164,6 +164,12 @@ export const agentlist = {
         },
         required: ["id", "object", "created", "model", "choices", "usage"],
       },
+      format: {
+        llmResponse: {
+          key: "choices.$0.message.content",
+          type: "string",
+        },
+      },
       stream: true,
     },
     {

diff --git a/src/utils/setAgentToLiteGraph.ts b/src/utils/setAgentToLiteGraph.ts
@@ -59,6 +59,13 @@ const outputs2outputs = (outputs: any) => {
   return jsonSchemaToIO(outputs, "Output");
 };
 
+const format2output = (format: any) => {
+  return Object.keys(format).map((property) => {
+    const data = format[property];
+    return [property, data.type];
+  });
+};
+
 const jsonSchemaToI2IOType = (inputs: any) => {
   if (!inputs) {
     return [];
@@ -68,6 +75,12 @@ const jsonSchemaToI2IOType = (inputs: any) => {
   }
   return [];
 };
+const format2output2 = (format: any) => {
+  return Object.keys(format).map((property) => {
+    const data = format[property];
+    return data.key
+  });
+};
 
 const setAgentToLiteGraph = (agents: AgentFunctionInfoDictionary) => {
   [
@@ -89,21 +102,22 @@ const setAgentToLiteGraph = (agents: AgentFunctionInfoDictionary) => {
   const lite2inputs: Record<string, string[]> = {};
   const lite2output: Record<string, string[]> = {};
 
-  Object.values(agents).map((agent) => {
+  // TODO remove any after add format to AgentFunctionInfo
+  Object.values(agents).map((agent: any) => {
     if (agent.category) {
-      agent.category.forEach((category) => {
+      agent.category.forEach((category: any) => {
         const name = agent.name.replace(/Agent$/, "");
         const nodeType = [category, name].join("/");
         lite2inputs[nodeType] = jsonSchemaToI2IOType(agent.inputs);
-        lite2output[nodeType] = jsonSchemaToI2IOType(agent.output);
+        lite2output[nodeType] = agent.format ? format2output2(agent.format) : jsonSchemaToI2IOType(agent.output);
         lite2agent[nodeType] = agent;
         LiteGraph.registerNodeType(
           nodeType,
           createAgentNode({
             name: name,
             category: category,
             inputs: inputs2inputs(agent.inputs),
-            outputs: outputs2outputs(agent.output),
+            outputs: agent.format ? format2output(agent.format) : outputs2outputs(agent.output),
           }),
         );
       });