Feature/task conversation memory

griptape/engines/extraction/base_extraction_engine.py

@@ -7,6 +7,7 @@
 from griptape.drivers import BasePromptDriver, OpenAiChatPromptDriver
 from griptape.rules import Ruleset
 from griptape.tokenizers import OpenAiTokenizer
+from griptape.utils import PromptStack
 
 
 @define
@@ -44,5 +45,11 @@ def min_response_tokens(self) -> int:
  )
 
  @abstractmethod
- def extract(self, text: str | ListArtifact, rulesets: Optional[list[Ruleset]] = None, **kwargs) -> ListArtifact:
+ def extract(
+ self,
+ text: str | ListArtifact,
+ rulesets: list[Ruleset] | None = None,
+ prompt_stack: PromptStack | None = None,
+ **kwargs,
+ ) -> ListArtifact:
  ...

griptape/engines/extraction/csv_extraction_engine.py

@@ -1,21 +1,24 @@
 from __future__ import annotations
-from typing import Optional
 import csv
 import io
 from attr import field, Factory, define
 from griptape.artifacts import TextArtifact, CsvRowArtifact, ListArtifact, ErrorArtifact
 from griptape.utils import PromptStack
 from griptape.engines import BaseExtractionEngine
 from griptape.utils import J2
-from griptape.rules import Ruleset, rule
+from griptape.rules import Ruleset
 
 
 @define
 class CsvExtractionEngine(BaseExtractionEngine):
  template_generator: J2 = field(default=Factory(lambda: J2("engines/extraction/csv_extraction.j2")), kw_only=True)
 
  def extract(
- self, text: str | ListArtifact, column_names: list[str], rulesets: Optional[Ruleset] = None
+ self,
+ text: str | ListArtifact,
+ column_names: list[str],
+ rulesets: Ruleset | None = None,
+ prompt_stack: PromptStack | None = None,
  ) -> ListArtifact | ErrorArtifact:
  try:
  return ListArtifact(
@@ -24,6 +27,7 @@ def extract(
  column_names,
  [],
  rulesets=rulesets,
+ prompt_stack=prompt_stack,
  ),
  item_separator="\n",
  )
@@ -44,7 +48,8 @@ def _extract_rec(
  artifacts: list[TextArtifact],
  column_names: list[str],
  rows: list[CsvRowArtifact],
- rulesets: Optional[Ruleset] = None,
+ rulesets: Ruleset | None = None,
+ prompt_stack: PromptStack | None = None,
  ) -> list[CsvRowArtifact]:
  artifacts_text = self.chunk_joiner.join([a.value for a in artifacts])
  full_text = self.template_generator.render(
@@ -54,14 +59,11 @@ def _extract_rec(
  )
 
  if self.prompt_driver.tokenizer.count_tokens_left(full_text) >= self.min_response_tokens:
- rows.extend(
- self.text_to_csv_rows(
- self.prompt_driver.run(
- PromptStack(inputs=[PromptStack.Input(full_text, role=PromptStack.USER_ROLE)])
- ).value,
- column_names,
- )
- )
+ if prompt_stack:
+ prompt_stack.add_user_input(full_text)
+ else:
+ prompt_stack = PromptStack(inputs=[PromptStack.Input(full_text, role=PromptStack.USER_ROLE)])
+ rows.extend(self.text_to_csv_rows(self.prompt_driver.run(prompt_stack=prompt_stack).value, column_names))
 
  return rows
  else:
@@ -72,13 +74,11 @@ def _extract_rec(
  rulesets=J2("rulesets/rulesets.j2").render(rulesets=rulesets),
  )
 
- rows.extend(
- self.text_to_csv_rows(
- self.prompt_driver.run(
- PromptStack(inputs=[PromptStack.Input(partial_text, role=PromptStack.USER_ROLE)])
- ).value,
- column_names,
- )
- )
+ if prompt_stack:
+ prompt_stack.add_user_input(partial_text)
+ else:
+ prompt_stack = PromptStack(inputs=[PromptStack.Input(partial_text, role=PromptStack.USER_ROLE)])
+
+ rows.extend(self.text_to_csv_rows(self.prompt_driver.run(prompt_stack=prompt_stack).value, column_names))
 
  return self._extract_rec(chunks[1:], column_names, rows, rulesets=rulesets)

griptape/engines/extraction/json_extraction_engine.py

@@ -14,7 +14,11 @@ class JsonExtractionEngine(BaseExtractionEngine):
  template_generator: J2 = field(default=Factory(lambda: J2("engines/extraction/json_extraction.j2")), kw_only=True)
 
  def extract(
- self, text: str | ListArtifact, template_schema: dict, rulesets: Optional[Ruleset] = None
+ self,
+ text: str | ListArtifact,
+ template_schema: dict,
+ rulesets: Ruleset | None = None,
+ prompt_stack: PromptStack | None = None,
  ) -> ListArtifact | ErrorArtifact:
  try:
  json_schema = json.dumps(template_schema)
@@ -25,6 +29,7 @@ def extract(
  json_schema,
  [],
  rulesets=rulesets,
+ prompt_stack=prompt_stack,
  ),
  item_separator="\n",
  )
@@ -39,7 +44,8 @@ def _extract_rec(
  artifacts: list[TextArtifact],
  json_template_schema: str,
  extractions: list[TextArtifact],
- rulesets: Optional[Ruleset] = None,
+ rulesets: Ruleset | None = None,
+ prompt_stack: PromptStack | None = None,
  ) -> list[TextArtifact]:
  artifacts_text = self.chunk_joiner.join([a.value for a in artifacts])
  full_text = self.template_generator.render(
@@ -49,6 +55,10 @@ def _extract_rec(
  )
 
  if self.prompt_driver.tokenizer.count_tokens_left(full_text) >= self.min_response_tokens:
+ if prompt_stack:
+ prompt_stack.add_user_input(full_text)
+ else:
+ prompt_stack = PromptStack(inputs=[PromptStack.Input(full_text, role=PromptStack.USER_ROLE)])
  extractions.extend(
  self.json_to_text_artifacts(
  self.prompt_driver.run(
@@ -66,6 +76,11 @@ def _extract_rec(
  rulesets=J2("rulesets/rulesets.j2").render(rulesets=rulesets),
  )
 
+ if prompt_stack:
+ prompt_stack.add_user_input(partial_text)
+ else:
+ prompt_stack = PromptStack(inputs=[PromptStack.Input(partial_text, role=PromptStack.USER_ROLE)])
+
  extractions.extend(
  self.json_to_text_artifacts(
  self.prompt_driver.run(

griptape/engines/query/base_query_engine.py

@@ -4,13 +4,18 @@
 from typing import Optional
 from griptape.artifacts import TextArtifact, ListArtifact
 from griptape.rules import Ruleset
+from griptape.utils import PromptStack
 
 
 @define
 class BaseQueryEngine(ABC):
  @abstractmethod
  def query(
- self, query: str, namespace: Optional[str] = None, rulesets: Optional[list[Ruleset]] = None
+ self,
+ query: str,
+ namespace: Optional[str] = None,
+ rulesets: Optional[list[Ruleset]] = None,
+ prompt_stack: Optional[PromptStack] = None,
  ) -> TextArtifact:
  ...
 

griptape/engines/query/vector_query_engine.py

@@ -1,5 +1,5 @@
 from __future__ import annotations
-from typing import TYPE_CHECKING, Optional
+from typing import TYPE_CHECKING
 from attr import define, field, Factory
 from griptape.artifacts import TextArtifact, BaseArtifact, ListArtifact
 from griptape.utils import PromptStack
@@ -26,10 +26,11 @@ class VectorQueryEngine(BaseQueryEngine):
  def query(
  self,
  query: str,
- namespace: Optional[str] = None,
- rulesets: Optional[list[Ruleset]] = None,
- metadata: Optional[str] = None,
- top_n: Optional[int] = None,
+ metadata: str | None = None,
+ top_n: int | None = None,
+ namespace: str | None = None,
+ rulesets: str | None = None,
+ prompt_stack: PromptStack | None = None,
  ) -> TextArtifact:
  tokenizer = self.prompt_driver.tokenizer
  result = self.vector_store_driver.query(query, top_n, namespace)
@@ -65,8 +66,11 @@ def query(
  )
 
  break
+ if prompt_stack is None:
+ prompt_stack = PromptStack()
+ prompt_stack.add_user_input(message)
 
- return self.prompt_driver.run(PromptStack(inputs=[PromptStack.Input(message, role=PromptStack.USER_ROLE)]))
+ return self.prompt_driver.run(prompt_stack=prompt_stack)
 
  def upsert_text_artifact(self, artifact: TextArtifact, namespace: Optional[str] = None) -> str:
  result = self.vector_store_driver.upsert_text_artifact(artifact, namespace=namespace)

griptape/engines/summary/base_summary_engine.py

@@ -3,13 +3,23 @@
 from attr import define
 from griptape.artifacts import TextArtifact, ListArtifact
 from griptape.rules import Ruleset
+from griptape.utils import PromptStack
 
 
 @define
 class BaseSummaryEngine(ABC):
- def summarize_text(self, text: str, rulesets: Optional[list[Ruleset]] = None) -> str:
- return self.summarize_artifacts(ListArtifact([TextArtifact(text)]), rulesets=rulesets).value
+ def summarize_text(
+ self, text: str, rulesets: Optional[list[Ruleset]] = None, prompt_stack: Optional[PromptStack] = None
+ ) -> str:
+ return self.summarize_artifacts(
+ ListArtifact([TextArtifact(text)]), rulesets=rulesets, prompt_stack=prompt_stack
+ ).value
 
  @abstractmethod
- def summarize_artifacts(self, artifacts: ListArtifact, rulesets: Optional[list[Ruleset]] = None) -> TextArtifact:
+ def summarize_artifacts(
+ self,
+ artifacts: ListArtifact,
+ rulesets: Optional[list[Ruleset]] = None,
+ prompt_stack: Optional[PromptStack] = None,
+ ) -> TextArtifact:
  ...

griptape/engines/summary/prompt_summary_engine.py

@@ -1,3 +1,4 @@
+from __future__ import annotations
 from typing import Optional
 from attr import define, Factory, field
 from griptape.artifacts import TextArtifact, BaseArtifact, ListArtifact
@@ -45,11 +46,17 @@ def min_response_tokens(self) -> int:
  - self.prompt_driver.tokenizer.max_tokens * self.max_token_multiplier
  )
 
- def summarize_artifacts(self, artifacts: ListArtifact, rulesets: Optional[Ruleset] = None) -> TextArtifact:
- return self.summarize_artifacts_rec(artifacts.value, None, rulesets=rulesets)
+ def summarize_artifacts(
+ self, artifacts: ListArtifact, rulesets: Ruleset | None = None, prompt_stack: PromptStack | None = None
+ ) -> TextArtifact:
+ return self.summarize_artifacts_rec(artifacts.value, None, rulesets=rulesets, prompt_stack=prompt_stack)
 
  def summarize_artifacts_rec(
- self, artifacts: list[BaseArtifact], summary: Optional[str], rulesets: Optional[Ruleset] = None
+ self,
+ artifacts: list[BaseArtifact],
+ summary: str | None,
+ rulesets: Ruleset | None = None,
+ prompt_stack: PromptStack | None = None,
  ) -> TextArtifact:
  artifacts_text = self.chunk_joiner.join([a.to_text() for a in artifacts])
 
@@ -58,20 +65,24 @@ def summarize_artifacts_rec(
  )
 
  if self.prompt_driver.tokenizer.count_tokens_left(full_text) >= self.min_response_tokens:
- return self.prompt_driver.run(
- PromptStack(inputs=[PromptStack.Input(full_text, role=PromptStack.USER_ROLE)])
- )
+ message_input = PromptStack.Input(full_text, role=PromptStack.USER_ROLE)
+ if prompt_stack is None:
+ prompt_stack = PromptStack(inputs=[message_input])
+ else:
+ prompt_stack.inputs.append(message_input)
+ return self.prompt_driver.run(prompt_stack=prompt_stack)
  else:
  chunks = self.chunker.chunk(artifacts_text)
 
  partial_text = self.template_generator.render(
  summary=summary, text=chunks[0].value, rulesets=J2("rulesets/rulesets.j2").render(rulesets=rulesets)
  )
+ message_input = PromptStack.Input(partial_text, role=PromptStack.USER_ROLE)
+ if prompt_stack is None:
+ prompt_stack = PromptStack(inputs=[message_input])
+ else:
+ prompt_stack.inputs.append(message_input)
 
  return self.summarize_artifacts_rec(
- chunks[1:],
- self.prompt_driver.run(
- PromptStack(inputs=[PromptStack.Input(partial_text, role=PromptStack.USER_ROLE)])
- ).value,
- rulesets=rulesets,
+ chunks[1:], self.prompt_driver.run(prompt_stack=prompt_stack).value, rulesets=rulesets
  )

griptape/tasks/extraction_task.py

@@ -1,13 +1,15 @@
 from attr import define, field
 from griptape.artifacts import ListArtifact
 from griptape.engines import BaseExtractionEngine
-from griptape.tasks import BaseTextInputTask
+from griptape.tasks import PromptTask
 
 
 @define
-class ExtractionTask(BaseTextInputTask):
+class ExtractionTask(PromptTask):
  extraction_engine: BaseExtractionEngine = field(kw_only=True)
  args: dict = field(kw_only=True)
 
  def run(self) -> ListArtifact:
- return self.extraction_engine.extract(self.input.to_text(), rulesets=self.all_rulesets, **self.args)
+ return self.extraction_engine.extract(
+ self.input.to_text(), rulesets=self.all_rulesets, prompt_stack=self.prompt_stack, **self.args
+ )

griptape/tasks/text_query_task.py

@@ -3,14 +3,16 @@
 from griptape.artifacts import TextArtifact
 from griptape.engines import BaseQueryEngine
 from griptape.loaders import TextLoader
-from griptape.tasks import BaseTextInputTask
+from griptape.tasks import PromptTask
 
 
 @define
-class TextQueryTask(BaseTextInputTask):
+class TextQueryTask(PromptTask):
  query_engine: BaseQueryEngine = field(kw_only=True)
  loader: TextLoader = field(default=Factory(lambda: TextLoader()), kw_only=True)
  namespace: Optional[str] = field(default=None, kw_only=True)
 
  def run(self) -> TextArtifact:
- return self.query_engine.query(self.input.to_text(), namespace=self.namespace, rulesets=self.all_rulesets)
+ return self.query_engine.query(
+ self.input.to_text(), namespace=self.namespace, rulesets=self.all_rulesets, prompt_stack=self.prompt_stack
+ )

griptape/tasks/text_summary_task.py

@@ -3,15 +3,19 @@
 from attr import define, field, Factory
 from griptape.artifacts import TextArtifact
 from griptape.engines import PromptSummaryEngine
-from griptape.tasks import BaseTextInputTask
+from griptape.tasks import PromptTask
 
 if TYPE_CHECKING:
  from griptape.engines import BaseSummaryEngine
 
 
 @define
-class TextSummaryTask(BaseTextInputTask):
+class TextSummaryTask(PromptTask):
  summary_engine: BaseSummaryEngine = field(kw_only=True, default=Factory(lambda: PromptSummaryEngine()))
 
  def run(self) -> TextArtifact:
- return TextArtifact(self.summary_engine.summarize_text(self.input.to_text(), rulesets=self.all_rulesets))
+ return TextArtifact(
+ self.summary_engine.summarize_text(
+ self.input.to_text(), rulesets=self.all_rulesets, prompt_stack=self.prompt_stack
+ )
+ )