victor7246 commited on
Commit
e76020d
·
verified ·
1 Parent(s): e79ed62

Update utils.py

Browse files
Files changed (1) hide show
  1. utils.py +4 -1
utils.py CHANGED
@@ -244,8 +244,11 @@ def clean_sql(s: str) -> str:
244
  if "LIMIT 1" in s:
245
  s = s.replace("LIMIT 1","").strip()
246
  s = s.replace("SELECT","SELECT TOP 1")
 
 
 
247
  return s
248
-
249
  class SQLDatabaseChainPatched(SQLDatabaseChain):
250
  intermediate_steps: List[Any] = Field(default_factory=list)
251
  llms: Dict[str, Any] = Field(default_factory=dict)
 
244
  if "LIMIT 1" in s:
245
  s = s.replace("LIMIT 1","").strip()
246
  s = s.replace("SELECT","SELECT TOP 1")
247
+ if s.endswith("TOP 1"):
248
+ s = s.replace("TOP 1","").strip()
249
+ s = s.replace("SELECT","SELECT TOP 1")
250
  return s
251
+
252
  class SQLDatabaseChainPatched(SQLDatabaseChain):
253
  intermediate_steps: List[Any] = Field(default_factory=list)
254
  llms: Dict[str, Any] = Field(default_factory=dict)