Merge pull request #867 from isamu/gemini_agent_stream

Gemini agent stream
receptron · Dec 30, 2024 · c941949 · c941949
2 parents e02347f + 028e044
commit c941949
Show file tree

Hide file tree

Showing 5 changed files with 64 additions and 8 deletions.
diff --git a/llm_agents/gemini_agent/lib/gemini_agent.d.ts b/llm_agents/gemini_agent/lib/gemini_agent.d.ts
@@ -9,6 +9,7 @@ type GeminiInputs = {
 } & GraphAILLMInputBase;
 type GeminiConfig = {
     apiKey?: string;
+    stream?: boolean;
 };
 type GeminiParams = GeminiInputs & GeminiConfig;
 export declare const geminiAgent: AgentFunction<GeminiParams, Record<string, any> | string, GeminiInputs, GeminiConfig>;

diff --git a/llm_agents/gemini_agent/lib/gemini_agent.js b/llm_agents/gemini_agent/lib/gemini_agent.js
@@ -4,9 +4,9 @@ exports.geminiAgent = void 0;
 const graphai_1 = require("graphai");
 const generative_ai_1 = require("@google/generative-ai");
 const llm_utils_1 = require("@graphai/llm_utils");
-const geminiAgent = async ({ params, namedInputs, config }) => {
+const geminiAgent = async ({ params, namedInputs, config, filterParams, }) => {
     const { model, system, temperature, max_tokens, tools, prompt, messages } = { ...params, ...namedInputs };
-    const { apiKey /* stream */ } = {
+    const { apiKey, stream } = {
         ...params,
         ...(config || {}),
     };
@@ -60,6 +60,20 @@ const geminiAgent = async ({ params, namedInputs, config }) => {
         }),
         generationConfig,
     });
+    if (stream) {
+        const result = await chat.sendMessageStream(lastMessage.content);
+        const chunks = [];
+        for await (const chunk of result.stream) {
+            const chunkText = chunk.text();
+            if (filterParams && filterParams.streamTokenCallback && chunkText) {
+                filterParams.streamTokenCallback(chunkText);
+            }
+            chunks.push(chunkText);
+        }
+        const text = chunks.join("");
+        const message = { role: "assistant", content: text };
+        return { choices: [{ message }], text, message };
+    }
     const result = await chat.sendMessage(lastMessage.content);
     const response = result.response;
     const text = response.text();
@@ -111,7 +125,7 @@ const geminiAgentInfo = {
     author: "Receptron team",
     repository: "https://github.com/receptron/graphai",
     license: "MIT",
-    // stream: true,
+    stream: true,
     npms: ["@anthropic-ai/sdk"],
     environmentVariables: ["GOOGLE_GENAI_API_KEY"],
 };

diff --git a/llm_agents/gemini_agent/package.json b/llm_agents/gemini_agent/package.json
@@ -1,6 +1,6 @@
 {
   "name": "@graphai/gemini_agent",
-  "version": "0.2.1",
+  "version": "0.2.2",
   "description": "Gemini agents for GraphAI.",
   "main": "lib/index.js",
   "files": [

diff --git a/llm_agents/gemini_agent/src/gemini_agent.ts b/llm_agents/gemini_agent/src/gemini_agent.ts
@@ -14,15 +14,20 @@ type GeminiInputs = {
 
 type GeminiConfig = {
   apiKey?: string;
-  // stream?: boolean;
+  stream?: boolean;
 };
 
 type GeminiParams = GeminiInputs & GeminiConfig;
 
-export const geminiAgent: AgentFunction<GeminiParams, Record<string, any> | string, GeminiInputs, GeminiConfig> = async ({ params, namedInputs, config }) => {
+export const geminiAgent: AgentFunction<GeminiParams, Record<string, any> | string, GeminiInputs, GeminiConfig> = async ({
+  params,
+  namedInputs,
+  config,
+  filterParams,
+}) => {
   const { model, system, temperature, max_tokens, tools, prompt, messages } = { ...params, ...namedInputs };
 
-  const { apiKey /* stream */ } = {
+  const { apiKey, stream } = {
     ...params,
     ...(config || {}),
   };
@@ -83,6 +88,21 @@ export const geminiAgent: AgentFunction<GeminiParams, Record<string, any> | stri
     generationConfig,
   });
 
+  if (stream) {
+    const result = await chat.sendMessageStream(lastMessage.content);
+    const chunks = [];
+    for await (const chunk of result.stream) {
+      const chunkText = chunk.text();
+      if (filterParams && filterParams.streamTokenCallback && chunkText) {
+        filterParams.streamTokenCallback(chunkText);
+      }
+      chunks.push(chunkText);
+    }
+    const text = chunks.join("");
+    const message: any = { role: "assistant", content: text };
+    return { choices: [{ message }], text, message };
+  }
+
   const result = await chat.sendMessage(lastMessage.content);
   const response = result.response;
   const text = response.text();
@@ -136,7 +156,7 @@ const geminiAgentInfo: AgentFunctionInfo = {
   author: "Receptron team",
   repository: "https://github.com/receptron/graphai",
   license: "MIT",
-  // stream: true,
+  stream: true,
   npms: ["@anthropic-ai/sdk"],
   environmentVariables: ["GOOGLE_GENAI_API_KEY"],
 };

diff --git a/llm_agents/gemini_agent/tests/run_gemini.ts b/llm_agents/gemini_agent/tests/run_gemini.ts
@@ -50,3 +50,24 @@ test("test gemini", async () => {
   }
   assert.deepStrictEqual(true, true);
 });
+
+test("test gemini stream", async () => {
+  const namedInputs = { prompt: ["tell me world history"] };
+  const params = { stream: true };
+  const res = (await geminiAgent({
+    namedInputs,
+    params,
+    filterParams: {
+      streamTokenCallback: (token: string) => {
+        console.log(token);
+      },
+    },
+    debugInfo: { verbose: false, nodeId: "test", retry: 5 },
+  })) as any;
+
+  if (res) {
+    console.log(res.choices[0].message["content"]);
+    console.log(res.text);
+  }
+  assert.deepStrictEqual(true, true);
+});