Update SQL templates (#12464)

templates/sql-llama2/sql_llama2/__init__.py

@@ -1,3 +1,3 @@
-from llama2.chain import chain
+from sql_llama2.chain import chain
 
 __all__ = ["chain"]

templates/sql-llama2/sql_llama2/chain.py

@@ -2,6 +2,7 @@ from pathlib import Path
 
 from langchain.llms import Replicate
 from langchain.prompts import ChatPromptTemplate
+from langchain.pydantic_v1 import BaseModel
 from langchain.schema.output_parser import StrOutputParser
 from langchain.schema.runnable import RunnablePassthrough
 from langchain.utilities import SQLDatabase
@@ -14,13 +15,11 @@ llm = Replicate(
     model_kwargs={"temperature": 0.01, "max_length": 500, "top_p": 1},
 )
 
-
 db_path = Path(__file__).parent / "nba_roster.db"
 rel = db_path.relative_to(Path.cwd())
 db_string = f"sqlite:///{rel}"
 db = SQLDatabase.from_uri(db_string, sample_rows_in_table_info=0)
 
-
 def get_schema(_):
     return db.get_table_info()
 
@@ -28,7 +27,6 @@ def get_schema(_):
 def run_query(query):
     return db.run(query)
 
-
 template_query = """Based on the table schema below, write a SQL query that would answer the user's question:
 {schema}
 
@@ -66,8 +64,12 @@ prompt_response = ChatPromptTemplate.from_messages(
     ]
 )
 
+# Supply the input types to the prompt 
+class InputType(BaseModel):
+    question: str
+
 chain = (
-    RunnablePassthrough.assign(query=sql_response)
+    RunnablePassthrough.assign(query=sql_response).with_types(input_type=InputType)
     | RunnablePassthrough.assign(
         schema=get_schema,
         response=lambda x: db.run(x["query"]),