Skip to content

Commit e5dca97

Browse files
authored
fix: Fix QueryOptions not applied to similarity search bug (#185)
* fix: Fix QueryOptions not applied to similarity search bug * syntax
1 parent 1489f81 commit e5dca97

File tree

2 files changed

+18
-4
lines changed

2 files changed

+18
-4
lines changed

src/langchain_google_cloud_sql_pg/engine.py

Lines changed: 13 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -327,7 +327,7 @@ async def _aexecute(self, query: str, params: Optional[dict] = None) -> None:
327327
await conn.commit()
328328

329329
async def _aexecute_outside_tx(self, query: str) -> None:
330-
"""Execute a SQL query."""
330+
"""Execute a SQL query in a new transaction."""
331331
async with self._engine.connect() as conn:
332332
await conn.execute(text("COMMIT"))
333333
await conn.execute(text(query))
@@ -343,6 +343,18 @@ async def _afetch(
343343

344344
return result_fetch
345345

346+
async def _afetch_with_query_options(
347+
self, query: str, query_options: str
348+
) -> Sequence[RowMapping]:
349+
"""Set temporary database flags and fetch results from a SQL query."""
350+
async with self._engine.connect() as conn:
351+
await conn.execute(text(query_options))
352+
result = await conn.execute(text(query))
353+
result_map = result.mappings()
354+
result_fetch = result_map.fetchall()
355+
356+
return result_fetch
357+
346358
def _execute(self, query: str, params: Optional[dict] = None) -> None:
347359
"""Execute a SQL query."""
348360
return self._run_as_sync(self._aexecute(query, params))

src/langchain_google_cloud_sql_pg/vectorstore.py

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -599,10 +599,12 @@ async def __query_collection(
599599
filter = f"WHERE {filter}" if filter else ""
600600
stmt = f"SELECT *, {search_function}({self.embedding_column}, '{embedding}') as distance FROM \"{self.table_name}\" {filter} ORDER BY {self.embedding_column} {operator} '{embedding}' LIMIT {k};"
601601
if self.index_query_options:
602-
await self.engine._aexecute(
603-
f"SET LOCAL {self.index_query_options.to_string()};"
602+
query_options_stmt = f"SET LOCAL {self.index_query_options.to_string()};"
603+
results = await self.engine._afetch_with_query_options(
604+
stmt, query_options_stmt
604605
)
605-
results = await self.engine._afetch(stmt)
606+
else:
607+
results = await self.engine._afetch(stmt)
606608
return results
607609

608610
def similarity_search(

0 commit comments

Comments
 (0)