kobai-sdk 0.2.6__py3-none-any.whl → 0.2.8rc1__py3-none-any.whl

kobai/ai_query.py

@@ -79,7 +79,13 @@ def followup_question(question, data, question_name, llm_config:llm_config, over
                 openai_api_version=llm_config.api_version,
                 temperature = llm_config.temperature,
                 max_tokens = llm_config.max_tokens,
-            )           
+            )   
+    else:
+        chat_model = ChatDatabricks(
+            endpoint = llm_config.endpoint,
+            temperature = llm_config.temperature,
+            max_tokens = llm_config.max_tokens,
+            )        
 
     if llm_config.use_simple_prompt: 
         prompt = PromptTemplate.from_template(SIMPLE_PROMPT_TEMPLATE)

kobai/ai_rag.py

@@ -0,0 +1,113 @@
+from kobai import tenant_client
+
+def generate_sentences(tc: tenant_client.TenantClient, replica_schema=None):
+    if tc.spark_client is None:
+        return None
+    
+    ss = tc.spark_client.spark_session
+
+    print("Getting Tenant Config")
+    tenant_json = tc.get_tenant_config()
+
+    concepts = __get_concept_metadata(tenant_json, tc.schema, tc.model_id)
+
+    print("Dropping and Recreating the RAG Table")
+    ss.sql(__create_rag_table_sql(tc.schema, tc.model_id))
+
+    print("Generating Extraction SQL")
+    sql_statements = []
+    sql_statements.extend(__generate_sentence_sql_concept_literals(concepts, tc.schema, tc.model_id))
+    sql_statements.extend(__generate_sentence_sql_concept_relations(concepts, tc.schema, tc.model_id))
+
+    print("Running the Extraction")
+    for sql_statement in sql_statements:
+        ss.sql(sql_statement)
+
+    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 __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)"
+    
+def __replicate_to_catalog_sql(base_schema, target_schema, model_id):
+    move_sql = f"INSERT INTO {target_schema}.rag_{model_id} (content, concept_id, type)"
+    move_sql += f" SELECT content, concept_id, type FROM {base_schema}.rag_{model_id}"
+    return move_sql
+
+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_from = f"{con['con_table_name']} cid"
+        for prop in con["properties"]:
+
+            sql_from += f" INNER JOIN {con['prop_table_name']} AS {prop['label']}"
+            sql_from += f" ON cid._conceptid = {prop['label']}._conceptid"
+            sql_from += f" AND {prop['label']}.type = 'l'"
+            sql_from += f" AND {prop['label']}.name = '{prop['name']}'"
+            
+            sql += f" || 'The {prop['label']} is ' || any_value({prop['label']}.value) IGNORE NULLS || '. '"
+            
+            full_sql = f"INSERT INTO {schema}.rag_{model_id} (content, concept_id, type)"
+            full_sql += f" SELECT {sql} content, cid._conceptid concept_id, 'c' type FROM {sql_from} GROUP BY cid._conceptid"
+            
+            statements.append(full_sql)
+            #test_df = spark.sql(full_sql)
+    return statements
+
+def __generate_sentence_sql_concept_relations(concepts, schema, model_id):
+    statements = []
+    for con in concepts:
+     
+        sql_from = f"{con['prop_table_name']} "
+        for rel in con["relations"]:
+
+            sql = f"'The {con['label']} identified by ' || split(_conceptid,'#')[1]"
+            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])) || '. '"
+
+            
+            full_sql = f"INSERT INTO {schema}.rag_{model_id} (content, concept_id, type)"
+            full_sql += f" SELECT {sql} content, _conceptid concept_id, 'e' type FROM {sql_from} GROUP BY _conceptid"
+
+            statements.append(full_sql)
+    return statements
+
+def __get_concept_metadata(tenant_json, schema, model_id):
+    target_concept_labels = {}
+    for d in tenant_json["domains"]:
+        for c in d["concepts"]:
+            target_concept_labels[c["uri"]] = d["name"] + " " + c["label"]
+
+    concepts = []
+
+    for d in tenant_json["domains"]:
+        for c in d["concepts"]:
+            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']}"
+                    })
+            con_rels = []
+            for rel in c["relations"]:
+                con_rels.append({
+                    "label": rel["label"],
+                    "name": f"{model_id}/{d['name']}/{c['label']}#{rel['label']}",
+                    "target_con_label": target_concept_labels[rel["relationTypeUri"]]
+                })
+            concepts.append({
+                "label": d["name"] + " " + c["label"],
+                #"id_column": d["name"] + "_" + c["label"],
+                "relations": con_rels,
+                "properties": con_props,
+                #"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",
+                })
+    return concepts

kobai/tenant_client.py

@@ -339,26 +339,26 @@ class TenantClient:
                     hasRels = True
                 if hasProps or hasRels:
                     hasEither = True
-                con_label = dom["name"] + "_" + con["name"]
+                con_label = dom["name"] + "_" + con["label"]
                 out_table = con["schema_table"].replace(".data_", ".genie_").replace("_np", "")
                 sql = "CREATE OR REPLACE VIEW " + out_table + " "
                 sql += "(" + con_label + " COMMENT '" + con["schema_id_sentence"] + "' "
                 if hasEither:
                     sql += ", "
-                from_sql = "(SELECT DISTINCT _conceptid, p1 FROM " + con["schema_table"] + ") AS " + dom["name"] + "_" + con["name"] + "_ID "
+                from_sql = "(SELECT DISTINCT _conceptid, p1 FROM " + con["schema_table"] + ") AS " + dom["name"] + "_" + con["label"] + "_ID "
                 as_sql = "SELECT DISTINCT " + con_label + "_ID._conceptid " + con_label
                 if hasEither:
                     as_sql += ", "
                 as_props = []
                 top_props = []
                 for prop in con["properties"]:
-                    prop_label = con_label + "_" + prop["name"]
+                    prop_label = con_label + "_" + prop["label"]
                     prop_name = self.model_id + "/" + prop["uri"].split("/")[-2] + "/" + prop["uri"].split("/")[-1]
                     from_sql += "LEFT JOIN " + con["schema_table"] + " AS " + prop_label + " ON " + prop_label + ".type='l' AND " + prop_label + ".name='" + prop_name + "' AND " + prop_label + ".scenario='' AND " + con_label + "_ID.p1=" + prop_label + ".p1 AND " + con_label + "_ID._conceptid=" + prop_label + "._conceptid "
                     as_props.append(prop_label + ".value " + prop_label)
                     top_props.append(prop_label + " COMMENT '" + prop["schema_sentence"] + "'")
                 for prop in con["relations"]:
-                    prop_label = con_label + "_" + prop["name"]
+                    prop_label = con_label + "_" + prop["label"]
                     prop_name = self.model_id + "/" + prop["uri"].split("/")[-2] + "/" + prop["uri"].split("/")[-1]
                     from_sql += "LEFT JOIN " + con["schema_table"] + " AS " + prop_label + " ON " + prop_label + ".type='r' AND " + prop_label + ".name='" + prop_name + "' AND " + prop_label + ".scenario='' AND " + con_label + "_ID.p1=" + prop_label + ".p1 AND " + con_label + "_ID._conceptid=" + prop_label + "._conceptid "
                     as_props.append(prop_label + ".value " + prop_label)
@@ -371,7 +371,7 @@ class TenantClient:
 
 
                 if not_concepts is not None:
-                    if con["name"] in not_concepts:
+                    if con["label"] in not_concepts:
                         continue
 
                 concept_added = False
@@ -388,7 +388,7 @@ class TenantClient:
                             concept_added = True
                             sql_list.append({"table": out_table, "sql": sql})
                 if concepts is not None:
-                    if con["name"] in concepts and not concept_added:
+                    if con["label"] in concepts and not concept_added:
                         if enforce_map and con["map_count"] > 0:
                             sql_list.append({"table": out_table, "sql": sql})
                         else:
@@ -441,10 +441,10 @@ class TenantClient:
 
         return ai_query.followup_question(followup_question, 
                                           data, 
-                                          question_name, 
+                                          question_name,
+                                          None, 
                                           override_model=override_model, 
-                                          use_simple_prompt=use_simple_prompt, 
-                                          debug=debug)
+                                          )
 
     def process_question_results(self, question_def):
 
@@ -621,7 +621,7 @@ class TenantClient:
 # Tenant Questions
 ########################################
 
-    def run_question_remote(self, question_id):
+    def run_question_remote(self, question_id, dynamic_filters: dict = None):
 
         """
         Returns JSON formatted result of Kobai question.
@@ -630,10 +630,13 @@ class TenantClient:
         question_id (int): Numeric identifier of Kobai question.
         """
 
-        uri = '/data-svcs/api/query/' + str(question_id) + '/execute?jsontype=tableau' #'/data-svcs/api/query/4518/solution/9/execute/tabular?'
+        uri = '/data-svcs/api/query/' + str(question_id) + '/execute?' #'/data-svcs/api/query/4518/solution/9/execute/tabular?'
 
         queryParams = {'jsontype': 'tableau'}
 
+        if bool(dynamic_filters):
+            queryParams.update(dynamic_filters)
+
         uri += urllib.parse.urlencode(queryParams)
 
         json={
@@ -643,7 +646,7 @@ class TenantClient:
 
         return response.json()
     
-    def run_question_remote_spark(self, question_id, schema=None):
+    def run_question_remote_spark(self, question_id, dynamic_filters: dict = None, schema=None):
 
         """
         Returns result of Kobai question in PySpark Dataframe.
@@ -655,7 +658,7 @@ class TenantClient:
         if not self.__spark_check_init_status():
             return None
 
-        question_data = self.run_question_remote(question_id)
+        question_data = self.run_question_remote(question_id, dynamic_filters)
 
         if question_data is None:
             return None
@@ -1015,7 +1018,7 @@ class TenantClient:
         datasource_label (string): Label of datasource to use.
         table_name (string): Name of table to use from specified datasource.
         """
-        
+        data_source_id = 0
         existing_datasource = self.list_data_sources()
         for d in existing_datasource["used"]:
             if datasource_label.lower() == d["name"].lower():

{kobai_sdk-0.2.6.dist-info → kobai_sdk-0.2.8rc1.dist-info}/METADATA

@@ -1,6 +1,6 @@
-Metadata-Version: 2.1
+Metadata-Version: 2.2
 Name: kobai-sdk
-Version: 0.2.6
+Version: 0.2.8rc1
 Summary: A package that enables interaction with a Kobai tenant.
 Author-email: Ryan Oattes <ryan@kobai.io>
 License:                                  Apache License

kobai_sdk-0.2.8rc1.dist-info/RECORD

@@ -0,0 +1,14 @@
+kobai/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+kobai/ai_query.py,sha256=fMTcfj-6Ma3FRB08VYEDj8PwOEOtFGsJHyQrha5yvPg,4512
+kobai/ai_rag.py,sha256=y_N7qVu8HfUHHZPIyQSO7L995RBeNtDhva7U5HBHSfY,5063
+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.8rc1.dist-info/LICENSE,sha256=QwcOLU5TJoTeUhuIXzhdCEEDDvorGiC6-3YTOl4TecE,11356
+kobai_sdk-0.2.8rc1.dist-info/METADATA,sha256=nZTb2svQk01wT32zBZDPKgeYnSAx22YER5YLHEIjoAQ,19167
+kobai_sdk-0.2.8rc1.dist-info/WHEEL,sha256=In9FTNxeP60KnTkGw7wk6mJPYd_dQSjEZmXdBdMCI-8,91
+kobai_sdk-0.2.8rc1.dist-info/top_level.txt,sha256=ns1El3BrTTHKvoAgU1XtiSaVIudYeCXbEEUVY8HFDZ4,6
+kobai_sdk-0.2.8rc1.dist-info/RECORD,,

{kobai_sdk-0.2.6.dist-info → kobai_sdk-0.2.8rc1.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (75.6.0)
+Generator: setuptools (75.8.0)
 Root-Is-Purelib: true
 Tag: py3-none-any
 

kobai/test.py

@@ -1,5 +0,0 @@
-import llm_config, ai_query
-
-llm_config = llm_config.LLMConfig(api_key="sV9LuoA5n0PwqggMXOYMhhZlt56FpgnMXFohimPhD7Ug3CnBLbO8JQQJ99ALACYeBjFXJ3w3AAABACOGZm8X", llm_provider="azure_openai")
-llm_config.get_azure_ad_token()
-ai_query.followup_question_1(question="abc", data={}, question_name="sample", llm_config=llm_config)

kobai_sdk-0.2.6.dist-info/RECORD

@@ -1,14 +0,0 @@
-kobai/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-kobai/ai_query.py,sha256=QciKN643VsoZ5rP9zzxshuSMyXouLfp46WR9g4et1-M,4309
-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=LhAZPMYbFdMnE16g9jqdc5w2qSmubLPnW9_VBo90LFw,41696
-kobai/test.py,sha256=LHCKClyFNPgLA1hF62RlDzilaFw2Nd3BWcaLtOc5xdc,320
-kobai_sdk-0.2.6.dist-info/LICENSE,sha256=QwcOLU5TJoTeUhuIXzhdCEEDDvorGiC6-3YTOl4TecE,11356
-kobai_sdk-0.2.6.dist-info/METADATA,sha256=q3tpys2ZjqVXDl441A4TA-0QE8dagHxZEp_U68H-ePk,19164
-kobai_sdk-0.2.6.dist-info/WHEEL,sha256=PZUExdf71Ui_so67QXpySuHtCi3-J3wvF4ORK6k_S8U,91
-kobai_sdk-0.2.6.dist-info/top_level.txt,sha256=ns1El3BrTTHKvoAgU1XtiSaVIudYeCXbEEUVY8HFDZ4,6
-kobai_sdk-0.2.6.dist-info/RECORD,,