Spaces:
Runtime error
Runtime error
"""Chain that interprets a prompt and executes bash code to perform bash operations.""" | |
from typing import Dict, List | |
from pydantic import BaseModel, Extra | |
from langchain.chains.base import Chain | |
from langchain.chains.llm import LLMChain | |
from langchain.chains.llm_bash.prompt import PROMPT | |
from langchain.prompts.base import BasePromptTemplate | |
from langchain.schema import BaseLanguageModel | |
from langchain.utilities.bash import BashProcess | |
class LLMBashChain(Chain, BaseModel): | |
"""Chain that interprets a prompt and executes bash code to perform bash operations. | |
Example: | |
.. code-block:: python | |
from langchain import LLMBashChain, OpenAI | |
llm_bash = LLMBashChain(llm=OpenAI()) | |
""" | |
llm: BaseLanguageModel | |
"""LLM wrapper to use.""" | |
input_key: str = "question" #: :meta private: | |
output_key: str = "answer" #: :meta private: | |
prompt: BasePromptTemplate = PROMPT | |
class Config: | |
"""Configuration for this pydantic object.""" | |
extra = Extra.forbid | |
arbitrary_types_allowed = True | |
def input_keys(self) -> List[str]: | |
"""Expect input key. | |
:meta private: | |
""" | |
return [self.input_key] | |
def output_keys(self) -> List[str]: | |
"""Expect output key. | |
:meta private: | |
""" | |
return [self.output_key] | |
def _call(self, inputs: Dict[str, str]) -> Dict[str, str]: | |
llm_executor = LLMChain(prompt=self.prompt, llm=self.llm) | |
bash_executor = BashProcess() | |
self.callback_manager.on_text(inputs[self.input_key], verbose=self.verbose) | |
t = llm_executor.predict(question=inputs[self.input_key]) | |
self.callback_manager.on_text(t, color="green", verbose=self.verbose) | |
t = t.strip() | |
if t.startswith("```bash"): | |
# Split the string into a list of substrings | |
command_list = t.split("\n") | |
print(command_list) | |
# Remove the first and last substrings | |
command_list = [s for s in command_list[1:-1]] | |
output = bash_executor.run(command_list) | |
self.callback_manager.on_text("\nAnswer: ", verbose=self.verbose) | |
self.callback_manager.on_text(output, color="yellow", verbose=self.verbose) | |
else: | |
raise ValueError(f"unknown format from LLM: {t}") | |
return {self.output_key: output} | |
def _chain_type(self) -> str: | |
return "llm_bash_chain" | |