kobai-sdk 0.3.2__py3-none-any.whl → 0.3.3__py3-none-any.whl

kobai/genie.py

@@ -0,0 +1,194 @@
+
+def get_genie_descriptions(solution_id, structure, schema):
+    for di, dom in enumerate(structure["domains"]):
+        structure["domains"][di]["label"] = dom["name"]
+
+    add_episteme_metadata(structure)
+    download_inherited_props(structure)
+    add_map_count(structure)
+    add_semantic_sentences(structure)
+    add_schema_sentences(structure, schema, solution_id)
+
+    #add map count
+
+    return structure
+
+#def get_genie_questions(solution_id, structure):
+
+#    graph_uri = structure["uri"]
+
+#    question_config = get_tenant_question_config(solution_id, structure)
+
+#    questions = []
+#    for q in question_config:
+#        q_name = q["description"]
+#        q_def = json.loads(q["definition"])
+
+#        if q["published"]:
+#            result = get_question_sql(q_def, graph_uri)
+#            if result is not None:
+#                questions.append({"name": q_name, "sql": result})
+
+#    return questions
+
+############################
+# Config Building
+############################
+
+def add_episteme_metadata(structure):
+    for dom in structure["domains"]:
+        dom["e_id"] = dom["label"]
+        for con in dom["concepts"]:
+            con["e_id"] = dom["label"] + "_" + con["label"]
+            
+            for prop in con["properties"]:
+                prop["e_id"] = dom["label"] + "_" + con["label"] + "_" + prop["label"]
+
+            for rel in con["relations"]:
+                rel["e_id"] = dom["label"] + "_" + con["label"] + "_" + rel["label"]
+                rel["e_target_id"] = rel["relationTypeUri"].split("/")[-1].replace("#", "_")
+
+def download_inherited_props(structure):
+    for dom in structure["domains"]:
+        if "concepts" in dom:
+            for con in dom["concepts"]:
+                recurse_parent_props(con["uri"], structure, con["properties"], con["relations"], visited=[])
+
+def recurse_parent_props(uri, structure, props, rels, visited=None):
+    visited.append(uri)
+    for dom in structure["domains"]:
+        for con in dom["concepts"]:
+            if con["uri"] == uri:
+                for icon in con["inheritedConcepts"]:
+                    for pdom in structure["domains"]:
+                        for pcon in pdom["concepts"]:
+                            if pcon["uri"] == icon:
+                                for pprop in pcon["properties"]:
+                                    prop_found = False
+                                    for pf in props:
+                                        if pf["uri"] == pprop["uri"]:
+                                            prop_found = True
+                                    if not prop_found:
+                                        props.append(pprop)
+                                for prel in pcon["relations"]:
+                                    rel_found = False
+                                    for rf in rels:
+                                        if rf["uri"] == prel["uri"]:
+                                            rel_found = True
+                                    if not rel_found:
+                                        rels.append(prel)
+                    if icon not in visited:
+                        recurse_parent_props(icon, structure, props, rels, visited)
+
+def add_map_count(structure):
+
+    for dom in structure["domains"]:
+        for con in dom["concepts"]:
+            map_count = 0
+            #for md in mapping_defs:
+            for md in structure["mappingDefs"]:
+                if con["uri"] == md["conceptTypeUri"]:
+                    map_count = map_count + 1
+            con["map_count"] = map_count
+
+############################
+# Question Config
+############################
+
+def get_tenant_question_config(solution_id, structure):
+
+    question_defs = []
+    for row in structure["queries"]:
+        print(row)
+        is_published = False
+        if row[4] is not None and row[4] != "":
+            is_published = True
+        question_def = {"question_id": row[1], "description": row[2], "definition": row[3], "published": is_published}
+        question_defs.append(question_def)
+    return question_defs
+
+############################
+# Sentences
+############################
+
+def add_semantic_sentences(structure):
+    for dom in structure["domains"]:
+        concept_list = []
+
+        for con in dom["concepts"]:
+            concept_sentence = ""
+            property_list = []
+            relation_list = []
+            concept_list.append(con["label"])
+
+            for prop in con["properties"]:
+                property_list.append(prop["label"])
+                property_sentence = "The " + prop["label"] + " for the " + dom["label"] + " " + con["label"] + "."
+                prop["sementic_sentence"] = property_sentence
+
+            for rel in con["relations"]:
+                relation_list.append(rel["label"])
+
+            concept_sentence = "The " + con["label"] + " concept contains details about " + smart_comma_formatting(property_list) + "."
+            con["semantic_sentence"] = concept_sentence
+
+        domain_sentence = "The " + dom["label"] + " domain contains concepts called " + smart_comma_formatting(concept_list) + "."
+        dom["semantic_sentence"] = domain_sentence
+
+def add_schema_sentences(structure, schema, tenant_id):
+    for dom in structure["domains"]:
+
+        for con in dom["concepts"]:
+            neighbours_list = []
+            concept_sentence = "The " + schema + ".data_" + tenant_id + "_" + dom["label"] + "_" + con["label"] + "_w table contains information about " + dom["label"] + " " + con["label"] + "s. "
+            concept_sentence += "This refers to a class " + con["label"] + " in a domain of similar classes called " + dom["label"] + ". "
+            concept_sentence += "It includes details such as "
+            property_list = []
+
+            if len(con["properties"]) > 0:
+                for prop in con["properties"]:
+                    property_list.append(prop["label"])
+                    property_sentence = "The " + prop["label"] + " for the " + dom["label"] + " " + con["label"] + "."
+                    prop["schema_sentence"] = property_sentence
+                concept_sentence += smart_comma_formatting(property_list) + ". "
+
+            if len(con["relations"]) > 0:
+                for rel in con["relations"]:
+                    rel_dom = rel["relationTypeUri"].split("/")[-1].split("#")[0]
+                    rel_con = rel["relationTypeUri"].split("/")[-1].split("#")[1]
+                    rel_table = schema + ".data_" + tenant_id + "_" + rel_dom + "_" + rel_con + "_w"
+                    relation_sentence = "A relationship called " + rel["label"] + " connecting " + dom["label"] + " " + con["label"] + "s to " + rel_dom + " " + rel_con + "s. "
+                    relation_sentence += "A key connecting this table to the unique identifier of the " + rel_table + " table. "
+                    rel["schema_sentence"] = relation_sentence
+                    neighbours_list.append(rel_dom + " " + rel_con)
+                    #structure["domains"][idom]["concepts"][icon]["relations"][irel].pop("target", None)
+            if len(neighbours_list) > 0:
+                concept_sentence += "For context, it is connected to neighbor classes like " + smart_comma_formatting(neighbours_list)
+
+            con["schema_sentence"] = concept_sentence
+            con["schema_id_sentence"] = "The unique identifier for each " + dom["label"] + " " + con["label"] + ". "
+            con["schema_table"] = schema + ".data_" + tenant_id + "_" + dom["label"] + "_" + con["label"] + "_np"
+
+def smart_comma_formatting(items):
+    if items is None:
+        return ""
+    match len(items):
+        case 0:
+            return ""
+        case 1:
+            return items[0]
+        case 2:
+            return items[0] + " and " + items[1]
+        case _:
+            return ", ".join(items[0: -1]) + " and " + items[-1]
+        
+def label_from_url(inpt):
+    try:
+        return inpt.split("#")[1].replace("_", " ")
+    except IndexError:
+        return inpt
+    
+############################
+# Mapping
+############################
+

kobai/tenant_client.py

@@ -11,6 +11,7 @@ from langchain_core.language_models.chat_models import BaseChatModel
 from langchain_core.embeddings import Embeddings
 
 from . import spark_client, databricks_client, ai_query, tenant_api, ai_rag
+from .genie import get_genie_descriptions
 
 class TenantClient:
 
@@ -228,7 +229,8 @@ class TenantClient:
             #print(t["sql"])
             try:
                 self.spark_client._SparkClient__run_sql(t["sql"])
-            except:
+            except Exception as e:
+                print("Error creating view.", e)
                 print(t["sql"])
         print("Updated " + str(len(tables)) + " views for Genie.")
 
@@ -265,7 +267,6 @@ class TenantClient:
         """
         Use the Databricks Client to create a Genie Data Room for this tenant.
         """
-
         data_rooms = self.databricks_client._DatabricksClient__api_get("/api/2.0/data-rooms")
         room_id = "-1"
         if data_rooms:
@@ -280,9 +281,10 @@ class TenantClient:
 
         for t in self.__get_view_sql(domains=domains, concepts=concepts, not_concepts=not_concepts, enforce_map=enforce_map):
             payload["table_identifiers"].append(t["table"])
+            print(t["table"])
         response = self.databricks_client._DatabricksClient__api_patch("/api/2.0/data-rooms/" + room_id, payload)
 
-        payload = {"title":"Notes","content":"When filtering for a named entity, use a like comparison instead of equality. All tables are denormalized, so columns may have repeated rows for the same primary identifier. You should handle this by putting each table in a subquery and using the DISTINCT keyword.","instruction_type":"TEXT_INSTRUCTION"}
+        payload = {"title":"Notes","content":"When filtering for a named entity, use a like comparison instead of equality. All tables are denormalized, so columns may have repeated rows for the same primary identifier. You should handle this by putting each table in a subquery and using the DISTINCT keyword. The first column in each view is a unique identifier that should only be used for joins, and never shown to a user. Find another column to identify the subject of the table.","instruction_type":"TEXT_INSTRUCTION"}
         instructions = self.databricks_client._DatabricksClient__api_get("/api/2.0/data-rooms/" + room_id + "/instructions")
         inst_id = "-1"
 
@@ -324,11 +326,10 @@ class TenantClient:
 ########################################
 
     def __get_descriptions(self):
-        
-        params={"schema": self.schema, "tenant_id": self.model_id}
-        response = self.api_client._TenantAPI__run_get("/episteme-svcs/api/descriptions", params=params)
-        return response.json()
 
+        tenant_config = self.get_tenant_config()
+        descriptions = get_genie_descriptions(self.model_id, tenant_config, self.schema)
+        return descriptions
 
     def __get_view_sql(self, domains=None, concepts=None, not_concepts=None, enforce_map=True):
         sql_list = []
@@ -404,16 +405,20 @@ class TenantClient:
 
     def __get_questions(self):
         
-        response = self.api_client._TenantAPI__run_get("/episteme-svcs/api/questions")
+        #response = self.api_client._TenantAPI__run_get("/episteme-svcs/api/questions")
+
+        #tenant_config = self.get_tenant_config()
+        #questions = get_genie_questions(self.id, tenant_config)
 
-        return_questions = []
-        for q in response.json():
-            sql = q["sql"]
-            sql = sql[2:-2]
-            sql = sql.replace(".data_", ".genie_").replace("_Literals", "").replace("_w", "")
-            return_questions.append({"name": q["name"], "sql": sql})
+        #return_questions = []
+        #for q in questions:
+        #    sql = q["sql"]
+        #    sql = sql[2:-2]
+        #    sql = sql.replace(".data_", ".genie_").replace("_Literals", "").replace("_w", "")
+        #    return_questions.append({"name": q["name"], "sql": sql})
 
-        return return_questions
+        #return return_questions
+        return []
     
 ########################################
 # RAG Functions

{kobai_sdk-0.3.2.dist-info → kobai_sdk-0.3.3.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: kobai-sdk
-Version: 0.3.2
+Version: 0.3.3
 Summary: A package that enables interaction with a Kobai tenant.
 Author-email: Ryan Oattes <ryan@kobai.io>
 License:                                  Apache License
@@ -292,6 +292,22 @@ k.use_browser_token(access_token="KOBAI_ACESS_TOKEN_FROM_BROWSER")
 #### Authentication via on-behalf-of flow
 The sample code demonstrating authentication via the on-behalf-of flow will be provided, if requested.
 
+3. Initialize a Spark client using your current `SparkSession`, and generate semantically-rich SQL views describing this Kobai tenant:
+
+```python
+k.spark_init_session(spark)
+k.spark_generate_genie_views()
+```
+
+4. Initialize a Databricks API client using your Notebook context, and create a Genie Data Rooms environment for this Kobai tenant.
+
+```python
+notebook_context = dbutils.notebook.entry_point.getDbutils().notebook().getContext()
+sql_warehouse = '8834d98a8agffa76'
+
+k.databricks_init_notebook(notebook_context, sql_warehouse)
+k.databricks_build_genie()
+```
 
 ## AI Functionality
 The Kobai SDK enables users to ask follow-up questions based on the results of previous queries. This functionality currently supports models hosted on Databricks and Azure OpenAI. 

{kobai_sdk-0.3.2.dist-info → kobai_sdk-0.3.3.dist-info}/RECORD

@@ -3,12 +3,13 @@ kobai/ai_query.py,sha256=FnXn2pabJpXfTUcJvieVkAgMAjSTH9u5SFR9SJUJ-Lk,9556
 kobai/ai_rag.py,sha256=XUq_SnJw17P53Zk75hHJgTryGjHEAyYPwC0r2WtuNp4,14627
 kobai/databricks_client.py,sha256=fyqqMly2Qm0r1AHWsQjkYeNsDdH0G1JSgTkF9KJ55qA,2118
 kobai/demo_tenant_client.py,sha256=wlNc-bdI2wotRXo8ppUOalv4hYdBlek_WzJNARZV-AE,9293
+kobai/genie.py,sha256=-EbEYpu9xj_3zIXaPdwbNJEAmoeM7nb9qK-h1f_STtM,8061
 kobai/ms_authenticate.py,sha256=rlmhtvAaSRBlYmvIBy5epMVa4MBGBLPaMwawu1T_xDQ,2252
 kobai/spark_client.py,sha256=opM_F-4Ut5Hq5zZjWMuLvUps9sDULvyPNZHXGL8dW1k,776
 kobai/tenant_api.py,sha256=Q5yuFd9_V4lo3LWzvYEEO3LpDRWFgQD4TlRPXDTGbiE,4368
-kobai/tenant_client.py,sha256=w83NmLuOEyJjOVUuLva2vbq0zpGFzhi9LdSq1pKClA8,38613
-kobai_sdk-0.3.2.dist-info/licenses/LICENSE,sha256=QwcOLU5TJoTeUhuIXzhdCEEDDvorGiC6-3YTOl4TecE,11356
-kobai_sdk-0.3.2.dist-info/METADATA,sha256=7WIGEJBGHn2QIsYPLdbsnkwvtjqx3RJHZJR3kl0gu_M,19304
-kobai_sdk-0.3.2.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-kobai_sdk-0.3.2.dist-info/top_level.txt,sha256=ns1El3BrTTHKvoAgU1XtiSaVIudYeCXbEEUVY8HFDZ4,6
-kobai_sdk-0.3.2.dist-info/RECORD,,
+kobai/tenant_client.py,sha256=QWITygInZBZCkW3M7xNQCkNfQoGpTvKW_BykG2-E7Is,39012
+kobai_sdk-0.3.3.dist-info/licenses/LICENSE,sha256=QwcOLU5TJoTeUhuIXzhdCEEDDvorGiC6-3YTOl4TecE,11356
+kobai_sdk-0.3.3.dist-info/METADATA,sha256=bKWV-R83S7Mz4N1Sqn7cqyJN6WlYGADGD3Zd0Why-Oc,19869
+kobai_sdk-0.3.3.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+kobai_sdk-0.3.3.dist-info/top_level.txt,sha256=ns1El3BrTTHKvoAgU1XtiSaVIudYeCXbEEUVY8HFDZ4,6
+kobai_sdk-0.3.3.dist-info/RECORD,,