60 lines
1.9 KiB
Python
60 lines
1.9 KiB
Python
|
"""Pydantic schemas for the GraphRAG plugin."""
|
||
|
|
|
||
|
|
from __future__ import annotations
|
||
|
|
|
||
|
|
from datetime import datetime
|
||
|
|
from typing import Any
|
||
|
|
|
||
|
|
from pydantic import BaseModel, Field
|
||
|
|
|
||
|
|
|
||
|
|
class RelationshipCreate(BaseModel):
|
||
|
|
source_type: str = Field(..., max_length=50, description="Entity type of the source (e.g. 'contact', 'company')")
|
||
|
|
source_id: str = Field(..., description="UUID of the source entity")
|
||
|
|
target_type: str = Field(..., max_length=50, description="Entity type of the target (e.g. 'email', 'task')")
|
||
|
|
target_id: str = Field(..., description="UUID of the target entity")
|
||
|
|
relationship_type: str = Field(..., max_length=50, description="Type of relationship (e.g. 'works_for', 'has_email')")
|
||
|
|
metadata: dict[str, Any] | None = Field(None, description="Arbitrary metadata about the relationship")
|
||
|
|
|
||
|
|
|
||
|
|
class RelationshipRead(BaseModel):
|
||
|
|
id: str
|
||
|
|
source_type: str
|
||
|
|
source_id: str
|
||
|
|
target_type: str
|
||
|
|
target_id: str
|
||
|
|
relationship_type: str
|
||
|
|
metadata: dict[str, Any] | None = None
|
||
|
|
owner_id: str | None = None
|
||
|
|
created_at: datetime | None = None
|
||
|
|
|
||
|
|
|
||
|
|
class TraverseRequest(BaseModel):
|
||
|
|
source_type: str = Field(..., max_length=50, description="Entity type of the starting node")
|
||
|
|
source_id: str = Field(..., description="UUID of the starting node")
|
||
|
|
max_hops: int = Field(3, ge=1, le=10, description="Maximum traversal depth")
|
||
|
|
relationship_types: list[str] | None = Field(None, description="Optional filter by relationship types")
|
||
|
|
|
||
|
|
|
||
|
|
class GraphNode(BaseModel):
|
||
|
|
entity_type: str
|
||
|
|
entity_id: str
|
||
|
|
depth: int
|
||
|
|
path: list[str] = []
|
||
|
|
|
||
|
|
|
||
|
|
class GraphEdge(BaseModel):
|
||
|
|
source_type: str
|
||
|
|
source_id: str
|
||
|
|
target_type: str
|
||
|
|
target_id: str
|
||
|
|
relationship_type: str
|
||
|
|
metadata: dict[str, Any] | None = None
|
||
|
|
|
||
|
|
|
||
|
|
class TraverseResponse(BaseModel):
|
||
|
|
nodes: list[GraphNode]
|
||
|
|
edges: list[GraphEdge]
|
||
|
|
total_nodes: int
|
||
|
|
total_edges: int
|