From 576609e66508e863c62972e4a5d4de2e905df656 Mon Sep 17 00:00:00 2001 From: Nick Furlotte Date: Thu, 2 Feb 2023 08:34:23 -0800 Subject: [PATCH] Update PAL to allow passing local and global context to PythonREPL (#774) Passing additional variables to the python environment can be useful for example if you want to generate code to analyze a dataset. I also added a tracker for the executed code - `code_history`. --- langchain/chains/pal/base.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/langchain/chains/pal/base.py b/langchain/chains/pal/base.py index ccecd9ec..36c0be62 100644 --- a/langchain/chains/pal/base.py +++ b/langchain/chains/pal/base.py @@ -4,7 +4,7 @@ As in https://arxiv.org/pdf/2211.10435.pdf. """ from __future__ import annotations -from typing import Any, Dict, List +from typing import Any, Dict, List, Optional from pydantic import BaseModel, Extra @@ -24,6 +24,8 @@ class PALChain(Chain, BaseModel): prompt: BasePromptTemplate stop: str = "\n\n" get_answer_expr: str = "print(solution())" + python_globals: Optional[Dict[str, Any]] = None + python_locals: Optional[Dict[str, Any]] = None output_key: str = "result" #: :meta private: class Config: @@ -54,7 +56,7 @@ class PALChain(Chain, BaseModel): self.callback_manager.on_text( code, color="green", end="\n", verbose=self.verbose ) - repl = PythonREPL() + repl = PythonREPL(_globals=self.python_globals, _locals=self.python_locals) res = repl.run(code + f"\n{self.get_answer_expr}") return {self.output_key: res.strip()}