wip for query builder

santiatpml · santiatpml · commit e2e7085ab44a · 2023-06-06T22:14:23.000-07:00
diff --git a/pgml-sdks/python/pgml/examples/pika_ex1.py b/pgml-sdks/python/pgml/examples/pika_ex1.py
@@ -0,0 +1,27 @@
+from pypika import Query, Table, AliasedQuery, Order
+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)
+
+table_embed = (
+    Query()
+    .with_(
+        Query()
+        .from_(embeddings_table)
+        .cross_join(query_table)
+        .on(embeddings_table.embedding == query_table.embedding).select('*'),
+        "cte",
+    )
+    .from_(AliasedQuery("cte"))
+    .select("score")
+)
+print(table_embed)
diff --git a/pgml-sdks/python/pgml/pgml/__init__.py b/pgml-sdks/python/pgml/pgml/__init__.py
@@ -4,4 +4,5 @@
     run_create_or_insert_statement,
     run_select_statement,
     run_drop_or_delete_statement,
-)
+)
+from .queries import Embed, CosineDistance
diff --git a/pgml-sdks/python/pgml/pgml/collection.py b/pgml-sdks/python/pgml/pgml/collection.py
@@ -20,6 +20,7 @@
     run_select_statement,
 )
 
+from .queries import Embed, CosineDistance
 from langchain.text_splitter import RecursiveCharacterTextSplitter
 from pypika.queries import Schema, Table, Query, QueryBuilder
 
@@ -849,6 +850,9 @@ def vector_recall(self,
                 return []
 
         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_select_statement = """
         WITH query_cte AS (
             SELECT pgml.embed(transformer => {model}, text => '{query_text}', kwargs => {model_params}) AS query_embedding
