configuring AST for possible better code generation

Author: Tejas Ganesh Naik

.github/workflows/sync-S3-KB.yml

@@ -1,7 +1,5 @@
 name: syncS3andKB
 on:
-  push:
-    branches: ["main"]
   workflow_dispatch:
     inputs:
       sdk_name:
@@ -41,6 +39,22 @@ jobs:
           role-to-assume: ${{ secrets.AWS_ASSUME_ROLE }}     # once merged, update trust policy of the role to point to main branch
           aws-region: us-west-2
 
+      - name: Setup Python (for tree-sitter processing)
+        if: ${{ env.sdk_name == 'python' }}
+        uses: actions/setup-python@v4
+        with:
+          python-version: '3.9'
+
+      - name: Install tree-sitter dependencies
+        if: ${{ env.sdk_name == 'python' }}
+        run: |
+          pip install tree-sitter tree-sitter-python boto3
+
+      - name: Process Python files with tree-sitter
+        if: ${{ env.sdk_name == 'python' }}
+        run: |
+          python premium-processor.py
+
       - name: Set SDK and language mapping for S3
         run: |
           if [ "$sdk_name" == "javascriptv3" ]; then
@@ -103,7 +117,7 @@ jobs:
         run: |
           for level in "basics" "feature-scenario" "complex-feature-scenario"; do
             if [ -d "./extracted_snippets/$level" ]; then
-              aws s3 sync "./extracted_snippets/$level/" "s3://$S3_LANGUAGE-premium-bucket/$level/" --delete
+              aws s3 sync "./extracted_snippets/$level/" "s3://ast-$S3_LANGUAGE-premium-bucket/$level/" --delete
               echo "Uploaded $level examples to S3"
             fi
           done

cpp/premium-ex.md

@@ -1,7 +1,7 @@
 //Add paths for premium examples here which will be consumed by workflow to update KB with quality code 
 
 ## basics:
-/example_code/s3/s3_getting_started_scenario.cpp
+/example_code/s3
 
 ## feature-scenario:
 /example_code/medical-imaging/imaging_set_and_frames_workflow

gov2/premium-ex.md

@@ -2,10 +2,11 @@
 //Don't forget to add new line at the end of this file 
 
 ## basics:
-/s3/actions/bucket_basics.go
+/s3
 
 ## feature-scenario:
-/s3/actions/bucket_basics.go
+/dynamodb/scenarios
 
 ## complex-feature-scenario:
-/dynamodb/scenarios/scenario_movie_table.go
+
+

premium-processor.py

@@ -0,0 +1,114 @@
+import os
+from pathlib import Path
+from typing import Generator, List
+from tree_sitter import Language, Parser, Tree, Node
+import tree_sitter_python
+
+def extract_paths_from_markdown(markdown_file: str, sdk_name: str) -> dict:
+    """Extract paths from premium-ex.md for given SDK"""
+    paths = {"basics": [], "feature-scenario": [], "complex-feature-scenario": []}
+    
+    with open(markdown_file, 'r') as f:
+        content = f.read()
+    
+    current_section = None
+    for line in content.split('\n'):
+        line = line.strip()
+        if line.startswith('## basics:'):
+            current_section = "basics"
+        elif line.startswith('## feature-scenario:'):
+            current_section = "feature-scenario"
+        elif line.startswith('## complex-feature-scenario:'):
+            current_section = "complex-feature-scenario"
+        elif line.startswith('##'):
+            current_section = None
+        elif current_section and line.startswith('/'):
+            paths[current_section].append(f"./{sdk_name}{line}")
+    
+    return paths
+
+def get_python_files(directory: str) -> List[str]:
+    """Get all Python files in directory recursively"""
+    python_files = []
+    for root, dirs, files in os.walk(directory):
+        for file in files:
+            if file.endswith('.py'):
+                python_files.append(os.path.join(root, file))
+    return python_files
+
+def extract_subtrees(tree: Tree) -> List[Node]:
+    """Extract terminal subtrees from AST"""
+    terminal = [
+        'function_definition', 'async_function_definition', 'class_definition',
+        'if_statement', 'while_statement', 'for_statement', 'try_statement',
+        'with_statement', 'import_statement', 'import_from_statement'
+    ]
+    
+    def extract_subtree(subtree_root):
+        queue = [subtree_root]
+        subtree_nodes = []
+        while queue:
+            current_node = queue.pop(0)
+            for child in current_node.children:
+                if str(child.type) not in ["\n"]:
+                    queue.append(child)
+                if str(child.type) in terminal:
+                    subtree_nodes.append(child)
+        return subtree_nodes
+
+    root = tree.root_node
+    all_subtrees = []
+    queue = [root]
+    while queue:
+        current_node = queue.pop(0)
+        if str(current_node.type) in terminal:
+            all_subtrees.append(current_node)
+        else:
+            subtree = extract_subtree(current_node)
+            all_subtrees.extend(subtree)
+            children = [x for x in current_node.children]
+            queue.extend(children)
+    return all_subtrees
+
+def process_python_file(file_path: str, s3_client, bucket_name: str, level: str):
+    """Process a single Python file and upload chunks to S3"""
+    try:
+        PY_LANGUAGE = Language(tree_sitter_python.language())
+        parser = Parser()
+        parser.language = PY_LANGUAGE
+        
+        code = Path(file_path).read_text()
+        tree = parser.parse(bytes(code, "utf8"))
+        
+        subtrees = extract_subtrees(tree)
+        
+        for i, subtree in enumerate(subtrees):
+            chunk_text = code[subtree.start_byte:subtree.end_byte]
+            file_key = f"{level}/{Path(file_path).stem}_chunk_{i}_{subtree.type}.py"
+            
+            s3_client.put_object(
+                Bucket=bucket_name,
+                Key=file_key,
+                Body=chunk_text,
+                ContentType='text/plain'
+            )
+            
+    except Exception as e:
+        print(f"Error processing {file_path}: {e}")
+
+def main():
+    sdk_name = os.environ.get('sdk_name', 'python')
+    markdown_file = f"./{sdk_name}/premium-ex.md"
+    paths = extract_paths_from_markdown(markdown_file, sdk_name)
+    
+    for level, directories in paths.items():
+        for directory in directories:
+            if os.path.exists(directory):
+                for root, dirs, files in os.walk(directory):
+                    for file in files:
+                        if file.endswith('.py'):
+                            py_file = os.path.join(root, file)
+                            process_python_file(py_file, level)
+
+if __name__ == "__main__":
+    main()