From 01cf28ad2f156fcce3ed02a968894fa38d42c916 Mon Sep 17 00:00:00 2001
From: jlonge4 <josh.longenecker1@gmail.com>
Date: Sat, 30 Mar 2024 15:23:10 -0400
Subject: [PATCH] add keyword test

---
 .../components/retrievers/pgvector/keyword_retriever.py     | 6 ++++++
 .../document_stores/pgvector/document_store.py              | 2 +-
 2 files changed, 7 insertions(+), 1 deletion(-)

diff --git a/integrations/pgvector/src/haystack_integrations/components/retrievers/pgvector/keyword_retriever.py b/integrations/pgvector/src/haystack_integrations/components/retrievers/pgvector/keyword_retriever.py
index ce0e8601a..856395b6e 100644
--- a/integrations/pgvector/src/haystack_integrations/components/retrievers/pgvector/keyword_retriever.py
+++ b/integrations/pgvector/src/haystack_integrations/components/retrievers/pgvector/keyword_retriever.py
@@ -58,6 +58,7 @@ def __init__(
         document_store: PgvectorDocumentStore,
         filters: Optional[Dict[str, Any]] = None,
         top_k: int = 10,
+        language: str = "english",
     ):
         """
         :param document_store: An instance of `PgvectorDocumentStore}.
@@ -74,6 +75,7 @@ def __init__(
         self.document_store = document_store
         self.filters = filters or {}
         self.top_k = top_k
+        self.language = language
 
     def to_dict(self) -> Dict[str, Any]:
         """
@@ -87,6 +89,7 @@ def to_dict(self) -> Dict[str, Any]:
             filters=self.filters,
             top_k=self.top_k,
             document_store=self.document_store.to_dict(),
+            language=self.language,
         )
 
     @classmethod
@@ -109,6 +112,7 @@ def run(
         user_query: str,
         filters: Optional[Dict[str, Any]] = None,
         top_k: Optional[int] = None,
+        language: Optional[str] = "english",
     ):
         """
         Retrieve documents from the `PgvectorDocumentStore`, based on their embeddings.
@@ -121,10 +125,12 @@ def run(
         """
         filters = filters or self.filters
         top_k = top_k or self.top_k
+        language = language or self.language
 
         docs = self.document_store._keyword_retrieval(
             user_query=user_query,
             filters=filters,
             top_k=top_k,
+            language=language,
         )
         return {"documents": docs}
diff --git a/integrations/pgvector/src/haystack_integrations/document_stores/pgvector/document_store.py b/integrations/pgvector/src/haystack_integrations/document_stores/pgvector/document_store.py
index a4aad1916..2dbe750f5 100644
--- a/integrations/pgvector/src/haystack_integrations/document_stores/pgvector/document_store.py
+++ b/integrations/pgvector/src/haystack_integrations/document_stores/pgvector/document_store.py
@@ -532,7 +532,7 @@ def _keyword_retrieval(
             ts_rank_cd(to_tsvector(%(language)s, content), query) DESC) AS rank
             FROM {table_name}, plainto_tsquery(%(language)s, %(query)s) query
             WHERE to_tsvector(%(language)s, content) @@ query"""
-        ).format(table_name=Identifier(self.table_name), language=language)
+        ).format(table_name=Identifier(self.table_name), language=language, query=user_query)
 
         sql_where_clause = SQL("")
         params = ()