"""RAG (Retrieval-Augmented Generation) system for Logos ADK."""
from __future__ import annotations

from dataclasses import dataclass, field
from typing import Any
from uuid import uuid4


class DocumentLoadError(Exception):
    """Raised when a document cannot be loaded."""


@dataclass
class Document:
    """A loaded document before chunking."""

    content: str
    metadata: dict[str, Any] = field(default_factory=dict)
    id: str = field(default_factory=lambda: uuid4().hex)


@dataclass
class Chunk:
    """A piece of a document with optional embedding."""

    content: str
    metadata: dict[str, Any]
    document_id: str
    embedding: list[float] | None = None
    id: str = field(default_factory=lambda: uuid4().hex)


from logos_adk.rag.knowledge_base import KnowledgeBase  # noqa: E402
from logos_adk.rag.loaders import (  # noqa: E402
    CSVLoader,
    DocumentLoader,
    HTMLLoader,
    JSONLoader,
    MarkdownLoader,
    TextLoader,
    URLLoader,
)
from logos_adk.rag.memory import RAGMemory  # noqa: E402
from logos_adk.rag.splitter import RecursiveCharacterSplitter, TextSplitter  # noqa: E402
from logos_adk.rag.tool import rag_tool  # noqa: E402
from logos_adk.rag.vector_store import InMemoryVectorStore, VectorStore  # noqa: E402

try:  # noqa: E402
    from logos_adk.rag.vector_postgresql import PostgreSQLVectorStore as PostgreSQLVectorStore
except ImportError:
    pass

__all__ = [
    "Chunk",
    "Document",
    "CSVLoader",
    "DocumentLoadError",
    "DocumentLoader",
    "HTMLLoader",
    "InMemoryVectorStore",
    "JSONLoader",
    "KnowledgeBase",
    "MarkdownLoader",
    "PostgreSQLVectorStore",
    "RAGMemory",
    "RecursiveCharacterSplitter",
    "TextLoader",
    "TextSplitter",
    "URLLoader",
    "VectorStore",
    "rag_tool",
]
