forked from continuedev/continue
-
Notifications
You must be signed in to change notification settings - Fork 1
/
DocsContextProvider.ts
170 lines (155 loc) · 5.52 KB
/
DocsContextProvider.ts
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
import fetch from "node-fetch";
import {
ContextItem,
ContextProviderDescription,
ContextProviderExtras,
ContextSubmenuItem,
LoadSubmenuItemsArgs,
} from "../../index.js";
import { DocsService } from "../../indexing/docs/DocsService.js";
import configs from "../../indexing/docs/preIndexedDocs.js";
import TransformersJsEmbeddingsProvider from "../../indexing/embeddings/TransformersJsEmbeddingsProvider.js";
import { BaseContextProvider } from "../index.js";
class DocsContextProvider extends BaseContextProvider {
static DEFAULT_N_RETRIEVE = 30;
static DEFAULT_N_FINAL = 15;
static description: ContextProviderDescription = {
title: "docs",
displayTitle: "Docs",
description: "Type to search docs",
type: "submenu",
};
private docsService: DocsService;
constructor(options: any) {
super(options);
this.docsService = DocsService.getInstance();
}
private async _getIconDataUrl(url: string): Promise<string | undefined> {
try {
const response = await fetch(url);
if (!response.headers.get("content-type")?.startsWith("image/")) {
console.log("Not an image: ", await response.text());
return undefined;
}
const buffer = await response.buffer();
const base64data = buffer.toString("base64");
return `data:${response.headers.get("content-type")};base64,${base64data}`;
} catch (e) {
console.log("E: ", e);
return undefined;
}
}
async getContextItems(
query: string,
extras: ContextProviderExtras,
): Promise<ContextItem[]> {
// Not supported in JetBrains IDEs right now
if ((await extras.ide.getIdeInfo()).ideType === "jetbrains") {
throw new Error(
"The @docs context provider is not currently supported in JetBrains IDEs. We'll have an update soon!",
);
}
const embeddingsProvider = new TransformersJsEmbeddingsProvider();
const [vector] = await embeddingsProvider.embed([extras.fullInput]);
let chunks = await this.docsService.retrieve(
query,
vector,
this.options?.nRetrieve ?? DocsContextProvider.DEFAULT_N_RETRIEVE,
embeddingsProvider.id,
);
if (extras.reranker) {
try {
const scores = await extras.reranker.rerank(extras.fullInput, chunks);
chunks.sort(
(a, b) => scores[chunks.indexOf(b)] - scores[chunks.indexOf(a)],
);
chunks = chunks.splice(
0,
this.options?.nFinal ?? DocsContextProvider.DEFAULT_N_FINAL,
);
} catch (e) {
console.warn(`Failed to rerank docs results: ${e}`);
chunks = chunks.splice(
0,
this.options?.nFinal ?? DocsContextProvider.DEFAULT_N_FINAL,
);
}
}
return [
...chunks
.map((chunk) => ({
name: chunk.filepath.includes("/tree/main") // For display of GitHub files
? chunk.filepath
.split("/")
.slice(1)
.join("/")
.split("/tree/main/")
.slice(1)
.join("/")
: chunk.otherMetadata?.title || chunk.filepath,
description: chunk.filepath, // new URL(chunk.filepath, query).toString(),
content: chunk.content,
}))
.reverse(),
{
name: "Instructions",
description: "Instructions",
content:
"Use the above documentation to answer the following question. You should not reference anything outside of what is shown, unless it is a commonly known concept. Reference URLs whenever possible using markdown formatting. If there isn't enough information to answer the question, suggest where the user might look to learn more.",
},
];
}
async loadSubmenuItems(
args: LoadSubmenuItemsArgs,
): Promise<ContextSubmenuItem[]> {
const docs = await this.docsService.list();
const submenuItems: ContextSubmenuItem[] = docs.map((doc) => ({
title: doc.title,
description: new URL(doc.baseUrl).hostname,
id: doc.baseUrl,
metadata: {
preIndexed: !!configs.find((config) => config.title === doc.title),
},
}));
submenuItems.push(
...configs
// After it's actually downloaded, we don't want to show twice
.filter(
(config) => !submenuItems.some((item) => item.id === config.startUrl),
)
.map((config) => ({
title: config.title,
description: new URL(config.startUrl).hostname,
id: config.startUrl,
metadata: {
preIndexed: true,
},
// iconUrl: config.faviconUrl,
})),
);
// Sort submenuItems such that the objects with titles which don't occur in configs occur first, and alphabetized
submenuItems.sort((a, b) => {
const aTitleInConfigs = a.metadata?.preIndexed;
const bTitleInConfigs = b.metadata?.preIndexed;
// Primary criterion: Items not in configs come first
if (!aTitleInConfigs && bTitleInConfigs) {
return -1;
} else if (aTitleInConfigs && !bTitleInConfigs) {
return 1;
} else {
// Secondary criterion: Alphabetical order when both items are in the same category
return a.title.toString().localeCompare(b.title.toString());
}
});
// const icons = await Promise.all(
// submenuItems.map(async (item) =>
// item.iconUrl ? this._getIconDataUrl(item.iconUrl) : undefined,
// ),
// );
// icons.forEach((icon, i) => {
// submenuItems[i].iconUrl = icon;
// });
return submenuItems;
}
}
export default DocsContextProvider;