aiecs 1.5.1__py3-none-any.whl

aiecs/domain/knowledge_graph/models/path_pattern.py

@@ -0,0 +1,173 @@
+"""
+PathPattern Domain Model
+
+Represents a pattern for graph traversal specifications.
+"""
+
+from typing import List, Optional, Set
+from pydantic import BaseModel, Field
+from enum import Enum
+
+
+class TraversalDirection(str, Enum):
+    """Direction for graph traversal"""
+
+    OUTGOING = "outgoing"
+    INCOMING = "incoming"
+    BOTH = "both"
+
+
+class PathPattern(BaseModel):
+    """
+    Path Pattern for Graph Traversal
+
+    Specifies constraints and preferences for graph traversal operations.
+    Used to control how the graph is explored and which paths are valid.
+
+    Attributes:
+        relation_types: Optional list of allowed relation types
+        entity_types: Optional list of allowed entity types
+        direction: Traversal direction (outgoing, incoming, both)
+        max_depth: Maximum path length
+        allow_cycles: Whether to allow revisiting nodes
+        required_relation_sequence: Optional sequence of relation types that must be followed
+
+    Example:
+        ```python
+        # Find paths following WORKS_FOR -> LOCATED_IN pattern
+        pattern = PathPattern(
+            relation_types=["WORKS_FOR", "LOCATED_IN"],
+            required_relation_sequence=["WORKS_FOR", "LOCATED_IN"],
+            max_depth=2,
+            allow_cycles=False
+        )
+        ```
+    """
+
+    relation_types: Optional[List[str]] = Field(
+        default=None,
+        description="Optional list of allowed relation types (None = all types allowed)",
+    )
+
+    entity_types: Optional[List[str]] = Field(
+        default=None,
+        description="Optional list of allowed entity types (None = all types allowed)",
+    )
+
+    direction: TraversalDirection = Field(
+        default=TraversalDirection.OUTGOING,
+        description="Direction for traversal",
+    )
+
+    max_depth: int = Field(
+        default=3,
+        ge=1,
+        le=10,
+        description="Maximum path length (number of hops)",
+    )
+
+    allow_cycles: bool = Field(
+        default=False, description="Whether to allow revisiting nodes (cycles)"
+    )
+
+    required_relation_sequence: Optional[List[str]] = Field(
+        default=None,
+        description="Optional sequence of relation types that must be followed in order",
+    )
+
+    min_path_length: int = Field(default=1, ge=1, description="Minimum path length to return")
+
+    excluded_entity_ids: Set[str] = Field(
+        default_factory=set,
+        description="Set of entity IDs to exclude from traversal",
+    )
+
+    def is_relation_allowed(self, relation_type: str, depth: int = 0) -> bool:
+        """
+        Check if a relation type is allowed at the current depth
+
+        Args:
+            relation_type: Relation type to check
+            depth: Current depth in the traversal (0-indexed)
+
+        Returns:
+            True if the relation is allowed
+        """
+        # Check required sequence if specified
+        if self.required_relation_sequence:
+            if depth >= len(self.required_relation_sequence):
+                return False
+            return relation_type == self.required_relation_sequence[depth]
+
+        # Check allowed types if specified
+        if self.relation_types:
+            return relation_type in self.relation_types
+
+        # All types allowed if no constraints
+        return True
+
+    def is_entity_allowed(self, entity_id: str, entity_type: str) -> bool:
+        """
+        Check if an entity is allowed in the path
+
+        Args:
+            entity_id: Entity ID to check
+            entity_type: Entity type to check
+
+        Returns:
+            True if the entity is allowed
+        """
+        # Check excluded entities
+        if entity_id in self.excluded_entity_ids:
+            return False
+
+        # Check allowed types if specified
+        if self.entity_types:
+            return entity_type in self.entity_types
+
+        # All entities allowed if no constraints
+        return True
+
+    def is_valid_path_length(self, length: int) -> bool:
+        """
+        Check if a path length is valid according to the pattern
+
+        Args:
+            length: Path length to check
+
+        Returns:
+            True if the length is valid
+        """
+        return self.min_path_length <= length <= self.max_depth
+
+    def should_continue_traversal(self, depth: int) -> bool:
+        """
+        Check if traversal should continue at the current depth
+
+        Args:
+            depth: Current depth in the traversal
+
+        Returns:
+            True if traversal should continue
+        """
+        return depth < self.max_depth
+
+    class Config:
+        use_enum_values = True
+
+    def __str__(self) -> str:
+        parts = [f"depth={self.max_depth}"]
+
+        if self.relation_types:
+            parts.append(f"relations={','.join(self.relation_types)}")
+
+        if self.entity_types:
+            parts.append(f"entities={','.join(self.entity_types)}")
+
+        if self.required_relation_sequence:
+            parts.append(f"sequence={' -> '.join(self.required_relation_sequence)}")
+
+        if self.allow_cycles:
+            parts.append("allow_cycles")
+
+        return f"PathPattern({', '.join(parts)})"

aiecs/domain/knowledge_graph/models/query.py

@@ -0,0 +1,272 @@
+"""
+Query Domain Models
+
+Models for specifying graph queries and their results.
+"""
+
+from typing import Any, List, Optional, Dict
+from enum import Enum
+from pydantic import BaseModel, Field
+from aiecs.domain.knowledge_graph.models.entity import Entity
+from aiecs.domain.knowledge_graph.models.relation import Relation
+from aiecs.domain.knowledge_graph.models.path import Path
+from aiecs.domain.knowledge_graph.models.path_pattern import PathPattern
+
+
+class QueryType(str, Enum):
+    """Types of graph queries"""
+
+    ENTITY_LOOKUP = "entity_lookup"
+    VECTOR_SEARCH = "vector_search"
+    TRAVERSAL = "traversal"
+    PATH_FINDING = "path_finding"
+    SUBGRAPH = "subgraph"
+    CUSTOM = "custom"
+
+
+class GraphQuery(BaseModel):
+    """
+    Graph Query Specification
+
+    Specifies a query to execute against the knowledge graph.
+
+    Attributes:
+        query_type: Type of query to execute
+        entity_id: Entity ID for entity lookup queries
+        entity_type: Filter by entity type
+        relation_type: Filter by relation type
+        embedding: Query embedding for vector search
+        properties: Property constraints for filtering
+        max_results: Maximum number of results to return
+        max_depth: Maximum traversal depth for path queries
+        score_threshold: Minimum score threshold for results
+
+    Example:
+        ```python
+        # Vector search query
+        query = GraphQuery(
+            query_type=QueryType.VECTOR_SEARCH,
+            embedding=[0.1, 0.2, ...],
+            entity_type="Document",
+            max_results=10,
+            score_threshold=0.7
+        )
+
+        # Traversal query
+        query = GraphQuery(
+            query_type=QueryType.TRAVERSAL,
+            entity_id="person_001",
+            relation_type="KNOWS",
+            max_depth=3
+        )
+        ```
+    """
+
+    query_type: QueryType = Field(..., description="Type of query to execute")
+
+    # Entity lookup
+    entity_id: Optional[str] = Field(
+        default=None, description="Entity ID for entity lookup queries"
+    )
+
+    # Filtering
+    entity_type: Optional[str] = Field(default=None, description="Filter results by entity type")
+
+    relation_type: Optional[str] = Field(
+        default=None, description="Filter by relation type (for traversal)"
+    )
+
+    # Vector search
+    embedding: Optional[List[float]] = Field(
+        default=None, description="Query embedding for vector search"
+    )
+
+    # Property constraints
+    properties: Dict[str, Any] = Field(
+        default_factory=dict, description="Property constraints for filtering"
+    )
+
+    # Result constraints
+    max_results: int = Field(default=10, ge=1, description="Maximum number of results to return")
+
+    max_depth: int = Field(
+        default=3,
+        ge=1,
+        le=10,
+        description="Maximum traversal depth for path queries",
+    )
+
+    score_threshold: float = Field(
+        default=0.0,
+        ge=0.0,
+        le=1.0,
+        description="Minimum score threshold for results",
+    )
+
+    # Source/target constraints for path finding
+    source_entity_id: Optional[str] = Field(
+        default=None, description="Source entity ID for path finding"
+    )
+
+    target_entity_id: Optional[str] = Field(
+        default=None, description="Target entity ID for path finding"
+    )
+
+    # Custom query parameters
+    custom_params: Dict[str, Any] = Field(
+        default_factory=dict, description="Additional custom query parameters"
+    )
+
+    # Custom pattern matching (Task 3.3)
+    pattern: Optional[PathPattern] = Field(
+        default=None,
+        description="Graph pattern for custom pattern matching queries",
+    )
+
+    patterns: Optional[List[PathPattern]] = Field(
+        default=None,
+        description="Multiple graph patterns for complex pattern matching",
+    )
+
+    optional_patterns: Optional[List[PathPattern]] = Field(
+        default=None, description="Optional patterns that may or may not match"
+    )
+
+    projection: Optional[List[str]] = Field(
+        default=None,
+        description="Fields to project in results (e.g., ['id', 'name', 'properties.age'])",
+    )
+
+    aggregations: Optional[Dict[str, str]] = Field(
+        default=None,
+        description="Aggregations to apply (e.g., {'count': 'COUNT', 'avg_age': 'AVG(properties.age)'})",
+    )
+
+    group_by: Optional[List[str]] = Field(
+        default=None, description="Fields to group by for aggregations"
+    )
+
+    class Config:
+        use_enum_values = True
+
+    def __str__(self) -> str:
+        parts = [f"GraphQuery(type={self.query_type})"]
+        if self.entity_id:
+            parts.append(f"entity_id={self.entity_id}")
+        if self.entity_type:
+            parts.append(f"entity_type={self.entity_type}")
+        if self.relation_type:
+            parts.append(f"relation_type={self.relation_type}")
+        return " ".join(parts)
+
+
+class GraphResult(BaseModel):
+    """
+    Graph Query Result
+
+    Contains the results of a graph query execution.
+
+    Attributes:
+        query: The original query that was executed
+        entities: List of entity results
+        relations: List of relation results (for traversal/subgraph queries)
+        paths: List of path results (for path-finding queries)
+        scores: Scores for each result (parallel to entities)
+        total_count: Total number of matching results (before max_results limit)
+        execution_time_ms: Query execution time in milliseconds
+
+    Example:
+        ```python
+        result = GraphResult(
+            query=query,
+            entities=[entity1, entity2],
+            scores=[0.95, 0.87],
+            total_count=2,
+            execution_time_ms=15.3
+        )
+        ```
+    """
+
+    query: GraphQuery = Field(..., description="The original query that was executed")
+
+    entities: List[Entity] = Field(default_factory=list, description="List of entity results")
+
+    relations: List[Relation] = Field(
+        default_factory=list,
+        description="List of relation results (for graph queries)",
+    )
+
+    paths: List[Path] = Field(
+        default_factory=list,
+        description="List of path results (for path-finding)",
+    )
+
+    scores: List[float] = Field(default_factory=list, description="Scores for each entity result")
+
+    total_count: int = Field(default=0, ge=0, description="Total matching results (before limit)")
+
+    execution_time_ms: Optional[float] = Field(
+        default=None,
+        ge=0.0,
+        description="Query execution time in milliseconds",
+    )
+
+    class Config:
+        arbitrary_types_allowed = True
+
+    @property
+    def has_results(self) -> bool:
+        """Check if result contains any data"""
+        return len(self.entities) > 0 or len(self.paths) > 0
+
+    @property
+    def entity_count(self) -> int:
+        """Get number of entities in result"""
+        return len(self.entities)
+
+    @property
+    def path_count(self) -> int:
+        """Get number of paths in result"""
+        return len(self.paths)
+
+    def get_top_entities(self, n: int = 5) -> List[Entity]:
+        """
+        Get top N entities by score
+
+        Args:
+            n: Number of entities to return
+
+        Returns:
+            Top N entities
+        """
+        if not self.scores:
+            return self.entities[:n]
+
+        # Sort by score (descending)
+        scored_entities = list(zip(self.entities, self.scores))
+        scored_entities.sort(key=lambda x: x[1], reverse=True)
+        return [entity for entity, _ in scored_entities[:n]]
+
+    def get_entity_ids(self) -> List[str]:
+        """
+        Get list of all entity IDs in result
+
+        Returns:
+            List of entity IDs
+        """
+        return [entity.id for entity in self.entities]
+
+    def __str__(self) -> str:
+        parts = [f"GraphResult(entities={self.entity_count}, paths={self.path_count})"]
+        if self.execution_time_ms:
+            parts.append(f"time={self.execution_time_ms:.2f}ms")
+        return " ".join(parts)
+
+    def __repr__(self) -> str:
+        return (
+            f"GraphResult("
+            f"entities={self.entity_count}, "
+            f"relations={len(self.relations)}, "
+            f"paths={self.path_count}, "
+            f"total_count={self.total_count})"
+        )

aiecs/domain/knowledge_graph/models/query_plan.py

@@ -0,0 +1,187 @@
+"""
+Query Plan Domain Models
+
+Models for query planning, decomposition, and optimization.
+"""
+
+from typing import Any, List, Dict, Set
+from enum import Enum
+from pydantic import BaseModel, Field
+from aiecs.domain.knowledge_graph.models.query import GraphQuery
+
+
+class QueryOperation(str, Enum):
+    """Types of query operations"""
+
+    ENTITY_LOOKUP = "entity_lookup"
+    VECTOR_SEARCH = "vector_search"
+    TRAVERSAL = "traversal"
+    FILTER = "filter"
+    JOIN = "join"
+    AGGREGATE = "aggregate"
+    RANK = "rank"
+
+
+class QueryStep(BaseModel):
+    """
+    A single step in a query execution plan
+
+    Represents one operation in a multi-step query plan.
+
+    Attributes:
+        step_id: Unique identifier for this step
+        operation: Type of operation to perform
+        query: Graph query specification for this step
+        depends_on: List of step IDs that must complete before this step
+        description: Human-readable description of what this step does
+        estimated_cost: Estimated computational cost (0-1, higher = more expensive)
+
+    Example:
+        ```python
+        step = QueryStep(
+            step_id="step_1",
+            operation=QueryOperation.VECTOR_SEARCH,
+            query=GraphQuery(
+                query_type=QueryType.VECTOR_SEARCH,
+                embedding=[0.1, 0.2, ...],
+                max_results=10
+            ),
+            description="Find semantically similar entities"
+        )
+        ```
+    """
+
+    step_id: str = Field(..., description="Unique identifier for this step")
+
+    operation: QueryOperation = Field(..., description="Type of operation to perform")
+
+    query: GraphQuery = Field(..., description="Graph query specification for this step")
+
+    depends_on: List[str] = Field(
+        default_factory=list,
+        description="List of step IDs that must complete before this step",
+    )
+
+    description: str = Field(..., description="Human-readable description of what this step does")
+
+    estimated_cost: float = Field(
+        default=0.5,
+        ge=0.0,
+        le=1.0,
+        description="Estimated computational cost (0-1, higher = more expensive)",
+    )
+
+    metadata: Dict[str, Any] = Field(
+        default_factory=dict, description="Additional metadata for this step"
+    )
+
+
+class QueryPlan(BaseModel):
+    """
+    Query Execution Plan
+
+    Represents a structured plan for executing a complex query as a series of steps.
+
+    Attributes:
+        plan_id: Unique identifier for this plan
+        original_query: The original natural language or complex query
+        steps: List of query steps to execute in order
+        total_estimated_cost: Total estimated cost of executing this plan
+        optimized: Whether this plan has been optimized
+        explanation: Human-readable explanation of the plan
+
+    Example:
+        ```python
+        plan = QueryPlan(
+            plan_id="plan_001",
+            original_query="Who works at companies that Alice knows people at?",
+            steps=[
+                QueryStep(step_id="step_1", ...),  # Find Alice
+                QueryStep(step_id="step_2", ...),  # Find people Alice knows
+                QueryStep(step_id="step_3", ...)   # Find their companies
+            ],
+            explanation="Multi-hop query to find companies through social connections"
+        )
+        ```
+    """
+
+    plan_id: str = Field(..., description="Unique identifier for this plan")
+
+    original_query: str = Field(..., description="The original natural language or complex query")
+
+    steps: List[QueryStep] = Field(..., description="List of query steps to execute in order")
+
+    total_estimated_cost: float = Field(
+        default=0.0, description="Total estimated cost of executing this plan"
+    )
+
+    optimized: bool = Field(default=False, description="Whether this plan has been optimized")
+
+    explanation: str = Field(default="", description="Human-readable explanation of the plan")
+
+    metadata: Dict[str, Any] = Field(
+        default_factory=dict, description="Additional metadata for this plan"
+    )
+
+    def calculate_total_cost(self) -> float:
+        """Calculate total estimated cost from all steps"""
+        return sum(step.estimated_cost for step in self.steps)
+
+    def get_executable_steps(self, completed_step_ids: set[str]) -> List[QueryStep]:
+        """
+        Get steps that can be executed given completed steps
+
+        Args:
+            completed_step_ids: Set of step IDs that have been completed
+
+        Returns:
+            List of steps whose dependencies are all satisfied
+        """
+        executable = []
+        for step in self.steps:
+            if step.step_id in completed_step_ids:
+                continue
+            if all(dep in completed_step_ids for dep in step.depends_on):
+                executable.append(step)
+        return executable
+
+    def get_execution_order(self) -> List[List[str]]:
+        """
+        Get optimal execution order for steps
+
+        Returns:
+            List of lists, where each inner list contains step IDs that can run in parallel
+
+        Example:
+            [[step_1], [step_2, step_3], [step_4]]
+            - step_1 runs first
+            - step_2 and step_3 can run in parallel after step_1
+            - step_4 runs after both step_2 and step_3 complete
+        """
+        completed: Set[str] = set()
+        execution_order: List[List[str]] = []
+
+        while len(completed) < len(self.steps):
+            # Get all steps that can run now
+            current_batch = []
+            for step in self.steps:
+                if step.step_id not in completed:
+                    if all(dep in completed for dep in step.depends_on):
+                        current_batch.append(step.step_id)
+
+            if not current_batch:
+                # Should not happen if dependencies are valid
+                break
+
+            execution_order.append(current_batch)
+            completed.update(current_batch)
+
+        return execution_order
+
+
+class OptimizationStrategy(str, Enum):
+    """Query optimization strategies"""
+
+    MINIMIZE_COST = "minimize_cost"  # Reorder to minimize total cost
+    MINIMIZE_LATENCY = "minimize_latency"  # Maximize parallelization
+    BALANCED = "balanced"  # Balance cost and latency