feat: replace sqlalchemy query by executing SQL

custom_pgvector.py

@@ -1,5 +1,5 @@
 from __future__ import annotations
-
+import pandas as pd
 import asyncio
 import contextlib
 import enum
@@ -28,6 +28,7 @@ from langchain.vectorstores.utils import maximal_marginal_relevance
 from pgvector.sqlalchemy import Vector
 from sqlalchemy import delete
 from sqlalchemy.orm import Session, declarative_base, relationship
+from sqlalchemy import text
 
 
 class DistanceStrategy(str, enum.Enum):
@@ -348,16 +349,19 @@ class CustomPGVector(VectorStore):
         docs = [
             (
                 Document(
-                    page_content=result.Article.abstract,
+                    page_content=result.abstract,
                     metadata={
-                        "id": result.Article.id,
-                        "title": result.Article.title,
-                        "doi": result.Article.doi,
+                        "id": result.id,
+                        "title": result.title,
+                        "authors": result.authors,
+                        "doi": result.doi,
+                        "keywords": results.keywords,
+                        "distance": results.distance,
                     },
                 ),
                 result.distance if self.embedding_function is not None else None,
             )
-            for result in results
+            for result in results.itertuples()
         ]
         return docs
 
@@ -369,16 +373,31 @@ class CustomPGVector(VectorStore):
     ) -> List[Any]:
         """Query the collection."""
         with Session(self._conn) as session:
-            results: List[Any] = (
-                session.query(
-                    self.EmbeddingStore,
-                    self.distance_strategy(embedding).label("distance"),  # type: ignore
+            results = session.execute(
+                text(
+                    f"""
+                    select
+                        a.id,
+                        a.title,
+                        a.doi,
+                        a.abstract,
+                        string_agg(distinct keyword."name", ',') as keywords, 
+                        string_agg(distinct author."name", ',') as authors,
+                        abstract_embedding <-> '{str(embedding)}' as distance
+                    from article a
+                    left join article_keyword ON article_keyword.article_id = a.id 
+                    left join keyword on article_keyword.keyword_id = keyword.id
+                    left join article_author ON article_author.article_id = a.id
+                    left join author on author.id = article_author.author_id  
+                    where abstract != 'NaN'
+                    GROUP BY a.id
+                    ORDER BY distance
+                    LIMIT {k};
+                """
                 )
-                .order_by(sqlalchemy.asc("distance"))
-                .limit(k)
-                .all()
             )
-        print(results)
+            results = results.fetchall()
+            results = pd.DataFrame(results, columns=["id", "title", "doi", "abstract", "keywords",  "authors", "distance"])
         return results
 
     def similarity_search_by_vector(