kobai-sdk 0.2.7__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

@@ -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):
 
@@ -1018,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.7.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.7
+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.7.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.7.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=R74NZpeo547vLpi_pt3QWT7I7phPy25kG_ZPuJhdEAg,41846
-kobai/test.py,sha256=LHCKClyFNPgLA1hF62RlDzilaFw2Nd3BWcaLtOc5xdc,320
-kobai_sdk-0.2.7.dist-info/LICENSE,sha256=QwcOLU5TJoTeUhuIXzhdCEEDDvorGiC6-3YTOl4TecE,11356
-kobai_sdk-0.2.7.dist-info/METADATA,sha256=EUiOF9YK0TElH6fZpb5125yy3TCFyyxJryhMgYOFuFo,19164
-kobai_sdk-0.2.7.dist-info/WHEEL,sha256=PZUExdf71Ui_so67QXpySuHtCi3-J3wvF4ORK6k_S8U,91
-kobai_sdk-0.2.7.dist-info/top_level.txt,sha256=ns1El3BrTTHKvoAgU1XtiSaVIudYeCXbEEUVY8HFDZ4,6
-kobai_sdk-0.2.7.dist-info/RECORD,,