Update app.py

app.py

@@ -4,69 +4,83 @@ import base64
 import streamlit.components.v1 as components
 from transformers import pipeline
 
-# --- Load summarizer model ---
+st.set_page_config(page_title="AR/VR Code Visualizer", layout="wide")
+st.title("👓 AR/VR Code Visualizer with Call Relationships")
+
 @st.cache_resource
 def load_model():
-    summarizer = pipeline("summarization", model="philschmid/bart-large-cnn-samsum")
-    return summarizer
+    return pipeline("summarization", model="philschmid/bart-large-cnn-samsum")
 
 summarizer = load_model()
 
-st.set_page_config(page_title="AR/VR Code Visualizer", layout="wide")
-st.title("👓 AR/VR Enhanced Code Visualizer")
-
-# --- Upload Python file ---
-uploaded_file = st.file_uploader("📁 Upload your Python file", type=["py"])
+uploaded_file = st.file_uploader("📁 Upload a Python file", type=["py"])
 
-# --- Parse code and show output ---
-if uploaded_file is not None:
-    code = uploaded_file.read().decode('utf-8')
-    st.subheader("📄 Uploaded Code")
-    st.code(code, language='python')
+if uploaded_file:
+    code = uploaded_file.read().decode("utf-8")
+    st.code(code, language="python")
 
     # Parse AST
     tree = ast.parse(code)
-    functions = [node.name for node in ast.walk(tree) if isinstance(node, ast.FunctionDef)]
-    classes = [node.name for node in ast.walk(tree) if isinstance(node, ast.ClassDef)]
-
-    st.subheader("🔍 Code Structure")
-    st.write(f"**Functions:** {functions or 'None'}")
-    st.write(f"**Classes:** {classes or 'None'}")
-
-    # Hugging Face Summarization
-    st.subheader("🧠 Code Explanation (AI)")
-    prompt = f"This code contains functions: {functions} and classes: {classes}. Explain what this code does."
-    try:
-        summary = summarizer(prompt, max_length=60, min_length=15, do_sample=False)
-        st.success(summary[0]['summary_text'])
-    except Exception as e:
-        st.error(f"Summarization failed: {e}")
-
-    # Build dynamic A-Frame HTML
-    def generate_aframe_html(functions, classes):
-        elements = []
-        x_pos = -3
-
-        for fn in functions:
-            elements.append(f"""
-                <a-box position="{x_pos} 1 -3" color="#FFC65D" depth="0.5" height="0.5" width="2"></a-box>
-                <a-text value="Function: {fn}" position="{x_pos} 1.7 -3" align="center" color="black"></a-text>
-            """)
-            x_pos += 3
-
-        for cls in classes:
-            elements.append(f"""
-                <a-box position="{x_pos} 1 -3" color="#F16775" depth="0.5" height="0.5" width="2"></a-box>
-                <a-text value="Class: {cls}" position="{x_pos} 1.7 -3" align="center" color="black"></a-text>
-            """)
-            x_pos += 3
-
-        html = f"""
+
+    # Find functions and their calls
+    class FunctionCallVisitor(ast.NodeVisitor):
+        def __init__(self):
+            self.calls = {}  # {caller: [callee1, callee2]}
+
+        def visit_FunctionDef(self, node):
+            caller = node.name
+            self.calls[caller] = []
+            for child in ast.walk(node):
+                if isinstance(child, ast.Call) and isinstance(child.func, ast.Name):
+                    self.calls[caller].append(child.func.id)
+            self.generic_visit(node)
+
+    visitor = FunctionCallVisitor()
+    visitor.visit(tree)
+    call_graph = visitor.calls
+    all_functions = list(call_graph.keys())
+
+    # Show extracted relationships
+    st.subheader("📊 Call Relationships")
+    for fn, callees in call_graph.items():
+        st.write(f"🔹 `{fn}` calls: {', '.join(callees) if callees else 'None'}")
+
+    # Summarize using Hugging Face
+    prompt = f"Explain the structure and purpose of the following functions and how they call each other: {call_graph}"
+    summary = summarizer(prompt, max_length=60, min_length=15, do_sample=False)
+    st.success(summary[0]['summary_text'])
+
+    # Generate A-Frame HTML
+    def generate_aframe(call_graph):
+        spacing = 3
+        boxes = []
+        lines = []
+        positions = {}
+
+        # Assign fixed positions on x-axis
+        for i, fn in enumerate(call_graph):
+            x = i * spacing
+            positions[fn] = (x, 1, -3)
+            boxes.append(f'''
+                <a-box position="{x} 1 -3" depth="0.5" height="0.5" width="2" color="#FFC65D"></a-box>
+                <a-text value="{fn}" position="{x} 1.8 -3" align="center" color="#000"></a-text>
+            ''')
+
+        # Draw arrows between caller → callee
+        for caller, callees in call_graph.items():
+            for callee in callees:
+                if callee in positions:
+                    x1, y1, z1 = positions[caller]
+                    x2, y2, z2 = positions[callee]
+                    lines.append(f'''
+                        <a-entity line="start: {x1} {y1} {z1}; end: {x2} {y2} {z2}; color: red"></a-entity>
+                    ''')
+
+        html = f'''
         <!DOCTYPE html>
         <html>
           <head>
             <script src="https://aframe.io/releases/1.3.0/aframe.min.js"></script>
-            <meta name="viewport" content="width=device-width, initial-scale=1">
           </head>
           <body>
             <a-scene background="color: #ECECEC">
@@ -74,22 +88,21 @@ if uploaded_file is not None:
                 <a-camera wasd-controls-enabled="true" look-controls-enabled="true"></a-camera>
               </a-entity>
               <a-light type="ambient" color="#FFF"></a-light>
-              <a-plane rotation="-90 0 0" width="30" height="30" color="#7BC8A4"></a-plane>
-              {''.join(elements)}
+              <a-plane rotation="-90 0 0" width="40" height="40" color="#7BC8A4"></a-plane>
+              {''.join(boxes)}
+              {''.join(lines)}
             </a-scene>
           </body>
         </html>
-        """
+        '''
         return html
 
-    # Generate and encode A-Frame scene as base64
-    aframe_html = generate_aframe_html(functions, classes)
-    encoded_html = base64.b64encode(aframe_html.encode()).decode()
-    data_url = f"data:text/html;base64,{encoded_html}"
+    aframe_html = generate_aframe(call_graph)
+    b64 = base64.b64encode(aframe_html.encode()).decode()
+    data_url = f"data:text/html;base64,{b64}"
 
-    st.subheader("🌐 AR/VR Visualization")
-    st.info("Below is an interactive 3D visualization of your code structure (functions & classes).")
-    components.iframe(data_url, height=500, scrolling=True)
+    st.subheader("🌐 3D Call Relationship Visualization")
+    components.iframe(data_url, height=500)
 
 else:
-    st.info("Please upload a Python file to get started.")
+    st.info("Upload a Python file to visualize function call relationships.")