@jpoly1219/context-extractor 0.1.2 → 0.2.1

package/README.md

@@ -70,7 +70,7 @@ dune build
 
 Ignore the wildcard build errors. The command is meant to setup the modules and imports.
 
-Almost there! Create a `credentials.json` file following the steps at the **credentials.json** section below in the README.
+Almost there! Create a `config.json` file following the steps at the **config.json** section below in the README.
 
 Finally, build and run.
 
@@ -88,16 +88,21 @@ node dist/runner.js
 3. Extract relevant headers.
 4. Optionally complete the hole with an LLM.
 
-This library exposes the method `extractContext`, which has the following definition:
+This library exposes two methods `extractContext` and `completeWithLLM`, which have the following definitions:
 
 ```ts
 const extractContext = async (
   language: Language,
   sketchPath: string,
   repoPath: string,
-  credentialsPath: string,
-  getCompletion: boolean
-): Promise<{ context: Context | null, completion: string | null }
+): Promise<Context | null>;
+
+const completeWithLLM = async (
+  ctx: Context,
+  language: Language,
+  sketchPath: string,
+  configPath: string
+): Promise<string>;
 
 enum Language {
   TypeScript,
@@ -105,22 +110,26 @@ enum Language {
 }
 
 interface Context {
-  hole: string,
-  relevantTypes: Map<string, string[]>,
-  relevantHeaders: Map<string, string[]>
+  holeType: string,
+  relevantTypes: Map<Filepath, RelevantType[]>,
+  relevantHeaders: Map<Filepath, RelevantHeader[]>
 }
+
+type Filepath = string;
+type RelevantType = string;
+type RelevantHeader = string;
 ```
 
-- `sketchPath` is the full path to your sketch file with the typed hole construct (`_()` for TypeScript, `_` for OCaml).
+- `sketchPath` is the full path to your sketch file with the typed hole construct (`_()` for TypeScript, `_` for OCaml). This is NOT prefixed with `file://`.
 - `repoPath` is the full path to your repository root.
-- `credentialsPath` is the full path to your `credentials.json`.
-- `getCompletion` is a flag to set if you want the LLM to complete the typed hole. This completion is saved in the `completion` field of the return result.
-- `null` values will only be set if something goes wrong internally. When `getCompletion` is set to false, the `completion` field's value will be an empty string.
+- `configPath` is the full path to your `config.json`.
+- `null` values will only be set if something goes wrong internally.
+- `ctx` is the result from `extractContext`.
 
-### credentials.json
+### config.json
 
 The extractor calls OpenAI for code completion.
-For this you need a `credentials.json` file that holds your specific OpenAI parameters.
+For this you need a `config.json` file that holds your specific OpenAI parameters.
 
 The json has the following format:
 
@@ -137,7 +146,7 @@ The json has the following format:
 }
 ```
 
-Internally, this is how fields above are populated when creating a new OpenAI client.
+Internally, this is how the credentials are populated when creating a new OpenAI client.
 
 ```ts
 const openai = new OpenAI({

package/dist/app.d.ts

@@ -0,0 +1,27 @@
+import { Language, Context } from "./types";
+export declare class App {
+    private language;
+    private languageDriver;
+    private languageServer;
+    private lspClient;
+    private sketchPath;
+    private repoPath;
+    private result;
+    private timeout;
+    constructor(language: Language, sketchPath: string, repoPath: string);
+    init(): Promise<void>;
+    run(): Promise<void>;
+    close(): void;
+    getSavedResult(): Context | null;
+}
+export declare class CompletionEngine {
+    private language;
+    private config;
+    private sketchPath;
+    constructor(language: Language, sketchPath: string, configPath: string);
+    completeWithLLM(context: Context): Promise<string>;
+    generateTypesAndHeadersPrompt(sketchFileContent: string, holeType: string, relevantTypes: string, relevantHeaders: string): {
+        role: string;
+        content: string;
+    }[];
+}

package/dist/app.js

@@ -0,0 +1,284 @@
+"use strict";
+var __createBinding = (this && this.__createBinding) || (Object.create ? (function(o, m, k, k2) {
+    if (k2 === undefined) k2 = k;
+    var desc = Object.getOwnPropertyDescriptor(m, k);
+    if (!desc || ("get" in desc ? !m.__esModule : desc.writable || desc.configurable)) {
+      desc = { enumerable: true, get: function() { return m[k]; } };
+    }
+    Object.defineProperty(o, k2, desc);
+}) : (function(o, m, k, k2) {
+    if (k2 === undefined) k2 = k;
+    o[k2] = m[k];
+}));
+var __setModuleDefault = (this && this.__setModuleDefault) || (Object.create ? (function(o, v) {
+    Object.defineProperty(o, "default", { enumerable: true, value: v });
+}) : function(o, v) {
+    o["default"] = v;
+});
+var __importStar = (this && this.__importStar) || function (mod) {
+    if (mod && mod.__esModule) return mod;
+    var result = {};
+    if (mod != null) for (var k in mod) if (k !== "default" && Object.prototype.hasOwnProperty.call(mod, k)) __createBinding(result, mod, k);
+    __setModuleDefault(result, mod);
+    return result;
+};
+var __importDefault = (this && this.__importDefault) || function (mod) {
+    return (mod && mod.__esModule) ? mod : { "default": mod };
+};
+Object.defineProperty(exports, "__esModule", { value: true });
+exports.CompletionEngine = exports.App = void 0;
+const path = __importStar(require("path"));
+const fs = __importStar(require("fs"));
+const child_process_1 = require("child_process");
+const child_process_2 = require("child_process");
+const openai_1 = __importDefault(require("openai"));
+const main_1 = require("../ts-lsp-client-dist/src/main");
+const types_1 = require("./types");
+// TODO: Bundle the drivers as barrel exports.
+const typescript_driver_1 = require("./typescript-driver");
+const ocaml_driver_1 = require("./ocaml-driver");
+const utils_1 = require("./utils");
+class App {
+    constructor(language, sketchPath, repoPath) {
+        // private result: {
+        //   hole: string;
+        //   relevantTypes: string[];
+        //   relevantHeaders: string[];
+        // } | null = null;
+        this.result = null;
+        // Optional timeout for forced termination
+        this.timeout = setTimeout(() => {
+            if (!this.languageServer.killed) {
+                console.log('Forcibly killing the process...');
+                this.languageServer.kill('SIGKILL');
+            }
+        }, 5000);
+        this.language = language;
+        this.sketchPath = sketchPath;
+        this.repoPath = repoPath;
+        const r = (() => {
+            switch (language) {
+                case types_1.Language.TypeScript: {
+                    this.languageDriver = new typescript_driver_1.TypeScriptDriver();
+                    return (0, child_process_1.spawn)("typescript-language-server", ["--stdio"], { stdio: ["pipe", "pipe", "pipe"] });
+                }
+                case types_1.Language.OCaml: {
+                    this.languageDriver = new ocaml_driver_1.OcamlDriver();
+                    try {
+                        (0, child_process_2.execSync)(`eval $(opam env --switch=. --set-switch)`, { shell: "/bin/bash" });
+                        // execSync("opam switch .", { shell: "/bin/bash" })
+                        const currDir = __dirname;
+                        process.chdir(path.dirname(sketchPath));
+                        // execSync("which dune", { shell: "/bin/bash" })
+                        (0, child_process_1.spawn)("dune", ["build", "-w"]);
+                        process.chdir(currDir);
+                    }
+                    catch (err) {
+                        console.log("ERROR:", err);
+                    }
+                    // TODO: Spawn a dune build -w on sketch directory.
+                    // try {
+                    //   execSync("which dune", { shell: "/bin/bash" })
+                    //   spawn("dune", ["build", "-w"]);
+                    // } catch (err) {
+                    //   console.log("ERROR:", err)
+                    // }
+                    // process.chdir(currDir);
+                    return (0, child_process_1.spawn)("ocamllsp", ["--stdio"]);
+                }
+            }
+        })();
+        const e = new main_1.JSONRPCEndpoint(r.stdin, r.stdout);
+        const c = new main_1.LspClient(e);
+        this.languageServer = r;
+        this.lspClient = c;
+        this.languageServer.on('close', (code) => {
+            if (code !== 0) {
+                console.log(`ls process exited with code ${code}`);
+            }
+        });
+        // Clear timeout once the process exits
+        this.languageServer.on('exit', () => {
+            clearTimeout(this.timeout);
+            console.log('Process terminated cleanly.');
+        });
+        // const logFile = fs.createWriteStream("log.txt");
+        // r.stdout.on('data', (d) => logFile.write(d));
+    }
+    async init() {
+        await this.languageDriver.init(this.lspClient, this.sketchPath);
+    }
+    async run() {
+        // const outputFile = fs.createWriteStream("output.txt");
+        try {
+            await this.init();
+            const holeContext = await this.languageDriver.getHoleContext(this.lspClient, this.sketchPath);
+            const relevantTypes = await this.languageDriver.extractRelevantTypes(this.lspClient, holeContext.fullHoverResult, holeContext.functionName, holeContext.range.start.line, holeContext.range.end.line, new Map(), holeContext.source);
+            // Postprocess the map.
+            if (this.language === types_1.Language.TypeScript) {
+                relevantTypes.delete("_()");
+                for (const [k, { typeSpan: v, sourceFile: src }] of relevantTypes.entries()) {
+                    relevantTypes.set(k, { typeSpan: v.slice(0, -1), sourceFile: src });
+                }
+            }
+            else if (this.language === types_1.Language.OCaml) {
+                relevantTypes.delete("_");
+            }
+            console.log(path.join(path.dirname(this.sketchPath), `sketch${path.extname(this.sketchPath)}`));
+            let repo = [];
+            if (this.language === types_1.Language.TypeScript) {
+                repo = (0, utils_1.getAllTSFiles)(this.repoPath);
+            }
+            else if (this.language === types_1.Language.OCaml) {
+                repo = (0, utils_1.getAllOCamlFiles)(this.repoPath);
+            }
+            const relevantHeaders = await this.languageDriver.extractRelevantHeaders(this.lspClient, repo, relevantTypes, holeContext.functionTypeSpan);
+            // Postprocess the map.
+            if (this.language === types_1.Language.TypeScript) {
+                relevantTypes.delete("");
+                for (const [k, { typeSpan: v, sourceFile: src }] of relevantTypes.entries()) {
+                    relevantTypes.set(k, { typeSpan: v + ";", sourceFile: src });
+                }
+                for (const obj of relevantHeaders) {
+                    obj.typeSpan += ";";
+                }
+            }
+            const relevantTypesToReturn = new Map();
+            relevantTypes.forEach(({ typeSpan: v, sourceFile: src }, _) => {
+                if (relevantTypesToReturn.has(src)) {
+                    const updated = relevantTypesToReturn.get(src);
+                    updated.push(v);
+                    relevantTypesToReturn.set(src, updated);
+                }
+                else {
+                    relevantTypesToReturn.set(src, [v]);
+                }
+            });
+            const relevantHeadersToReturn = new Map();
+            relevantHeaders.forEach(({ typeSpan: v, sourceFile: src }) => {
+                if (relevantHeadersToReturn.has(src)) {
+                    const updated = relevantHeadersToReturn.get(src);
+                    if (!updated.includes(v)) {
+                        updated.push(v);
+                    }
+                    relevantHeadersToReturn.set(src, updated);
+                }
+                else {
+                    relevantHeadersToReturn.set(src, [v]);
+                }
+            });
+            this.result = {
+                holeType: holeContext.functionTypeSpan,
+                relevantTypes: relevantTypesToReturn,
+                relevantHeaders: relevantHeadersToReturn
+            };
+        }
+        catch (err) {
+            console.error("Error during execution:", err);
+            throw err;
+        }
+        finally {
+            // outputFile.end();
+        }
+    }
+    close() {
+        // TODO:
+        try {
+            this.lspClient.exit();
+        }
+        catch (err) {
+            console.log(err);
+        }
+    }
+    getSavedResult() {
+        return this.result;
+    }
+}
+exports.App = App;
+class CompletionEngine {
+    constructor(language, sketchPath, configPath) {
+        this.language = language;
+        this.config = JSON.parse(fs.readFileSync(configPath, "utf8"));
+        this.sketchPath = sketchPath;
+    }
+    async completeWithLLM(context) {
+        let joinedTypes = "";
+        let joinedHeaders = "";
+        context.relevantTypes.forEach((v, _) => {
+            joinedTypes = joinedTypes + v.join("\n") + "\n";
+        });
+        context.relevantHeaders.forEach((v, _) => {
+            joinedHeaders = joinedHeaders + v.join("\n") + "\n";
+        });
+        // Create a prompt.
+        const prompt = this.generateTypesAndHeadersPrompt(
+        // fs.readFileSync(path.join(targetDirectoryPath, "sketch.ts"), "utf8"),
+        fs.readFileSync(this.sketchPath, "utf8"), context.holeType, joinedTypes, joinedHeaders);
+        // Call the LLM to get completion results back.
+        const apiBase = this.config.apiBase;
+        const deployment = this.config.deployment;
+        const model = this.config.gptModel;
+        const apiVersion = this.config.apiVersion;
+        const apiKey = this.config.apiKey;
+        const openai = new openai_1.default({
+            apiKey,
+            baseURL: `${apiBase}/openai/deployments/${deployment}`,
+            defaultQuery: { "api-version": apiVersion },
+            defaultHeaders: { "api-key": apiKey }
+        });
+        const llmResult = await openai.chat.completions.create({
+            model,
+            messages: prompt,
+            temperature: this.config.temperature
+        });
+        return llmResult.choices[0].message.content;
+    }
+    generateTypesAndHeadersPrompt(sketchFileContent, holeType, relevantTypes, relevantHeaders) {
+        let holeConstruct = "";
+        switch (this.language) {
+            case types_1.Language.TypeScript: {
+                holeConstruct = "_()";
+            }
+            case types_1.Language.OCaml: {
+                holeConstruct = "_";
+            }
+        }
+        const prompt = [{
+                role: "system",
+                content: [
+                    "CODE COMPLETION INSTRUCTIONS:",
+                    `- Reply with a functional, idiomatic replacement for the program hole marked '${holeConstruct}' in the provided TypeScript program sketch`,
+                    `- Reply only with a single replacement term for the unqiue distinguished hole marked '${holeConstruct}'`,
+                    "Reply only with code",
+                    "- DO NOT include the program sketch in your reply",
+                    "- DO NOT include a period at the end of your response and DO NOT use markdown",
+                    "- DO NOT include a type signature for the program hole, as this is redundant and is already in the provided program sketch"
+                ].join("\n"),
+            }];
+        let userPrompt = {
+            role: "user",
+            content: ""
+        };
+        if (relevantTypes) {
+            userPrompt.content +=
+                `# The expected type of the goal completion is ${holeType} #
+
+# The following type definitions are likely relevant: #
+${relevantTypes}
+
+      `;
+        }
+        if (relevantHeaders) {
+            userPrompt.content += `
+# Consider using these variables relevant to the expected type: #
+${relevantHeaders}
+
+      `;
+        }
+        userPrompt.content += `# Program Sketch to be completed: #\n${(0, utils_1.removeLines)(sketchFileContent).join("\n")}`;
+        prompt.push(userPrompt);
+        return prompt;
+    }
+    ;
+}
+exports.CompletionEngine = CompletionEngine;

package/dist/codeql.d.ts

@@ -0,0 +1,17 @@
+import { relevantTypeObject, varsObject, typesObject } from "./types";
+declare const createDatabaseWithCodeQL: (pathToCodeQL: string, targetPath: string) => string;
+declare const extractHoleType: (pathToCodeQL: string, pathToQuery: string, pathToDatabase: string, outDir: string) => typesObject;
+declare const extractRelevantTypesWithCodeQL: (pathToCodeQL: string, pathToQuery: string, pathToDatabase: string, outDir: string) => Map<string, relevantTypeObject>;
+declare const extractHeadersWithCodeQL: (pathToCodeQL: string, pathToQuery: string, pathToDatabase: string, outDir: string) => Map<string, varsObject>;
+declare const extractTypesAndLocations: (pathToCodeQL: string, pathToQuery: string, pathToDatabase: string, outDir: string) => {
+    locationToType: Map<string, string[]>;
+    typeToLocation: Map<string, string>;
+};
+declare const extractRelevantContextWithCodeQL: (pathToCodeQL: string, pathToQuery: string, pathToDatabase: string, outDir: string, headers: Map<string, varsObject>, relevantTypes: Map<string, relevantTypeObject>) => Set<string>;
+declare const getRelevantHeaders: (pathToCodeQL: string, pathToQuery: string, pathToDatabase: string, outDir: string, headers: Map<string, varsObject>, holeType: typesObject) => Set<string>;
+declare const getRelevantHeaders3: (pathToCodeQL: string, pathToQuery: string, pathToDatabase: string, outDir: string, headers: Map<string, varsObject>, holeType: typesObject, relevantTypes: Map<string, relevantTypeObject>) => Set<string>;
+declare const getRelevantHeaders4: (pathToCodeQL: string, pathToQueryDir: string, pathToDatabase: string, outDir: string, headers: Map<string, varsObject>, holeType: typesObject, relevantTypes: Map<string, relevantTypeObject>, knownTypeLocations: {
+    locationToType: Map<string, string[]>;
+    typeToLocation: Map<string, string>;
+}) => Set<string>;
+export { createDatabaseWithCodeQL, extractHoleType, extractRelevantTypesWithCodeQL, extractHeadersWithCodeQL, extractRelevantContextWithCodeQL, extractTypesAndLocations, getRelevantHeaders, getRelevantHeaders3, getRelevantHeaders4 };