example for query builder

santiadavani · santiadavani · commit d157f3449798 · 2023-06-07T18:09:00.000-07:00
diff --git a/pgml-sdks/python/pgml/examples/question_answering.py b/pgml-sdks/python/pgml/examples/question_answering.py
@@ -33,7 +33,7 @@
 
 start = time()
 query = "Who won 20 grammy awards?"
-results = collection.vector_search(query, top_k=5, metadata_filter={"title" : "Beyoncé"})
+results = collection.vector_search(query, top_k=5, metadata_filter={"title": "Beyoncé"})
 _end = time()
 console.print("\nResults for '%s'" % (query), style="bold")
 console.print(results)
diff --git a/pgml-sdks/python/pgml/examples/question_answering_query_builder.py b/pgml-sdks/python/pgml/examples/question_answering_query_builder.py
@@ -0,0 +1,49 @@
+from pgml import Database
+import os
+import json
+from datasets import load_dataset
+from time import time
+from dotenv import load_dotenv
+from rich.console import Console
+from pypika import Table
+
+load_dotenv()
+console = Console()
+
+local_pgml = "postgres://postgres@127.0.0.1:5433/pgml_development"
+
+conninfo = os.environ.get("PGML_CONNECTION", local_pgml)
+db = Database(conninfo)
+
+collection_name = "squad_collection"
+collection = db.create_or_get_collection(collection_name)
+
+
+data = load_dataset("squad", split="train")
+data = data.to_pandas()
+data = data.drop_duplicates(subset=["context"])
+
+documents = [
+    {"id": r["id"], "text": r["context"], "title": r["title"]}
+    for r in data.to_dict(orient="records")
+]
+
+collection.upsert_documents(documents[:200])
+collection.generate_chunks()
+collection.generate_embeddings()
+
+start = time()
+query = "Who won 20 grammy awards?"
+documents_table = Table("documents", schema=collection_name)
+sql_query = (
+    collection.vector_recall(query)
+    .where(documents_table.metadata.contains({"title": "Beyoncé"}))
+    .limit(5)
+)
+results = collection.execute(sql_query)
+_end = time()
+console.print("\nResults for '%s'" % (query), style="bold")
+console.print(results)
+console.print("Query time = %0.3f" % (_end - start))
+
+# db.archive_collection(collection_name)
diff --git a/pgml-sdks/python/pgml/pgml/collection.py b/pgml-sdks/python/pgml/pgml/collection.py
@@ -876,11 +876,14 @@ def vector_recall(
             .from_(embeddings_table)
             .select(
                 "chunk_id",
-                CosineDistance(
-                    embeddings_table.embedding, Cast(query_cte.embedding, "vector")
+                (
+                    1.0
+                    - CosineDistance(
+                        embeddings_table.embedding, Cast(query_cte.embedding, "vector")
+                    )
                 ).as_("score"),
             )
-            .inner_join(AliasedQuery("query_cte"))
+            .join(AliasedQuery("query_cte"))
             .cross()
         )
 
diff --git a/pgml-sdks/python/pgml/tests/test_collection.py b/pgml-sdks/python/pgml/tests/test_collection.py
@@ -174,8 +174,10 @@ def test_vector_recall(self):
         query = (
             self.collection.vector_recall("product is abc")
             .where(documents_table.metadata.contains({"source": "amazon"}))
-            .where(Cast(documents_table.metadata.get_json_value("reviews"),'INTEGER') < 45)
-            .limit(10)
+            .where(
+                Cast(documents_table.metadata.get_json_value("reviews"), "INTEGER") < 45
+            )
+            .limit(2)
         )
         results = self.collection.execute(query)
         assert results[0]["metadata"]["user"] == "John Doe"
