diff --git a/package.json b/package.json index efcde345630494fa3472842133e233bcec0d1d2b..4666b5004436ff2f27fbf7b9f288e357e3ce1602 100644 --- a/package.json +++ b/package.json @@ -48,6 +48,7 @@ "picocolors": "1.0.0", "prompts": "2.1.0", "rimraf": "^5.0.5", + "smol-toml": "^1.1.3", "tar": "6.1.15", "terminal-link": "^3.0.0", "update-check": "1.5.4", diff --git a/templates/index.ts b/templates/index.ts index 1165c9467008b9ada4ebad0f7a3747bf40d9d38e..4508b24b1636735ff2b5ace778c7c1834a3b4446 100644 --- a/templates/index.ts +++ b/templates/index.ts @@ -1,3 +1,4 @@ +import { parse, stringify } from "smol-toml"; import { copy } from "../helpers/copy"; import { callPackageManager } from "../helpers/install"; @@ -313,6 +314,62 @@ const installTSTemplate = async ({ await callPackageManager(packageManager, isOnline); }; +interface IDependencyItem { + name: string; + version: string; +} + +const getPythonAddOnDependencies = (vectorDb?: TemplateVectorDB) => { + const addOnDependencies: IDependencyItem[] = []; + + switch (vectorDb) { + case "mongo": { + addOnDependencies.push({ + name: "pymongo", + version: "^4.6.1", + }); + break; + } + } + + return addOnDependencies; +}; + +const preparePythonDependencies = async ( + root: string, + addOnDependencies: IDependencyItem[], +) => { + if (addOnDependencies.length === 0) return; + + const FILENAME = "pyproject.toml"; + try { + // Parse toml file + const file = path.join(root, FILENAME); + const fileContent = await fs.readFile(file, "utf8"); + const fileParsed = parse(fileContent); + + // Modify toml dependencies + const tool = fileParsed.tool as any; + const dependencies = tool.poetry.dependencies as any; + for (const dependency of addOnDependencies) { + dependencies[dependency.name] = dependency.version; + } + + // Write toml file + const newFileContent = stringify(fileParsed); + await fs.writeFile(file, newFileContent); + + const dependenciesString = addOnDependencies.map((d) => d.name).join(", "); + console.log(`\nAdded ${dependenciesString} to ${cyan(FILENAME)}\n`); + } catch (error) { + console.log( + `Error when preparing ${FILENAME} file for Python template\n`, + error, + ); + console.log(error); + } +}; + const installPythonTemplate = async ({ root, template, @@ -359,6 +416,9 @@ const installPythonTemplate = async ({ }); } + const addOnDependencies = getPythonAddOnDependencies(vectorDb); + await preparePythonDependencies(root, addOnDependencies); + console.log( "\nPython project, dependencies won't be installed automatically.\n", ); diff --git a/templates/types/simple/fastapi/pyproject.toml b/templates/types/simple/fastapi/pyproject.toml index a563ffbfd3bce508bf9fc287dc11940404f4ab85..f9bb9605b78f53417fd9f61e9e03812fdc21047e 100644 --- a/templates/types/simple/fastapi/pyproject.toml +++ b/templates/types/simple/fastapi/pyproject.toml @@ -12,7 +12,6 @@ uvicorn = { extras = ["standard"], version = "^0.23.2" } llama-index = "^0.9.19" pypdf = "^3.17.0" python-dotenv = "^1.0.0" -pymongo = "^4.6.1" [build-system] diff --git a/templates/types/streaming/fastapi/pyproject.toml b/templates/types/streaming/fastapi/pyproject.toml index a563ffbfd3bce508bf9fc287dc11940404f4ab85..f9bb9605b78f53417fd9f61e9e03812fdc21047e 100644 --- a/templates/types/streaming/fastapi/pyproject.toml +++ b/templates/types/streaming/fastapi/pyproject.toml @@ -12,7 +12,6 @@ uvicorn = { extras = ["standard"], version = "^0.23.2" } llama-index = "^0.9.19" pypdf = "^3.17.0" python-dotenv = "^1.0.0" -pymongo = "^4.6.1" [build-system]