Querybuilder for vector search prototyped

santiadavani · santiadavani · commit adf64d2da82b · 2023-06-07T12:54:35.000-07:00
diff --git a/pgml-sdks/python/pgml/examples/pika_ex1.py b/pgml-sdks/python/pgml/examples/pika_ex1.py
@@ -1,27 +1,39 @@
-from pypika import Query, Table, AliasedQuery, Order
+from pypika import Query, Table, AliasedQuery, Order, Field
+from pypika.functions import Cast
 from pgml.queries import Embed, CosineDistance
 
-embeddings_table = Table("embeddings_table")
-chunks_table = Table("chunks_table")
-documents_table = Table("documents_table")
-
-query_embed = Query().select(Embed(transformer="instructxl", text="hello"))
-print(query_embed)
-
-query_table = AliasedQuery("query_cte")
-query_cte = Query().with_(query_embed, "query_cte").from_(query_table).select('*')
-print(query_cte)
+embeddings_table = Table("test_collection_1.embeddings_d2beb7")
+chunks_table = Table("test_collection_1.chunks")
+documents_table = Table("test_collection_1.documents")
 
+model = "intfloat/e5-small"
+text = "hello world"
+query_embed = Query().select(Embed(transformer=model, text=text))
+query_cte = AliasedQuery("query_cte")
+cte = AliasedQuery("cte")
 table_embed = (
     Query()
-    .with_(
-        Query()
-        .from_(embeddings_table)
-        .cross_join(query_table)
-        .on(embeddings_table.embedding == query_table.embedding).select('*'),
-        "cte",
+    .from_(embeddings_table)
+    .select(
+        "chunk_id",
+        CosineDistance(
+            embeddings_table.embedding, Cast(query_cte.embedding, "vector")
+        ).as_("score"),
     )
-    .from_(AliasedQuery("cte"))
-    .select("score")
+    .inner_join(AliasedQuery("query_cte"))
+    .on(Field(1) == Field(1))
+)
+
+query_cte = (
+    Query()
+    .with_(query_embed, "query_cte")
+    .with_(table_embed, "cte")
+    .from_("cte")
+    .select(cte.score, chunks_table.chunk, documents_table.metadata).orderby(cte.score, order=Order.desc)
+    .inner_join(chunks_table)
+    .on(chunks_table.id == cte.chunk_id)
+    .inner_join(documents_table)
+    .on(documents_table.id == chunks_table.document_id)
 )
-print(table_embed)
+print(query_cte.get_sql().replace('"', ""))
+
diff --git a/pgml-sdks/python/pgml/pgml/__init__.py b/pgml-sdks/python/pgml/pgml/__init__.py
@@ -5,4 +5,4 @@
     run_select_statement,
     run_drop_or_delete_statement,
 )
-from .queries import Embed, CosineDistance
+from .queries import Embed, CosineDistance
diff --git a/pgml-sdks/python/pgml/pgml/collection.py b/pgml-sdks/python/pgml/pgml/collection.py
@@ -805,23 +805,23 @@ def vector_search(
 
     def execute(self, sql_statement: QueryBuilder) -> List[Dict[str, Any]]:
         conn = self.pool.getconn()
-        results = run_select_statement(conn,sql_statement.get_sql().replace("\"",""))
+        results = run_select_statement(conn, sql_statement.get_sql().replace('"', ""))
         self.pool.putconn(conn)
         return results
 
-    def vector_recall(self,        
+    def vector_recall(
+        self,
         query: str,
         query_parameters: Optional[Dict[str, Any]] = {},
         top_k: int = 5,
         model_id: int = 1,
-        splitter_id: int = 1) -> List[Dict[str, Any]]:
-        
-
+        splitter_id: int = 1,
+    ) -> List[Dict[str, Any]]:
         if model_id in self._cache_model_names.keys():
             model = self._cache_model_names[model_id]
         else:
             models = Table(self.models_table)
-            q = Query.from_(models).select('name').where(models.id == model_id)
+            q = Query.from_(models).select("name").where(models.id == model_id)
             results = self.execute(q)
             model = results[0]["name"]
             self._cache_model_names[model_id] = model
@@ -835,7 +835,12 @@ def vector_recall(self,
 
         if not embeddings_table:
             transforms_table = Table(self.transforms_table)
-            q = Query.from_(transforms_table).select('table_name').where(transforms_table.model_id == model_id).where(transforms_table.splitter_id == splitter_id)
+            q = (
+                Query.from_(transforms_table)
+                .select("table_name")
+                .where(transforms_table.model_id == model_id)
+                .where(transforms_table.splitter_id == splitter_id)
+            )
             embedding_table_results = self.execute(q)
             if embedding_table_results:
                 embeddings_table = embedding_table_results[0]["table_name"]
@@ -851,8 +856,17 @@ def vector_recall(self,
 
         conn = self.pool.getconn()
 
-        cte_query = Query.select(Embed(transformer=model,text=query,parameters=query_parameters)).with_()
-        table_embedding = Query.from_(embeddings_table).select('chunk_id',CosineDistance(embeddings_table.embedding,query_embedding.cosine)).cross_join(query_embedding)
+        cte_query = Query.select(
+            Embed(transformer=model, text=query, parameters=query_parameters)
+        ).with_()
+        table_embedding = (
+            Query.from_(embeddings_table)
+            .select(
+                "chunk_id",
+                CosineDistance(embeddings_table.embedding, query_embedding.cosine),
+            )
+            .cross_join(query_embedding)
+        )
         cte_select_statement = """
         WITH query_cte AS (
             SELECT pgml.embed(transformer => {model}, text => '{query_text}', kwargs => {model_params}) AS query_embedding
@@ -883,4 +897,4 @@ def vector_recall(self,
         )
         self.pool.putconn(conn)
 
-        return search_results
+        return search_results
diff --git a/pgml-sdks/python/pgml/pgml/queries.py b/pgml-sdks/python/pgml/pgml/queries.py
@@ -3,10 +3,20 @@
 from typing import Dict
 from pypika import JSON, Array
 
+
 class Embed(Function):
-    def __init__(self, transformer: str, text: str, parameters: Dict[str,Any] = {}, alias: str = "embedding") -> None:
-        super(Embed,self).__init__('pgml.embed', transformer, text, JSON(parameters), alias=alias)
+    def __init__(
+        self,
+        transformer: str,
+        text: str,
+        parameters: Dict[str, Any] = {},
+        alias: str = "embedding",
+    ) -> None:
+        super(Embed, self).__init__(
+            "pgml.embed", transformer, text, JSON(parameters), alias=alias
+        )
+
 
 class CosineDistance(Function):
     def __init__(self, lhs: Array, rhs: Array, alias: str = "cosine") -> None:
-        super(CosineDistance,self).__init__('cosine_distance', lhs, rhs, alias=alias)
+        super(CosineDistance, self).__init__("cosine_distance", lhs, rhs, alias=alias)

Original file line number	Diff line number	Diff line change
`@@ -5,4 +5,4 @@`
`5`	`5`	`run_select_statement,`
`6`	`6`	`run_drop_or_delete_statement,`
`7`	`7`	`)`
`8`		`-from .queries import Embed, CosineDistance`
	`8`	`+from .queries import Embed, CosineDistance`