kobai-sdk 0.2.7__tar.gz → 0.2.8rc1__tar.gz

{kobai_sdk-0.2.7/kobai_sdk.egg-info → kobai_sdk-0.2.8rc1}/PKG-INFO

@@ -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.7 → kobai_sdk-0.2.8rc1}/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_sdk-0.2.8rc1/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_sdk-0.2.7 → kobai_sdk-0.2.8rc1}/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 → kobai_sdk-0.2.8rc1/kobai_sdk.egg-info}/PKG-INFO

@@ -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.7 → kobai_sdk-0.2.8rc1}/kobai_sdk.egg-info/SOURCES.txt

@@ -4,13 +4,13 @@ README.md
 pyproject.toml
 kobai/__init__.py
 kobai/ai_query.py
+kobai/ai_rag.py
 kobai/databricks_client.py
 kobai/demo_tenant_client.py
 kobai/llm_config.py
 kobai/spark_client.py
 kobai/tenant_api.py
 kobai/tenant_client.py
-kobai/test.py
 kobai_sdk.egg-info/PKG-INFO
 kobai_sdk.egg-info/SOURCES.txt
 kobai_sdk.egg-info/dependency_links.txt

{kobai_sdk-0.2.7 → kobai_sdk-0.2.8rc1}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "kobai-sdk"
-version = "0.2.7"
+version = "0.2.8rc1"
 description = "A package that enables interaction with a Kobai tenant."
 readme = "README.md"
 authors = [{ name = "Ryan Oattes", email = "ryan@kobai.io" }]

kobai_sdk-0.2.7/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)