Team Ai
Modelpublic

niishantth/codegraph

sourceHugging Faceupdated 6mo agoView on Hugging Face
0likes
query_engine.py168 linesDownload Raw Back to codegraph
1"""Query engine for error mapping and graph navigation."""2 3from __future__ import annotations4 5import re6from pathlib import Path7from typing import Any, Optional8 9from codegraph.models import Entity, EntityType, ErrorMapping, RelationType10from codegraph.graph_builder import GraphBuilder11 12 13class QueryEngine:14    """High-level query interface over the knowledge graph."""15 16    def __init__(self, builder: GraphBuilder) -> None:17        self.builder = builder18 19    # ------------------------------------------------------------------20    # Basic lookups21    # ------------------------------------------------------------------22 23    def find_entity(self, name: str) -> list[Entity]:24        return self.builder.query_by_name(name)25 26    def find_by_file(self, file_path: str) -> list[Entity]:27        return self.builder.query_by_file(file_path)28 29    def find_by_line(self, file_path: str, line: int) -> Optional[Entity]:30        return self.builder.query_by_line(file_path, line)31 32    # ------------------------------------------------------------------33    # Error mapping34    # ------------------------------------------------------------------35 36    def map_error(self, error_message: str, file_path: Optional[str] = None, line: Optional[int] = None, column: Optional[int] = None) -> ErrorMapping:37        """Map a runtime/compile-time error to the knowledge graph."""38        mapping = ErrorMapping(error_message=error_message, error_type=self._classify_error(error_message), file_path=file_path, line=line, column=column)39 40        # Try to extract file/line from the message if not provided41        extracted = self._extract_location(error_message)42        if extracted:43            mapping.file_path = mapping.file_path or extracted.get("file_path")44            mapping.line = mapping.line or extracted.get("line")45            mapping.column = mapping.column or extracted.get("column")46 47        # Match to entity48        if mapping.file_path and mapping.line:49            entity = self.find_by_line(mapping.file_path, mapping.line)50            mapping.matched_entity = entity51            if entity:52                mapping.related_entities = self._related_entities(entity)53                mapping.suggested_fix = self._suggest_fix(mapping)54        elif mapping.file_path:55            # No line – match to file entity56            file_ents = self.find_by_file(mapping.file_path)57            if file_ents:58                mapping.matched_entity = file_ents[0]59 60        return mapping61 62    def _classify_error(self, message: str) -> str:63        msg = message.lower()64        patterns = {65            "syntax_error": ["syntaxerror", "unexpected token", "invalid syntax"],66            "name_error": ["nameerror", "is not defined", "referenceerror", "cannot find name"],67            "type_error": ["typeerror", "unsupported operand", "cannot read property", "attributeerror"],68            "import_error": ["importerror", "modulenotfounderror", "cannot find module", "package not found"],69            "index_error": ["indexerror", "out of range", "arrayindexoutofbounds"],70            "key_error": ["keyerror", "map has no entry"],71            "value_error": ["valueerror", "invalid literal"],72            "runtime_error": ["runtimeerror", "panic", "exception"],73        }74        for err_type, keywords in patterns.items():75            if any(k in msg for k in keywords):76                return err_type77        return "unknown"78 79    def _extract_location(self, message: str) -> dict[str, Any]:80        """Try to parse file path and line number from common stack-trace formats."""81        # Python traceback: File "/path/to/file.py", line 4282        m = re.search(r'File "([^"]+)", line (\d+)', message)83        if m:84            return {"file_path": m.group(1), "line": int(m.group(2))}85        # Generic: path/to/file.py:42:10 or path\to\file.py:4286        m = re.search(r'([\w/\\.-]+\.(?:py|js|ts|go|java)):(\d+)(?::(\d+))?', message)87        if m:88            return {"file_path": m.group(1), "line": int(m.group(2)), "column": int(m.group(3)) if m.group(3) else None}89        return {}90 91    def _related_entities(self, entity: Entity) -> list[Entity]:92        """Return entities directly connected to the given entity."""93        related: list[Entity] = []94        g = self.builder.graph95        for neighbor in g.predecessors(entity.id):96            related.append(self.builder.entity_map[neighbor])97        for neighbor in g.successors(entity.id):98            related.append(self.builder.entity_map[neighbor])99        return related100 101    def _suggest_fix(self, mapping: ErrorMapping) -> Optional[str]:102        err_type = mapping.error_type103        entity = mapping.matched_entity104        if not entity:105            return None106        if err_type == "name_error":107            return f"Check spelling of '{entity.name}' or ensure the variable/function is defined before use."108        if err_type == "import_error":109            return f"Ensure the module imported in '{entity.file_path}' is installed and accessible."110        if err_type == "type_error":111            return f"Review type expectations around '{entity.name}' in {entity.file_path}:{entity.span.start.line if entity.span else '?'}."112        if err_type == "index_error":113            return f"Add bounds checking before indexing in '{entity.name}'."114        if err_type == "key_error":115            return f"Use .get() or check key existence before accessing dictionary in '{entity.name}'."116        if err_type == "attribute_error":117            return f"Verify the object has attribute '{entity.name}' or check for None."118        return None119 120    # ------------------------------------------------------------------121    # Graph navigation122    # ------------------------------------------------------------------123 124    def callers_of(self, entity_id: str) -> list[Entity]:125        """Return entities that call the given entity."""126        g = self.builder.graph127        return [self.builder.entity_map[n] for n in g.predecessors(entity_id) if g.edges[n, entity_id].get("type") == RelationType.CALLS.value]128 129    def callees_of(self, entity_id: str) -> list[Entity]:130        """Return entities called by the given entity."""131        g = self.builder.graph132        return [self.builder.entity_map[n] for n in g.successors(entity_id) if g.edges[entity_id, n].get("type") == RelationType.CALLS.value]133 134    def inheritance_chain(self, entity_id: str) -> list[Entity]:135        """Return inheritance ancestors for a class."""136        chain: list[Entity] = []137        current = entity_id138        g = self.builder.graph139        visited: set[str] = set()140        while current not in visited:141            visited.add(current)142            found = False143            for n in g.successors(current):144                if g.edges[current, n].get("type") == RelationType.INHERITS.value:145                    chain.append(self.builder.entity_map[n])146                    current = n147                    found = True148                    break149            if not found:150                break151        return chain152 153    def file_outline(self, file_path: str) -> list[Entity]:154        """Return top-level entities in a file ordered by line number."""155        ents = self.find_by_file(file_path)156        top = [e for e in ents if e.type != EntityType.FILE and "parent_id" not in e.meta]157        top.sort(key=lambda e: e.span.start.line if e.span else 0)158        return top159 160    def search(self, query: str) -> list[Entity]:161        """Fuzzy search over entity names and file paths."""162        q = query.lower()163        results: list[Entity] = []164        for e in self.builder.entity_map.values():165            if q in e.name.lower() or q in e.file_path.lower():166                results.append(e)167        return results168