11import { createOpenAI } from "@ai-sdk/openai"
2+ import { createOpenRouter } from "@openrouter/ai-sdk-provider"
23import FirecrawlApp , { SearchResponse } from "@mendable/firecrawl-js"
34import { generateObject , LanguageModel , Message , streamText } from "ai"
45import { z } from "zod"
56import pLimit from "p-limit"
67
78import { ExtensionMessage } from "../../shared/ExtensionMessage"
9+ import { ResearchTaskPayload } from "../../shared/WebviewMessage"
810import { ClineProvider } from "../../core/webview/ClineProvider"
911
1012import {
@@ -20,6 +22,14 @@ import {
2022import { truncatePrompt , trimPrompt } from "./utils/prompt"
2123
2224export class DeepResearchService {
25+ public readonly providerId : string
26+ public readonly providerApiKey : string
27+ public readonly firecrawlApiKey : string
28+ public readonly modelId : string
29+ public readonly breadth : number
30+ public readonly depth : number
31+ public readonly concurrency : number
32+
2333 private providerRef : WeakRef < ClineProvider >
2434 private firecrawl : FirecrawlApp
2535 private model : LanguageModel
@@ -30,22 +40,28 @@ export class DeepResearchService {
3040 private messages : Message [ ] = [ ]
3141
3242 constructor (
43+ { providerId, providerApiKey, firecrawlApiKey, modelId, breadth, depth } : ResearchTaskPayload [ "session" ] ,
3344 clineProvider : ClineProvider ,
34- public readonly modelId : string ,
35- public readonly breadth : number ,
36- public readonly depth : number ,
37- public readonly concurrency : number ,
38- public readonly firecrawlApiKey : string ,
39- public readonly openaiApiKey : string ,
4045 ) {
46+ this . providerId = providerId
47+ this . providerApiKey = providerApiKey
48+ this . firecrawlApiKey = firecrawlApiKey
49+ this . modelId = modelId
50+ this . breadth = breadth
51+ this . depth = depth
52+ this . concurrency = 2
53+
4154 this . providerRef = new WeakRef ( clineProvider )
4255
4356 this . firecrawl = new FirecrawlApp ( { apiKey : firecrawlApiKey } )
4457
45- this . model = createOpenAI ( { apiKey : openaiApiKey } ) ( modelId , {
46- // reasoningEffort: "medium",
47- structuredOutputs : true ,
48- } )
58+ if ( providerId === "openai-native" ) {
59+ const openai = createOpenAI ( { apiKey : providerApiKey } )
60+ this . model = openai ( modelId , { structuredOutputs : true } )
61+ } else {
62+ const openrouter = createOpenRouter ( { apiKey : providerApiKey } )
63+ this . model = openrouter ( modelId )
64+ }
4965 }
5066
5167 /**
@@ -72,34 +88,34 @@ export class DeepResearchService {
7288 }
7389
7490 private chatSystemPrompt ( ) {
75- // return trimPrompt(`
76- // You are an expert research assistant helping to explain and clarify research findings. Follow these guidelines:
77-
78- // - You always answer the with markdown formatting. You will be penalized if you do not answer with markdown when it would be possible.
79- // - The markdown formatting you support: headings, bold, italic, links, tables, lists, code blocks, and blockquotes.
80- // - You do not support images and never include images. You will be penalized if you render images.
81- // - You also support Mermaid formatting. You will be penalized if you do not render Mermaid diagrams when it would be possible.
82- // - The Mermaid diagrams you support: sequenceDiagram, flowChart, classDiagram, stateDiagram, erDiagram, gantt, journey, gitGraph, pie.
83- // - Reference specific findings from the research when answering.
84- // - Be precise and detailed in explanations.
85- // - If asked about something outside the research scope, acknowledge this and stick to what was actually researched.
86- // - Feel free to make connections between different parts of the research.
87- // - When speculating or making inferences beyond the direct research, clearly label these as such.
88- // - If asked about sources, refer to the URLs provided in the research.
89- // - Maintain a professional, analytical tone.
90- // - Never include images in responses.
91- // `)
92-
9391 return trimPrompt ( `
94- You are a general answering assistant that can comply with any request.
92+ You are an expert research assistant helping to explain and clarify research findings. Follow these guidelines:
93+
94+ - You always answer the with markdown formatting. You will be penalized if you do not answer with markdown when it would be possible.
95+ - The markdown formatting you support: headings, bold, italic, links, tables, lists, code blocks, and blockquotes.
96+ - You do not support images and never include images. You will be penalized if you render images.
97+ - You also support Mermaid formatting. You will be penalized if you do not render Mermaid diagrams when it would be possible.
98+ - The Mermaid diagrams you support: sequenceDiagram, flowChart, classDiagram, stateDiagram, erDiagram, gantt, journey, gitGraph, pie.
99+ - Reference specific findings from the research when answering.
100+ - Be precise and detailed in explanations.
101+ - If asked about something outside the research scope, acknowledge this and stick to what was actually researched.
102+ - Feel free to make connections between different parts of the research.
103+ - When speculating or making inferences beyond the direct research, clearly label these as such.
104+ - If asked about sources, refer to the URLs provided in the research.
105+ - Maintain a professional, analytical tone.
106+ - Never include images in responses.
107+ ` )
108+
109+ // return trimPrompt(`
110+ // You are a general answering assistant that can comply with any request.
95111
96- You always answer the with markdown formatting. You will be penalized if you do not answer with markdown when it would be possible.
97- The markdown formatting you support: headings, bold, italic, links, tables, lists, code blocks, and blockquotes.
98- You do not support images and never include images. You will be penalized if you render images.
112+ // You always answer the with markdown formatting. You will be penalized if you do not answer with markdown when it would be possible.
113+ // The markdown formatting you support: headings, bold, italic, links, tables, lists, code blocks, and blockquotes.
114+ // You do not support images and never include images. You will be penalized if you render images.
99115
100- You also support Mermaid formatting. You will be penalized if you do not render Mermaid diagrams when it would be possible.
101- The Mermaid diagrams you support: sequenceDiagram, flowChart, classDiagram, stateDiagram, erDiagram, gantt, journey, gitGraph, pie.
102- `)
116+ // You also support Mermaid formatting. You will be penalized if you do not render Mermaid diagrams when it would be possible.
117+ // The Mermaid diagrams you support: sequenceDiagram, flowChart, classDiagram, stateDiagram, erDiagram, gantt, journey, gitGraph, pie.
118+ // `)
103119 }
104120
105121 /**
@@ -476,13 +492,49 @@ export class DeepResearchService {
476492
477493 this . inquiry . query = query
478494
495+ const onProgressUpdated = ( ) => {
496+ const { expectedQueries, completedQueries } = this . progress
497+ this . progress . progressPercentage = Math . round ( ( completedQueries / expectedQueries ) * 100 )
498+ this . postMessage ( { type : "research.progress" , text : JSON . stringify ( this . progress ) } )
499+ }
500+
501+ const onGeneratedQueries = ( queries : ResearchQuery [ ] ) =>
502+ this . postMessage ( {
503+ type : "research.output" ,
504+ text : JSON . stringify ( {
505+ content : `Generated ${ queries . length } ${ queries . map ( ( { query } ) => `- ${ query } ) . join ( "\n" ) } ,
506+ annotations : [
507+ {
508+ type : "badge" ,
509+ data : { label : "Idea" , variant : "outline" } ,
510+ } ,
511+ ] ,
512+ } ) ,
513+ } )
514+
515+ const onExtractedLearnings = ( learnings : ResearchLearnings & { urls : string [ ] } ) =>
516+ this . postMessage ( {
517+ type : "research.output" ,
518+ text : JSON . stringify ( {
519+ content : `Extracted ${ learnings . learnings . length } ${ learnings . urls . length } ${ learnings . urls . map ( ( url ) => `- ${ url } ) . join ( "\n" ) } ,
520+ annotations : [
521+ {
522+ type : "badge" ,
523+ data : { label : "Learning" , variant : "outline" } ,
524+ } ,
525+ ] ,
526+ } ) ,
527+ } )
528+
479529 // Calculate total expected queries across all depth levels.
480530 // At each level, the breadth is halved, so level 1 has full breadth,
481531 // level 2 has breadth/2, level 3 has breadth/4, etc.
482532 for ( let i = this . depth ; i > 0 ; i -- ) {
483533 this . progress . expectedQueries += Math . ceil ( this . breadth / Math . pow ( 2 , this . depth - i ) )
484534 }
485535
536+ onProgressUpdated ( )
537+
486538 const { learnings, visitedUrls } = await this . withLoading (
487539 ( ) =>
488540 this . deepResearch ( {
@@ -491,37 +543,9 @@ export class DeepResearchService {
491543 depth : this . depth ,
492544 learnings : [ ] ,
493545 visitedUrls : [ ] ,
494- onProgressUpdated : ( ) => {
495- const { expectedQueries, completedQueries } = this . progress
496- this . progress . progressPercentage = Math . round ( ( completedQueries / expectedQueries ) * 100 )
497- this . postMessage ( { type : "research.progress" , text : JSON . stringify ( this . progress ) } )
498- } ,
499- onGeneratedQueries : ( queries ) =>
500- this . postMessage ( {
501- type : "research.output" ,
502- text : JSON . stringify ( {
503- content : `Generated ${ queries . length } ${ queries . map ( ( { query } ) => `- ${ query } ) . join ( "\n" ) } ,
504- annotations : [
505- {
506- type : "badge" ,
507- data : { label : "Idea" , variant : "outline" } ,
508- } ,
509- ] ,
510- } ) ,
511- } ) ,
512- onExtractedLearnings : ( learnings ) =>
513- this . postMessage ( {
514- type : "research.output" ,
515- text : JSON . stringify ( {
516- content : `Extracted ${ learnings . learnings . length } ${ learnings . urls . length } ${ learnings . urls . map ( ( url ) => `- ${ url } ) . join ( "\n" ) } ,
517- annotations : [
518- {
519- type : "badge" ,
520- data : { label : "Learning" , variant : "outline" } ,
521- } ,
522- ] ,
523- } ) ,
524- } ) ,
546+ onProgressUpdated,
547+ onGeneratedQueries,
548+ onExtractedLearnings,
525549 } ) ,
526550 "Researching..." ,
527551 )
0 commit comments