Skip to content

Commit 6da0cfe

Browse files
anishnagErick Friis
andauthored
experimental[patch]: SmartLLMChain Output Key Customization (#14466)
**Description** The `SmartLLMChain` was was fixed to output key "resolution". Unfortunately, this prevents the ability to use multiple `SmartLLMChain` in a `SequentialChain` because of colliding output keys. This change simply gives the option the customize the output key to allow for sequential chaining. The default behavior is the same as the current behavior. Now, it's possible to do the following: ``` from langchain.chat_models import ChatOpenAI from langchain.prompts import PromptTemplate from langchain_experimental.smart_llm import SmartLLMChain from langchain.chains import SequentialChain joke_prompt = PromptTemplate( input_variables=["content"], template="Tell me a joke about {content}.", ) review_prompt = PromptTemplate( input_variables=["scale", "joke"], template="Rate the following joke from 1 to {scale}: {joke}" ) llm = ChatOpenAI(temperature=0.9, model_name="gpt-4-32k") joke_chain = SmartLLMChain(llm=llm, prompt=joke_prompt, output_key="joke") review_chain = SmartLLMChain(llm=llm, prompt=review_prompt, output_key="review") chain = SequentialChain( chains=[joke_chain, review_chain], input_variables=["content", "scale"], output_variables=["review"], verbose=True ) response = chain.run({"content": "chickens", "scale": "10"}) print(response) ``` --------- Co-authored-by: Erick Friis <erick@langchain.dev>
1 parent 0797358 commit 6da0cfe

1 file changed

Lines changed: 5 additions & 4 deletions

File tree

  • libs/experimental/langchain_experimental/smart_llm

libs/experimental/langchain_experimental/smart_llm/base.py

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -66,6 +66,7 @@ def resolve_prompt_inputs(self) -> Dict[str, Any]:
6666

6767
prompt: BasePromptTemplate
6868
"""Prompt object to use."""
69+
output_key: str = "resolution"
6970
ideation_llm: Optional[BaseLanguageModel] = None
7071
"""LLM to use in ideation step. If None given, 'llm' will be used."""
7172
critique_llm: Optional[BaseLanguageModel] = None
@@ -132,8 +133,8 @@ def input_keys(self) -> List[str]:
132133
def output_keys(self) -> List[str]:
133134
"""Defines the output keys."""
134135
if self.return_intermediate_steps:
135-
return ["ideas", "critique", "resolution"]
136-
return ["resolution"]
136+
return ["ideas", "critique", self.output_key]
137+
return [self.output_key]
137138

138139
def prep_prompts(
139140
self,
@@ -169,8 +170,8 @@ def _call(
169170
self.history.critique = critique
170171
resolution = self._resolve(stop, run_manager)
171172
if self.return_intermediate_steps:
172-
return {"ideas": ideas, "critique": critique, "resolution": resolution}
173-
return {"resolution": resolution}
173+
return {"ideas": ideas, "critique": critique, self.output_key: resolution}
174+
return {self.output_key: resolution}
174175

175176
def _get_text_from_llm_result(self, result: LLMResult, step: str) -> str:
176177
"""Between steps, only the LLM result text is passed, not the LLMResult object.

0 commit comments

Comments
 (0)