aiecs 1.5.1__py3-none-any.whl

aiecs/application/knowledge_graph/search/hybrid_search.py

@@ -0,0 +1,423 @@
+"""
+Hybrid Search Strategy
+
+Combines vector similarity search with graph structure traversal
+to provide enhanced search results.
+"""
+
+from typing import List, Optional, Dict, Tuple
+from enum import Enum
+from pydantic import BaseModel, Field
+from aiecs.domain.knowledge_graph.models.entity import Entity
+from aiecs.domain.knowledge_graph.models.path import Path
+from aiecs.infrastructure.graph_storage.base import GraphStore
+
+
+class SearchMode(str, Enum):
+    """Search mode for hybrid search"""
+
+    VECTOR_ONLY = "vector_only"
+    GRAPH_ONLY = "graph_only"
+    HYBRID = "hybrid"
+
+
+class HybridSearchConfig(BaseModel):
+    """
+    Configuration for hybrid search
+
+    Attributes:
+        mode: Search mode (vector_only, graph_only, hybrid)
+        vector_weight: Weight for vector similarity scores (0.0-1.0)
+        graph_weight: Weight for graph structure scores (0.0-1.0)
+        max_results: Maximum number of results to return
+        vector_threshold: Minimum similarity threshold for vector search
+        max_graph_depth: Maximum depth for graph traversal
+        expand_results: Whether to expand vector results with graph neighbors
+        min_combined_score: Minimum combined score threshold
+    """
+
+    mode: SearchMode = Field(default=SearchMode.HYBRID, description="Search mode")
+
+    vector_weight: float = Field(
+        default=0.6,
+        ge=0.0,
+        le=1.0,
+        description="Weight for vector similarity scores",
+    )
+
+    graph_weight: float = Field(
+        default=0.4,
+        ge=0.0,
+        le=1.0,
+        description="Weight for graph structure scores",
+    )
+
+    max_results: int = Field(default=10, ge=1, description="Maximum number of results")
+
+    vector_threshold: float = Field(
+        default=0.0,
+        ge=0.0,
+        le=1.0,
+        description="Minimum similarity threshold for vector search",
+    )
+
+    max_graph_depth: int = Field(
+        default=2, ge=1, le=5, description="Maximum depth for graph traversal"
+    )
+
+    expand_results: bool = Field(
+        default=True,
+        description="Whether to expand vector results with graph neighbors",
+    )
+
+    min_combined_score: float = Field(
+        default=0.0,
+        ge=0.0,
+        le=1.0,
+        description="Minimum combined score threshold",
+    )
+
+    entity_type_filter: Optional[str] = Field(
+        default=None, description="Optional entity type filter"
+    )
+
+    class Config:
+        use_enum_values = True
+
+
+class HybridSearchStrategy:
+    """
+    Hybrid Search Strategy
+
+    Combines vector similarity search with graph structure traversal
+    to provide enhanced search results that leverage both semantic
+    similarity and structural relationships.
+
+    Search Modes:
+    - VECTOR_ONLY: Pure vector similarity search
+    - GRAPH_ONLY: Pure graph traversal from seed entities
+    - HYBRID: Combines both approaches with weighted scoring
+
+    Example:
+        ```python
+        strategy = HybridSearchStrategy(graph_store)
+
+        config = HybridSearchConfig(
+            mode=SearchMode.HYBRID,
+            vector_weight=0.6,
+            graph_weight=0.4,
+            max_results=10,
+            expand_results=True
+        )
+
+        results = await strategy.search(
+            query_embedding=[0.1, 0.2, ...],
+            config=config
+        )
+
+        for entity, score in results:
+            print(f"{entity.id}: {score:.3f}")
+        ```
+    """
+
+    def __init__(self, graph_store: GraphStore):
+        """
+        Initialize hybrid search strategy
+
+        Args:
+            graph_store: Graph storage backend
+        """
+        self.graph_store = graph_store
+
+    async def search(
+        self,
+        query_embedding: List[float],
+        config: Optional[HybridSearchConfig] = None,
+        seed_entity_ids: Optional[List[str]] = None,
+    ) -> List[Tuple[Entity, float]]:
+        """
+        Perform hybrid search
+
+        Args:
+            query_embedding: Query vector embedding
+            config: Search configuration (uses defaults if None)
+            seed_entity_ids: Optional seed entities for graph traversal
+
+        Returns:
+            List of (entity, score) tuples sorted by score descending
+        """
+        if config is None:
+            config = HybridSearchConfig()
+
+        if config.mode == SearchMode.VECTOR_ONLY:
+            return await self._vector_search(query_embedding, config)
+        elif config.mode == SearchMode.GRAPH_ONLY:
+            if not seed_entity_ids:
+                # If no seeds provided, do vector search to find seeds
+                vector_results = await self._vector_search(query_embedding, config, max_results=5)
+                seed_entity_ids = [entity.id for entity, _ in vector_results]
+            return await self._graph_search(seed_entity_ids, config)
+        else:  # HYBRID
+            return await self._hybrid_search(query_embedding, config, seed_entity_ids)
+
+    async def _vector_search(
+        self,
+        query_embedding: List[float],
+        config: HybridSearchConfig,
+        max_results: Optional[int] = None,
+    ) -> List[Tuple[Entity, float]]:
+        """
+        Perform vector similarity search
+
+        Args:
+            query_embedding: Query vector
+            config: Search configuration
+            max_results: Optional override for max results
+
+        Returns:
+            List of (entity, score) tuples
+        """
+        results = await self.graph_store.vector_search(
+            query_embedding=query_embedding,
+            entity_type=config.entity_type_filter,
+            max_results=max_results or config.max_results,
+            score_threshold=config.vector_threshold,
+        )
+
+        return results
+
+    async def _graph_search(
+        self, seed_entity_ids: List[str], config: HybridSearchConfig
+    ) -> List[Tuple[Entity, float]]:
+        """
+        Perform graph structure search from seed entities
+
+        Args:
+            seed_entity_ids: Starting entities for traversal
+            config: Search configuration
+
+        Returns:
+            List of (entity, score) tuples
+        """
+        # Collect entities from graph traversal
+        entity_scores: Dict[str, float] = {}
+
+        for seed_id in seed_entity_ids:
+            # Get neighbors at different depths
+            current_entities = {seed_id}
+            visited = set()
+
+            for depth in range(config.max_graph_depth):
+                next_entities = set()
+
+                for entity_id in current_entities:
+                    if entity_id in visited:
+                        continue
+
+                    visited.add(entity_id)
+
+                    # Score decreases with depth
+                    depth_score = 1.0 / (depth + 1)
+
+                    # Update score (take max if entity seen from multiple
+                    # paths)
+                    if entity_id not in entity_scores:
+                        entity_scores[entity_id] = depth_score
+                    else:
+                        entity_scores[entity_id] = max(entity_scores[entity_id], depth_score)
+
+                    # Get neighbors for next depth
+                    neighbors = await self.graph_store.get_neighbors(
+                        entity_id, direction="outgoing"
+                    )
+
+                    for neighbor in neighbors:
+                        if neighbor.id not in visited:
+                            next_entities.add(neighbor.id)
+
+                current_entities = next_entities
+
+                if not current_entities:
+                    break
+
+        # Retrieve entities and create result list
+        results = []
+        for entity_id, score in entity_scores.items():
+            entity = await self.graph_store.get_entity(entity_id)
+            if entity:
+                # Apply entity type filter if specified
+                if config.entity_type_filter:
+                    if entity.entity_type == config.entity_type_filter:
+                        results.append((entity, score))
+                else:
+                    results.append((entity, score))
+
+        # Sort by score descending
+        results.sort(key=lambda x: x[1], reverse=True)
+
+        # Return top results
+        return results[: config.max_results]
+
+    async def _hybrid_search(
+        self,
+        query_embedding: List[float],
+        config: HybridSearchConfig,
+        seed_entity_ids: Optional[List[str]] = None,
+    ) -> List[Tuple[Entity, float]]:
+        """
+        Perform hybrid search combining vector and graph
+
+        Args:
+            query_embedding: Query vector
+            config: Search configuration
+            seed_entity_ids: Optional seed entities
+
+        Returns:
+            List of (entity, score) tuples with combined scores
+        """
+        # Step 1: Vector search
+        vector_results = await self._vector_search(
+            query_embedding,
+            config,
+            max_results=config.max_results * 2,  # Get more for expansion
+        )
+
+        # Create score dictionaries
+        vector_scores: Dict[str, float] = {entity.id: score for entity, score in vector_results}
+
+        # Step 2: Graph expansion (if enabled)
+        graph_scores: Dict[str, float] = {}
+
+        if config.expand_results:
+            # Use top vector results as seeds
+            seeds = seed_entity_ids or [entity.id for entity, _ in vector_results[:5]]
+
+            graph_results = await self._graph_search(seeds, config)
+            graph_scores = {entity.id: score for entity, score in graph_results}
+
+        # Step 3: Combine scores
+        combined_scores = await self._combine_scores(vector_scores, graph_scores, config)
+
+        # Step 4: Retrieve entities and create results
+        results = []
+        for entity_id, combined_score in combined_scores.items():
+            # Apply minimum score threshold
+            if combined_score < config.min_combined_score:
+                continue
+
+            entity = await self.graph_store.get_entity(entity_id)
+            if entity:
+                results.append((entity, combined_score))
+
+        # Sort by combined score descending
+        results.sort(key=lambda x: x[1], reverse=True)
+
+        # Return top results
+        return results[: config.max_results]
+
+    async def _combine_scores(
+        self,
+        vector_scores: Dict[str, float],
+        graph_scores: Dict[str, float],
+        config: HybridSearchConfig,
+    ) -> Dict[str, float]:
+        """
+        Combine vector and graph scores with weighted averaging
+
+        Args:
+            vector_scores: Entity ID to vector similarity score
+            graph_scores: Entity ID to graph structure score
+            config: Search configuration
+
+        Returns:
+            Combined scores dictionary
+        """
+        # Normalize weights
+        total_weight = config.vector_weight + config.graph_weight
+        if total_weight == 0:
+            total_weight = 1.0
+
+        norm_vector_weight = config.vector_weight / total_weight
+        norm_graph_weight = config.graph_weight / total_weight
+
+        # Get all entity IDs
+        all_entity_ids = set(vector_scores.keys()) | set(graph_scores.keys())
+
+        # Combine scores
+        combined: Dict[str, float] = {}
+
+        for entity_id in all_entity_ids:
+            v_score = vector_scores.get(entity_id, 0.0)
+            g_score = graph_scores.get(entity_id, 0.0)
+
+            # Weighted combination
+            combined[entity_id] = v_score * norm_vector_weight + g_score * norm_graph_weight
+
+        return combined
+
+    async def search_with_expansion(
+        self,
+        query_embedding: List[float],
+        config: Optional[HybridSearchConfig] = None,
+        include_paths: bool = False,
+    ) -> Tuple[List[Tuple[Entity, float]], Optional[List[Path]]]:
+        """
+        Search with result expansion and optional path tracking
+
+        Args:
+            query_embedding: Query vector
+            config: Search configuration
+            include_paths: Whether to include paths to results
+
+        Returns:
+            Tuple of (results, paths) where paths is None if not requested
+        """
+        if config is None:
+            config = HybridSearchConfig()
+
+        # Perform search
+        results = await self.search(query_embedding, config)
+
+        paths = None
+        if include_paths and config.expand_results:
+            # Find paths from top vector results to expanded results
+            paths = await self._find_result_paths(results, config)
+
+        return results, paths
+
+    async def _find_result_paths(
+        self, results: List[Tuple[Entity, float]], config: HybridSearchConfig
+    ) -> List[Path]:
+        """
+        Find paths between top results
+
+        Args:
+            results: Search results
+            config: Search configuration
+
+        Returns:
+            List of paths connecting results
+        """
+        if len(results) < 2:
+            return []
+
+        paths = []
+
+        # Find paths between top results
+        for i in range(min(3, len(results))):
+            source_id = results[i][0].id
+
+            for j in range(i + 1, min(i + 4, len(results))):
+                target_id = results[j][0].id
+
+                # Find paths between these entities
+                found_paths = await self.graph_store.find_paths(
+                    source_entity_id=source_id,
+                    target_entity_id=target_id,
+                    max_depth=config.max_graph_depth,
+                    max_paths=2,
+                )
+
+                paths.extend(found_paths)
+
+        return paths

aiecs/application/knowledge_graph/search/reranker.py

@@ -0,0 +1,295 @@
+"""
+Result Reranking Framework
+
+Pluggable reranking strategies for improving search result relevance.
+"""
+
+from abc import ABC, abstractmethod
+from typing import List, Dict, Optional, Tuple
+from enum import Enum
+
+from aiecs.domain.knowledge_graph.models.entity import Entity
+
+
+class ScoreCombinationMethod(str, Enum):
+    """Methods for combining scores from multiple reranking strategies"""
+
+    WEIGHTED_AVERAGE = "weighted_average"
+    RRF = "rrf"  # Reciprocal Rank Fusion
+    MAX = "max"
+    MIN = "min"
+
+
+class RerankerStrategy(ABC):
+    """
+    Abstract base class for reranking strategies
+
+    Each strategy computes a relevance score for entities given a query.
+    Strategies can be combined using different combination methods.
+
+    Example:
+        ```python
+        class TextSimilarityReranker(RerankerStrategy):
+            async def score(
+                self,
+                query: str,
+                entities: List[Entity]
+            ) -> List[float]:
+                # Compute BM25 scores
+                return scores
+        ```
+    """
+
+    @property
+    @abstractmethod
+    def name(self) -> str:
+        """Strategy name for identification"""
+
+    @abstractmethod
+    async def score(self, query: str, entities: List[Entity], **kwargs) -> List[float]:
+        """
+        Compute relevance scores for entities
+
+        Args:
+            query: Query text or context
+            entities: List of entities to score
+            **kwargs: Strategy-specific parameters
+
+        Returns:
+            List of scores (one per entity), same order as entities
+            Scores should be in range [0.0, 1.0] for best results
+        """
+
+
+def normalize_scores(scores: List[float], method: str = "min_max") -> List[float]:
+    """
+    Normalize scores to [0.0, 1.0] range
+
+    Args:
+        scores: Raw scores to normalize
+        method: Normalization method ("min_max", "z_score", "softmax")
+
+    Returns:
+        Normalized scores in [0.0, 1.0] range
+    """
+    if not scores:
+        return []
+
+    if method == "min_max":
+        min_score = min(scores)
+        max_score = max(scores)
+        if max_score == min_score:
+            return [1.0] * len(scores)
+        return [(s - min_score) / (max_score - min_score) for s in scores]
+
+    elif method == "z_score":
+        import statistics
+
+        if len(scores) < 2:
+            return [1.0] * len(scores)
+        mean = statistics.mean(scores)
+        stdev = statistics.stdev(scores) if len(scores) > 1 else 1.0
+        if stdev == 0:
+            return [1.0] * len(scores)
+        # Normalize to [0, 1] using sigmoid
+        normalized = [(s - mean) / stdev for s in scores]
+        import math
+
+        return [1 / (1 + math.exp(-n)) for n in normalized]
+
+    elif method == "softmax":
+        import math
+
+        # Shift to avoid overflow
+        max_score = max(scores)
+        exp_scores = [math.exp(s - max_score) for s in scores]
+        sum_exp = sum(exp_scores)
+        if sum_exp == 0:
+            return [1.0 / len(scores)] * len(scores)
+        return [e / sum_exp for e in exp_scores]
+
+    else:
+        raise ValueError(f"Unknown normalization method: {method}")
+
+
+def combine_scores(
+    score_dicts: List[Dict[str, float]],
+    method: ScoreCombinationMethod = ScoreCombinationMethod.WEIGHTED_AVERAGE,
+    weights: Optional[Dict[str, float]] = None,
+) -> Dict[str, float]:
+    """
+    Combine scores from multiple strategies
+
+    Args:
+        score_dicts: List of {entity_id: score} dictionaries from each strategy
+        method: Combination method
+        weights: Optional weights for each strategy (for weighted_average)
+
+    Returns:
+        Combined scores as {entity_id: combined_score}
+    """
+    if not score_dicts:
+        return {}
+
+    # Collect all entity IDs
+    all_entity_ids = set()
+    for score_dict in score_dicts:
+        all_entity_ids.update(score_dict.keys())
+
+    if method == ScoreCombinationMethod.WEIGHTED_AVERAGE:
+        if weights is None:
+            # Equal weights
+            weight = 1.0 / len(score_dicts)
+            weights = {f"strategy_{i}": weight for i in range(len(score_dicts))}
+
+        combined = {}
+        for entity_id in all_entity_ids:
+            weighted_sum = 0.0
+            total_weight = 0.0
+            for i, score_dict in enumerate(score_dicts):
+                strategy_name = f"strategy_{i}"
+                weight = weights.get(strategy_name, 1.0 / len(score_dicts))
+                score = score_dict.get(entity_id, 0.0)
+                weighted_sum += weight * score
+                total_weight += weight
+            combined[entity_id] = weighted_sum / total_weight if total_weight > 0 else 0.0
+        return combined
+
+    elif method == ScoreCombinationMethod.RRF:
+        # Reciprocal Rank Fusion
+        k = 60  # RRF constant
+        combined = {}
+        for entity_id in all_entity_ids:
+            rrf_score = 0.0
+            for score_dict in score_dicts:
+                if entity_id in score_dict:
+                    # Get rank (1-indexed, higher score = lower rank)
+                    scores = sorted(score_dict.values(), reverse=True)
+                    rank = scores.index(score_dict[entity_id]) + 1
+                    rrf_score += 1.0 / (k + rank)
+            combined[entity_id] = rrf_score
+        return combined
+
+    elif method == ScoreCombinationMethod.MAX:
+        combined = {}
+        for entity_id in all_entity_ids:
+            combined[entity_id] = max(score_dict.get(entity_id, 0.0) for score_dict in score_dicts)
+        return combined
+
+    elif method == ScoreCombinationMethod.MIN:
+        combined = {}
+        for entity_id in all_entity_ids:
+            combined[entity_id] = min(score_dict.get(entity_id, 1.0) for score_dict in score_dicts)
+        return combined
+
+    else:
+        raise ValueError(f"Unknown combination method: {method}")
+
+
+class ResultReranker:
+    """
+    Result Reranker orchestrator
+
+    Combines multiple reranking strategies to improve search result relevance.
+
+    Example:
+        ```python
+        # Create strategies
+        text_reranker = TextSimilarityReranker()
+        semantic_reranker = SemanticReranker()
+
+        # Create reranker
+        reranker = ResultReranker(
+            strategies=[text_reranker, semantic_reranker],
+            combination_method=ScoreCombinationMethod.WEIGHTED_AVERAGE,
+            weights={"text": 0.6, "semantic": 0.4}
+        )
+
+        # Rerank results
+        reranked = await reranker.rerank(
+            query="machine learning",
+            entities=search_results
+        )
+        ```
+    """
+
+    def __init__(
+        self,
+        strategies: List[RerankerStrategy],
+        combination_method: ScoreCombinationMethod = ScoreCombinationMethod.WEIGHTED_AVERAGE,
+        weights: Optional[Dict[str, float]] = None,
+        normalize_scores: bool = True,
+        normalization_method: str = "min_max",
+    ):
+        """
+        Initialize ResultReranker
+
+        Args:
+            strategies: List of reranking strategies
+            combination_method: Method for combining scores
+            weights: Optional weights for strategies (for weighted_average)
+            normalize_scores: Whether to normalize scores before combining
+            normalization_method: Normalization method ("min_max", "z_score", "softmax")
+        """
+        if not strategies:
+            raise ValueError("At least one strategy is required")
+
+        self.strategies = strategies
+        self.combination_method = combination_method
+        self.weights = weights or {}
+        self.normalize_scores = normalize_scores
+        self.normalization_method = normalization_method
+
+    async def rerank(
+        self,
+        query: str,
+        entities: List[Entity],
+        top_k: Optional[int] = None,
+        **kwargs,
+    ) -> List[Tuple[Entity, float]]:
+        """
+        Rerank entities using all strategies
+
+        Args:
+            query: Query text or context
+            entities: List of entities to rerank
+            top_k: Optional limit on number of results
+            **kwargs: Additional parameters passed to strategies
+
+        Returns:
+            List of (entity, combined_score) tuples, sorted by score descending
+        """
+        if not entities:
+            return []
+
+        # Get scores from each strategy
+        strategy_scores = []
+        for strategy in self.strategies:
+            scores = await strategy.score(query, entities, **kwargs)
+
+            # Normalize if requested
+            if self.normalize_scores:
+                scores = normalize_scores(scores, self.normalization_method)
+
+            # Convert to entity_id -> score dictionary
+            score_dict = {entity.id: score for entity, score in zip(entities, scores)}
+            strategy_scores.append(score_dict)
+
+        # Combine scores
+        combined_scores = combine_scores(
+            strategy_scores,
+            method=self.combination_method,
+            weights=self.weights,
+        )
+
+        # Create (entity, score) tuples
+        reranked = [(entity, combined_scores.get(entity.id, 0.0)) for entity in entities]
+
+        # Sort by score descending
+        reranked.sort(key=lambda x: x[1], reverse=True)
+
+        # Apply top_k limit
+        if top_k is not None:
+            reranked = reranked[:top_k]
+
+        return reranked