memorygraphMCP 0.11.7__py3-none-any.whl

memorygraph/intelligence/entity_extraction.py

@@ -0,0 +1,386 @@
+"""
+Entity Extraction - Automatic entity identification and linking.
+
+This module extracts entities from memory content using regex patterns
+and optional NLP models. Supports file paths, functions, classes, errors,
+technologies, concepts, and more.
+"""
+
+import re
+import logging
+from enum import Enum
+from typing import Optional
+from datetime import datetime
+from pydantic import BaseModel, Field
+
+logger = logging.getLogger(__name__)
+
+
+class EntityType(Enum):
+    """Types of entities that can be extracted from memory content."""
+
+    FILE = "file"  # /path/to/file.py, file.txt
+    FUNCTION = "function"  # function_name(), methodName()
+    CLASS = "class"  # ClassName, ComponentName
+    ERROR = "error"  # ErrorType, Exception, error codes
+    TECHNOLOGY = "technology"  # Python, React, PostgreSQL
+    CONCEPT = "concept"  # authentication, caching, CORS
+    PERSON = "person"  # @username, developer names
+    PROJECT = "project"  # project/repo names
+    COMMAND = "command"  # CLI commands
+    PACKAGE = "package"  # npm/pip package names
+    URL = "url"  # HTTP(S) URLs
+    VARIABLE = "variable"  # variable_name, CONSTANT_NAME
+
+
+class Entity(BaseModel):
+    """Represents an extracted entity."""
+
+    text: str = Field(..., description="The extracted entity text")
+    entity_type: EntityType = Field(..., description="Type of the entity")
+    confidence: float = Field(default=1.0, ge=0.0, le=1.0, description="Extraction confidence")
+    context: Optional[str] = Field(None, description="Surrounding context")
+    start_pos: Optional[int] = Field(None, description="Start position in text")
+    end_pos: Optional[int] = Field(None, description="End position in text")
+
+
+class EntityExtractor:
+    """Extracts entities from text using regex patterns."""
+
+    # Regex patterns for different entity types
+    PATTERNS = {
+        EntityType.FILE: [
+            # Absolute paths: /path/to/file.py
+            r"(?:/[\w\-./]+)",
+            # Relative paths with extension: src/file.py
+            r"(?:[\w\-./]+\.[\w]+)",
+            # Windows paths: C:\path\to\file.py
+            r"(?:[A-Z]:\\[\w\-\\./]+)",
+        ],
+        EntityType.FUNCTION: [
+            # function_name()
+            r"\b([a-z_]\w*)\(\)",
+            # methodName()
+            r"\b([a-z]\w*[A-Z]\w*)\(\)",
+        ],
+        EntityType.CLASS: [
+            # ClassName, Handler, Service, Manager, etc.
+            r"\b([A-Z][\w]*(?:Class|Handler|Service|Manager|Controller|Provider|Factory|Builder|Strategy|Adapter|Facade|Proxy|Decorator|Observer|Singleton|Component|Module|Store|Action|Reducer|Hook|Context))\b",
+            # Generic PascalCase
+            r"\b([A-Z][a-z]+(?:[A-Z][a-z]+)+)\b",
+        ],
+        EntityType.ERROR: [
+            # *Error, *Exception
+            r"\b(\w*(?:Error|Exception))\b",
+            # HTTP status codes
+            r"\b([45]\d{2})\b",
+            # Error codes like ERR_*, E_*
+            r"\b(E(?:RR)?_[\w_]+)\b",
+        ],
+        EntityType.TECHNOLOGY: [
+            # Programming languages
+            r"\b(Python|JavaScript|TypeScript|Java|Kotlin|Swift|Go|Rust|C\+\+|C#|Ruby|PHP|Scala|Haskell|Elixir|Clojure|Erlang)\b",
+            # Frameworks
+            r"\b(React|Vue|Angular|Django|Flask|FastAPI|Express|Spring|Rails|Laravel|Symfony|Nest\.?js|Next\.?js|Nuxt\.?js|Svelte|Solid)\b",
+            # Databases
+            r"\b(PostgreSQL|MySQL|MongoDB|Redis|Neo4j|Memgraph|SQLite|DynamoDB|Cassandra|CouchDB|Elasticsearch|MariaDB|Oracle|MSSQL)\b",
+            # Cloud/Infrastructure
+            r"\b(AWS|Azure|GCP|Docker|Kubernetes|Terraform|Ansible|Jenkins|GitHub|GitLab|CircleCI|Travis)\b",
+        ],
+        EntityType.CONCEPT: [
+            # Common programming concepts
+            r"\b(authentication|authorization|caching|logging|testing|debugging|deployment|migration|refactoring|optimization|validation|serialization|deserialization|encryption|decryption|compression|decompression)\b",
+            # Architecture patterns
+            r"\b(MVC|MVVM|MVP|REST|GraphQL|gRPC|microservices|monolith|serverless|event-driven|CQRS|DDD|hexagonal|clean architecture)\b",
+            # Security concepts
+            r"\b(CORS|XSS|CSRF|SQL injection|JWT|OAuth|SAML|TLS|SSL|HTTPS|firewall|WAF)\b",
+        ],
+        EntityType.COMMAND: [
+            # Commands in backticks or quotes
+            r"`([^`]+)`",
+            r'"([^"]+)"' + r'\s*(?:command|cmd|run|exec)',
+        ],
+        EntityType.PACKAGE: [
+            # npm/pip packages
+            r"\b((?:@[\w\-]+\/)?[\w\-]+)\b(?=\s*(?:package|library|module|dependency))",
+            # Common package patterns
+            r"\b(react-\w+|vue-\w+|@types/\w+|webpack-\w+|babel-\w+|eslint-\w+|pytest-\w+)\b",
+        ],
+        EntityType.URL: [
+            # HTTP(S) URLs
+            r"https?://[\w\-./]+(?:\?[\w\-=&]*)?",
+        ],
+        EntityType.VARIABLE: [
+            # CONSTANT_NAME
+            r"\b([A-Z][A-Z0-9_]{2,})\b",
+            # snake_case
+            r"\b([a-z_]\w*[a-z]\w*)\b(?=\s*[:=])",
+        ],
+    }
+
+    def __init__(self, enable_nlp: bool = False):
+        """
+        Initialize the entity extractor.
+
+        Args:
+            enable_nlp: Enable NLP-based extraction (requires spaCy). Default: False
+        """
+        self.enable_nlp = enable_nlp
+        self.nlp_model = None
+
+        if enable_nlp:
+            try:
+                import spacy  # type: ignore
+
+                self.nlp_model = spacy.load("en_core_web_sm")
+                logger.info("NLP entity extraction enabled")
+            except (ImportError, OSError):
+                logger.warning(
+                    "spaCy not available, falling back to regex-only extraction. "
+                    "Install with: pip install spacy && python -m spacy download en_core_web_sm"
+                )
+                self.enable_nlp = False
+
+    def extract(self, text: str, min_confidence: float = 0.5) -> list[Entity]:
+        """
+        Extract entities from text.
+
+        Args:
+            text: Text to extract entities from
+            min_confidence: Minimum confidence threshold (0.0-1.0)
+
+        Returns:
+            List of extracted entities
+        """
+        entities: list[Entity] = []
+
+        # Extract using regex patterns
+        entities.extend(self._extract_with_regex(text))
+
+        # Extract using NLP if enabled
+        if self.enable_nlp and self.nlp_model:
+            entities.extend(self._extract_with_nlp(text))
+
+        # Deduplicate and filter by confidence
+        entities = self._deduplicate(entities)
+        entities = [e for e in entities if e.confidence >= min_confidence]
+
+        return entities
+
+    def _extract_with_regex(self, text: str) -> list[Entity]:
+        """Extract entities using regex patterns."""
+        entities: list[Entity] = []
+
+        for entity_type, patterns in self.PATTERNS.items():
+            for pattern in patterns:
+                for match in re.finditer(pattern, text, re.IGNORECASE):
+                    entity_text = match.group(1) if match.groups() else match.group(0)
+
+                    # Skip very short or very long matches
+                    if len(entity_text) < 2 or len(entity_text) > 100:
+                        continue
+
+                    # Calculate confidence based on pattern specificity
+                    confidence = self._calculate_confidence(entity_type, entity_text, text)
+
+                    # Extract context (50 chars before and after)
+                    start = max(0, match.start() - 50)
+                    end = min(len(text), match.end() + 50)
+                    context = text[start:end]
+
+                    entities.append(
+                        Entity(
+                            text=entity_text,
+                            entity_type=entity_type,
+                            confidence=confidence,
+                            context=context,
+                            start_pos=match.start(),
+                            end_pos=match.end(),
+                        )
+                    )
+
+        return entities
+
+    def _extract_with_nlp(self, text: str) -> list[Entity]:
+        """Extract entities using NLP (spaCy)."""
+        entities: list[Entity] = []
+
+        if not self.nlp_model:
+            return entities
+
+        doc = self.nlp_model(text)
+
+        # Map spaCy entity types to our EntityType
+        nlp_type_mapping = {
+            "PERSON": EntityType.PERSON,
+            "ORG": EntityType.PROJECT,  # Organizations often map to projects
+            "PRODUCT": EntityType.TECHNOLOGY,
+            "GPE": EntityType.CONCEPT,  # Geopolitical entities as concepts
+        }
+
+        for ent in doc.ents:
+            if ent.label_ in nlp_type_mapping:
+                entities.append(
+                    Entity(
+                        text=ent.text,
+                        entity_type=nlp_type_mapping[ent.label_],
+                        confidence=0.8,  # NLP confidence is generally high
+                        context=ent.sent.text if ent.sent else None,
+                        start_pos=ent.start_char,
+                        end_pos=ent.end_char,
+                    )
+                )
+
+        return entities
+
+    def _calculate_confidence(self, entity_type: EntityType, text: str, full_text: str) -> float:
+        """
+        Calculate extraction confidence based on entity type and context.
+
+        Args:
+            entity_type: Type of entity
+            text: Extracted entity text
+            full_text: Full text being analyzed
+
+        Returns:
+            Confidence score (0.0-1.0)
+        """
+        confidence = 0.7  # Base confidence
+
+        # Boost confidence for specific patterns
+        if entity_type == EntityType.FILE:
+            if text.endswith((".py", ".js", ".ts", ".jsx", ".tsx", ".md", ".txt", ".json", ".yaml", ".yml")):
+                confidence = 0.95
+            elif "/" in text or "\\" in text:
+                confidence = 0.85
+
+        elif entity_type == EntityType.FUNCTION:
+            if "()" in text:
+                confidence = 0.9
+
+        elif entity_type == EntityType.CLASS:
+            # Higher confidence for known suffixes
+            if any(text.endswith(suffix) for suffix in ["Handler", "Service", "Manager", "Controller"]):
+                confidence = 0.95
+            else:
+                confidence = 0.75
+
+        elif entity_type == EntityType.ERROR:
+            if text.endswith(("Error", "Exception")):
+                confidence = 0.95
+            elif re.match(r"[45]\d{2}", text):  # HTTP status codes
+                confidence = 0.9
+
+        elif entity_type == EntityType.TECHNOLOGY:
+            # Known technologies have high confidence
+            confidence = 0.95
+
+        elif entity_type == EntityType.URL:
+            confidence = 0.99
+
+        elif entity_type == EntityType.COMMAND:
+            # Commands in backticks are very reliable
+            confidence = 0.9
+
+        return min(confidence, 1.0)
+
+    def _deduplicate(self, entities: list[Entity]) -> list[Entity]:
+        """Remove duplicate entities, keeping highest confidence."""
+        seen: dict[tuple[str, EntityType], Entity] = {}
+
+        for entity in entities:
+            key = (entity.text.lower(), entity.entity_type)
+
+            if key not in seen or entity.confidence > seen[key].confidence:
+                seen[key] = entity
+
+        return list(seen.values())
+
+
+# Singleton instance for convenience
+_default_extractor = EntityExtractor()
+
+
+def extract_entities(text: str, min_confidence: float = 0.5) -> list[Entity]:
+    """
+    Extract entities from text using the default extractor.
+
+    Args:
+        text: Text to extract entities from
+        min_confidence: Minimum confidence threshold (0.0-1.0)
+
+    Returns:
+        List of extracted entities
+
+    Example:
+        >>> entities = extract_entities("Fixed authentication bug in src/auth.py")
+        >>> for entity in entities:
+        ...     print(f"{entity.entity_type.value}: {entity.text}")
+        file: src/auth.py
+        concept: authentication
+    """
+    return _default_extractor.extract(text, min_confidence)
+
+
+async def link_entities(
+    backend,
+    memory_id: str,
+    entities: list[Entity],
+) -> list[str]:
+    """
+    Link extracted entities to a memory by creating entity nodes and MENTIONS relationships.
+
+    Args:
+        backend: Database backend instance
+        memory_id: ID of the memory to link entities to
+        entities: List of entities to link
+
+    Returns:
+        List of created entity IDs
+
+    Example:
+        >>> entities = extract_entities("Fixed React hooks issue")
+        >>> entity_ids = await link_entities(backend, memory_id, entities)
+    """
+    entity_ids: list[str] = []
+
+    for entity in entities:
+        # Create or find entity node
+        query = """
+        MERGE (e:Entity {text: $text, type: $type})
+        ON CREATE SET
+            e.id = randomUUID(),
+            e.created_at = datetime(),
+            e.occurrence_count = 1
+        ON MATCH SET
+            e.occurrence_count = e.occurrence_count + 1,
+            e.last_seen = datetime()
+        WITH e
+        MATCH (m:Memory {id: $memory_id})
+        MERGE (m)-[r:MENTIONS]->(e)
+        ON CREATE SET
+            r.confidence = $confidence,
+            r.created_at = datetime()
+        RETURN e.id as entity_id
+        """
+
+        params = {
+            "text": entity.text,
+            "type": entity.entity_type.value,
+            "memory_id": memory_id,
+            "confidence": entity.confidence,
+        }
+
+        try:
+            result = await backend.execute_query(query, params)
+            if result:
+                entity_ids.append(result[0]["entity_id"])
+                logger.debug(
+                    f"Linked entity '{entity.text}' ({entity.entity_type.value}) "
+                    f"to memory {memory_id}"
+                )
+        except Exception as e:
+            logger.error(f"Failed to link entity '{entity.text}': {e}")
+            continue
+
+    return entity_ids