feat: print only first title && abstract

custom_pgvector.py

@@ -16,7 +16,7 @@ from langchain.vectorstores.base import VectorStore
 from pgvector.sqlalchemy import Vector
 from sqlalchemy import delete, text
 from sqlalchemy.orm import Session, declarative_base
-
+from utils import str_to_list
 
 class DistanceStrategy(str, enum.Enum):
     """Enumerator of the Distance strategies."""
@@ -316,9 +316,9 @@ class CustomPGVector(VectorStore):
                 Document(
                     page_content=json.dumps(
                         {
-                            "abstract": result["abstract"],
+                            "abstract": result["abstract"][0],
                             "id": result["id"],
-                            "title": result["title"],
+                            "title": result["title"][0],
                             "authors": result["authors"],
                             "doi": result["doi"],
                             # "halID": result["halID"],
@@ -376,6 +376,8 @@ class CustomPGVector(VectorStore):
                     "distance",
                 ],
             )
+            results["abstract"] = results["abstract"].apply(str_to_list)
+            results["title"] = results["title"].apply(str_to_list)
         results = results.to_dict(orient="records")
         return results
 

utils.py

@@ -0,0 +1,20 @@
+import re
+
+
+def str_to_list(str_input: str) -> list[str]:
+    if isinstance(str_input, list):
+        return str_input
+
+    splits = re.split(r"', '|\", \"|', \"|\", '", str_input)
+    splits = [
+        split.removeprefix("[")
+        .removesuffix("]")
+        .removeprefix("(")
+        .removesuffix(")")
+        .removeprefix("'")
+        .removesuffix("'")
+        .removeprefix('"')
+        .removesuffix('"')
+        for split in splits
+    ]
+    return splits