make chat context specific. create multi-context dir

Author: vysakh0

README.md

@@ -4,14 +4,10 @@ Data Neuron is a powerful AI-driven data framework to create and maintain AI DAT
 
 Supports SQLite, PostgreSQL, MySQL, MSSQL, CSV files(through duckdb). Works with major LLMs like Claude (default), OpenAI, LLAMA etc(through groq, nvidia, ..), OLLAMA.
 
-
 https://github.com/user-attachments/assets/2301e7cd-a895-4b9b-8a8f-2f30c3e02e16
 
-
 https://github.com/user-attachments/assets/5f6ed1f2-58cd-4e75-ae0b-745f7177a746
 
-
-
 ### The framework:
 
 <img width="621" alt="Screenshot 2024-07-25 at 11 30 35 PM" src="https://github.com/user-attachments/assets/09353e34-a0f7-4650-b477-746eaf10c354">
@@ -102,15 +98,24 @@ pip install "dataneuron[mysql]"
    dnn --init
    ```
 
-   This will create YAML files in the `context/` directory which will be your semantic layer for your data.
+   This will prompt for a context name, you can give `product_analytics` or `customer_success` or any and it will then create YAML files in the `context/<contextname>` directory which will be your semantic layer for your data.
    You will be told to select couple of tables, so that it can be auto-labelled which you can edit later.
 
 3. Or start an interactive chat session:
 
    ```
-   dnn --chat
+   dnn --chat <context_name>
    ```
 
+   eg:
+
+   ```
+   dnn --chat product_analytics
+   ```
+
+   You can chat with the semantic layer that you have created. And you will also be able to save the metric
+   to a dashboard, this will get created under `dashboards/<dashname>.yml`
+
 4. You can generate reports with image as input for your dashboards. You need to have `wkhtmltopdf` in your system.
    For mac
 

src/data_neuron/chat_cmd/main.py

@@ -34,12 +34,23 @@ def list_dashboards():
     return [f.split('.')[0] for f in os.listdir(dashboards_dir) if f.endswith('.yml')]
 
 
-def start_chat():
+def start_chat(context_name=None):
     print_header("Starting DATA neuron chat session...")
     print_warning(
-        "Please remember your messages and result wll be sent to LLM for API call.")
-    print_info("🗄️ Fetching the context from your context folder\n")
-    context = load_context()
+        "Please remember your messages and result will be sent to LLM for API call.")
+
+    if context_name is None:
+        context_name = styled_prompt(
+            "Please enter the name of the context you want to use")
+
+    context_dir = os.path.join('context', context_name)
+    if not os.path.exists(context_dir):
+        print_warning(
+            f"Context '{context_name}' does not exist. Please create it using 'dnn --init' first.")
+        return
+
+    print_info(f"🗄️ Fetching the context from {context_dir}\n")
+    context = load_context(context_name)
     print_success("Context is loaded!\n")
 
     chat_history = []
@@ -83,7 +94,7 @@ def process_with_llm(query: str, context: dict, chat_history: list):
     print_header("DATA neuron is thinking...")
     print_prompt("🤖 Sending request to LLM\n")
     print_info("Rephrasing the question suited for db")
-    changed_query = process_query(query)
+    changed_query = process_query(query, context)
     prompt = sql_query_prompt(changed_query, context)
     system_prompt = "You are a helpful assistant that generates SQL queries based on natural language questions and maintains context throughout the conversation."
 

src/data_neuron/cli.py

@@ -14,7 +14,7 @@
 # @click.option('--ask', help='The question you want to ask about the database.')
 @click.option('--db-init', type=click.Choice(['sqlite', 'mysql', 'mssql', 'postgres', 'csv', 'clickhouse']), help='Initialize database configuration.')
 # New option for chat
-@click.option('--chat', is_flag=True, help='Start an interactive chat session.')
+@click.option('--chat', help='Start an interactive chat session.', type=str)
 @click.option('--version', is_flag=True, help='Show the version of the tool')
 # New option
 @click.option('--report', is_flag=True, help='Generate a dashboard report.')
@@ -26,7 +26,7 @@ def cli(init, db_init, chat, version, report):
     elif db_init:
         init_database_config(db_init)
     elif chat:
-        start_chat()  # New function call for chat functionality
+        start_chat(chat)  # New function call for chat functionality
     elif version:
         click.echo(f"Data Neuron CLI version {VERSION}")
         return

src/data_neuron/context_init_cmd/main.py

@@ -1,4 +1,5 @@
 import os
+import click
 from .table_operations import get_table_list, choose_tables, generate_yaml_for_table
 from .yaml_generator import generate_definitions_and_relationships
 from ..utils.print import print_header, print_info, print_success, print_warning, print_prompt
@@ -7,18 +8,29 @@
 
 @handle_database_errors
 def init_context():
-    print_header("Setting up the context(semantic layer) for your database..")
+    print_header(
+        "Setting up a new context(semantic layer) for your database..")
+
+    context_name = click.prompt(
+        "Please enter a name for this context eg: customer_success or product_analytics", type=str)
+    context_dir = os.path.join('context', context_name)
+
+    if os.path.exists(context_dir):
+        overwrite = click.confirm(
+            f"Context '{context_name}' already exists. Do you want to overwrite it?", default=False)
+        if not overwrite:
+            print_info("Context initialization cancelled.")
+            return
+
     print_warning(
-        "This will create a folder context in the current directory. And override if there is an existing file")
+        f"This will create a folder {context_dir} in the current directory. Existing files will be overwritten.")
     print_prompt("You can edit it anytime...")
     print_warning("Please choose a set of 10 or lesser tables..\n")
     print_info("🗄️ Fetching tables from the database")
-
     db_type, all_tables = get_table_list()
-
     chosen_tables = choose_tables(all_tables)
 
-    os.makedirs('context/tables', exist_ok=True)
+    os.makedirs(os.path.join(context_dir, 'tables'), exist_ok=True)
 
     for table_info in chosen_tables:
         if isinstance(table_info, dict):
@@ -29,8 +41,10 @@ def init_context():
             table = table_info
 
         yaml_content = generate_yaml_for_table(schema, table)
-        with open(f'context/tables/{schema}___{table}.yaml', 'w') as f:
+
+        with open(os.path.join(context_dir, 'tables', f'{schema}__{table}.yaml'), 'w') as f:
             f.write(yaml_content)
+
         print("\n")
         print_success(f"Generated YAML for table: {schema}.{table}")
 
@@ -39,15 +53,15 @@ def init_context():
         definitions_yaml, relationships_yaml = generate_definitions_and_relationships(
             chosen_tables, db_type)
 
-        with open('context/definitions.yaml', 'w') as f:
+        with open(os.path.join(context_dir, 'definitions.yaml'), 'w') as f:
             f.write(definitions_yaml)
         print_success("Generated definitions.yaml")
 
-        with open('context/relationships.yaml', 'w') as f:
+        with open(os.path.join(context_dir, 'relationships.yaml'), 'w') as f:
             f.write(relationships_yaml)
         print_success("Generated relationships.yaml")
 
-        print_success("Initialization complete!")
+        print_success(f"Initialization complete for context: {context_name}")
     except Exception as e:
         print_warning(f"An error occurred: {str(e)}")
 

src/data_neuron/context_loader.py

@@ -5,16 +5,18 @@
 CONFIG_PATH = 'database.yaml'
 
 
-def load_context():
+def load_context(context_name):
     context = {
         'tables': {},
         'relationships': {},
         'global_definitions': {},
         'database': {}
     }
 
+    context_dir = os.path.join('context', context_name)
+
     # Load table-specific context
-    tables_path = os.path.join('context', 'tables')
+    tables_path = os.path.join(context_dir, 'tables')
     for filename in os.listdir(tables_path):
         if filename.endswith('.yaml'):
             file_path = os.path.join(tables_path, filename)
@@ -33,13 +35,13 @@ def load_context():
                             f"Warning: Invalid or empty YAML content in {filename}. Skipping this table.")
 
     # Load relationships
-    relationships_path = os.path.join('context', 'relationships.yaml')
+    relationships_path = os.path.join(context_dir, 'relationships.yaml')
     if os.path.exists(relationships_path):
         with open(relationships_path, 'r') as f:
             context['relationships'] = yaml.safe_load(f)
 
     # Load global definitions
-    definitions_path = os.path.join('context', 'definitions.yaml')
+    definitions_path = os.path.join(context_dir, 'definitions.yaml')
     if os.path.exists(definitions_path):
         with open(definitions_path, 'r') as f:
             context['global_definitions'] = yaml.safe_load(f)

src/data_neuron/query_refiner.py

@@ -8,8 +8,8 @@
 
 
 class LLMQueryRefiner:
-    def __init__(self):
-        self.context = load_context()
+    def __init__(self, context):
+        self.context = context
         self.db = DatabaseFactory.get_database()
         self.db_helper = DatabaseHelper(
             self.context['database'], self.db)
@@ -142,8 +142,8 @@ def further_refine_query(self, query: str, refined_entities: List[Dict]) -> str:
         return query
 
 
-def process_query(user_query: str) -> str:
-    refiner = LLMQueryRefiner()
+def process_query(user_query: str, context: Dict) -> str:
+    refiner = LLMQueryRefiner(context)
     refined_query, changes, refined_entities, invalid_entities = refiner.refine_query(
         user_query)