Merge branch 'schemabased'

Author: vysakh0

src/data_neuron/api/claude_api.py

@@ -78,7 +78,6 @@ def call_claude_vision_api_with_pagination(query: str, image_path: str, include_
     api_key = get_api_key()
     headers = get_headers(api_key)
 
-    print("cluade...vision")
     mime_type, image_data = convert_to_base64(image_path)
 
     full_response = ""

src/data_neuron/context_init_cmd/main.py

@@ -1,5 +1,4 @@
 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
@@ -9,40 +8,48 @@
 @handle_database_errors
 def init_context():
     print_header("Setting up the context(semantic layer) for your database..")
-
     print_warning(
         "This will create a folder context in the current directory. And override if there is an existing file")
-
     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)
 
-    for table in chosen_tables:
-        yaml_content = generate_yaml_for_table(table)
-        with open(f'context/tables/{table}.yaml', 'w') as f:
+    for table_info in chosen_tables:
+        if isinstance(table_info, dict):
+            schema = table_info['schema']
+            table = table_info['table']
+        else:
+            schema = 'main'  # default schema
+            table = table_info
+
+        yaml_content = generate_yaml_for_table(schema, table)
+        with open(f'context/tables/{schema}___{table}.yaml', 'w') as f:
             f.write(yaml_content)
         print("\n")
-        print_success(f"Generated YAML for table: {table}")
+        print_success(f"Generated YAML for table: {schema}.{table}")
 
     print_info("Generating definitions and relationships...")
-    definitions_yaml, relationships_yaml = generate_definitions_and_relationships(
-        chosen_tables, db_type)
+    try:
+        definitions_yaml, relationships_yaml = generate_definitions_and_relationships(
+            chosen_tables, db_type)
 
-    with open('context/definitions.yaml', 'w') as f:
-        f.write(definitions_yaml)
-    print_success("Generated definitions.yaml")
+        with open('context/definitions.yaml', 'w') as f:
+            f.write(definitions_yaml)
+        print_success("Generated definitions.yaml")
 
-    with open('context/relationships.yaml', 'w') as f:
-        f.write(relationships_yaml)
-    print_success("Generated relationships.yaml")
+        with open('context/relationships.yaml', 'w') as f:
+            f.write(relationships_yaml)
+        print_success("Generated relationships.yaml")
 
-    print_success("Initialization complete!")
+        print_success("Initialization complete!")
+    except Exception as e:
+        print_warning(f"An error occurred: {str(e)}")
 
 
 if __name__ == '__main__':

src/data_neuron/context_init_cmd/table_operations.py

@@ -1,12 +1,10 @@
-
 import click
 from ..db_operations.factory import DatabaseFactory
 
 
 def get_table_list():
     db = DatabaseFactory.get_database()
-    res = db.get_table_list()
-    return db.db_type, res
+    return db.db_type, db.get_table_list()
 
 
 def choose_tables(all_tables):
@@ -15,8 +13,8 @@ def choose_tables(all_tables):
         batch = all_tables[i:i+10]
         click.echo(
             "Choose tables (enter numbers separated by commas, or 'all' for all, 'skip' for next batch, 'done' to finish):")
-        for idx, table in enumerate(batch, start=1):
-            click.echo(f"{idx}. {table}")
+        for idx, table_info in enumerate(batch, start=1):
+            click.echo(f"{idx}. {table_info['schema']}.{table_info['table']}")
         choice = click.prompt("Your choice").lower()
         if choice == 'all':
             chosen_tables.extend(batch)
@@ -31,9 +29,9 @@ def choose_tables(all_tables):
     return chosen_tables
 
 
-def generate_yaml_for_table(table_name):
+def generate_yaml_for_table(schema, table):
     db = DatabaseFactory.get_database()
-    table_info = db.get_table_info(table_name)
+    table_info = db.get_table_info(schema, table)
 
     from ..prompts.yaml_generation_prompt import table_yaml_prompt
     from ..api.main import stream_neuron_api

src/data_neuron/context_init_cmd/yaml_generator.py

@@ -3,12 +3,22 @@
 
 
 def generate_definitions_and_relationships(tables, db_type):
-    prompt = definitions_relationships_prompt(tables, db_type)
+    # Check if tables is a list of dictionaries or a list of strings
+    if tables and isinstance(tables[0], dict):
+        table_names = [
+            f"{table['schema']}.{table['table']}" for table in tables]
+    elif tables and isinstance(tables[0], str):
+        table_names = tables
+    else:
+        raise ValueError("Unexpected table format")
+
+    prompt = definitions_relationships_prompt(table_names, db_type)
     system_prompt = "You are a helpful assistant that generates YAML content for database definitions and relationships."
 
     yaml_content = ""
     for chunk in stream_neuron_api(prompt, instruction_prompt=system_prompt):
         yaml_content += chunk
 
+    # Split the YAML content into definitions and relationships
     definitions_yaml, relationships_yaml = yaml_content.split('---')
     return definitions_yaml.strip(), relationships_yaml.strip()

src/data_neuron/context_loader.py

@@ -17,17 +17,32 @@ def load_context():
     tables_path = os.path.join('context', 'tables')
     for filename in os.listdir(tables_path):
         if filename.endswith('.yaml'):
-            with open(os.path.join(tables_path, filename), 'r') as f:
-                table_data = yaml.safe_load(f)
-                context['tables'][table_data['table_name']] = table_data
+            file_path = os.path.join(tables_path, filename)
+            if os.path.getsize(file_path) > 0:  # Check if file is not empty
+                with open(file_path, 'r') as f:
+                    table_data = yaml.safe_load(f)
+                    if table_data and isinstance(table_data, dict):
+                        full_name = table_data.get('full_name')
+                        if full_name:
+                            context['tables'][full_name] = table_data
+                        else:
+                            print(
+                                f"Warning: 'full_name' not found in {filename}. Skipping this table.")
+                    else:
+                        print(
+                            f"Warning: Invalid or empty YAML content in {filename}. Skipping this table.")
 
     # Load relationships
-    with open(os.path.join('context', 'relationships.yaml'), 'r') as f:
-        context['relationships'] = yaml.safe_load(f)
+    relationships_path = os.path.join('context', 'relationships.yaml')
+    if os.path.exists(relationships_path):
+        with open(relationships_path, 'r') as f:
+            context['relationships'] = yaml.safe_load(f)
 
     # Load global definitions
-    with open(os.path.join('context', 'definitions.yaml'), 'r') as f:
-        context['global_definitions'] = yaml.safe_load(f)
+    definitions_path = os.path.join('context', 'definitions.yaml')
+    if os.path.exists(definitions_path):
+        with open(definitions_path, 'r') as f:
+            context['global_definitions'] = yaml.safe_load(f)
 
     # Load database configuration
     if not os.path.exists(CONFIG_PATH):
@@ -41,7 +56,7 @@ def load_context():
         if not db_config:
             raise ConfigurationError(
                 "No database configuration found in the YAML file.")
-        # careful this loads into context,
+        # IMPORTANT not to send db_config itself
         context['database'] = db_config.get('name')
     except yaml.YAMLError as e:
         raise ConfigurationError(f"Error parsing YAML configuration: {str(e)}")

src/data_neuron/db_operations/base.py

@@ -11,18 +11,13 @@ def get_table_list(self):
         pass
 
     @abstractmethod
-    def get_table_info(self, table_name):
+    def get_table_info(self, schema, table_name):
         pass
 
     @abstractmethod
     def execute_query(self, query: str) -> str:
         pass
 
-    @abstractmethod
-    def get_schema_info(self) -> str:
-        pass
-
-    @abstractmethod
     def execute_query_with_column_names(self, query: str) -> Tuple[List[Tuple], List[str]]:
         pass
 

src/data_neuron/db_operations/clickhouse.py

@@ -1,6 +1,6 @@
 from .base import DatabaseOperations
 from .exceptions import ConnectionError, OperationError
-from typing import List, Tuple
+from typing import List, Tuple, Dict, Any
 
 
 class ClickHouseOperations(DatabaseOperations):
@@ -27,26 +27,29 @@ def _get_connection(self):
             raise ConnectionError(
                 f"Failed to connect to ClickHouse database: {str(e)}")
 
-    def get_table_list(self) -> List[str]:
+    def get_table_list(self) -> List[Dict[str, str]]:
         try:
             client = self._get_connection()
             result = client.query("SHOW TABLES")
-            return [row[0] for row in result.result_rows]
+            return [{"schema": self.conn_params.get('database', 'default'), "table": table[0]} for table in result.result_rows]
         except Exception as e:
             raise OperationError(f"Failed to get table list: {str(e)}")
 
-    def get_table_info(self, table_name: str) -> dict:
+    def get_table_info(self, schema: str, table: str) -> Dict[str, Any]:
         try:
             client = self._get_connection()
-            result = client.query(f"DESCRIBE {table_name}")
+            result = client.query(f"DESCRIBE {schema}.{table}")
             return {
-                'table_name': table_name,
+                'schema': schema,
+                'table_name': table,
+                'full_name': f"{schema}.{table}",
                 'columns': [
                     {
                         'name': col[0],
                         'type': col[1],
                         'nullable': col[5] == 'YES',
-                        'primary_key': col[6] == 'true'
+                        # ClickHouse uses 'key' to indicate primary key columns
+                        'primary_key': col[2] == 'key'
                     } for col in result.result_rows
                 ]
             }
@@ -61,24 +64,10 @@ def execute_query_with_column_names(self, query: str) -> Tuple[List[Tuple], List
         except Exception as e:
             raise OperationError(f"Failed to execute query: {str(e)}")
 
-    def execute_query(self, query: str) -> str:
+    def execute_query(self, query: str) -> List[Tuple]:
         try:
             client = self._get_connection()
             result = client.query(query)
             return result.result_rows
         except Exception as e:
             raise OperationError(f"Failed to execute query: {str(e)}")
-
-    def get_schema_info(self) -> str:
-        try:
-            client = self._get_connection()
-            tables = self.get_table_list()
-            schema_info = []
-            for table in tables:
-                schema_info.append(f"\nTable: {table}")
-                columns = client.query(f"DESCRIBE {table}")
-                for column in columns.result_rows:
-                    schema_info.append(f"  {column[0]} ({column[1]})")
-            return "\n".join(schema_info)
-        except Exception as e:
-            raise OperationError(f"Failed to get schema info: {str(e)}")

src/data_neuron/db_operations/database_helpers.py

@@ -14,19 +14,6 @@ def quote_identifier(self, identifier: str) -> str:
         else:  # postgres and sqlite
             return f'"{identifier}"'
 
-    def get_fully_qualified_table_name(self, schema_name: str, table_name: str) -> str:
-        quoted_table = self.quote_identifier(table_name)
-        if self.database == 'postgres':
-            if schema_name and schema_name != 'public':
-                return f'{self.quote_identifier(schema_name)}.{quoted_table}'
-        elif self.database == 'mssql':
-            if schema_name and schema_name != 'dbo':
-                return f'{self.quote_identifier(schema_name)}.{quoted_table}'
-        elif self.database == 'mysql':
-            if schema_name:
-                return f'{self.quote_identifier(schema_name)}.{quoted_table}'
-        return quoted_table
-
     def get_pattern_match_clause(self, column: str, value: str) -> str:
         if self.database == 'mssql':
             return f"LOWER({column}) LIKE '%' + LOWER('{value}') + '%'"
@@ -37,10 +24,7 @@ def execute_query(self, query: str) -> List[tuple]:
         return self.db.execute_query(query)
 
 
-def top_few_records(db_helper: DatabaseHelper, column_name: str, schema_name: str, table_name: str, potential_value: str = None, limit: int = 10) -> str:
-    quoted_table = db_helper.get_fully_qualified_table_name(
-        schema_name, table_name)
-
+def top_few_records(db_helper: DatabaseHelper, column_name: str, table_name: str, potential_value: str = None, limit: int = 10) -> str:
     # Handle the asterisk separately
     if column_name == '*':
         select_clause = '*'
@@ -55,20 +39,20 @@ def top_few_records(db_helper: DatabaseHelper, column_name: str, schema_name: st
         if column_name == '*':
             # When column is '*', we can't use it in the WHERE clause
             # Instead, we'll search across all columns
-            where_clause = f"WHERE EXISTS (SELECT 1 FROM {quoted_table} FOR JSON PATH) WHERE JSON_VALUE(BulkColumn, '$.*') LIKE '%{potential_value}%'"
+            where_clause = f"WHERE EXISTS (SELECT 1 FROM {table_name} FOR JSON PATH) WHERE JSON_VALUE(BulkColumn, '$.*') LIKE '%{potential_value}%'"
         else:
             where_clause = f"WHERE {db_helper.get_pattern_match_clause(quoted_column, potential_value)}"
 
     if db_helper.database == 'mssql':
         query = f"""
         SELECT {distinct_clause} TOP {limit}  {select_clause}
-        FROM {quoted_table}
+        FROM {table_name}
         {where_clause}
         """
     else:  # postgres, mysql, and sqlite
         query = f"""
         SELECT {distinct_clause}{select_clause}
-        FROM {quoted_table}
+        FROM {table_name}
         {where_clause}
         LIMIT {limit}
         """