antfly-sdk 0.0.2.dev2__py3-none-any.whl

antfly/client_generated/models/graph_query.py

@@ -0,0 +1,192 @@
+from __future__ import annotations
+
+from collections.abc import Mapping
+from typing import TYPE_CHECKING, Any, TypeVar, cast
+
+from attrs import define as _attrs_define
+from attrs import field as _attrs_field
+
+from ..models.graph_query_type import GraphQueryType
+from ..types import UNSET, Unset
+
+if TYPE_CHECKING:
+    from ..models.graph_node_selector import GraphNodeSelector
+    from ..models.graph_query_params import GraphQueryParams
+    from ..models.pattern_step import PatternStep
+
+
+T = TypeVar("T", bound="GraphQuery")
+
+
+@_attrs_define
+class GraphQuery:
+    """Declarative graph query to execute after full-text/vector searches
+
+    Attributes:
+        type_ (GraphQueryType): Type of graph query to execute
+        index_name (str): Graph index name (must be graph type)
+        start_nodes (GraphNodeSelector | Unset): Defines how to select start/target nodes for graph queries
+        target_nodes (GraphNodeSelector | Unset): Defines how to select start/target nodes for graph queries
+        params (GraphQueryParams | Unset): Parameters for graph traversal and pathfinding
+        pattern (list[PatternStep] | Unset): Pattern steps for pattern query type
+        return_aliases (list[str] | Unset): Which aliases to return from pattern query (empty = all)
+        include_documents (bool | Unset): Fetch full documents for graph results
+        include_edges (bool | Unset): Include edge details for each node
+        fields (list[str] | Unset): Which fields to return from documents
+    """
+
+    type_: GraphQueryType
+    index_name: str
+    start_nodes: GraphNodeSelector | Unset = UNSET
+    target_nodes: GraphNodeSelector | Unset = UNSET
+    params: GraphQueryParams | Unset = UNSET
+    pattern: list[PatternStep] | Unset = UNSET
+    return_aliases: list[str] | Unset = UNSET
+    include_documents: bool | Unset = UNSET
+    include_edges: bool | Unset = UNSET
+    fields: list[str] | Unset = UNSET
+    additional_properties: dict[str, Any] = _attrs_field(init=False, factory=dict)
+
+    def to_dict(self) -> dict[str, Any]:
+        type_ = self.type_.value
+
+        index_name = self.index_name
+
+        start_nodes: dict[str, Any] | Unset = UNSET
+        if not isinstance(self.start_nodes, Unset):
+            start_nodes = self.start_nodes.to_dict()
+
+        target_nodes: dict[str, Any] | Unset = UNSET
+        if not isinstance(self.target_nodes, Unset):
+            target_nodes = self.target_nodes.to_dict()
+
+        params: dict[str, Any] | Unset = UNSET
+        if not isinstance(self.params, Unset):
+            params = self.params.to_dict()
+
+        pattern: list[dict[str, Any]] | Unset = UNSET
+        if not isinstance(self.pattern, Unset):
+            pattern = []
+            for pattern_item_data in self.pattern:
+                pattern_item = pattern_item_data.to_dict()
+                pattern.append(pattern_item)
+
+        return_aliases: list[str] | Unset = UNSET
+        if not isinstance(self.return_aliases, Unset):
+            return_aliases = self.return_aliases
+
+        include_documents = self.include_documents
+
+        include_edges = self.include_edges
+
+        fields: list[str] | Unset = UNSET
+        if not isinstance(self.fields, Unset):
+            fields = self.fields
+
+        field_dict: dict[str, Any] = {}
+        field_dict.update(self.additional_properties)
+        field_dict.update(
+            {
+                "type": type_,
+                "index_name": index_name,
+            }
+        )
+        if start_nodes is not UNSET:
+            field_dict["start_nodes"] = start_nodes
+        if target_nodes is not UNSET:
+            field_dict["target_nodes"] = target_nodes
+        if params is not UNSET:
+            field_dict["params"] = params
+        if pattern is not UNSET:
+            field_dict["pattern"] = pattern
+        if return_aliases is not UNSET:
+            field_dict["return_aliases"] = return_aliases
+        if include_documents is not UNSET:
+            field_dict["include_documents"] = include_documents
+        if include_edges is not UNSET:
+            field_dict["include_edges"] = include_edges
+        if fields is not UNSET:
+            field_dict["fields"] = fields
+
+        return field_dict
+
+    @classmethod
+    def from_dict(cls: type[T], src_dict: Mapping[str, Any]) -> T:
+        from ..models.graph_node_selector import GraphNodeSelector
+        from ..models.graph_query_params import GraphQueryParams
+        from ..models.pattern_step import PatternStep
+
+        d = dict(src_dict)
+        type_ = GraphQueryType(d.pop("type"))
+
+        index_name = d.pop("index_name")
+
+        _start_nodes = d.pop("start_nodes", UNSET)
+        start_nodes: GraphNodeSelector | Unset
+        if isinstance(_start_nodes, Unset):
+            start_nodes = UNSET
+        else:
+            start_nodes = GraphNodeSelector.from_dict(_start_nodes)
+
+        _target_nodes = d.pop("target_nodes", UNSET)
+        target_nodes: GraphNodeSelector | Unset
+        if isinstance(_target_nodes, Unset):
+            target_nodes = UNSET
+        else:
+            target_nodes = GraphNodeSelector.from_dict(_target_nodes)
+
+        _params = d.pop("params", UNSET)
+        params: GraphQueryParams | Unset
+        if isinstance(_params, Unset):
+            params = UNSET
+        else:
+            params = GraphQueryParams.from_dict(_params)
+
+        _pattern = d.pop("pattern", UNSET)
+        pattern: list[PatternStep] | Unset = UNSET
+        if _pattern is not UNSET:
+            pattern = []
+            for pattern_item_data in _pattern:
+                pattern_item = PatternStep.from_dict(pattern_item_data)
+
+                pattern.append(pattern_item)
+
+        return_aliases = cast(list[str], d.pop("return_aliases", UNSET))
+
+        include_documents = d.pop("include_documents", UNSET)
+
+        include_edges = d.pop("include_edges", UNSET)
+
+        fields = cast(list[str], d.pop("fields", UNSET))
+
+        graph_query = cls(
+            type_=type_,
+            index_name=index_name,
+            start_nodes=start_nodes,
+            target_nodes=target_nodes,
+            params=params,
+            pattern=pattern,
+            return_aliases=return_aliases,
+            include_documents=include_documents,
+            include_edges=include_edges,
+            fields=fields,
+        )
+
+        graph_query.additional_properties = d
+        return graph_query
+
+    @property
+    def additional_keys(self) -> list[str]:
+        return list(self.additional_properties.keys())
+
+    def __getitem__(self, key: str) -> Any:
+        return self.additional_properties[key]
+
+    def __setitem__(self, key: str, value: Any) -> None:
+        self.additional_properties[key] = value
+
+    def __delitem__(self, key: str) -> None:
+        del self.additional_properties[key]
+
+    def __contains__(self, key: str) -> bool:
+        return key in self.additional_properties

antfly/client_generated/models/graph_query_params.py

@@ -0,0 +1,216 @@
+from __future__ import annotations
+
+from collections.abc import Mapping
+from typing import TYPE_CHECKING, Any, TypeVar, cast
+
+from attrs import define as _attrs_define
+from attrs import field as _attrs_field
+
+from ..models.edge_direction import EdgeDirection
+from ..models.path_weight_mode import PathWeightMode
+from ..types import UNSET, Unset
+
+if TYPE_CHECKING:
+    from ..models.graph_query_params_algorithm_params import GraphQueryParamsAlgorithmParams
+    from ..models.node_filter import NodeFilter
+
+
+T = TypeVar("T", bound="GraphQueryParams")
+
+
+@_attrs_define
+class GraphQueryParams:
+    """Parameters for graph traversal and pathfinding
+
+    Attributes:
+        edge_types (list[str] | Unset): Filter by edge types
+        direction (EdgeDirection | Unset): Direction of edges to query:
+            - out: Outgoing edges from the node
+            - in: Incoming edges to the node
+            - both: Both outgoing and incoming edges
+        max_depth (int | Unset): Maximum traversal depth
+        min_weight (float | Unset): Minimum edge weight
+        max_weight (float | Unset): Maximum edge weight
+        max_results (int | Unset): Maximum number of results (traversal)
+        deduplicate_nodes (bool | Unset): Remove duplicate nodes (traversal)
+        include_paths (bool | Unset): Include path information (traversal)
+        weight_mode (PathWeightMode | Unset): Path weighting algorithm for pathfinding:
+            - min_hops: Minimize number of edges
+            - min_weight: Minimize sum of edge weights
+            - max_weight: Maximize product of edge weights
+        k (int | Unset): Number of paths to find (k-shortest-paths)
+        node_filter (NodeFilter | Unset): Filter nodes during graph traversal using existing query primitives
+        algorithm (str | Unset): Graph algorithm to run (e.g., 'pagerank', 'betweenness')
+        algorithm_params (GraphQueryParamsAlgorithmParams | Unset): Parameters for the graph algorithm
+    """
+
+    edge_types: list[str] | Unset = UNSET
+    direction: EdgeDirection | Unset = UNSET
+    max_depth: int | Unset = UNSET
+    min_weight: float | Unset = UNSET
+    max_weight: float | Unset = UNSET
+    max_results: int | Unset = UNSET
+    deduplicate_nodes: bool | Unset = UNSET
+    include_paths: bool | Unset = UNSET
+    weight_mode: PathWeightMode | Unset = UNSET
+    k: int | Unset = UNSET
+    node_filter: NodeFilter | Unset = UNSET
+    algorithm: str | Unset = UNSET
+    algorithm_params: GraphQueryParamsAlgorithmParams | Unset = UNSET
+    additional_properties: dict[str, Any] = _attrs_field(init=False, factory=dict)
+
+    def to_dict(self) -> dict[str, Any]:
+        edge_types: list[str] | Unset = UNSET
+        if not isinstance(self.edge_types, Unset):
+            edge_types = self.edge_types
+
+        direction: str | Unset = UNSET
+        if not isinstance(self.direction, Unset):
+            direction = self.direction.value
+
+        max_depth = self.max_depth
+
+        min_weight = self.min_weight
+
+        max_weight = self.max_weight
+
+        max_results = self.max_results
+
+        deduplicate_nodes = self.deduplicate_nodes
+
+        include_paths = self.include_paths
+
+        weight_mode: str | Unset = UNSET
+        if not isinstance(self.weight_mode, Unset):
+            weight_mode = self.weight_mode.value
+
+        k = self.k
+
+        node_filter: dict[str, Any] | Unset = UNSET
+        if not isinstance(self.node_filter, Unset):
+            node_filter = self.node_filter.to_dict()
+
+        algorithm = self.algorithm
+
+        algorithm_params: dict[str, Any] | Unset = UNSET
+        if not isinstance(self.algorithm_params, Unset):
+            algorithm_params = self.algorithm_params.to_dict()
+
+        field_dict: dict[str, Any] = {}
+        field_dict.update(self.additional_properties)
+        field_dict.update({})
+        if edge_types is not UNSET:
+            field_dict["edge_types"] = edge_types
+        if direction is not UNSET:
+            field_dict["direction"] = direction
+        if max_depth is not UNSET:
+            field_dict["max_depth"] = max_depth
+        if min_weight is not UNSET:
+            field_dict["min_weight"] = min_weight
+        if max_weight is not UNSET:
+            field_dict["max_weight"] = max_weight
+        if max_results is not UNSET:
+            field_dict["max_results"] = max_results
+        if deduplicate_nodes is not UNSET:
+            field_dict["deduplicate_nodes"] = deduplicate_nodes
+        if include_paths is not UNSET:
+            field_dict["include_paths"] = include_paths
+        if weight_mode is not UNSET:
+            field_dict["weight_mode"] = weight_mode
+        if k is not UNSET:
+            field_dict["k"] = k
+        if node_filter is not UNSET:
+            field_dict["node_filter"] = node_filter
+        if algorithm is not UNSET:
+            field_dict["algorithm"] = algorithm
+        if algorithm_params is not UNSET:
+            field_dict["algorithm_params"] = algorithm_params
+
+        return field_dict
+
+    @classmethod
+    def from_dict(cls: type[T], src_dict: Mapping[str, Any]) -> T:
+        from ..models.graph_query_params_algorithm_params import GraphQueryParamsAlgorithmParams
+        from ..models.node_filter import NodeFilter
+
+        d = dict(src_dict)
+        edge_types = cast(list[str], d.pop("edge_types", UNSET))
+
+        _direction = d.pop("direction", UNSET)
+        direction: EdgeDirection | Unset
+        if isinstance(_direction, Unset):
+            direction = UNSET
+        else:
+            direction = EdgeDirection(_direction)
+
+        max_depth = d.pop("max_depth", UNSET)
+
+        min_weight = d.pop("min_weight", UNSET)
+
+        max_weight = d.pop("max_weight", UNSET)
+
+        max_results = d.pop("max_results", UNSET)
+
+        deduplicate_nodes = d.pop("deduplicate_nodes", UNSET)
+
+        include_paths = d.pop("include_paths", UNSET)
+
+        _weight_mode = d.pop("weight_mode", UNSET)
+        weight_mode: PathWeightMode | Unset
+        if isinstance(_weight_mode, Unset):
+            weight_mode = UNSET
+        else:
+            weight_mode = PathWeightMode(_weight_mode)
+
+        k = d.pop("k", UNSET)
+
+        _node_filter = d.pop("node_filter", UNSET)
+        node_filter: NodeFilter | Unset
+        if isinstance(_node_filter, Unset):
+            node_filter = UNSET
+        else:
+            node_filter = NodeFilter.from_dict(_node_filter)
+
+        algorithm = d.pop("algorithm", UNSET)
+
+        _algorithm_params = d.pop("algorithm_params", UNSET)
+        algorithm_params: GraphQueryParamsAlgorithmParams | Unset
+        if isinstance(_algorithm_params, Unset):
+            algorithm_params = UNSET
+        else:
+            algorithm_params = GraphQueryParamsAlgorithmParams.from_dict(_algorithm_params)
+
+        graph_query_params = cls(
+            edge_types=edge_types,
+            direction=direction,
+            max_depth=max_depth,
+            min_weight=min_weight,
+            max_weight=max_weight,
+            max_results=max_results,
+            deduplicate_nodes=deduplicate_nodes,
+            include_paths=include_paths,
+            weight_mode=weight_mode,
+            k=k,
+            node_filter=node_filter,
+            algorithm=algorithm,
+            algorithm_params=algorithm_params,
+        )
+
+        graph_query_params.additional_properties = d
+        return graph_query_params
+
+    @property
+    def additional_keys(self) -> list[str]:
+        return list(self.additional_properties.keys())
+
+    def __getitem__(self, key: str) -> Any:
+        return self.additional_properties[key]
+
+    def __setitem__(self, key: str, value: Any) -> None:
+        self.additional_properties[key] = value
+
+    def __delitem__(self, key: str) -> None:
+        del self.additional_properties[key]
+
+    def __contains__(self, key: str) -> bool:
+        return key in self.additional_properties

antfly/client_generated/models/graph_query_params_algorithm_params.py

@@ -0,0 +1,47 @@
+from __future__ import annotations
+
+from collections.abc import Mapping
+from typing import Any, TypeVar
+
+from attrs import define as _attrs_define
+from attrs import field as _attrs_field
+
+T = TypeVar("T", bound="GraphQueryParamsAlgorithmParams")
+
+
+@_attrs_define
+class GraphQueryParamsAlgorithmParams:
+    """Parameters for the graph algorithm"""
+
+    additional_properties: dict[str, Any] = _attrs_field(init=False, factory=dict)
+
+    def to_dict(self) -> dict[str, Any]:
+
+        field_dict: dict[str, Any] = {}
+        field_dict.update(self.additional_properties)
+
+        return field_dict
+
+    @classmethod
+    def from_dict(cls: type[T], src_dict: Mapping[str, Any]) -> T:
+        d = dict(src_dict)
+        graph_query_params_algorithm_params = cls()
+
+        graph_query_params_algorithm_params.additional_properties = d
+        return graph_query_params_algorithm_params
+
+    @property
+    def additional_keys(self) -> list[str]:
+        return list(self.additional_properties.keys())
+
+    def __getitem__(self, key: str) -> Any:
+        return self.additional_properties[key]
+
+    def __setitem__(self, key: str, value: Any) -> None:
+        self.additional_properties[key] = value
+
+    def __delitem__(self, key: str) -> None:
+        del self.additional_properties[key]
+
+    def __contains__(self, key: str) -> bool:
+        return key in self.additional_properties

antfly/client_generated/models/graph_query_result.py

@@ -0,0 +1,155 @@
+from __future__ import annotations
+
+from collections.abc import Mapping
+from typing import TYPE_CHECKING, Any, TypeVar
+
+from attrs import define as _attrs_define
+from attrs import field as _attrs_field
+
+from ..models.graph_query_type import GraphQueryType
+from ..types import UNSET, Unset
+
+if TYPE_CHECKING:
+    from ..models.graph_result_node import GraphResultNode
+    from ..models.path import Path
+    from ..models.pattern_match import PatternMatch
+
+
+T = TypeVar("T", bound="GraphQueryResult")
+
+
+@_attrs_define
+class GraphQueryResult:
+    """Results of a graph query
+
+    Attributes:
+        type_ (GraphQueryType): Type of graph query to execute
+        total (int): Total number of results
+        nodes (list[GraphResultNode] | Unset): Result nodes
+        paths (list[Path] | Unset): Result paths (for pathfinding queries)
+        matches (list[PatternMatch] | Unset): Pattern matches (for pattern queries)
+        took (int | Unset): Query execution time
+    """
+
+    type_: GraphQueryType
+    total: int
+    nodes: list[GraphResultNode] | Unset = UNSET
+    paths: list[Path] | Unset = UNSET
+    matches: list[PatternMatch] | Unset = UNSET
+    took: int | Unset = UNSET
+    additional_properties: dict[str, Any] = _attrs_field(init=False, factory=dict)
+
+    def to_dict(self) -> dict[str, Any]:
+        type_ = self.type_.value
+
+        total = self.total
+
+        nodes: list[dict[str, Any]] | Unset = UNSET
+        if not isinstance(self.nodes, Unset):
+            nodes = []
+            for nodes_item_data in self.nodes:
+                nodes_item = nodes_item_data.to_dict()
+                nodes.append(nodes_item)
+
+        paths: list[dict[str, Any]] | Unset = UNSET
+        if not isinstance(self.paths, Unset):
+            paths = []
+            for paths_item_data in self.paths:
+                paths_item = paths_item_data.to_dict()
+                paths.append(paths_item)
+
+        matches: list[dict[str, Any]] | Unset = UNSET
+        if not isinstance(self.matches, Unset):
+            matches = []
+            for matches_item_data in self.matches:
+                matches_item = matches_item_data.to_dict()
+                matches.append(matches_item)
+
+        took = self.took
+
+        field_dict: dict[str, Any] = {}
+        field_dict.update(self.additional_properties)
+        field_dict.update(
+            {
+                "type": type_,
+                "total": total,
+            }
+        )
+        if nodes is not UNSET:
+            field_dict["nodes"] = nodes
+        if paths is not UNSET:
+            field_dict["paths"] = paths
+        if matches is not UNSET:
+            field_dict["matches"] = matches
+        if took is not UNSET:
+            field_dict["took"] = took
+
+        return field_dict
+
+    @classmethod
+    def from_dict(cls: type[T], src_dict: Mapping[str, Any]) -> T:
+        from ..models.graph_result_node import GraphResultNode
+        from ..models.path import Path
+        from ..models.pattern_match import PatternMatch
+
+        d = dict(src_dict)
+        type_ = GraphQueryType(d.pop("type"))
+
+        total = d.pop("total")
+
+        _nodes = d.pop("nodes", UNSET)
+        nodes: list[GraphResultNode] | Unset = UNSET
+        if _nodes is not UNSET:
+            nodes = []
+            for nodes_item_data in _nodes:
+                nodes_item = GraphResultNode.from_dict(nodes_item_data)
+
+                nodes.append(nodes_item)
+
+        _paths = d.pop("paths", UNSET)
+        paths: list[Path] | Unset = UNSET
+        if _paths is not UNSET:
+            paths = []
+            for paths_item_data in _paths:
+                paths_item = Path.from_dict(paths_item_data)
+
+                paths.append(paths_item)
+
+        _matches = d.pop("matches", UNSET)
+        matches: list[PatternMatch] | Unset = UNSET
+        if _matches is not UNSET:
+            matches = []
+            for matches_item_data in _matches:
+                matches_item = PatternMatch.from_dict(matches_item_data)
+
+                matches.append(matches_item)
+
+        took = d.pop("took", UNSET)
+
+        graph_query_result = cls(
+            type_=type_,
+            total=total,
+            nodes=nodes,
+            paths=paths,
+            matches=matches,
+            took=took,
+        )
+
+        graph_query_result.additional_properties = d
+        return graph_query_result
+
+    @property
+    def additional_keys(self) -> list[str]:
+        return list(self.additional_properties.keys())
+
+    def __getitem__(self, key: str) -> Any:
+        return self.additional_properties[key]
+
+    def __setitem__(self, key: str, value: Any) -> None:
+        self.additional_properties[key] = value
+
+    def __delitem__(self, key: str) -> None:
+        del self.additional_properties[key]
+
+    def __contains__(self, key: str) -> bool:
+        return key in self.additional_properties

antfly/client_generated/models/graph_query_type.py

@@ -0,0 +1,12 @@
+from enum import Enum
+
+
+class GraphQueryType(str, Enum):
+    K_SHORTEST_PATHS = "k_shortest_paths"
+    NEIGHBORS = "neighbors"
+    PATTERN = "pattern"
+    SHORTEST_PATH = "shortest_path"
+    TRAVERSE = "traverse"
+
+    def __str__(self) -> str:
+        return str(self.value)