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 | |