fmt

docs/docs/use_cases/sql/large_db.ipynb

@@ -192,7 +192,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 49,
+   "execution_count": 59,
    "id": "0ccb0bf5-c580-428f-9cde-a58772ae784e",
    "metadata": {},
    "outputs": [
@@ -202,7 +202,7 @@
        "[Table(name='Music')]"
       ]
      },
-     "execution_count": 49,
+     "execution_count": 59,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -219,7 +219,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 50,
+   "execution_count": 60,
    "id": "ae4899fc-6f8a-4b10-983c-9e3fef4a7bb9",
    "metadata": {},
    "outputs": [
@@ -229,7 +229,7 @@
        "['Album', 'Artist', 'Genre', 'MediaType', 'Playlist', 'PlaylistTrack', 'Track']"
       ]
      },
-     "execution_count": 50,
+     "execution_count": 60,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -258,7 +258,7 @@
     "    return tables\n",
     "\n",
     "\n",
-    "table_chain = category_chain | get_tables\n",
+    "table_chain = category_chain | get_tables  # noqa\n",
     "table_chain.invoke({\"input\": \"What are all the genres of Alanis Morisette songs\"})"
    ]
   },

libs/langchain/langchain/chains/sql_database/query.py

@@ -1,10 +1,10 @@
-from typing import List, Optional, TypedDict, Union
+from typing import Any, Dict, List, Optional, TypedDict, Union
 
 from langchain_community.utilities.sql_database import SQLDatabase
 from langchain_core.language_models import BaseLanguageModel
 from langchain_core.output_parsers import StrOutputParser
 from langchain_core.prompts import BasePromptTemplate
-from langchain_core.runnables import Runnable, RunnableParallel, RunnablePassthrough
+from langchain_core.runnables import Runnable, RunnablePassthrough
 
 from langchain.chains.sql_database.prompt import PROMPT, SQL_PROMPTS
 
@@ -31,7 +31,7 @@ def create_sql_query_chain(
     db: SQLDatabase,
     prompt: Optional[BasePromptTemplate] = None,
     k: int = 5,
-) -> Runnable[Union[SQLInput, SQLInputWithTables], str]:
+) -> Runnable[Union[SQLInput, SQLInputWithTables, Dict[str, Any]], str]:
     """Create a chain that generates SQL queries.
 
     *Security Note*: This chain generates SQL queries for the given database.
@@ -128,8 +128,14 @@ def create_sql_query_chain(
         ),
     }
     return (
-        RunnablePassthrough.assign(**inputs)
-        | (lambda x: {k: v for k, v in x.items() if k not in ("question", "table_names_to_use")})
+        RunnablePassthrough.assign(**inputs)  # type: ignore
+        | (
+            lambda x: {
+                k: v
+                for k, v in x.items()
+                if k not in ("question", "table_names_to_use")
+            }
+        )
         | prompt_to_use.partial(top_k=str(k))
         | llm.bind(stop=["\nSQLResult:"])
         | StrOutputParser()