Update app.py

app.py

@@ -1,18 +1,8 @@
 #!/usr/bin/env python3
 """
-Code Flow Analyzer with Gradio Interface - Colab Compatible
-A single-file application that uses LangChain agents to analyze code structure
+Code Flow Analyzer with Gradio Interface - Hugging Face Spaces & Colab Compatible
+A single-file application that uses LangChain agents with the Gemini model to analyze code structure
 and generate Mermaid.js flowchart diagrams through a web interface.
-
-For Google Colab:
-1. Run this cell to install dependencies and start the app
-2. Enter your GROQ_API_KEY when prompted
-3. The Gradio interface will appear inline
-
-For local usage:
-1. Set your API key: export GROQ_API_KEY="your-key-here"
-2. Run: python code_flow_analyzer.py
-3. Open the provided URL in your browser
 """
 
 import ast
@@ -30,21 +20,22 @@ try:
     print("🟢 Running in Google Colab")
 except ImportError:
     IN_COLAB = False
-    print("🟡 Running locally")
+    print("🟡 Running locally or in Hugging Face Spaces")
 
 # Install dependencies if in Colab
 if IN_COLAB:
     print("📦 Installing dependencies...")
-    os.system("pip install -q gradio langchain langgraph langchain-groq")
+    os.system("pip install -q gradio langchain langgraph langchain-google-genai")
     print("✅ Dependencies installed")
 
 import gradio as gr
 from langchain.chat_models import init_chat_model
+from langchain_google_genai import ChatGoogleGenerativeAI
 from langchain.tools import tool
 from langgraph.prebuilt import create_react_agent
 from langgraph.checkpoint.memory import MemorySaver
 
-# Sample code examples
+# Sample code examples (unchanged)
 SAMPLE_PYTHON = '''def main():
     user_input = get_user_input()
     if user_input:
@@ -147,29 +138,29 @@ SAMPLE_JAVA = '''public class Calculator {
     }
 }'''
 
-# Handle API key setup
+# --- Gemini API Key Setup ---
 def setup_api_key():
-    """Setup API key for both Colab and local environments"""
-    api_key = os.getenv("GROQ_API_KEY")
+    """Setup API key for Colab, Hugging Face Spaces, and local environments"""
+    api_key = os.getenv("GOOGLE_API_KEY")
 
     if not api_key:
         if IN_COLAB:
-            print("🔑 Please enter your Groq API key:")
-            print("   Get a free key from: https://console.groq.com/")
-            api_key = getpass.getpass("GROQ_API_KEY: ")
+            print("🔑 Please enter your Google API key:")
+            print("   Get a key from: https://aistudio.google.com/app/apikey")
+            api_key = getpass.getpass("GOOGLE_API_KEY: ")
             if api_key:
-                os.environ["GROQ_API_KEY"] = api_key
+                os.environ["GOOGLE_API_KEY"] = api_key
                 print("✅ API key set successfully")
             else:
                 print("⚠️ No API key provided - agent features will be disabled")
         else:
-            print("⚠️ GROQ_API_KEY not found in environment variables")
-            print("   Get a free API key from: https://console.groq.com/")
-            print("   Set it with: export GROQ_API_KEY='your-key-here'")
+            print("⚠️ GOOGLE_API_KEY not found in environment variables")
+            print("   Set it with: export GOOGLE_API_KEY='your-key-here'")
+            print("   In Hugging Face Spaces, use the 'Secrets' tab to set the key.")
     else:
-        print("✅ Groq API key found")
+        print("✅ Google API key found")
 
-    return api_key or os.getenv("GROQ_API_KEY")
+    return api_key or os.getenv("GOOGLE_API_KEY")
 
 # Setup API key
 api_key = setup_api_key()
@@ -181,15 +172,16 @@ agent_executor = None
 
 if api_key:
     try:
-        model = init_chat_model("groq:qwen2.5-32b-instruct")
+        model = ChatGoogleGenerativeAI(model="gemini-pro", temperature=0)
+        print("✅ Gemini model initialized successfully: gemini-pro")
         memory = MemorySaver()
-        print("✅ Groq model initialized successfully")
     except Exception as e:
-        print(f"❌ Could not initialize LangChain model: {e}")
-        print("   Please check your API key and internet connection")
+        print(f"❌ Could not initialize Gemini model: {e}")
+        print("   Please check your API key and internet connection.")
         model = None
         memory = None
 
+# --- Tool Definitions (unchanged) ---
 @tool
 def analyze_code_structure(source_code: str) -> Dict[str, Any]:
     """
@@ -481,7 +473,7 @@ def analyze_code_with_agent(source_code: str, language: str = "auto") -> Tuple[s
         return "", "No code provided", [], 0, "Please enter some source code to analyze"
 
     if not agent_executor:
-        return "", "Agent not available", [], 0, "❌ LangChain agent not initialized. Please check your GROQ_API_KEY"
+        return "", "Agent not available", [], 0, "❌ LangChain agent not initialized. Please check your GOOGLE_API_KEY"
 
     try:
         # Detect language if auto
@@ -497,62 +489,46 @@ def analyze_code_with_agent(source_code: str, language: str = "auto") -> Tuple[s
             else:
                 language = "Unknown"
 
-        config = {"configurable": {"thread_id": f"session_{hash(source_code) % 10000}"}}
+        config = {
+            "configurable": {"thread_id": f"session_{hash(source_code) % 10000}"},
+            "recursion_limit": 100
+        }
 
+        # Refined prompt for better tool use
         prompt = f"""
-        You are a code analysis expert. Analyze the following {language} source code and create a visual flow diagram.
-
-        Code to analyze:
+        You are a code analysis expert. Analyze the following {language} source code.
+        Your task is to:
+        1. Use the 'analyze_code_structure' tool with the full source code provided below.
+        2. Use the 'generate_mermaid_diagram' tool with the output of the first tool.
+        3. Use the 'calculate_complexity_score' tool with the output of the first tool.
+        4. Provide a brief, human-readable summary of the analysis, including the generated Mermaid diagram, complexity score, and a list of functions found.
+        5. Present the final result in a clear, easy-to-read format.
+        
+        Source Code to Analyze:
         ```{language.lower()}
         {source_code}
         ```
-
-        Please follow these steps exactly:
-        1. First use the analyze_code_structure tool to understand the code structure
-        2. Then use the generate_mermaid_diagram tool to create a flowchart
-        3. Use the calculate_complexity_score tool to assess code complexity
-        4. Provide a brief summary of what the code does and its structure
-
-        Important guidelines:
-        - Focus on the main flow between functions and important control structures
-        - Keep the diagram clean and readable
-        - Identify all functions and their relationships
-        - Note any loops, conditionals, and function calls
-        - Provide insights about code organization and complexity
-
-        Please be thorough in your analysis and make sure to use all three tools in order.
         """
 
-        # Execute the agent
         result = agent_executor.invoke(
             {"messages": [{"role": "user", "content": prompt}]},
             config
         )
 
         if result and "messages" in result:
-            response = result["messages"][-1].content
-
+            response_content = result["messages"][-1].content
+            
             # Extract Mermaid diagram
-            mermaid_match = re.search(r'```mermaid\n(.*?)\n```', response, re.DOTALL)
-            if not mermaid_match:
-                mermaid_match = re.search(r'flowchart TD.*?(?=\n\n|\Z)', response, re.DOTALL)
-
-            mermaid_diagram = mermaid_match.group(1) if mermaid_match else mermaid_match.group(0) if mermaid_match else ""
-
-            # If no mermaid found in response, try to generate directly
-            if not mermaid_diagram:
-                print("⚠️ No Mermaid diagram found in response, generating fallback...")
-                # Fallback: analyze directly
-                analysis_result = analyze_code_structure.invoke({"source_code": source_code})
-                mermaid_diagram = generate_mermaid_diagram.invoke({"analysis_data": analysis_result})
+            mermaid_match = re.search(r'```mermaid\n(.*?)\n```', response_content, re.DOTALL)
+            mermaid_diagram = mermaid_match.group(1) if mermaid_match else ""
 
             # Extract complexity score
-            complexity_match = re.search(r'complexity.*?(\d+)', response, re.IGNORECASE)
-            complexity_score = int(complexity_match.group(1)) if complexity_match else 1
+            complexity_match = re.search(r'complexity.*?(\d+)', response_content, re.IGNORECASE)
+            complexity_score = int(complexity_match.group(1)) if complexity_match else 0
 
             # Extract functions
             functions_found = []
-            func_matches = re.findall(r'Functions found:.*?([^\n]+)', response, re.IGNORECASE)
+            func_matches = re.findall(r'Functions found:.*?([^\n]+)', response_content, re.IGNORECASE)
             if func_matches:
                 functions_found = [f.strip() for f in func_matches[0].split(',')]
             else:
@@ -561,17 +537,26 @@ def analyze_code_with_agent(source_code: str, language: str = "auto") -> Tuple[s
                 functions_found = [f["name"] for f in analysis_result.get("functions", [])]
 
             # Clean up the response for summary
-            summary = re.sub(r'```mermaid.*?```', '', response, flags=re.DOTALL)
+            summary = re.sub(r'```mermaid.*?```', '', response_content, flags=re.DOTALL)
             summary = re.sub(r'flowchart TD.*?(?=\n\n|\Z)', '', summary, flags=re.DOTALL)
             summary = summary.strip()
 
+            if not mermaid_diagram and not summary:
+                # Last resort fallback if agent fails entirely
+                analysis_result = analyze_code_structure.invoke({"source_code": source_code})
+                mermaid_diagram = generate_mermaid_diagram.invoke({"analysis_data": analysis_result})
+                complexity_score = calculate_complexity_score.invoke({"analysis_data": analysis_result})
+                functions_found = [f["name"] for f in analysis_result.get("functions", [])]
+                summary = "Agent failed to provide a detailed summary, but a fallback analysis was successful."
+
             return mermaid_diagram, summary, functions_found, complexity_score, ""
 
     except Exception as e:
         error_msg = f"❌ Analysis failed: {str(e)}"
-        print(f"Error details: {traceback.format_exc()}")  # For debugging
+        print(f"Error details: {traceback.format_exc()}")
         return "", "", [], 0, error_msg
 
+# --- Gradio Interface Setup (unchanged) ---
 def create_gradio_interface():
     """Create and configure the Gradio interface"""
 
@@ -633,7 +618,10 @@ def create_gradio_interface():
         """)
 
         # API Status
-        api_status = "🟢 Groq LangChain Agent Ready (Qwen 2.5-32B)" if agent_executor else "🔴 Agent Not Available (Check GROQ_API_KEY)"
+        model_info = ""
+        if agent_executor and model:
+            model_info = " (Gemini LLM)"
+        api_status = f"🟢 Gemini LangChain Agent Ready{model_info}" if agent_executor else "🔴 Agent Not Available (Check GOOGLE_API_KEY)"
         gr.Markdown(f"**Status:** {api_status}")
 
         with gr.Row():
@@ -728,16 +716,16 @@ def create_gradio_interface():
         )
 
         # Footer
-        environment_info = "Google Colab" if IN_COLAB else "Local Environment"
+        environment_info = "Google Colab" if IN_COLAB else "Hugging Face Spaces or Local Environment"
         gr.Markdown(f"""
         ---
         **🛠️ Running in:** {environment_info}
 
-        **📦 Dependencies:** gradio, langchain, langgraph, langchain-groq
+        **📦 Dependencies:** gradio, langchain, langgraph, langchain-google-genai
 
-        **🔧 Powered by:** LangChain Agents, Groq API (Qwen 2.5-32B), Mermaid.js, Gradio
+        **🔧 Powered by:** LangChain Agents, Google Gemini, Mermaid.js, Gradio
 
-        **🆓 Get Groq API Key:** [console.groq.com](https://console.groq.com/) (Free tier available!)
+        **🆓 Get Google API Key:** [aistudio.google.com/app/apikey](https://aistudio.google.com/app/apikey)
         """)
 
     return interface
@@ -746,37 +734,29 @@ def main():
     """Main function to run the application"""
     print("🔄 Code Flow Analyzer with Gradio")
     print("=" * 50)
-    print(f"🌐 Environment: {'Google Colab' if IN_COLAB else 'Local'}")
+    print(f"🌐 Environment: {'Google Colab' if IN_COLAB else 'Hugging Face Spaces or Local'}")
 
     if agent_executor:
         print("✅ LangChain agent ready")
     else:
         print("❌ LangChain agent not available")
         if IN_COLAB:
-            print("   💡 Restart this cell and enter your GROQ_API_KEY when prompted")
+            print("   💡 Restart this cell and enter your GOOGLE_API_KEY when prompted")
+        else:
+            print("   💡 Please set your GOOGLE_API_KEY as an environment variable or secret")
 
     print("\n🚀 Starting Gradio interface...")
 
     # Create and launch the interface
     interface = create_gradio_interface()
 
-    # Launch configuration for Colab vs local
-    if IN_COLAB:
-        # For Colab, use public sharing and specific settings
-        interface.launch(
-            share=True,  # Enable public sharing for Colab
-            debug=False,
-            height=600,
-            show_error=True
-        )
-    else:
-        # For local usage
-        interface.launch(
-            server_name="0.0.0.0",
-            server_port=7860,
-            share=False,  # Disable sharing for local use
-            debug=False
-        )
+    # Launch configuration for Colab vs local/Spaces
+    interface.launch(
+        share=True if IN_COLAB else False,
+        debug=False,
+        height=600,
+        show_error=True
+    )
 
 # Auto-run if in Colab or when script is executed directly
 if __name__ == "__main__" or IN_COLAB: