fix: adding better validation

This commit is contained in:
Gabriel Almeida 2023-04-05 22:33:11 -03:00
commit 8dbd3adecd
2 changed files with 12 additions and 2 deletions

View file

@ -1,3 +1,4 @@
from langflow.graph.utils import extract_input_variables_from_prompt
from pydantic import BaseModel, validator
@ -25,3 +26,12 @@ class CodeValidationResponse(BaseModel):
class PromptValidationResponse(BaseModel):
input_variables: list
def validate_prompt(template: str):
input_variables = extract_input_variables_from_prompt(template)
if invalid := [variable for variable in input_variables if " " in variable]:
raise ValueError(
f"Invalid input variables: {invalid}. Please remove spaces from input variables"
)
return PromptValidationResponse(input_variables=input_variables)

View file

@ -5,6 +5,7 @@ from langflow.api.base import (
CodeValidationResponse,
Prompt,
PromptValidationResponse,
validate_prompt,
)
from langflow.graph.utils import extract_input_variables_from_prompt
from langflow.utils.logger import logger
@ -29,8 +30,7 @@ def post_validate_code(code: Code):
@router.post("/prompt", status_code=200, response_model=PromptValidationResponse)
def post_validate_prompt(prompt: Prompt):
try:
input_variables = extract_input_variables_from_prompt(prompt.template)
return PromptValidationResponse(input_variables=input_variables)
return validate_prompt(prompt.template)
except Exception as e:
logger.exception(e)
return HTTPException(status_code=500, detail=str(e))