fix(code-interpreter): use llama for preprocessing

Signed-off-by: Jan Pokorný <[email protected]>
i-am-bee · Oct 23, 2024 · 599d1f2 · 599d1f2
1 parent d6d3ed8
commit 599d1f2
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 5 deletions.
diff --git a/src/runs/execution/factory.ts b/src/runs/execution/factory.ts
@@ -177,7 +177,7 @@ export function createCodeLLM(backend: LLMBackend = LLM_BACKEND) {
       });
       return new IBMvLLM({
         client: vllmClient,
-        modelId: 'ibm/granite-34b-code-instruct',
+        modelId: 'meta-llama/llama-3-1-70b-instruct',
         parameters: {
           method: 'GREEDY',
           stopping: { include_stop_sequence: false, max_new_tokens: MAX_NEW_TOKENS }
@@ -188,7 +188,7 @@ export function createCodeLLM(backend: LLMBackend = LLM_BACKEND) {
       bamClient ??= new BAMClient({ apiKey: BAM_API_KEY ?? undefined });
       return new BAMLLM({
         client: bamClient,
-        modelId: 'ibm/granite-34b-code-instruct',
+        modelId: 'meta-llama/llama-3-1-70b-instruct',
         parameters: {
           decoding_method: 'greedy',
           include_stop_sequence: false,
@@ -200,7 +200,7 @@ export function createCodeLLM(backend: LLMBackend = LLM_BACKEND) {
       if (!WATSONX_API_KEY) throw new Error('Missing WATSONX_API_KEY');
       if (!WATSONX_PROJECT_ID) throw new Error('Missing WATSONX_PROJECT_ID');
       return new WatsonXLLM({
-        modelId: 'ibm/granite-34b-code-instruct',
+        modelId: 'meta-llama/llama-3-1-70b-instruct',
         apiKey: WATSONX_API_KEY,
         projectId: WATSONX_PROJECT_ID,
         region: WATSONX_REGION ?? undefined,

diff --git a/src/runs/execution/tools/helpers.ts b/src/runs/execution/tools/helpers.ts
@@ -30,9 +30,10 @@ import { SimilarityToolOutput } from 'bee-agent-framework/tools/similarity';
 import { GoogleSearchTool } from 'bee-agent-framework/tools/search/googleSearch';
 import { uniqueBy } from 'remeda';
 import { z } from 'zod';
-import { PromptTemplate } from 'bee-agent-framework/template';
+import { LLMChatTemplates } from 'bee-agent-framework/adapters/shared/llmChatTemplates';
 import { DuckDuckGoSearchTool } from 'bee-agent-framework/tools/search/duckDuckGoSearch';
 import { SearchToolOptions, SearchToolOutput } from 'bee-agent-framework/tools/search/base';
+import { PromptTemplate } from 'bee-agent-framework/template';
 
 import { AgentContext } from '../execute.js';
 import { getRunVectorStores } from '../helpers.js';
@@ -185,7 +186,24 @@ export async function getTools(run: LoadedRun, context: AgentContext): Promise<F
                   llm: codeLLM,
                   promptTemplate: new PromptTemplate({
                     schema: z.object({ input: z.string() }),
-                    template: `Your task is to fix provided Python Code that may or may not contain a syntax error.\nIMPORTANT: the output must not contain any additional comments or explanation.\n\nInput:\n{{input}}\n\nOutput:\n`
+                    template: LLMChatTemplates.get('llama3.1').template.render({
+                      messages: [
+                        {
+                          system: [
+                            `Your task is to fix the provided code that may or may not contain a syntax error.\nIMPORTANT: the output must not contain any additional comments or explanation.`
+                          ],
+                          user: [],
+                          assistant: [],
+                          ipython: []
+                        },
+                        {
+                          system: [],
+                          user: [`{{input}}`],
+                          assistant: [],
+                          ipython: []
+                        }
+                      ]
+                    })
                   })
                 }
               : undefined