kobai-sdk 0.2.8rc2__py3-none-any.whl → 0.2.8rc3__py3-none-any.whl

kobai/ai_rag.py

@@ -1,15 +1,12 @@
 from kobai import tenant_client
 
 from pyspark.sql.types import StructType, StructField, StringType, ArrayType, FloatType, IntegerType
-#from delta.tables import *
-from sentence_transformers import SentenceTransformer, util
-#from deltalake import DeltaTable
+from pyspark.sql import functions as F
+from sentence_transformers import SentenceTransformer
 from delta import DeltaTable
 from typing import Union
 from langchain_core.language_models.chat_models import BaseChatModel
 from langchain_core.embeddings import Embeddings
-from langchain_core.documents import Document
-#from databricks_langchain import DatabricksEmbeddings, ChatDatabricks
 from langchain_community.document_loaders import PySparkDataFrameLoader
 from langchain import hub
 from langchain_core.output_parsers import StrOutputParser
@@ -18,14 +15,16 @@ from langchain_core.output_parsers import StrOutputParser
 
 
 
-def generate_sentences(tc: tenant_client.TenantClient, replica_schema=None, concept_white_list=None):
+def generate_sentences(tc: tenant_client.TenantClient, replica_schema=None, concept_white_list=None, use_questions=False):
 
     """
     Extract Semantic Data from Graph to Delta Table
 
     Parameters:
     tc (TenantClient): The Kobai tenant_client instance instantiated via the SDK.
-    replica_schema (str): An alternate schema (catalog.database) to create the Delta table. Useful when the base Kobai schema is not on a Unity Catalog.
+    replica_schema (str) OPTIONAL: An alternate schema (catalog.database) to create the Delta table. Useful when the base Kobai schema is not on a Unity Catalog.
+    concept_white_list ([str]) OPTIONAL: A list of Domain and Concept names for extraction.
+    use_questions (bool) OPTIONAL: Extract facts from published Kobai questions. 
     """
     
     if tc.spark_client is None:
@@ -37,8 +36,6 @@ def generate_sentences(tc: tenant_client.TenantClient, replica_schema=None, conc
     tenant_json = tc.get_tenant_config()
 
     concepts = __get_concept_metadata(tenant_json, tc.schema, tc.model_id, concept_white_list)
-    print(concepts)
-    print("")
 
     print("Dropping and Recreating the RAG Table")
     ss.sql(__create_rag_table_sql(tc.schema, tc.model_id))
@@ -50,17 +47,77 @@ def generate_sentences(tc: tenant_client.TenantClient, replica_schema=None, conc
 
     print("Running the Extraction")
     for sql_statement in sql_statements:
-        print(sql_statement)
-        print("")
         ss.sql(sql_statement)
 
+    if use_questions:
+        __generate_sentences_from_questions(tc)
+
     if replica_schema is not None:
         print("Replicating Schema")
         ss.sql(__create_rag_table_sql(replica_schema, tc.model_id))
         ss.sql(__replicate_to_catalog_sql(tc.schema, replica_schema, tc.model_id))
 
+def __generate_sentences_from_questions(tc: tenant_client.TenantClient):
+    ss = tc.spark_client.spark_session
+
+    print("Getting Question Data")
+
+    tenant_json = tc.get_tenant_config()
+
+    published_queries = []
+    for p in tenant_json["publishedAPIs"]:
+        published_queries.append(p["queryId"])
+
+    question_names = {}
+    for q in tenant_json["queries"]:
+        if q["id"] in published_queries:
+            question_names[q["id"]] = q["description"]
+
+    schemaV = StructType([
+            StructField("sentence",StringType(),True),
+            StructField("query_id", StringType(), True)
+        ])
+
+    sentences = []
+    for p in published_queries:
+        output = tc.run_question_remote(p)
+        for r in output:
+            sentence = f"For {question_names[p]}: "
+            for c in r:
+                sentence += f"The {c.replace('_', ' ')} is {r[c]}. "
+            sentences.append([sentence, p])
+
+
+    sentences_df = ss.createDataFrame(sentences, schemaV)
+    sentences_df = sentences_df.select(
+        F.col("sentence").alias("sentence"),
+        F.col("query_id").alias("concept_id"),
+        F.lit("q").alias("type"),
+    )
+
+    schema = tc.schema
+        
+    view_name = f"rag_{tc.model_id}_question_sentences"
+    sentences_df.createOrReplaceTempView(view_name)
+
+    full_sql = f"INSERT INTO {schema}.rag_{tc.model_id} (content, concept_id, type)"
+    full_sql += f" SELECT sentence, concept_id, type FROM {view_name}"
+    
+    ss.sql(full_sql)
+
+
+
 def encode_to_delta_local(tc: tenant_client.TenantClient, st_model: SentenceTransformer, replica_schema=None):
     
+    """
+    Encode Semantic Data to Vectors in Delta Table
+
+    Parameters:
+    tc (TenantClient): The Kobai tenant_client instance instantiated via the SDK.
+    st_model (SentenceTransformer): A sentence_transformers model to use for encoding.
+    replica_schema (str) OPTIONAL: An alternate schema (catalog.database) to create the Delta table. Useful when the base Kobai schema is not on a Unity Catalog.
+    """
+
     if tc.spark_client is None:
         return None
     
@@ -76,7 +133,6 @@ def encode_to_delta_local(tc: tenant_client.TenantClient, st_model: SentenceTran
     num_records = sentences_df.count()
     query_batch_size = 100000
 
-    #pool = model.start_multi_process_pool()
 
     for x in range(0, num_records, query_batch_size):
         print(f"Running Batch Starting at {x}")
@@ -84,14 +140,8 @@ def encode_to_delta_local(tc: tenant_client.TenantClient, st_model: SentenceTran
         sentences_df = ss.sql(sentences_sql)
         content_list = [r["content"] for r in sentences_df.collect()]
         id_list = [r["id"] for r in sentences_df.collect()]
-        #num_records_batch = len(content_list)
-        #print("Done Getting Data")
 
-        
         vector_list = st_model.encode(content_list, normalize_embeddings=True, show_progress_bar=True)
-        #vector_list = model.encode_multi_process(content_list, pool)
-
-        #print("Done Encoding")
 
         schemaV = StructType([
             StructField("id",IntegerType(),True),
@@ -123,6 +173,18 @@ def encode_to_delta_local(tc: tenant_client.TenantClient, st_model: SentenceTran
 
 def rag_delta(tc: tenant_client.TenantClient, emb_model: Union[SentenceTransformer, Embeddings], chat_model: BaseChatModel, question, k=5, replica_schema=None):
 
+    """
+    Run a RAG query using vectors in Delta table.
+
+    Parameters:
+    tc (TenantClient): The Kobai tenant_client instance instantiated via the SDK.
+    emb_model (UNION[SentenceTransformer, Embeddings]): A sentence_transformers or langchain embedding model to use for encoding the query.
+    chat_model (BaseChatModel): A langchain chat model to use in the RAG pipeline.
+    question (str): The user's query.
+    k (int) OPTIONAL: The number of RAG documents to retrieve.
+    replica_schema (str) OPTIONAL: An alternate schema (catalog.database) to create the Delta table. Useful when the base Kobai schema is not on a Unity Catalog.
+    """
+        
     schema = tc.schema
     if replica_schema is not None:
         schema = replica_schema
@@ -145,11 +207,8 @@ def rag_delta(tc: tenant_client.TenantClient, emb_model: Union[SentenceTransform
         print("Invalid Chat Model Type")
         return None
 
-    #print(vector_list)
     vector_list = [str(x) for x in vector_list]
-    #print(vector_list)
     vector_sql = ", ".join(vector_list)
-    #print(vector_sql)
 
     results = ss.sql(f"""
             SELECT content, reduce(zip_with(vector, cast(array({vector_sql}) as array<float>), (x,y) -> x*y), float(0.0), (acc,x) -> acc + x) score
@@ -162,8 +221,6 @@ def rag_delta(tc: tenant_client.TenantClient, emb_model: Union[SentenceTransform
     documents = loader.load()
     docs_content = "\n\n".join(doc.page_content for doc in documents)
 
-    #print(docs_content)
-
     prompt = hub.pull("rlm/rag-prompt")
 
     output_parser = StrOutputParser()
@@ -179,34 +236,6 @@ def rag_delta(tc: tenant_client.TenantClient, emb_model: Union[SentenceTransform
 
     return response
 
-def dep_rag_delta(tc: tenant_client.TenantClient, st_model: SentenceTransformer, question, k=5, replica_schema=None):
-
-    schema = tc.schema
-    if replica_schema is not None:
-        schema = replica_schema
-
-    if tc.spark_client is None:
-        return None
-    
-    ss = tc.spark_client.spark_session
-
-    vector_list = st_model.encode(question, normalize_embeddings=True).tolist()
-
-    #print(vector_list)
-    vector_list = [str(x) for x in vector_list]
-    #print(vector_list)
-    vector_sql = ", ".join(vector_list)
-    #print(vector_sql)
-
-    results = ss.sql(f"""
-            SELECT content, reduce(zip_with(vector, cast(array({vector_sql}) as array<float>), (x,y) -> x*y), float(0.0), (acc,x) -> acc + x) score
-            FROM {schema}.rag_{tc.model_id}
-            ORDER BY score DESC
-            LIMIT {k}
-            """)
-    
-    return results
-
 def __create_rag_table_sql(schema, model_id):
     return f"CREATE OR REPLACE TABLE {schema}.rag_{model_id} (id BIGINT GENERATED BY DEFAULT AS IDENTITY, content STRING, type string, concept_id string, vector ARRAY<FLOAT>) TBLPROPERTIES (delta.enableChangeDataFeed = true)"
     
@@ -219,10 +248,8 @@ def __generate_sentence_sql_concept_literals(concepts, schema, model_id):
     statements = []
     for con in concepts:
         sql = f"'This is a {con['label']}. '"
-        #sql += " || 'It is identified by ' || split(cid._conceptid,'#')[1] || '. '"
         sql += " || 'It is identified by ' || cid._plain_conceptid || '. '"
             
-        #sql_from = f"{con['con_table_name']} cid"
         sql_from = f"(SELECT _conceptid, _plain_conceptid FROM {con['prop_table_name']} GROUP BY _conceptid, _plain_conceptid) cid"
         for prop in con["properties"]:
 
@@ -237,7 +264,6 @@ def __generate_sentence_sql_concept_literals(concepts, schema, model_id):
         full_sql += f" SELECT {sql} content, cid._conceptid concept_id, 'c' type FROM {sql_from} GROUP BY cid._conceptid, cid._plain_conceptid"
         
         statements.append(full_sql)
-            #test_df = spark.sql(full_sql)
     return statements
 
 def __generate_sentence_sql_concept_relations(concepts, schema, model_id):
@@ -252,7 +278,6 @@ def __generate_sentence_sql_concept_relations(concepts, schema, model_id):
 
             sql = f"'The {con['label']} identified by ' || rel._plain_conceptid"
             sql += f" || ' has a relationship called {rel['label']} that connects it to one or more {rel['target_con_label']} identified by '"
-            #sql += " || concat_ws(', ', array_agg(split(value, '#')[1])) || '. '"
             sql += " || concat_ws(', ', array_agg(cid._plain_conceptid)) || '. '"
 
             
@@ -274,15 +299,11 @@ def __get_concept_metadata(tenant_json, schema, model_id, whitelist):
             }
     
     concepts = []
-    #parents = {}
     for d in tenant_json["domains"]:
         for c in d["concepts"]:
-            #if whitelist is not None and d["name"] + " " + c["label"] not in whitelist:
-            #    continue
             con_props = []
             for col in c["properties"]:
                 con_props.append({
-                    #"col_name": d["name"] + "_" + c["label"] + "_" + col["label"],
                     "label": col["label"],
                     "name": f"{model_id}/{d['name']}/{c['label']}#{col['label']}"
                     })
@@ -302,13 +323,9 @@ def __get_concept_metadata(tenant_json, schema, model_id, whitelist):
             concepts.append({
                 "uri": c["uri"],
                 "label": d["name"] + " " + c["label"],
-                #"id_column": d["name"] + "_" + c["label"],
                 "relations": con_rels,
                 "properties": con_props,
                 "parents": con_parents,
-                #"table_name": "data_" + k.model_id + "_" + d["name"] + "_" + c["label"] + "_w",
-                #"prop_table_name": f"{schema}.data_{model_id}_{d['name']}_{c['label']}_np",
-                #"con_table_name": f"{schema}.data_{model_id}_{d['name']}_{c['label']}_c",
                 "prop_table_name": target_table_names[c["uri"]]["prop"],
                 "con_table_name": target_table_names[c["uri"]]["con"]
                 })
@@ -319,7 +336,6 @@ def __get_concept_metadata(tenant_json, schema, model_id, whitelist):
                 for a in concepts:
                     if a["uri"] == p:
                         concepts[ci]["properties"].extend(a["properties"])
-                        #concepts[ci]["properties"] = list(set(concepts[ci]["properties"]))
 
     out_concepts = []
     for c in concepts:

{kobai_sdk-0.2.8rc2.dist-info → kobai_sdk-0.2.8rc3.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: kobai-sdk
-Version: 0.2.8rc2
+Version: 0.2.8rc3
 Summary: A package that enables interaction with a Kobai tenant.
 Author-email: Ryan Oattes <ryan@kobai.io>
 License:                                  Apache License
@@ -211,7 +211,7 @@ Classifier: Intended Audience :: Developers
 Classifier: License :: OSI Approved :: Apache Software License
 Classifier: Programming Language :: Python
 Classifier: Programming Language :: Python :: 3
-Requires-Python: >=3.9
+Requires-Python: >=3.11
 Description-Content-Type: text/markdown
 License-File: LICENSE
 Requires-Dist: pyspark

{kobai_sdk-0.2.8rc2.dist-info → kobai_sdk-0.2.8rc3.dist-info}/RECORD

@@ -1,14 +1,14 @@
 kobai/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 kobai/ai_query.py,sha256=fMTcfj-6Ma3FRB08VYEDj8PwOEOtFGsJHyQrha5yvPg,4512
-kobai/ai_rag.py,sha256=KbIlrbOX-0hbt7HaOh7nyzIrROGotGt0ghQSlzN6ZUA,13096
+kobai/ai_rag.py,sha256=TtUbUcSN9mIsauGyS_nw8j58T9jEd4OFiAwNvzo-rr8,13593
 kobai/databricks_client.py,sha256=fyqqMly2Qm0r1AHWsQjkYeNsDdH0G1JSgTkF9KJ55qA,2118
 kobai/demo_tenant_client.py,sha256=wlNc-bdI2wotRXo8ppUOalv4hYdBlek_WzJNARZV-AE,9293
 kobai/llm_config.py,sha256=ZFx81cUAOHYZgRoTkTY-utQYaWYlmR8773ZJpj74C1A,1900
 kobai/spark_client.py,sha256=opM_F-4Ut5Hq5zZjWMuLvUps9sDULvyPNZHXGL8dW1k,776
 kobai/tenant_api.py,sha256=9U6UbxpaAb-kpbuADXx3kbkNKaOzYy0I-GGwbpiCCOk,4212
 kobai/tenant_client.py,sha256=AyJ5R2oukEv3q1dcItpojvTUVp5-gwUKvyGjofjBKyc,41821
-kobai_sdk-0.2.8rc2.dist-info/LICENSE,sha256=QwcOLU5TJoTeUhuIXzhdCEEDDvorGiC6-3YTOl4TecE,11356
-kobai_sdk-0.2.8rc2.dist-info/METADATA,sha256=FiYuYjOjY5Hf5X58Cgv3Qhu_KnIWkeeCNmyzD8k-r4A,19204
-kobai_sdk-0.2.8rc2.dist-info/WHEEL,sha256=In9FTNxeP60KnTkGw7wk6mJPYd_dQSjEZmXdBdMCI-8,91
-kobai_sdk-0.2.8rc2.dist-info/top_level.txt,sha256=ns1El3BrTTHKvoAgU1XtiSaVIudYeCXbEEUVY8HFDZ4,6
-kobai_sdk-0.2.8rc2.dist-info/RECORD,,
+kobai_sdk-0.2.8rc3.dist-info/LICENSE,sha256=QwcOLU5TJoTeUhuIXzhdCEEDDvorGiC6-3YTOl4TecE,11356
+kobai_sdk-0.2.8rc3.dist-info/METADATA,sha256=f75oEdxRWLrr0bVmH1OvIlvc0KS9TrpNTh65eTlKX6k,19205
+kobai_sdk-0.2.8rc3.dist-info/WHEEL,sha256=In9FTNxeP60KnTkGw7wk6mJPYd_dQSjEZmXdBdMCI-8,91
+kobai_sdk-0.2.8rc3.dist-info/top_level.txt,sha256=ns1El3BrTTHKvoAgU1XtiSaVIudYeCXbEEUVY8HFDZ4,6
+kobai_sdk-0.2.8rc3.dist-info/RECORD,,