Spaces:
Running
Running
from langchain_core.messages import BaseMessage | |
from langchain_core.prompts import PromptTemplate | |
from langflow.custom import CustomComponent | |
from langflow.field_typing import LanguageModel, Text | |
class ShouldRunNextComponent(CustomComponent): | |
display_name = "Should Run Next" | |
description = "Determines if a vertex is runnable." | |
name = "ShouldRunNext" | |
def build(self, llm: LanguageModel, question: str, context: str, retries: int = 3) -> Text: | |
template = ( | |
"Given the following question and the context below, answer with a yes or no.\n\n" | |
"{error_message}\n\n" | |
"Question: {question}\n\n" # noqa: RUF100, RUF027 | |
"Context: {context}\n\n" # noqa: RUF100, RUF027 | |
"Answer:" | |
) | |
prompt = PromptTemplate.from_template(template) | |
chain = prompt | llm | |
error_message = "" | |
for _i in range(retries): | |
result = chain.invoke( | |
{"question": question, "context": context, "error_message": error_message}, | |
config={"callbacks": self.get_langchain_callbacks()}, | |
) | |
if isinstance(result, BaseMessage): | |
content = result.content | |
elif isinstance(result, str): | |
content = result | |
if isinstance(content, str) and content.lower().strip() in {"yes", "no"}: | |
break | |
condition = str(content).lower().strip() == "yes" | |
self.status = f"Should Run Next: {condition}" | |
if condition is False: | |
self.stop() | |
return context | |