From e86a0328a8c748ff9a5089c4ebd6c1f5faed332f Mon Sep 17 00:00:00 2001
From: David Zhang <dzz0615@gmail.com>
Date: Fri, 7 Mar 2025 14:23:20 -0800
Subject: [PATCH] x

---
 src/ai/providers.ts | 15 +++++++++++++--
 src/run.ts          |  5 ++++-
 2 files changed, 17 insertions(+), 3 deletions(-)

diff --git a/src/ai/providers.ts b/src/ai/providers.ts
index 89a6477..332176a 100644
--- a/src/ai/providers.ts
+++ b/src/ai/providers.ts
@@ -1,6 +1,10 @@
 import { createFireworks } from '@ai-sdk/fireworks';
 import { createOpenAI } from '@ai-sdk/openai';
-import { LanguageModelV1 } from 'ai';
+import {
+  extractReasoningMiddleware,
+  LanguageModelV1,
+  wrapLanguageModel,
+} from 'ai';
 import { getEncoding } from 'js-tiktoken';
 
 import { RecursiveCharacterTextSplitter } from './text-splitter';
@@ -30,7 +34,14 @@ const o3MiniModel = openai?.('o3-mini', {
   structuredOutputs: true,
 });
 
-const deepSeekR1Model = fireworks?.('accounts/fireworks/models/deepseek-r1');
+const deepSeekR1Model = fireworks
+  ? wrapLanguageModel({
+      model: fireworks(
+        'accounts/fireworks/models/deepseek-r1',
+      ) as LanguageModelV1,
+      middleware: extractReasoningMiddleware({ tagName: 'think' }),
+    })
+  : undefined;
 
 export function getModel(): LanguageModelV1 {
   if (customModel) {
diff --git a/src/run.ts b/src/run.ts
index 27329d9..fa40b74 100644
--- a/src/run.ts
+++ b/src/run.ts
@@ -1,6 +1,7 @@
 import * as fs from 'fs/promises';
 import * as readline from 'readline';
 
+import { getModel } from './ai/providers';
 import {
   deepResearch,
   writeFinalAnswer,
@@ -29,6 +30,8 @@ function askQuestion(query: string): Promise<string> {
 
 // run the agent
 async function run() {
+  console.log('Using model: ', getModel().modelId);
+
   // Get initial query
   const initialQuery = await askQuestion('What would you like to research? ');
 
@@ -47,7 +50,7 @@ async function run() {
     ) || 2;
   const isReport =
     (await askQuestion(
-      'Do you want to generate a long report or a specific answer? (report/answer)',
+      'Do you want to generate a long report or a specific answer? (report/answer, default report)',
     )) === 'report';
 
   log(`Creating research plan...`);