codebase-brain 0.1.0__py3-none-any.whl

brain_parser/codebase_walker.py

@@ -0,0 +1,101 @@
+import os
+import json
+from brain_parser.universal_parser import parse_file
+
+SKIP_FOLDERS = {
+    'lib', '.idea', '__pycache__',
+    'codebase_brain.egg-info', 'node_modules',
+    '.git', 'venv', '.env', 'tests', 'temp'
+}
+SKIP_FILES = {
+    'brain.json', 'brain_map.html'
+}
+def should_skip(path):
+    parts = path.replace('\\', '/').split('/')
+    for part in parts:
+        if part in SKIP_FOLDERS:
+            return True
+    return False
+def walk_codebase(root_path):
+    # load existing brain to check hashes
+    existing_brain = load_brain() or {}
+    brain = {}
+
+    for folder, subfolders, files in os.walk(root_path):
+        subfolders[:] = [s for s in subfolders if s not in SKIP_FOLDERS]
+
+        for file in files:
+            if file in SKIP_FILES:
+                continue
+
+            full_path = os.path.join(folder, file)
+
+            # read content to generate hash
+            try:
+                with open(full_path, 'r', encoding='utf-8', errors='ignore') as f:
+                    content = f.read()
+            except:
+                continue
+
+            from brain_parser.universal_parser import get_file_hash
+            current_hash = get_file_hash(content)
+
+            # if file exists in brain and hash matches → skip re-parsing
+            if full_path in existing_brain:
+                if existing_brain[full_path].get('hash') == current_hash:
+                    brain[full_path] = existing_brain[full_path]
+                    continue
+
+            # file is new or changed → parse and summarize
+            result = parse_file(full_path)
+            if result is not None and result['language'] != 'unknown':
+                brain[full_path] = result
+
+    return brain
+
+BRAIN_PATH = os.path.join(os.path.dirname(os.path.abspath(__file__)), '..', 'brain.json')
+
+def save_brain(brain, output_path=BRAIN_PATH):
+    with open(output_path, "w") as f:
+        json.dump(brain, f, indent=4)
+    print(f"Brain saved to {output_path}")
+
+def load_brain(input_path=BRAIN_PATH):
+    if not os.path.exists(input_path):
+        print("No brain found. Run walker first.")
+        return None
+    with open(input_path, "r") as f:
+        return json.load(f)
+
+
+def summarize_high_risk_files(brain, risk):
+    from brain_parser.universal_parser import summarize_file
+
+    summarized = 0
+    for filepath, data in brain.items():
+        file_risk = risk.get(filepath, 'LOW')
+        content = data.get('content', '')
+
+        if file_risk == 'HIGH' and content:
+            print(f"Summarizing HIGH risk file: {filepath}")
+            data['summary'] = summarize_file(filepath, content, data.get('language', ''))
+            data.pop('content', None)
+            summarized += 1
+        else:
+            # low/medium risk - generate simple summary from structure
+            functions = [f['name'] for f in data.get('functions', [])]
+            classes = [c['name'] for c in data.get('classes', [])]
+            imports = [i['name'] for i in data.get('imports', [])]
+            data[
+                'summary'] = f"File with {len(functions)} functions: {', '.join(functions[:5])}. Classes: {', '.join(classes[:3])}. Imports: {', '.join(imports[:5])}"
+            data.pop('content', None)
+
+    print(f"Summarized {summarized} HIGH risk files with Groq. Rest used structure only.")
+    return brain
+
+if __name__ == "__main__":
+    brain = walk_codebase(".")
+    save_brain(brain)
+
+    loaded = load_brain()
+    print(f"Brain loaded with {len(loaded)} files")

brain_parser/file_watcher.py

@@ -0,0 +1,81 @@
+import time
+from watchdog.observers import Observer
+from watchdog.events import FileSystemEventHandler
+from brain_parser.universal_parser import parse_file
+from brain_parser.codebase_walker import walk_codebase, save_brain, load_brain
+
+SKIP_WATCH = {'.git', '.idea', '__pycache__', 'node_modules', 'lib'}
+
+class BrainEventHandler(FileSystemEventHandler):
+    def __init__(self):
+        self.brain = load_brain() or {}
+
+
+    def should_skip(self,path):
+        # skip git, ide, temp files
+        if path.endswith('~'):
+            return True
+        for skip in SKIP_WATCH:
+            if skip in path.replace('\\', '/'):
+                return True
+        return False
+
+    def on_modified(self, event):
+        if event.is_directory:
+            return
+        if self.should_skip(event.src_path):
+            return
+        # skip brain output files
+
+        if event.src_path.endswith('brain.json') or event.src_path.endswith('brain_map.html'):
+            return
+        print(f"Modified: {event.src_path}")
+        result = parse_file(event.src_path)
+        if result and result['language'] != 'unknown':
+            self.brain[event.src_path] = result
+            save_brain(self.brain)
+
+    def on_created(self, event):
+        if self.should_skip(event.src_path):
+            return
+        if not event.is_directory:
+            print(f"New file: {event.src_path}")
+            result = parse_file(event.src_path)
+            if result and result['language'] != 'unknown':
+                self.brain[event.src_path] = result
+                save_brain(self.brain)
+
+    def on_deleted(self, event):
+        if self.should_skip(event.src_path):
+            return
+        if not event.is_directory:
+            print(f"Deleted: {event.src_path}")
+            if event.src_path in self.brain:
+                del self.brain[event.src_path]
+                save_brain(self.brain)
+
+
+def start_watching(path="."):
+    brain = load_brain() or {}
+    if not brain:
+        print("No brain found. Building first...")
+        brain = walk_codebase(path)
+        save_brain(brain)
+
+    event_handler = BrainEventHandler()
+    observer = Observer()
+    observer.schedule(event_handler, path, recursive=True)
+    observer.start()
+    print(f"Brain is watching: {path}")
+
+    try:
+        while True:
+            time.sleep(1)
+    except KeyboardInterrupt:
+        observer.stop()
+        print("Brain stopped.")
+    observer.join()
+
+
+if __name__ == "__main__":
+    start_watching(".")

brain_parser/graph_builder.py

@@ -0,0 +1,190 @@
+import networkx as nx
+import matplotlib.pyplot as plt
+from brain_parser.codebase_walker import load_brain
+from pyvis.network import Network
+
+
+def build_graph(brain):
+    G = nx.DiGraph()
+
+    for filepath, data in brain.items():
+        G.add_node(filepath)
+
+        for imp in data.get('imports', []):
+            raw = imp.get('name', '')
+
+            # extract module from "from x import y" or "import x"
+            if raw.startswith('from '):
+                module = raw.split('from ')[1].split(' import')[0].strip()
+            elif raw.startswith('import '):
+                module = raw.replace('import ', '').split(' as ')[0].strip()
+            else:
+                continue
+
+            # convert module path to file path format
+            # tests.circular_test.file_b → tests/circular_test/file_b
+            module_as_path = module.replace('.', '/').lower()
+
+            for other_file in brain.keys():
+                other_normalized = other_file.replace('\\', '/').lower()
+                # match whole filename not partial
+                filename = other_normalized.split('/')[-1].replace('.py', '').replace('.js', '').replace('.java',
+                                                                                                         '').replace(
+                    '.ts', '')
+                if module_as_path == filename or other_normalized.endswith(module_as_path + '.py'):
+                    G.add_edge(filepath, other_file)
+
+    return G
+
+
+def visualize_graph(G):
+    risk = calculate_risk(G)
+
+    color_map = {
+        "HIGH": "red",
+        "MEDIUM": "orange",
+        "LOW": "lightgreen"
+    }
+
+    node_colors = [color_map[risk[node]] for node in G.nodes()]
+
+    plt.figure(figsize=(10, 7))
+    pos = nx.spring_layout(G)
+
+    nx.draw(G, pos,
+            with_labels=True,
+            node_color=node_colors,
+            node_size=2000,
+            font_size=8,
+            arrows=True,
+            edge_color='gray',
+            font_weight='bold'
+            )
+
+    plt.title("Live Codebase Brain - Risk Map")
+    plt.savefig("brain_graph.png")
+    plt.show()
+
+
+def calculate_risk(G):
+    risk = {}
+
+    for node in G.nodes():
+        in_degree = G.in_degree(node)
+
+        if in_degree >= 2:
+            risk[node] = "HIGH"
+        elif in_degree == 1:
+            risk[node] = "MEDIUM"
+        else:
+            risk[node] = "LOW"
+
+    return risk
+
+def visualize_interactive(G):
+    risk = calculate_risk(G)\
+
+    color_map = {
+        "HIGH": "#ff4444",
+        "MEDIUM": "#ffa500",
+        "LOW": "#44cc44"
+    }
+
+    net = Network(
+        height="750px",
+        width="100%",
+        bgcolor="#1a1a2e",
+        font_color="white",
+        directed=True
+    )
+
+    for node in G.nodes():
+        risk_level = risk[node]
+        color = color_map[risk_level]
+        label = node.split("\\")[-1].split("/")[-1]
+
+        net.add_node(
+            node,
+            label=label,
+            color=color,
+            size=20 + (G.in_degree(node) * 3),
+            title=f"File: {label}\nRisk: {risk_level}\nConnections: {G.in_degree(node)}"
+        )
+
+    for edge in G.edges():
+        net.add_edge(edge[0], edge[1])
+
+    net.set_options("""
+    {
+      "nodes": {
+        "borderWidth": 2,
+        "shadow": true
+      },
+      "edges": {
+        "smooth": {
+          "type": "dynamic"
+        },
+        "shadow": true
+      },
+      "interaction": {
+        "hover": true,
+        "navigationButtons": true,
+        "hideEdgesOnDrag": true
+      },
+      "physics": {
+        "stabilization": true,
+        "barnesHut": {
+          "gravitationalConstant": -8000,
+          "springLength": 200,
+          "springConstant": 0.02
+        }
+      }
+    }
+    """)
+    net.save_graph("brain_map.html")
+    print(f"Nodes in graph: {len(G.nodes())}")
+    print(f"Edges in graph: {len(G.edges())}")
+    import webbrowser
+    webbrowser.open("brain_map.html")
+    print("Interactive graph saved to brain_map.html")
+
+
+def get_impact(G, filepath):
+    import networkx as nx
+    import os
+
+    # Normalize incoming path
+    target_norm = os.path.abspath(filepath).replace('\\', '/')
+    target_rel = filepath.replace('\\', '/').lstrip('./')
+
+    # Find matching node — handles both absolute and relative keys
+    matched = None
+
+    for node in G.nodes():
+        node_norm = os.path.abspath(node).replace('\\', '/')
+        if node_norm == target_norm:
+            matched = node
+            break
+
+    if not matched:
+        return None
+
+    direct = list(G.predecessors(matched))
+    all_affected = nx.ancestors(G, matched)
+    indirect = [f for f in all_affected if f not in direct]
+
+    return {
+        'target': matched,
+        'direct': direct,
+        'indirect': indirect,
+        'total_affected': len(direct) + len(indirect)
+    }
+
+if __name__ == "__main__":
+    brain = load_brain()
+    G = build_graph(brain)
+    visualize_interactive(G)
+
+    print("\nRisk Report:")
+    for node, level in calculate_risk(G).items():
+        print(f"  {node.split(chr(92))[-1]}: {level}")

brain_parser/llm_router.py

@@ -0,0 +1,75 @@
+import os
+from dotenv import load_dotenv
+
+load_dotenv()
+
+
+def call_llm(prompt, max_tokens=500):
+    provider = os.getenv("LLM_PROVIDER", "groq").lower()
+    api_key = os.getenv("LLM_API_KEY") or os.getenv("GROQ_API_KEY")
+    model = os.getenv("LLM_MODEL") or os.getenv("GROQ_MODEL", "llama-3.3-70b-versatile")
+
+    if provider == "groq":
+        return _call_groq(api_key, model, prompt, max_tokens)
+    elif provider == "openai":
+        return _call_openai(api_key, model, prompt, max_tokens)
+    elif provider == "anthropic":
+        return _call_anthropic(api_key, model, prompt, max_tokens)
+    elif provider == "google":
+        return _call_google(api_key, model, prompt, max_tokens)
+    elif provider == "ollama":
+        return _call_ollama(model, prompt, max_tokens)
+    else:
+        return _call_groq(api_key, model, prompt, max_tokens)
+
+
+def _call_groq(api_key, model, prompt, max_tokens):
+    from groq import Groq
+    client = Groq(api_key=api_key)
+    response = client.chat.completions.create(
+        model=model,
+        messages=[{"role": "user", "content": prompt}],
+        max_tokens=max_tokens,
+        temperature=0
+    )
+    return response.choices[0].message.content
+
+
+def _call_openai(api_key, model, prompt, max_tokens):
+    from openai import OpenAI
+    client = OpenAI(api_key=api_key)
+    response = client.chat.completions.create(
+        model=model,
+        messages=[{"role": "user", "content": prompt}],
+        max_tokens=max_tokens,
+        temperature=0
+    )
+    return response.choices[0].message.content
+
+
+def _call_anthropic(api_key, model, prompt, max_tokens):
+    import anthropic
+    client = anthropic.Anthropic(api_key=api_key)
+    response = client.messages.create(
+        model=model,
+        max_tokens=max_tokens,
+        messages=[{"role": "user", "content": prompt}]
+    )
+    return response.content[0].text
+
+
+def _call_google(api_key, model, prompt, max_tokens):
+    import google.generativeai as genai
+    genai.configure(api_key=api_key)
+    gemini = genai.GenerativeModel(model)
+    response = gemini.generate_content(prompt)
+    return response.text
+
+
+def _call_ollama(model, prompt, max_tokens):
+    import requests
+    response = requests.post(
+        "http://localhost:11434/api/generate",
+        json={"model": model, "prompt": prompt, "stream": False}
+    )
+    return response.json().get("response", "")

brain_parser/query_engine.py

@@ -0,0 +1,127 @@
+import json
+import os
+from brain_parser.codebase_walker import load_brain
+from dotenv import load_dotenv
+
+
+def get_client():
+    from brain_parser.llm_router import call_llm
+    # returns a wrapper so existing code doesn't break
+    load_dotenv()
+    model = os.getenv("LLM_MODEL") or os.getenv("GROQ_MODEL", "llama-3.3-70b-versatile")
+    return None, model
+
+
+def find_relevant_context(brain, question):
+    from brain_parser.graph_builder import build_graph, calculate_risk
+
+    # build graph and get risk scores
+    G = build_graph(brain)
+    risk = calculate_risk(G)
+
+    scores = {}
+    question_lower = question.lower()
+    question_words = set(question_lower.split())
+
+    for filepath, data in brain.items():
+        score = 0
+
+        filename = filepath.lower()
+        for word in question_words:
+            if word in filename:
+                score += 3
+
+        for func in data.get('functions', []):
+            if func['name'].lower() in question_lower:
+                score += 3
+
+        for cls in data.get('classes', []):
+            if cls['name'].lower() in question_lower:
+                score += 3
+
+        for imp in data.get('imports', []):
+            name = imp.get('name', '').lower()
+            if any(word in name for word in question_words):
+                score += 2
+
+        summary = data.get('summary', '').lower()
+        for word in question_words:
+            if word in summary:
+                score += 1
+
+        # boost score for HIGH risk files when question is about risk
+        if risk.get(filepath) == 'HIGH':
+            score += 2
+
+        scores[filepath] = score
+
+    sorted_files = sorted(scores.items(), key=lambda x: x[1], reverse=True)
+
+    # build context with risk scores included
+    top_5 = {}
+    for f, s in sorted_files[:5]:
+        if f in brain:
+            file_data = dict(brain[f])
+            file_data['risk'] = risk.get(f, 'LOW')
+            file_data['dependents'] = len(list(G.predecessors(f)))
+            file_data.pop('content', None)  # strip raw content — too many tokens
+            top_5[f] = file_data
+
+    if not top_5:
+        for f, s in sorted_files[:5]:
+            if f in brain:
+                file_data = dict(brain[f])
+                file_data['risk'] = risk.get(f, 'LOW')
+                file_data['dependents'] = len(list(G.predecessors(f)))
+                file_data.pop('content', None)
+                top_5[f] = file_data
+
+    return top_5
+
+
+def ask_brain(question):
+    from brain_parser.llm_router import call_llm
+
+    brain = load_brain()
+    if not brain:
+        return "Brain is empty. Run walker first."
+
+    context = find_relevant_context(brain, question)
+    if len(context) > 5:
+        context = dict(list(context.items())[:5])
+    context_str = json.dumps(context, indent=2)
+
+    prompt = f"""You are an intelligent codebase brain.
+You have deep knowledge of this codebase structure.
+Here is what you know:
+
+{context_str}
+
+Answer this question precisely and clearly:
+{question}"""
+
+    return call_llm(prompt, max_tokens=500)
+
+def calculate_risk(G):
+    risk = {}
+
+    for node in G.nodes():
+        in_degree = G.in_degree(node)
+
+        if in_degree >= 2:
+            risk[node] = "HIGH"
+        elif in_degree == 1:
+            risk[node] = "MEDIUM"
+        else:
+            risk[node] = "LOW"
+
+    return risk
+
+
+if __name__ == "__main__":
+    while True:
+        question = input("\nAsk your brain: ")
+        if question == "exit":
+            break
+        answer = ask_brain(question)
+        print(f"\nBrain: {answer}")