from typing import Any, Dict, Optional

from pydantic import BaseModel

from solutions.graph.models import GraphId


class RunCypherParams(BaseModel):
    graph_id: GraphId
    query: str
    params: Optional[Dict[str, Any]] = None


class RunLlmCypherParams(BaseModel):
    graph_id: GraphId
    query: str


class ComputeAdjacentNodeGroupInfoParams(BaseModel):
    node_group: str
    node_id: Any


class GetAdjacentNodesParams(BaseModel):
    source_node_group: str
    target_node_group: Optional[str] = None
    node_id: Any
