jyapayne/langflow
1
0
Fork 0
Code Issues Pull requests Projects Releases 1 Packages Wiki Activity Actions

Refactor SQLGeneratorComponent to handle prompt template

Browse source
This commit is contained in:
Gabriel Luiz Freitas Almeida 2024-02-26 20:29:12 -03:00
commit 67aca6dd36
1 changed files with 25 additions and 7 deletions
Download patch file Download diff file Expand all files Collapse all files

32
src/backend/langflow/components/chains/SQLGenerator.py
View file

@ -32,21 +32,39 @@ class SQLGeneratorComponent(CustomComponent):
db: SQLDatabase, db: SQLDatabase,
llm: BaseLanguageModel, llm: BaseLanguageModel,
top_k: int = 5, top_k: int = 5,
prompt: Optional[PromptTemplate] = None, prompt: Optional[Text] = None,
) -> Text: ) -> Text:
if prompt:
prompt_template = PromptTemplate.from_template(template=prompt)
else:
prompt_template = None
if top_k > 0: if top_k > 0:
kwargs = { kwargs = {
"k": top_k, "k": top_k,
} }
if not prompt: if not prompt_template:
sql_query_chain = create_sql_query_chain(llm=llm, db=db, **kwargs) sql_query_chain = create_sql_query_chain(llm=llm, db=db, **kwargs)
else: else:
template = prompt.template if hasattr(prompt, "template") else prompt template = (
prompt_template.template
if hasattr(prompt, "template")
else prompt_template
)
# Check if {question} is in the prompt # Check if {question} is in the prompt
if "{question}" not in template or "question" not in template.input_variables: if (
raise ValueError("Prompt must contain `{question}` to be used with Natural Language to SQL.") "{question}" not in template
sql_query_chain = create_sql_query_chain(llm=llm, db=db, prompt=prompt, **kwargs) or "question" not in template.input_variables
query_writer = sql_query_chain | {"query": lambda x: x.replace("SQLQuery:", "").strip()} ):
raise ValueError(
"Prompt must contain `{question}` to be used with Natural Language to SQL."
)
sql_query_chain = create_sql_query_chain(
llm=llm, db=db, prompt=prompt_template, **kwargs
)
query_writer = sql_query_chain | {
"query": lambda x: x.replace("SQLQuery:", "").strip()
}
response = query_writer.invoke({"question": inputs}) response = query_writer.invoke({"question": inputs})
query = response.get("query") query = response.get("query")
self.status = query self.status = query