@@ -3,14 +3,84 @@ import { detectLanguage } from './processors/detectLanguage';
33import { fileHeaders } from './processors/fileHeaders' ;
44import { languages } from './processors/languages' ;
55
6+ var decoder = new TextDecoder ( "utf8" ) ;
7+
8+ function getNotebookDocument ( document :vscode . TextDocument ) :vscode . NotebookDocument | undefined {
9+ return vscode . workspace . notebookDocuments
10+ . find ( x => x . uri . path === document . uri . path ) ;
11+ }
12+
613export async function preparePrompt ( document :vscode . TextDocument , position :vscode . Position , context :vscode . InlineCompletionContext ) {
714
815// Load document text
16+ console . log ( document ) ;
917let text = document . getText ( ) ;
1018let offset = document . offsetAt ( position ) ;
1119let prefix = text . slice ( 0 , offset ) ;
1220let suffix :string = text . slice ( offset ) ;
1321
22+ // If this is a notebook, add the surrounding cells to the prefix and suffix
23+ let notebookDocument = getNotebookDocument ( document ) ;
24+ let language = detectLanguage ( document . uri . fsPath , document . languageId ) ;
25+ let commentStart :string | undefined = undefined ;
26+ if ( language ) {
27+ commentStart = languages [ language ] . comment ?. start ;
28+ }
29+
30+ if ( notebookDocument ) {
31+ let beforeCurrentCell = true ;
32+
33+ let prefixCells = "" ;
34+ let suffixCells = "" ;
35+
36+ notebookDocument . getCells ( ) . forEach ( ( cell ) => {
37+ let out = "" ;
38+
39+ if ( cell . document . uri . fragment === document . uri . fragment ) {
40+ beforeCurrentCell = false ; // switch to suffix mode
41+ return ;
42+ }
43+
44+ // add the markdown cell output to the prompt as a comment
45+ if ( cell . kind === vscode . NotebookCellKind . Markup && commentStart ) {
46+ for ( const line of cell . document . getText ( ) . split ( '\n' ) ) {
47+ out += `\n${ commentStart } ${ line } ;
48+ }
49+ } else {
50+ out += cell . document . getText ( ) ;
51+ }
52+
53+ // if there is any outputs add them to the prompt as a comment
54+ if ( cell . kind === vscode . NotebookCellKind . Code && commentStart ) {
55+ console . log ( cell . outputs ) ;
56+ let cellOutputs = cell . outputs
57+ . map ( x => x . items
58+ . filter ( x => x . mime === 'text/plain' )
59+ . map ( x => decoder . decode ( x . data ) )
60+ . map ( x => x . slice ( 0 , 256 ) . split ( '\n' ) ) ) // limit to 256 characters
61+ . flat ( 3 ) ;
62+
63+ if ( cellOutputs . length > 0 ) {
64+ out += `\n${ commentStart } ;
65+ for ( const line of cellOutputs ) {
66+ out += `\n${ commentStart } ${ line } ;
67+ }
68+ }
69+ }
70+
71+ // update the prefix/suffix
72+ if ( beforeCurrentCell ) {
73+ prefixCells += out ;
74+ } else {
75+ suffixCells += out ;
76+ }
77+
78+ } ) ;
79+
80+ prefix = prefixCells + prefix ;
81+ suffix = suffix + suffixCells ;
82+ }
83+
1484// Trim suffix
1585// If suffix is too small it is safe to assume that it could be ignored which would allow us to use
1686// more powerful completition instead of in middle one
@@ -22,7 +92,6 @@ export async function preparePrompt(document: vscode.TextDocument, position: vsc
2292// NOTE: Most networks don't have a concept of filenames and expected language, but we expect that some files in training set has something in title that
2393// would indicate filename and language
2494// NOTE: If we can't detect language, we could ignore this since the number of languages that need detection is limited
25- let language = detectLanguage ( document . uri . fsPath , document . languageId ) ;
2695if ( language ) {
2796prefix = fileHeaders ( prefix , document . uri . fsPath , languages [ language ] ) ;
2897}