"""Knowledge graph type definitions."""

from __future__ import annotations

from datetime import datetime
from typing import Any

from pydantic import BaseModel, Field


class KnowledgeEntity(BaseModel):
    """An entity in the knowledge graph."""

    id: str
    name: str
    entity_type: str | None = None
    properties: dict[str, Any] | None = None
    created_at: datetime | None = None


class KnowledgeRelationship(BaseModel):
    """A relationship between entities."""

    source_id: str
    target_id: str
    relationship_type: str
    properties: dict[str, Any] | None = None
    created_at: datetime | None = None
    valid_from: datetime | None = None
    valid_to: datetime | None = None


class KnowledgeQueryResponse(BaseModel):
    """Response from a knowledge graph query."""

    entities: list[KnowledgeEntity] = Field(default_factory=list)
    relationships: list[KnowledgeRelationship] = Field(default_factory=list)
    total: int = 0


class KnowledgeIngestResponse(BaseModel):
    """Response from knowledge ingestion."""

    entities_created: int = 0
    relationships_created: int = 0
    status: str = "completed"
