code-graph-builder 0.2.0__py3-none-any.whl

code_graph_builder/tools/semantic_search.py

@@ -0,0 +1,387 @@
+"""Semantic search tools for code using embeddings.
+
+This module provides tools for semantic code search using vector embeddings.
+Integrates with both Memgraph and Kuzu backends for graph data retrieval.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from pathlib import Path
+from typing import TYPE_CHECKING, Protocol, runtime_checkable
+
+from loguru import logger
+
+from ..embeddings.qwen3_embedder import BaseEmbedder
+from ..embeddings.vector_store import SearchResult, VectorStore
+
+if TYPE_CHECKING:
+    from ..services import IngestorProtocol, QueryProtocol
+    from ..types import ResultRow
+
+
+@dataclass
+class SemanticSearchResult:
+    """Result from semantic code search.
+
+    Attributes:
+        node_id: Node identifier in graph database
+        qualified_name: Fully qualified name of code entity
+        name: Simple name of the entity
+        type: Entity type (Function, Class, Method, etc.)
+        score: Similarity score (0-1)
+        source_code: Source code if available
+        file_path: File path if available
+        start_line: Start line number
+        end_line: End line number
+    """
+
+    node_id: int
+    qualified_name: str
+    name: str
+    type: str
+    score: float
+    source_code: str | None = None
+    file_path: str | None = None
+    start_line: int | None = None
+    end_line: int | None = None
+
+
+@runtime_checkable
+class GraphServiceProtocol(Protocol):
+    """Protocol for graph service operations needed by semantic search."""
+
+    def fetch_all(self, query: str, params: dict | None = None) -> list[ResultRow]: ...
+
+    def query(self, cypher: str, params: dict | None = None) -> list[ResultRow]: ...
+
+
+class SemanticSearchService:
+    """Service for semantic code search.
+
+    Combines vector similarity search with graph database queries
+    to provide rich semantic search capabilities.
+
+    Example:
+        >>> from code_graph_builder.embeddings import create_embedder, create_vector_store
+        >>> from code_graph_builder.services import MemgraphIngestor
+        >>> from code_graph_builder.tools.semantic_search import SemanticSearchService
+        >>>
+        >>> embedder = create_embedder()
+        >>> vector_store = create_vector_store(backend="memory", dimension=1536)
+        >>>
+        >>> with MemgraphIngestor("localhost", 7687) as ingestor:
+        ...     service = SemanticSearchService(
+        ...         embedder=embedder,
+        ...         vector_store=vector_store,
+        ...         graph_service=ingestor
+        ...     )
+        ...     results = service.search("recursive fibonacci implementation", top_k=5)
+    """
+
+    def __init__(
+        self,
+        embedder: BaseEmbedder,
+        vector_store: VectorStore,
+        graph_service: GraphServiceProtocol | None = None,
+    ):
+        """Initialize semantic search service.
+
+        Args:
+            embedder: Embedder for generating query embeddings
+            vector_store: Vector store for similarity search
+            graph_service: Optional graph service for retrieving full node data
+        """
+        self.embedder = embedder
+        self.vector_store = vector_store
+        self.graph_service = graph_service
+
+    def search(
+        self,
+        query: str,
+        top_k: int = 5,
+        entity_types: list[str] | None = None,
+    ) -> list[SemanticSearchResult]:
+        """Search for code semantically similar to the query.
+
+        Args:
+            query: Natural language query describing what to find
+            top_k: Number of results to return
+            entity_types: Optional filter for entity types (Function, Class, etc.)
+
+        Returns:
+            List of semantic search results
+        """
+        try:
+            # Generate query embedding
+            query_embedding = self.embedder.embed_query(query)
+
+            # Search vector store
+            filter_metadata = None
+            if entity_types:
+                filter_metadata = {"type": entity_types[0]} if len(entity_types) == 1 else None
+
+            vector_results = self.vector_store.search_similar(
+                query_embedding, top_k=top_k, filter_metadata=filter_metadata
+            )
+
+            if not vector_results:
+                return []
+
+            # Enrich with graph data if available
+            if self.graph_service:
+                return self._enrich_results_from_graph(vector_results)
+            else:
+                return self._convert_results(vector_results)
+
+        except Exception as e:
+            logger.error(f"Semantic search failed: {e}")
+            return []
+
+    def _convert_results(self, vector_results: list[SearchResult]) -> list[SemanticSearchResult]:
+        """Convert vector search results to semantic search results."""
+        results: list[SemanticSearchResult] = []
+
+        for vr in vector_results:
+            # Extract name from qualified name
+            name = vr.qualified_name.split(".")[-1] if "." in vr.qualified_name else vr.qualified_name
+
+            results.append(
+                SemanticSearchResult(
+                    node_id=vr.node_id,
+                    qualified_name=vr.qualified_name,
+                    name=name,
+                    type="Unknown",
+                    score=vr.score,
+                )
+            )
+
+        return results
+
+    def _enrich_results_from_graph(
+        self, vector_results: list[SearchResult]
+    ) -> list[SemanticSearchResult]:
+        """Enrich vector search results with data from graph database."""
+        if not self.graph_service:
+            return self._convert_results(vector_results)
+
+        qnames = [vr.qualified_name for vr in vector_results]
+        query = self._build_nodes_query(qnames)
+
+        try:
+            graph_results = self.graph_service.fetch_all(query, {"qnames": qnames})
+            graph_data_map = {
+                row.get("qualified_name", ""): row for row in graph_results
+            }
+        except Exception as e:
+            logger.warning(f"Failed to enrich results from graph: {e}")
+            return self._convert_results(vector_results)
+
+        results: list[SemanticSearchResult] = []
+        for vr in vector_results:
+            graph_data = graph_data_map.get(vr.qualified_name, {})
+            name = graph_data.get("name") or (
+                vr.qualified_name.split(".")[-1]
+                if "." in vr.qualified_name
+                else vr.qualified_name
+            )
+            results.append(
+                SemanticSearchResult(
+                    node_id=vr.node_id,
+                    qualified_name=vr.qualified_name,
+                    name=name,
+                    type=graph_data.get("type", "Unknown"),
+                    score=vr.score,
+                    source_code=graph_data.get("source_code"),
+                    file_path=graph_data.get("path") or None,
+                    start_line=graph_data.get("start_line"),
+                    end_line=graph_data.get("end_line"),
+                )
+            )
+        return results
+
+    def _build_nodes_query(self, qualified_names: list[str]) -> str:
+        """Build Cypher query to fetch node details by qualified names."""
+        return """
+            MATCH (m:Module)-[:DEFINES]->(f:Function)
+            WHERE f.qualified_name IN $qnames
+            RETURN DISTINCT f.qualified_name AS qualified_name,
+                   f.name AS name,
+                   m.path AS path,
+                   f.start_line AS start_line,
+                   f.end_line AS end_line
+        """
+
+    def get_source_code(self, node_id: int) -> str | None:
+        """Get source code for a specific node by ID.
+
+        Args:
+            node_id: Node identifier
+
+        Returns:
+            Source code string or None if not found
+        """
+        if not self.graph_service:
+            return None
+
+        query = """
+            MATCH (n)
+            WHERE n.node_id = $node_id
+               OR n.id = $node_id
+               OR id(n) = $node_id
+            RETURN n.source_code AS source_code,
+                   n.path AS path,
+                   n.start_line AS start_line,
+                   n.end_line AS end_line
+        """
+
+        try:
+            results = self.graph_service.fetch_all(query, {"node_id": node_id})
+            if results:
+                return str(results[0].get("source_code", "")) or None
+        except Exception as e:
+            logger.warning(f"Failed to get source code for node {node_id}: {e}")
+
+        return None
+
+    def get_source_from_file(
+        self,
+        file_path: str,
+        start_line: int,
+        end_line: int,
+        repo_path: Path | None = None,
+    ) -> str | None:
+        """Extract source code from file.
+
+        Args:
+            file_path: Path to the file
+            start_line: Start line (1-indexed)
+            end_line: End line (inclusive)
+            repo_path: Repository root path for resolving relative paths
+
+        Returns:
+            Source code string or None if extraction fails
+        """
+        try:
+            path = Path(file_path)
+            if repo_path and not path.is_absolute():
+                path = repo_path / path
+
+            if not path.exists():
+                return None
+
+            with open(path, "r", encoding="utf-8", errors="ignore") as f:
+                lines = f.readlines()
+
+            # Adjust for 1-indexed lines
+            start_idx = max(0, start_line - 1)
+            end_idx = min(len(lines), end_line)
+
+            return "".join(lines[start_idx:end_idx])
+
+        except Exception as e:
+            logger.debug(f"Failed to extract source from {file_path}: {e}")
+            return None
+
+
+def create_semantic_search_service(
+    embedder: BaseEmbedder,
+    vector_store: VectorStore,
+    graph_service: GraphServiceProtocol | None = None,
+) -> SemanticSearchService:
+    """Factory function to create semantic search service.
+
+    Args:
+        embedder: Embedder instance
+        vector_store: Vector store instance
+        graph_service: Optional graph service for data enrichment
+
+    Returns:
+        Configured SemanticSearchService
+    """
+    return SemanticSearchService(
+        embedder=embedder,
+        vector_store=vector_store,
+        graph_service=graph_service,
+    )
+
+
+# Convenience functions for direct use
+
+
+def semantic_code_search(
+    query: str,
+    embedder: BaseEmbedder,
+    vector_store: VectorStore,
+    graph_service: GraphServiceProtocol | None = None,
+    top_k: int = 5,
+) -> list[SemanticSearchResult]:
+    """Perform semantic code search.
+
+    Convenience function for one-off searches without creating a service.
+
+    Args:
+        query: Natural language query
+        embedder: Embedder for query encoding
+        vector_store: Vector store to search
+        graph_service: Optional graph service for enrichment
+        top_k: Number of results
+
+    Returns:
+        List of search results
+    """
+    service = SemanticSearchService(
+        embedder=embedder,
+        vector_store=vector_store,
+        graph_service=graph_service,
+    )
+    return service.search(query, top_k=top_k)
+
+
+def get_function_source_by_node_id(
+    node_id: int,
+    graph_service: GraphServiceProtocol,
+    repo_path: Path | None = None,
+) -> str | None:
+    """Get function source code by node ID.
+
+    Args:
+        node_id: Node identifier
+        graph_service: Graph service to query
+        repo_path: Repository path for file resolution
+
+    Returns:
+        Source code or None
+    """
+    service = SemanticSearchService(
+        embedder=None,  # Not needed for this operation
+        vector_store=None,
+        graph_service=graph_service,
+    )
+
+    # Try to get from graph first
+    source = service.get_source_code(node_id)
+    if source:
+        return source
+
+    # Try to get from file
+    query = """
+        MATCH (n)
+        WHERE n.node_id = $node_id OR n.id = $node_id OR id(n) = $node_id
+        RETURN n.path AS path, n.start_line AS start_line, n.end_line AS end_line
+    """
+
+    try:
+        results = graph_service.fetch_all(query, {"node_id": node_id})
+        if results and repo_path:
+            row = results[0]
+            return service.get_source_from_file(
+                str(row.get("path", "")),
+                int(row.get("start_line", 0)) if row.get("start_line") else 0,
+                int(row.get("end_line", 0)) if row.get("end_line") else 0,
+                repo_path,
+            )
+    except Exception as e:
+        logger.warning(f"Failed to get source for node {node_id}: {e}")
+
+    return None

code_graph_builder/types.py

@@ -0,0 +1,333 @@
+"""Code Graph Builder - Type Definitions."""
+
+from __future__ import annotations
+
+from collections.abc import ItemsView, KeysView, Sequence
+from dataclasses import dataclass, field
+from enum import StrEnum
+from pathlib import Path
+from typing import TYPE_CHECKING, Callable, NamedTuple, Protocol, TypedDict
+
+from .constants import NodeLabel, RelationshipType, SupportedLanguage
+
+if TYPE_CHECKING:
+    from tree_sitter import Language, Node, Parser, Query
+
+    from .models import LanguageSpec
+
+
+# Basic type aliases
+PropertyValue = str | int | float | bool | list[str] | None
+PropertyDict = dict[str, PropertyValue]
+
+ResultScalar = str | int | float | bool | None
+ResultValue = ResultScalar | list[ResultScalar] | dict[str, ResultScalar]
+ResultRow = dict[str, ResultValue]
+
+
+# Node and relationship types
+class NodeType(StrEnum):
+    FUNCTION = "Function"
+    METHOD = "Method"
+    CLASS = "Class"
+    MODULE = "Module"
+    INTERFACE = "Interface"
+    PACKAGE = "Package"
+    ENUM = "Enum"
+    TYPE = "Type"
+    UNION = "Union"
+
+
+# Type aliases for function registry
+SimpleName = str
+QualifiedName = str
+SimpleNameLookup = dict[SimpleName, set[QualifiedName]]
+TrieNode = dict[str, "TrieNode | QualifiedName | NodeType"]
+FunctionRegistry = dict[QualifiedName, NodeType]
+
+
+# AST types (use string literal for forward reference since Node is TYPE_CHECKING only)
+ASTNode = "Node"
+
+
+# Graph data types
+class GraphMetadata(TypedDict):
+    total_nodes: int
+    total_relationships: int
+    exported_at: str
+
+
+class NodeData(TypedDict):
+    node_id: int
+    labels: list[str]
+    properties: dict[str, PropertyValue]
+
+
+class RelationshipData(TypedDict):
+    from_id: int
+    to_id: int
+    type: str
+    properties: dict[str, PropertyValue]
+
+
+class GraphData(TypedDict):
+    nodes: list[NodeData] | list[ResultRow]
+    relationships: list[RelationshipData] | list[ResultRow]
+    metadata: GraphMetadata
+
+
+class GraphSummary(TypedDict):
+    total_nodes: int
+    total_relationships: int
+    node_labels: dict[str, int]
+    relationship_types: dict[str, int]
+    metadata: GraphMetadata
+
+
+# Batch types
+class NodeBatchRow(TypedDict):
+    id: PropertyValue
+    props: PropertyDict
+
+
+class RelBatchRow(TypedDict):
+    from_val: PropertyValue
+    to_val: PropertyValue
+    props: PropertyDict
+
+
+BatchParams = NodeBatchRow | RelBatchRow | PropertyDict
+
+
+class BatchWrapper(TypedDict):
+    batch: Sequence[BatchParams]
+
+
+# Language query types
+class LanguageQueries(TypedDict, total=False):
+    functions: Query | None
+    classes: Query | None
+    calls: Query | None
+    imports: Query | None
+    locals: Query | None
+    typedefs: Query | None
+    macros: Query | None
+    config: LanguageSpec
+    language: Language
+    parser: Parser
+
+
+# Function match types
+class FunctionMatch(TypedDict):
+    node: Node
+    simple_name: str
+    qualified_name: str
+    parent_class: str | None
+    line_number: int
+
+
+# Embedding query result
+class EmbeddingQueryResult(TypedDict):
+    node_id: int
+    qualified_name: str
+    start_line: int | None
+    end_line: int | None
+    path: str | None
+
+
+# Build result
+@dataclass
+class BuildResult:
+    """Result of building a code graph."""
+
+    project_name: str
+    nodes_created: int
+    relationships_created: int
+    functions_found: int
+    classes_found: int
+    files_processed: int
+    errors: list[str]
+
+    def __post_init__(self) -> None:
+        if self.errors is None:
+            self.errors = []
+
+
+# Protocols
+class FunctionRegistryTrieProtocol(Protocol):
+    def __contains__(self, qualified_name: QualifiedName) -> bool: ...
+    def __getitem__(self, qualified_name: QualifiedName) -> NodeType: ...
+    def __setitem__(self, qualified_name: QualifiedName, func_type: NodeType) -> None: ...
+    def get(self, qualified_name: QualifiedName, default: NodeType | None = None) -> NodeType | None: ...
+    def keys(self) -> KeysView[QualifiedName]: ...
+    def items(self) -> ItemsView[QualifiedName, NodeType]: ...
+    def find_with_prefix(self, prefix: str) -> list[tuple[QualifiedName, NodeType]]: ...
+    def find_ending_with(self, suffix: str) -> list[QualifiedName]: ...
+
+
+class ASTCacheProtocol(Protocol):
+    def __setitem__(self, key: Path, value: tuple[Node, SupportedLanguage]) -> None: ...
+    def __getitem__(self, key: Path) -> tuple[Node, SupportedLanguage]: ...
+    def __delitem__(self, key: Path) -> None: ...
+    def __contains__(self, key: Path) -> bool: ...
+    def items(self) -> ItemsView[Path, tuple[Node, SupportedLanguage]]: ...
+
+
+class ColumnDescriptor(Protocol):
+    @property
+    def name(self) -> str: ...
+
+
+class CursorProtocol(Protocol):
+    def execute(
+        self,
+        query: str,
+        params: PropertyDict | Sequence[BatchParams] | BatchWrapper | None = None,
+    ) -> None: ...
+    def close(self) -> None: ...
+    @property
+    def description(self) -> Sequence[ColumnDescriptor] | None: ...
+    def fetchall(self) -> list[tuple[PropertyValue, ...]]: ...
+
+
+# Node identifier type
+NodeIdentifier = tuple[NodeLabel | str, str, str | None]
+
+
+# Language import type
+class LanguageImport(NamedTuple):
+    lang_key: SupportedLanguage
+    module_path: str
+    attr_name: str
+    submodule_name: SupportedLanguage
+
+
+# Language loader type (use string literal for forward reference since Language is TYPE_CHECKING only)
+LanguageLoader = "Callable[[], Language] | None"
+
+
+# Graph node for search results
+@dataclass
+class GraphNode:
+    """A node in the code graph with full information.
+
+    Attributes:
+        node_id: Unique node identifier
+        labels: Node labels (e.g., ["Function", "Method"])
+        qualified_name: Fully qualified name
+        name: Simple name
+        path: File path
+        start_line: Start line number
+        end_line: End line number
+        docstring: Documentation string
+        properties: Additional properties
+    """
+
+    node_id: int
+    labels: list[str]
+    qualified_name: str
+    name: str
+    path: str | None = None
+    start_line: int | None = None
+    end_line: int | None = None
+    docstring: str | None = None
+    properties: dict[str, PropertyValue] = field(default_factory=dict)
+
+
+# Semantic search result with graph node
+@dataclass
+class SemanticSearchResult:
+    """Result from semantic search with graph node information.
+
+    Attributes:
+        node: The graph node
+        score: Similarity score (0-1)
+        embedding: The embedding vector (optional)
+    """
+
+    node: GraphNode
+    score: float
+    embedding: list[float] | None = None
+
+
+# Graph service protocol
+class GraphServiceProtocol(Protocol):
+    """Protocol for graph database services."""
+
+    def get_node_by_id(self, node_id: int) -> GraphNode | None:
+        """Get a node by its ID.
+
+        Args:
+            node_id: Node identifier
+
+        Returns:
+            GraphNode if found, None otherwise
+        """
+        ...
+
+    def get_nodes_by_ids(self, node_ids: list[int]) -> list[GraphNode]:
+        """Get multiple nodes by their IDs.
+
+        Args:
+            node_ids: List of node identifiers
+
+        Returns:
+            List of GraphNode objects
+        """
+        ...
+
+    def search_nodes(
+        self,
+        query: str,
+        label: str | None = None,
+        limit: int = 10,
+    ) -> list[GraphNode]:
+        """Search nodes by name or qualified name.
+
+        Args:
+            query: Search query string
+            label: Optional node label filter
+            limit: Maximum number of results
+
+        Returns:
+            List of matching GraphNode objects
+        """
+        ...
+
+    def get_node_relationships(
+        self,
+        node_id: int,
+        rel_type: str | None = None,
+        direction: str = "both",
+    ) -> list[tuple[GraphNode, str, str]]:
+        """Get relationships for a node.
+
+        Args:
+            node_id: Node identifier
+            rel_type: Optional relationship type filter
+            direction: "out", "in", or "both"
+
+        Returns:
+            List of (related_node, relationship_type, direction) tuples
+        """
+        ...
+
+    def execute_query(
+        self,
+        query: str,
+        params: PropertyDict | None = None,
+    ) -> list[ResultRow]:
+        """Execute a Cypher query.
+
+        Args:
+            query: Cypher query string
+            params: Optional query parameters
+
+        Returns:
+            List of result rows
+        """
+        ...
+
+    def close(self) -> None:
+        """Close the service connection."""
+        ...