search.py

from __future__ import annotations

import json
import uuid
from dataclasses import dataclass, field
from typing import Any, Optional

from dnotifier import DNotifier


@dataclass
class KnowledgeDoc:
    title: str
    content: str
    type: str = "docs"
    metadata: dict[str, Any] = field(default_factory=dict)
    record_id: Optional[str] = None


async def add_knowledge_base(
    notifier: DNotifier, sender_id: str, doc: KnowledgeDoc
) -> Any:
    record_id = (doc.record_id or "").strip() or str(uuid.uuid4())
    content = f"# {doc.title}\n\n{doc.content}".strip()
    return await notifier.add_document(
        sender_id=sender_id,
        record_id=record_id,
        content=content,
        doc_type=doc.type or "docs",
        metadata={
            "title": doc.title,
            "source": "semantic-search",
            **doc.metadata,
        },
    )


SAMPLE_SEARCH_DOCS = [
    KnowledgeDoc(
        record_id="search-vector-overview",
        title="Vector search overview",
        type="guide",
        content=(
            "Semantic search finds documents by meaning, not exact keywords. Embeddings map "
            "text into vectors; similar meaning yields higher similarity scores. Use "
            "minSimilarity to filter weak matches."
        ),
    ),
    KnowledgeDoc(
        record_id="search-indexing",
        title="Indexing documents",
        type="howto",
        content=(
            "Index with addDocument. Provide a stable recordId when you may update the same "
            "doc later. Keep documents focused on one topic for better retrieval precision."
        ),
    ),
    KnowledgeDoc(
        record_id="search-filters",
        title="Search parameters",
        type="guide",
        content=(
            "search accepts query, limit, minSimilarity, and filterbySource. Start with "
            "limit 5 and minSimilarity 0.2, then tighten. filterbySource matches document "
            "metadata.source when set."
        ),
    ),
    KnowledgeDoc(
        record_id="search-vs-rag",
        title="Search vs RAG chat",
        type="faq",
        content=(
            "Semantic search returns ranked chunks. RAG chat (sendAI with useKnowledgeBase) "
            "answers in natural language using retrieved context. Use search when you need "
            "raw hits; use RAG when you need an answer."
        ),
    ),
]


async def seed_search_docs(notifier: DNotifier, sender_id: str) -> None:
    for doc in SAMPLE_SEARCH_DOCS:
        result = await add_knowledge_base(notifier, sender_id, doc)
        print(f"Indexed: {doc.title}", json.dumps(result, default=str))


async def semantic_search(
    notifier: DNotifier,
    sender_id: str,
    query: str,
    *,
    limit: int = 5,
    min_similarity: float = 0.2,
    filterby_source: Optional[str] = None,
) -> Any:
    return await notifier.search(
        sender_id=sender_id,
        query=query,
        limit=limit,
        min_similarity=min_similarity,
        filterby_source=filterby_source,
    )


async def list_docs(notifier: DNotifier, sender_id: str) -> Any:
    return await notifier.list_documents(sender_id=sender_id, limit=50)