From 6d22fa2a50f33d81b1e452e17c5200f4c27620db Mon Sep 17 00:00:00 2001
From: Marcus Schiesser <mail@marcusschiesser.de>
Date: Tue, 19 Nov 2024 00:34:26 +0700
Subject: [PATCH] feat: Get PromptTemplate template variables at run-time
(#1502)
---
.changeset/rotten-fans-mix.md | 5 +++++
packages/core/src/prompts/base.ts | 13 +++++++++++++
packages/core/tests/prompts.test.ts | 11 +++++++++++
3 files changed, 29 insertions(+)
create mode 100644 .changeset/rotten-fans-mix.md
diff --git a/.changeset/rotten-fans-mix.md b/.changeset/rotten-fans-mix.md
new file mode 100644
index 000000000..58e4f519c
--- /dev/null
+++ b/.changeset/rotten-fans-mix.md
@@ -0,0 +1,5 @@
+---
+"@llamaindex/core": patch
+---
+
+Get PromptTemplate template variables at run-time
diff --git a/packages/core/src/prompts/base.ts b/packages/core/src/prompts/base.ts
index 2aacdf340..b17b11316 100644
--- a/packages/core/src/prompts/base.ts
+++ b/packages/core/src/prompts/base.ts
@@ -32,6 +32,10 @@ export abstract class BasePromptTemplate<
const Vars extends readonly string[] = string[],
> {
metadata: Metadata = {};
+ /**
+ * Set of template variables used in the prompt template. Used for type hints only.
+ * To get the list of template variables used in the prompt at run-time, use the `vars` method.
+ */
templateVars: Set<string> = new Set();
options: Partial<Record<TemplatesVar[number] | (string & {}), string>> = {};
outputParser: BaseOutputParser | undefined;
@@ -223,4 +227,13 @@ export class PromptTemplate<
get template(): Template {
return this.#template;
}
+
+ /**
+ * Returns all the template variables used in the prompt template.
+ */
+ vars(): string[] {
+ const template = this.template;
+ const matches = template.match(/\{([^}]+)\}/g) || [];
+ return [...new Set(matches.map((match) => match.slice(1, -1)))];
+ }
}
diff --git a/packages/core/tests/prompts.test.ts b/packages/core/tests/prompts.test.ts
index d0a839f3a..48b9bac02 100644
--- a/packages/core/tests/prompts.test.ts
+++ b/packages/core/tests/prompts.test.ts
@@ -158,4 +158,15 @@ describe("PromptTemplate", () => {
const formatted = prompt.format({ text: "world", foo: "bar" });
expect(formatted).toBe("hello world bar\noutput_instruction");
});
+
+ test("should return all unique template vars of a prompt", () => {
+ const prompt = new PromptTemplate({
+ template: "Hello {name}! Your age is {age}. Nice to meet you {name}!",
+ });
+
+ const vars = prompt.vars();
+ expect(vars).toHaveLength(2);
+ expect(vars).toContain("name");
+ expect(vars).toContain("age");
+ });
});
--
GitLab