lumen-py 1.0.0__tar.gz

lumen_py-1.0.0/PKG-INFO

@@ -0,0 +1,15 @@
+Metadata-Version: 2.4
+Name: lumen-py
+Version: 1.0.0
+Summary: An AI-powered, local, Object-Oriented coding mentor.
+Author: Bivo
+Requires-Dist: torch
+Requires-Dist: transformers
+Requires-Dist: datasets
+Requires-Dist: typer
+Requires-Dist: rich
+Requires-Dist: pyperclip
+Requires-Dist: requests
+Dynamic: author
+Dynamic: requires-dist
+Dynamic: summary

lumen_py-1.0.0/README.md

@@ -0,0 +1,66 @@
+# Lumen-Py 🧠
+
+An AI-powered, local, Object-Oriented coding mentor designed to fight "AI Brain Rot." 
+
+Instead of writing the code for you and making you dependent, Lumen-Py uses a locally fine-tuned PyTorch Deep Learning model to analyze your architectural maturity, and a local Socratic LLM to guide you to the answers yourself.
+
+## 🌟 What It Offers (Features)
+* **System-Level Architecture Scanning:** Point Lumen at an entire project directory. It will crawl your files, map your dependencies, and critique your overall system design (e.g., pointing out tightly coupled classes or poor separation of concerns).
+* **Deep Learning Code Classifier:** Uses a locally fine-tuned Hugging Face `CodeBERT` model to mathematically classify your Python scripts as "Junior" or "Senior" architecture.
+* **Socratic Sledgehammer:** Integrates with local LLMs via Ollama to review code and ask targeted, Socratic questions instead of just printing the solution.
+* **Real-Time Token Streaming:** A seamless, matrix-style terminal UI that streams the AI's responses character-by-character, complete with live Markdown rendering.
+* **Mermaid.js Visualizer:** Automatically translates architectural concepts into flowchart code that copies straight to your clipboard.
+
+## ⚙️ Installation & Setup
+
+1. **Clone the repository:**
+```bash
+git clone https://github.com/Bivo2004/Lumen-Py.git
+cd Lumen-Py
+```
+
+2. **Install the dependencies:**
+```bash
+pip install torch transformers datasets typer rich pyperclip
+```
+
+3. **Install Ollama:**
+Ensure you have [Ollama](https://ollama.com/) installed and running locally with your preferred model (the default for this tool is `llama3`).
+
+## 🧠 Training the Brain (Required First Step)
+To keep the GitHub repository fast and lightweight, the compiled neural network weights (`lumen_brain.pth`) are ignored by Git. **You must train the PyTorch brain locally before running a code review.**
+
+Run the training pipeline. This script downloads a dataset, tokenizes it via CodeBERT, and trains the PyTorch model directly on your machine:
+```bash
+python train_deep_model.py
+```
+*Note: This will take a few moments and will generate a local `lumen_brain.pth` file in your directory. Do not delete this file.*
+
+## 🚀 How to Use Lumen-Py
+
+Lumen-Py is operated entirely through the terminal using its custom CLI.
+
+### 1. The Interactive Mentor
+Boot up the terminal chat interface for general Socratic mentoring or to brainstorm architectures.
+```bash
+python main.py start
+```
+*(Pro-tip: Inside the chat, try typing `diagram: how a REST API connects to a database` to see the Mermaid visualizer automatically build a flowchart and copy it to your clipboard!)*
+
+### 2. Single-File Socratic Review
+Point Lumen at a specific file on your machine. It will analyze the architectural maturity of that script and initiate a targeted Socratic review.
+```bash
+python main.py review path/to/your/script.py
+```
+
+### 3. System-Level Architecture Review
+Point Lumen at an entire project folder. It will scan all `.py` files, ignore your virtual environments, and critique how your files and classes interact with one another.
+```bash
+python main.py review path/to/your/project_folder/
+```
+
+## 🏗️ Architecture & Stack
+* **Language:** Python 3
+* **Deep Learning:** PyTorch, Hugging Face Transformers (`microsoft/codebert-base`), Datasets
+* **CLI & UI:** Typer, Rich
+* **LLM Integration:** Ollama

lumen_py-1.0.0/deep_mapper.py

@@ -0,0 +1,138 @@
+import torch
+import torch.nn as nn
+import torch.optim as optim
+from transformers import RobertaTokenizer, RobertaModel
+import os
+
+class CodeLevelNN(nn.Module):
+    """A PyTorch Feedforward Neural Network to classify architectural maturity."""
+    def __init__(self, input_dim: int = 768, hidden_dim: int = 128):
+        super(CodeLevelNN, self).__init__()
+        self.network = nn.Sequential(
+            nn.Linear(input_dim, hidden_dim),
+            nn.ReLU(),           
+            
+            nn.Dropout(0.2),     
+            nn.Linear(hidden_dim, 2) 
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.network(x)
+
+class DeepKnowledgeMapper:
+    """Handles the PyTorch classifier and the CodeBERT embedding pipeline."""
+    def __init__(self):
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        self.model = CodeLevelNN(input_dim=768).to(self.device)
+        self.tokenizer = RobertaTokenizer.from_pretrained("microsoft/codebert-base")
+        self.codebert = RobertaModel.from_pretrained("microsoft/codebert-base").to(self.device)
+        self.weights_path = "lumen_brain.pth"
+        
+        
+        for param in self.codebert.parameters():
+            param.requires_grad = False
+            
+        
+        for param in self.codebert.encoder.layer[-2:].parameters():
+            param.requires_grad = True
+
+        if os.path.exists(self.weights_path):
+            checkpoint = torch.load(self.weights_path, map_location=self.device, weights_only=False)
+            if isinstance(checkpoint, dict) and 'model_state_dict' in checkpoint:
+                self.model.load_state_dict(checkpoint['model_state_dict'])
+                self.codebert.load_state_dict(checkpoint['codebert_state_dict'])
+            else:
+                self.model.load_state_dict(checkpoint)
+            self._is_trained = True
+        else:
+            self._is_trained = False
+
+    def get_embedding(self, code_snippet: str) -> torch.Tensor:
+        """Translates raw Python code into a 768-number semantic vector."""
+        tokens = self.tokenizer(code_snippet, return_tensors="pt", truncation=True, max_length=512)
+        input_ids = tokens['input_ids'].to(self.device)
+        attention_mask = tokens['attention_mask'].to(self.device)
+
+        
+        outputs = self.codebert(input_ids=input_ids, attention_mask=attention_mask)
+            
+        return outputs.last_hidden_state[:, 0, :]
+
+    def train(self, code_samples: list[str], labels: list[str], epochs: int = 25, batch_size: int = 32):
+        """Trains end-to-end, updating the un-frozen CodeBERT layers and Neural Network."""
+        print(f"\nStep 1: Tokenizing {len(code_samples)} scripts...")
+        
+        label_map = {"Junior": 0, "Senior": 1}
+        
+        # Tokenize everything and keep on CPU to save VRAM
+        tokens = self.tokenizer(code_samples, padding=True, truncation=True, max_length=512, return_tensors="pt")
+        X_input_ids = tokens['input_ids']
+        X_attention_mask = tokens['attention_mask']
+        Y = torch.tensor([label_map[l] for l in labels], dtype=torch.long)
+
+        print("\nStep 2: Training the Neural Network alongside CodeBERT...")
+        
+        trainable_params = list(self.model.parameters()) + \
+                           list(filter(lambda p: p.requires_grad, self.codebert.parameters()))
+                           
+        # We drop the learning rate (e.g., 1e-4) to make sure we don't destroy CodeBERT's weights
+        optimizer = optim.Adam(trainable_params, lr=1e-4)
+        criterion = nn.CrossEntropyLoss()
+        
+        self.model.train() 
+        self.codebert.train() 
+        
+        dataset_size = len(code_samples)
+
+        for epoch in range(epochs):
+            permutation = torch.randperm(dataset_size)
+            epoch_loss = 0
+            steps = 0
+
+            # Process in batches
+            for i in range(0, dataset_size, batch_size):
+                indices = permutation[i:i+batch_size]
+                
+                # Move batches to device during loop execution
+                batch_input_ids = X_input_ids[indices].to(self.device)
+                batch_attention_mask = X_attention_mask[indices].to(self.device)
+                batch_y = Y[indices].to(self.device)
+
+                optimizer.zero_grad()
+                
+                # Forward pass through CodeBERT then MLP
+                outputs = self.codebert(input_ids=batch_input_ids, attention_mask=batch_attention_mask)
+                cls_tokens = outputs.last_hidden_state[:, 0, :]
+                logits = self.model(cls_tokens)
+                
+                loss = criterion(logits, batch_y)
+                loss.backward()
+                optimizer.step()
+                
+                epoch_loss += loss.item()
+                steps += 1
+
+            if (epoch + 1) % 5 == 0 or epoch == 0:
+                print(f"Epoch {epoch+1:02d}/{epochs} | Avg Loss: {epoch_loss/steps:.4f}")
+
+        # Save both state_dicts
+        torch.save({
+            'model_state_dict': self.model.state_dict(),
+            'codebert_state_dict': self.codebert.state_dict(),
+        }, self.weights_path)
+        self._is_trained = True
+        print(f"\n✅ Neural Network and CodeBERT weights saved to {self.weights_path}!")
+
+    def predict(self, code_snippet: str) -> str:
+        """Predicts the level of a new piece of code."""
+        if not self._is_trained:
+            return "Unknown (Model untrained)"
+            
+        self.model.eval()
+        self.codebert.eval()
+        with torch.no_grad():
+            embedding = self.get_embedding(code_snippet)
+            output = self.model(embedding)
+            prediction_idx = torch.argmax(output, dim=1).item()
+
+        return "Senior" if prediction_idx == 1 else "Junior"

lumen_py-1.0.0/llm_handler.py

@@ -0,0 +1,26 @@
+import requests
+import json
+
+class LocalAIHandler:
+    def __init__(self, model_name: str = "llama3"):
+        self.model_name = model_name
+        self.url = "http://localhost:11434/api/generate"
+
+    def generate_response(self, prompt: str) -> str:
+        """The old method: waits for the full response."""
+        payload = {"model": self.model_name, "prompt": prompt, "stream": False}
+        response = requests.post(self.url, json=payload)
+        return response.json().get("response", "")
+
+    def generate_stream(self, prompt: str):
+        """The new method: streams the response token by token."""
+        payload = {"model": self.model_name, "prompt": prompt, "stream": True}
+        try:
+            with requests.post(self.url, json=payload, stream=True) as response:
+                response.raise_for_status()
+                for line in response.iter_lines():
+                    if line:
+                        data = json.loads(line.decode('utf-8'))
+                        yield data.get("response", "")
+        except Exception as e:
+            yield f"\n[Error communicating with local AI: {e}]"

lumen_py-1.0.0/lumen_py.egg-info/PKG-INFO

@@ -0,0 +1,15 @@
+Metadata-Version: 2.4
+Name: lumen-py
+Version: 1.0.0
+Summary: An AI-powered, local, Object-Oriented coding mentor.
+Author: Bivo
+Requires-Dist: torch
+Requires-Dist: transformers
+Requires-Dist: datasets
+Requires-Dist: typer
+Requires-Dist: rich
+Requires-Dist: pyperclip
+Requires-Dist: requests
+Dynamic: author
+Dynamic: requires-dist
+Dynamic: summary

lumen_py-1.0.0/lumen_py.egg-info/SOURCES.txt

@@ -0,0 +1,14 @@
+README.md
+deep_mapper.py
+llm_handler.py
+main.py
+setup.py
+socratic_engine.py
+train_deep_model.py
+visualizer.py
+lumen_py.egg-info/PKG-INFO
+lumen_py.egg-info/SOURCES.txt
+lumen_py.egg-info/dependency_links.txt
+lumen_py.egg-info/entry_points.txt
+lumen_py.egg-info/requires.txt
+lumen_py.egg-info/top_level.txt

lumen_py-1.0.0/lumen_py.egg-info/dependency_links.txt

@@ -0,0 +1 @@
+

lumen_py-1.0.0/lumen_py.egg-info/entry_points.txt

@@ -0,0 +1,2 @@
+[console_scripts]
+lumen = main:app

lumen_py-1.0.0/lumen_py.egg-info/requires.txt

@@ -0,0 +1,7 @@
+torch
+transformers
+datasets
+typer
+rich
+pyperclip
+requests

lumen_py-1.0.0/lumen_py.egg-info/top_level.txt

@@ -0,0 +1,6 @@
+deep_mapper
+llm_handler
+main
+socratic_engine
+train_deep_model
+visualizer

lumen_py-1.0.0/main.py

@@ -0,0 +1,138 @@
+import typer
+import pyperclip
+from rich.console import Console
+from rich.markdown import Markdown
+from rich.panel import Panel
+from pathlib import Path
+from rich.live import Live
+from llm_handler import LocalAIHandler
+from socratic_engine import SocraticEngine, SocraticStrategy
+from visualizer import MermaidVisualizer
+from deep_mapper import DeepKnowledgeMapper  # <-- Bring in the brain!
+
+app = typer.Typer()
+console = Console()
+
+@app.command()
+def start(model: str = typer.Option("llama3", help="The Ollama model to use.")):
+    """
+    Start the Lumen-Py AI Mentor terminal session.
+    """
+    console.print(Panel.fit("[bold green]Welcome to Lumen-Py[/bold green]\nYour Anti-Dependency AI Mentor.", border_style="green"))
+    console.print("[dim]Type 'exit' to quit, or 'diagram: <concept>' to generate an architecture map.[/dim]\n")
+
+    with console.status("[bold cyan]Booting neural engines...", spinner="dots"):
+        handler = LocalAIHandler(model_name=model)
+        strategy = SocraticStrategy()
+        engine = SocraticEngine(ai_handler=handler, strategy=strategy)
+        visualizer = MermaidVisualizer(ai_handler=handler)
+
+    while True:
+        try:
+            user_input = console.input("\n[bold blue]You:[/bold blue] ")
+            
+            if user_input.lower() in ['exit', 'quit']:
+                console.print("[bold green]Lumen:[/bold green] Keep coding cleanly. Goodbye!")
+                break
+                
+            if "diagram" in user_input.lower():
+                concept = user_input.lower().replace("diagram", "").strip()
+                with console.status("[bold magenta]Drawing on the whiteboard...", spinner="bouncingBar"):
+                    mermaid_code = visualizer.generate_class_diagram(concept)
+                
+                pyperclip.copy(mermaid_code)
+                console.print("\n[bold magenta]Lumen (Diagram):[/bold magenta]")
+                console.print(Panel(mermaid_code, title="Mermaid.js", border_style="magenta"))
+                console.print("[bold green]✅ Raw code automatically copied to your clipboard![/bold green]")
+                console.print("[dim]Just go to https://mermaid.live and press Ctrl+V / Cmd+V[/dim]")
+                continue
+
+            console.print("\n[bold green]Lumen:[/bold green]")
+            full_response = ""
+            # The Live display updates the Markdown rendering in real-time
+            with Live(Markdown(""), console=console, refresh_per_second=15) as live:
+                for chunk in engine.ask_tutor_stream(user_input):
+                    full_response += chunk
+                    live.update(Markdown(full_response))
+
+        except KeyboardInterrupt:
+            console.print("\n[bold green]Lumen:[/bold green] Session terminated. Goodbye!")
+            break
+
+
+@app.command()
+def review(target_path: Path, model: str = typer.Option("llama3", help="The Ollama model to use.")):
+    """
+    Review a single Python file OR an entire project directory.
+    """
+    if not target_path.exists():
+        console.print(f"[bold red]Error:[/bold red] Could not find {target_path}")
+        raise typer.Exit()
+
+    # --- THE SYSTEM SCANNER ---
+    code_content = ""
+    file_count = 0
+    
+    with console.status("[bold cyan]Scanning file system...", spinner="dots"):
+        if target_path.is_file():
+            # Single file logic
+            with open(target_path, "r", encoding="utf-8") as f:
+                code_content = f"--- FILE: {target_path.name} ---\n{f.read()}\n\n"
+            file_count = 1
+            console.print(Panel.fit(f"[bold green]Lumen Code Reviewer[/bold green]\nAnalyzing File: {target_path.name}", border_style="green"))
+            
+        elif target_path.is_dir():
+            # Directory logic: Crawl the folder
+            console.print(Panel.fit(f"[bold green]Lumen System Architect[/bold green]\nAnalyzing Project: {target_path.name}/", border_style="green"))
+            
+            # Define folders to ignore
+            ignore_dirs = {".venv", "venv", "env", "__pycache__", ".git", "build", "dist"}
+            
+            for filepath in target_path.rglob("*.py"):
+                # Skip ignored directories
+                if any(ignored in filepath.parts for ignored in ignore_dirs):
+                    continue
+                    
+                with open(filepath, "r", encoding="utf-8", errors="ignore") as f:
+                    code_content += f"--- FILE: {filepath.name} ---\n{f.read()}\n\n"
+                file_count += 1
+                
+        if file_count == 0:
+            console.print("[bold red]Error:[/bold red] No valid Python files found.")
+            raise typer.Exit()
+
+    console.print(f"[dim]Successfully loaded {file_count} file(s) into memory.[/dim]\n")
+
+    # --- THE REVIEW PIPELINE ---
+    with console.status("[bold cyan]Booting neural engines and loading CodeBERT...", spinner="dots"):
+        handler = LocalAIHandler(model_name=model)
+        mapper = DeepKnowledgeMapper()
+        
+    with console.status("[bold yellow]Analyzing system architecture...", spinner="dots"):
+        
+        level = mapper.predict(code_content[:2000]) 
+        
+    console.print(f"[bold blue]Overall Architectural Maturity:[/bold blue] {level} Level\n")
+
+    strategy = SocraticStrategy()
+    engine = SocraticEngine(ai_handler=handler, strategy=strategy)
+    
+    hidden_prompt = (
+        f"I have provided the code for a Python project below. Your ML model classified the overall architecture "
+        f"as {level}-level. Review it as a strict Socratic software architect. Do NOT rewrite "
+        f"the code. Look at how the files interact. Ask me exactly ONE targeted question about my project's "
+        f"design patterns, separation of concerns, or scalability.\n\n"
+        f"PROJECT CODE:\n{code_content}"
+    )
+    
+    
+    console.print("\n[bold green]Lumen:[/bold green]")
+    full_response = ""
+    
+    # The Live display updates the Markdown rendering in real-time
+    with Live(Markdown(""), console=console, refresh_per_second=15) as live:
+        for chunk in engine.ask_tutor_stream(hidden_prompt):
+            full_response += chunk
+            live.update(Markdown(full_response))
+if __name__ == "__main__":
+    app()

lumen_py-1.0.0/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

lumen_py-1.0.0/setup.py

@@ -0,0 +1,26 @@
+from setuptools import setup
+
+setup(
+    name="lumen-py",
+    version="1.0.0",
+    description="An AI-powered, local, Object-Oriented coding mentor.",
+    author="Bivo",
+    # We list the individual files that make up your tool
+    py_modules=["main", "deep_mapper", "llm_handler", "socratic_engine", "visualizer", "train_deep_model"],
+    # These are the libraries pip will automatically install for the user
+    install_requires=[
+        "torch",
+        "transformers",
+        "datasets",
+        "typer",
+        "rich",
+        "pyperclip",
+        "requests"
+    ],
+    # This creates the global terminal command "lumen"
+    entry_points={
+        "console_scripts": [
+            "lumen=main:app", 
+        ],
+    },
+)

lumen_py-1.0.0/socratic_engine.py

@@ -0,0 +1,49 @@
+from abc import ABC, abstractmethod
+from llm_handler import LocalAIHandler
+
+# --- 1. The Strategy Interface ---
+class TutorStrategy(ABC):
+    """Abstract base class defining a tutoring strategy."""
+    
+    @abstractmethod
+    def get_system_prompt(self) -> str:
+        pass
+
+# --- 2. Concrete Strategy ---
+class SocraticStrategy(TutorStrategy):
+    """A strict strategy that refuses to write code and only asks guiding questions."""
+    
+    def get_system_prompt(self) -> str:
+        return (
+            "You are Lumen, an expert Python Architect and strict Socratic tutor. "
+            "YOUR PRIME DIRECTIVE: NEVER write code for the user. EVER. "
+            "If the user asks for code, refuse politely. Instead, identify the core "
+            "concept they are struggling with and ask ONE guiding question to help them "
+            "figure it out themselves. Keep your responses short, analytical, and encouraging."
+        )
+
+# --- 3. The Context (Engine) ---
+class SocraticEngine:
+    """
+    The main engine that uses a TutorStrategy and the LocalAIHandler 
+    to process user requests.
+    """
+    def __init__(self, ai_handler: LocalAIHandler, strategy: TutorStrategy):
+        self.ai_handler = ai_handler
+        self.strategy = strategy
+
+    def set_strategy(self, strategy: TutorStrategy):
+        """Allows swapping the tutoring style at runtime."""
+        self.strategy = strategy
+
+    def ask_tutor(self, user_input: str) -> str:
+        """Processes the user input through the current strategy."""
+        system_prompt = self.strategy.get_system_prompt()
+        return self.ai_handler.generate_response(
+            prompt=user_input, 
+            system_prompt=system_prompt
+        )
+    def ask_tutor_stream(self, prompt: str):
+        """Passes the system prompt and yields the streamed response."""
+        full_prompt = f"{self.strategy.get_system_prompt()}\n\nUser: {prompt}\nLumen:"
+        return self.ai_handler.generate_stream(full_prompt)

lumen_py-1.0.0/train_deep_model.py

@@ -0,0 +1,58 @@
+from datasets import load_dataset
+from deep_mapper import DeepKnowledgeMapper
+
+def heuristic_labeler(code_str: str) -> str:
+    """Labels the raw training data."""
+    if 'class ' in code_str or '->' in code_str or 'yield ' in code_str or '@' in code_str:
+        return "Senior"
+    return "Junior"
+
+def run_deep_training():
+    print("Downloading dataset from Hugging Face...")
+    # Pull a larger chunk so we have enough data to mine for "Senior" examples
+    dataset = load_dataset("flytech/python-codes-25k", split="train[:5000]")
+    
+    X_train = []
+    y_train = []
+    
+    junior_count = 0
+    senior_count = 0
+    max_per_class = 250 # We want exactly 250 of each (500 total)
+    
+    print("Mining and balancing data...")
+    for item in dataset:
+        code = item['output']
+        label = heuristic_labeler(code)
+        
+        # Only add to our training list if we haven't hit our cap for that class
+        if label == "Senior" and senior_count < max_per_class:
+            X_train.append(code)
+            y_train.append(label)
+            senior_count += 1
+        elif label == "Junior" and junior_count < max_per_class:
+            X_train.append(code)
+            y_train.append(label)
+            junior_count += 1
+            
+        # Stop searching once we have a perfectly balanced dataset
+        if junior_count == max_per_class and senior_count == max_per_class:
+            break
+
+    print(f"Dataset Balanced: {junior_count} Junior | {senior_count} Senior")
+    
+    mapper = DeepKnowledgeMapper()
+    
+    print("\nStarting Deep Learning Pipeline...")
+    
+    mapper.train(X_train, y_train, epochs=20)
+    
+    print("\n--- Live Inference Test ---")
+    junior_script = "x = 10\ny = 20\nprint(x + y)"
+    senior_script = "class Server:\n    def __init__(self):\n        self.active = True"
+    
+    print(f"Junior Test -> Classified as: {mapper.predict(junior_script)}")
+    print(f"Senior Test -> Classified as: {mapper.predict(senior_script)}")
+    print("---------------------------")
+
+if __name__ == "__main__":
+    run_deep_training()

lumen_py-1.0.0/visualizer.py

@@ -0,0 +1,38 @@
+from llm_handler import LocalAIHandler
+
+class MermaidVisualizer:
+    """
+    Translates architectural concepts or Python code into Mermaid.js diagram strings.
+    """
+    def __init__(self, ai_handler: LocalAIHandler):
+        self.ai_handler = ai_handler
+        self._system_prompt = (
+            "You are an expert Software Architect. Your ONLY job is to generate "
+            "valid Mermaid.js diagram code. "
+            "DO NOT include any conversational text, explanations, or Markdown formatting. "
+            "Return strictly the raw Mermaid syntax.\n\n"
+            "CRITICAL SYNTAX RULES:\n"
+            "1. For Concepts/Flows: Use `graph TD`. Define nodes: `A[\"Name\"]` and links: `A -->|label| B`.\n"
+            "2. For OOP Classes: Use `classDiagram`. Define classes: `class Car { +start() }` and links: `Car --> Engine : contains`.\n"
+            "3. NEVER mix flowchart syntax (`graph`) with OOP syntax (`classDiagram`)."
+        )
+
+    def generate_class_diagram(self, concept_or_code: str) -> str:
+        """
+        Generates a UML class diagram based on a concept or code snippet.
+        """
+        prompt = (
+            f"Generate the most appropriate Mermaid diagram (either `graph TD` for concepts or `classDiagram` for strict code) "
+            f"for the following:\n\n{concept_or_code}"
+        )
+        
+        # We pass our strict system prompt to the handler
+        raw_mermaid = self.ai_handler.generate_response(
+            prompt=prompt, 
+            system_prompt=self._system_prompt
+        )
+        
+        # Safety filter to remove accidental markdown blocks
+        clean_mermaid = raw_mermaid.replace("```mermaid", "").replace("```", "").strip()
+        
+        return clean_mermaid