emdash-core 0.1.7__py3-none-any.whl

emdash_core/embeddings/service.py

@@ -0,0 +1,215 @@
+"""Embedding service with multi-provider support."""
+
+from typing import Optional, Union
+
+from .models import EmbeddingModel
+from .registry import get_provider, get_available_model, ProviderRegistry
+from .providers.base import EmbeddingProvider
+from ..core.config import get_config
+from ..core.models import PullRequestEntity, FunctionEntity, ClassEntity
+from ..utils.logger import log
+
+
+class EmbeddingService:
+    """
+    Unified embedding service with multi-provider support.
+
+    Supports OpenAI, Fireworks AI, and other providers through a registry.
+    Falls back to the first available provider if none is specified.
+
+    Usage:
+        # Use default model (OpenAI text-embedding-3-small)
+        service = EmbeddingService()
+
+        # Use specific model
+        service = EmbeddingService(model=EmbeddingModel.FIREWORKS_NOMIC_EMBED_V1_5)
+
+        # Use model from string
+        service = EmbeddingService(model="fireworks:nomic-ai/nomic-embed-text-v1.5")
+
+        # Generate embeddings
+        embeddings = service.embed_texts(["hello", "world"])
+    """
+
+    def __init__(
+        self,
+        model: Optional[Union[EmbeddingModel, str]] = None,
+        provider: Optional[EmbeddingProvider] = None,
+    ):
+        """
+        Initialize embedding service.
+
+        Args:
+            model: Embedding model to use. Can be EmbeddingModel enum or string.
+                  If None, uses first available model (OpenAI > Fireworks).
+            provider: Pre-configured provider instance. If provided, model is ignored.
+        """
+        if provider is not None:
+            self._provider = provider
+            self._model = provider.model
+        elif model is not None:
+            # Parse model from string if needed
+            if isinstance(model, str):
+                parsed = EmbeddingModel.from_string(model)
+                if parsed is None:
+                    raise ValueError(f"Unknown embedding model: {model}")
+                model = parsed
+            self._model = model
+            self._provider = None  # Lazy-load
+        else:
+            # Auto-select first available model
+            self._model = get_available_model()
+            self._provider = None
+
+    @property
+    def model(self) -> Optional[EmbeddingModel]:
+        """Get the embedding model."""
+        return self._model
+
+    @property
+    def provider(self) -> Optional[EmbeddingProvider]:
+        """Get the provider, lazy-loading if needed."""
+        if self._provider is None and self._model is not None:
+            self._provider = get_provider(self._model)
+        return self._provider
+
+    @property
+    def is_available(self) -> bool:
+        """Check if embedding service is available."""
+        if self._model is None:
+            return False
+        try:
+            return self.provider is not None and self.provider.is_available
+        except Exception:
+            return False
+
+    @property
+    def dimensions(self) -> int:
+        """Get embedding dimensions for the current model."""
+        if self._model is None:
+            return 1536  # Default OpenAI dimensions
+        return self._model.dimensions
+
+    def embed_texts(self, texts: list[str]) -> list[Optional[list[float]]]:
+        """
+        Generate embeddings for multiple texts.
+
+        Args:
+            texts: List of text strings to embed
+
+        Returns:
+            List of embedding vectors. None for failed embeddings.
+        """
+        if not texts:
+            return []
+
+        if not self.is_available:
+            log.warning("No embedding provider available")
+            return [None] * len(texts)
+
+        return self.provider.embed_texts(texts)
+
+    def embed_text(self, text: str) -> Optional[list[float]]:
+        """
+        Generate embedding for a single text.
+
+        Args:
+            text: Text string to embed
+
+        Returns:
+            Embedding vector or None if failed
+        """
+        if not text:
+            return None
+        embeddings = self.embed_texts([text])
+        return embeddings[0] if embeddings else None
+
+    def embed_query(self, query: str) -> Optional[list[float]]:
+        """
+        Generate embedding for a search query.
+
+        Args:
+            query: Search query string
+
+        Returns:
+            Embedding vector or None if failed
+        """
+        if not query:
+            return None
+
+        if not self.is_available:
+            log.warning("No embedding provider available")
+            return None
+
+        return self.provider.embed_query(query)
+
+    def embed_pr(self, pr: PullRequestEntity) -> Optional[list[float]]:
+        """
+        Generate embedding for a PR (title + description).
+
+        Args:
+            pr: PullRequestEntity
+
+        Returns:
+            Embedding vector or None
+        """
+        text = f"{pr.title}\n\n{pr.description or ''}"
+        return self.embed_text(text)
+
+    def embed_function(self, func: FunctionEntity) -> Optional[list[float]]:
+        """
+        Generate embedding for a function (name + docstring).
+
+        Args:
+            func: FunctionEntity
+
+        Returns:
+            Embedding vector or None
+        """
+        text = f"{func.name}\n\n{func.docstring or ''}"
+        return self.embed_text(text)
+
+    def embed_class(self, cls: ClassEntity) -> Optional[list[float]]:
+        """
+        Generate embedding for a class (name + docstring).
+
+        Args:
+            cls: ClassEntity
+
+        Returns:
+            Embedding vector or None
+        """
+        text = f"{cls.name}\n\n{cls.docstring or ''}"
+        return self.embed_text(text)
+
+    @staticmethod
+    def list_models() -> list[dict]:
+        """List all available embedding models."""
+        return EmbeddingModel.list_all()
+
+    @staticmethod
+    def list_available_providers() -> list[str]:
+        """List providers that are configured and available."""
+        return [
+            name
+            for name in ProviderRegistry.list_providers()
+            if ProviderRegistry.is_provider_available(name)
+        ]
+
+    @staticmethod
+    def get_model_info(model: Union[EmbeddingModel, str]) -> Optional[dict]:
+        """Get information about a specific model."""
+        if isinstance(model, str):
+            model = EmbeddingModel.from_string(model)
+            if model is None:
+                return None
+
+        return {
+            "name": model.name,
+            "provider": model.provider,
+            "model_id": model.model_id,
+            "dimensions": model.dimensions,
+            "max_tokens": model.spec.max_tokens,
+            "batch_size": model.spec.batch_size,
+            "description": model.spec.description,
+        }

emdash_core/graph/__init__.py

@@ -0,0 +1,26 @@
+"""Graph module for Kuzu database operations."""
+
+from .connection import (
+    KuzuConnection,
+    get_connection,
+    set_connection,
+    close_connection,
+)
+from .schema import SchemaManager, initialize_database
+from .builder import GraphBuilder
+from .writer import GraphWriter
+
+__all__ = [
+    # Connection
+    "KuzuConnection",
+    "get_connection",
+    "set_connection",
+    "close_connection",
+    # Schema
+    "SchemaManager",
+    "initialize_database",
+    # Builder
+    "GraphBuilder",
+    # Writer
+    "GraphWriter",
+]

emdash_core/graph/builder.py

@@ -0,0 +1,134 @@
+"""Graph construction - coordinates entity writing to Kuzu."""
+
+from ..core.models import CodebaseEntities, GitData
+from .connection import KuzuConnection
+from .writer import GraphWriter
+from ..utils.logger import log
+
+
+class GraphBuilder:
+    """Builds the Kuzu knowledge graph from extracted entities."""
+
+    def __init__(self, connection: KuzuConnection, batch_size: int = 1000):
+        """Initialize graph builder.
+
+        Args:
+            connection: Neo4j connection
+            batch_size: Batch size for write operations
+        """
+        self.connection = connection
+        self.batch_size = batch_size
+
+    def build_code_graph(self, entities: CodebaseEntities):
+        """Build Layer A: Code structure graph.
+
+        Args:
+            entities: Codebase entities to write
+        """
+        log.info("Building code structure graph (Layer A)...")
+
+        # Use connection directly for Kuzu (no session needed)
+        writer = GraphWriter(self.connection, self.batch_size)
+
+        # Pass 1: Create nodes
+        log.info("Pass 1: Creating nodes...")
+        writer.write_files(entities.files)
+        writer.write_modules(entities.modules)
+        writer.write_classes(entities.classes)
+        writer.write_functions(entities.functions)
+
+        # Pass 2: Create relationships
+        log.info("Pass 2: Creating relationships...")
+        writer.write_imports(entities.imports)
+        writer.write_inheritance(entities.classes)
+        writer.write_calls(entities.functions)
+
+        log.info("Code graph construction complete")
+
+    def build_git_graph(self, git_data: GitData):
+        """Build Layer B: Git history graph.
+
+        Args:
+            git_data: Git data to write
+        """
+        log.info("Building git history graph (Layer B)...")
+
+        # Use connection directly for Kuzu (no session needed)
+        writer = GraphWriter(self.connection, self.batch_size)
+
+        # Create nodes
+        writer.write_authors(git_data.authors)
+        writer.write_commits(git_data.commits)
+
+        # Create relationships
+        writer.write_commit_authorship(git_data.commits)
+        writer.write_file_modifications(git_data.modifications)
+
+        log.info("Git graph construction complete")
+
+    def delete_files(self, file_paths: list[str]) -> int:
+        """Remove files and their associated entities from graph.
+
+        Deletes:
+        1. Classes and Functions belonging to these files
+        2. Relationships (CONTAINS_CLASS, CONTAINS_FUNCTION, IMPORTS, etc.)
+        3. The File nodes themselves
+
+        Args:
+            file_paths: List of file paths to remove
+
+        Returns:
+            Number of files deleted
+        """
+        if not file_paths:
+            return 0
+
+        log.info(f"Deleting {len(file_paths)} files from graph...")
+
+        deleted_count = 0
+        for file_path in file_paths:
+            try:
+                # Delete classes belonging to this file
+                self.connection.execute_write(
+                    "MATCH (c:Class) WHERE c.file_path = $path DETACH DELETE c",
+                    {"path": file_path}
+                )
+
+                # Delete functions belonging to this file
+                self.connection.execute_write(
+                    "MATCH (f:Function) WHERE f.file_path = $path DETACH DELETE f",
+                    {"path": file_path}
+                )
+
+                # Delete the file node itself (and its relationships)
+                self.connection.execute_write(
+                    "MATCH (f:File {path: $path}) DETACH DELETE f",
+                    {"path": file_path}
+                )
+
+                deleted_count += 1
+            except Exception as e:
+                log.warning(f"Failed to delete file {file_path}: {e}")
+
+        log.info(f"Deleted {deleted_count} files from graph")
+        return deleted_count
+
+    def clear_repository_data(self, repo_url: str):
+        """Clear all data for a specific repository.
+
+        Args:
+            repo_url: Repository URL
+        """
+        log.warning(f"Clearing data for repository: {repo_url}")
+
+        # Use connection directly for Kuzu (no session needed)
+        # Delete all nodes related to this repository
+        # This is a simplified version - in production you'd want more sophisticated cleanup
+        query = """
+        MATCH (r:Repository {url: $url})
+        OPTIONAL MATCH (r)-[*]-(n)
+        DETACH DELETE r, n
+        """
+        self.connection.execute_write(query, {"url": repo_url})
+
+        log.info("Repository data cleared")