app/services/ingestion.py

from pathlib import Path

from app.core.config import settings
from app.core.models import Document, IngestionResult, SourceType
from app.extractors.arxiv import extract_arxiv
from app.extractors.medium import extract_medium
from app.extractors.pdf import extract_pdf
from app.services.chat import NvidiaChatClient
from app.services.chunking import chunk_document
from app.services.embeddings import get_embedding_client
from app.services.vector_store import QdrantVectorStore
from app.utils.source_detection import detect_source


EXPORT_DIR = Path("data/exports")


def extract_document(
    url: str | None = None,
    pdf_path: str | None = None,
) -> Document:
    source_type = detect_source(url, pdf_path)
    if source_type == SourceType.PDF:
        return extract_pdf(str(pdf_path))
    if source_type == SourceType.ARXIV:
        return extract_arxiv(str(url))
    if source_type == SourceType.MEDIUM:
        return extract_medium(str(url))
    raise ValueError(f"Unsupported source type: {source_type}")


def save_markdown(document: Document, chunks_count: int) -> Path:
    EXPORT_DIR.mkdir(parents=True, exist_ok=True)
    safe_title = "".join(char if char.isalnum() or char in "-_" else "_" for char in document.title)[:80]
    path = EXPORT_DIR / f"{safe_title or document.source_type.value}.md"
    metadata_lines = "\n".join(f"- {key}: {value}" for key, value in document.metadata.items())
    path.write_text(
        "\n".join(
            [
                f"# {document.title}",
                "",
                f"- Source type: {document.source_type.value}",
                f"- Source: {document.source}",
                f"- Chunks uploaded: {chunks_count}",
                metadata_lines,
                "",
                "## Extracted Text",
                "",
                document.text,
            ]
        ),
        encoding="utf-8",
    )
    return path


def ingest_source(
    url: str | None,
    pdf_path: str | None,
    chunk_size: int | None = None,
    chunk_overlap: int | None = None,
    collection_name: str | None = None,
) -> IngestionResult:
    document = extract_document(url=url, pdf_path=pdf_path)
    chunks = chunk_document(
        document,
        chunk_size=chunk_size or settings.CHUNK_SIZE,
        overlap=chunk_overlap or settings.CHUNK_OVERLAP,
    )
    embeddings = get_embedding_client().embed_texts([chunk.text for chunk in chunks])
    store = QdrantVectorStore(collection_name=collection_name)
    store.upsert_chunks(chunks, embeddings)
    export_path = save_markdown(document, len(chunks))
    return IngestionResult(
        document=document,
        chunks=chunks,
        collection_name=store.collection_name,
        export_path=export_path,
    )


def search_knowledge_base(query: str, limit: int = 5, collection_name: str | None = None):
    query_text = query.strip()
    if not query_text:
        raise ValueError("Enter a query to search.")
    embedding = get_embedding_client().embed_texts([query_text])[0]
    return QdrantVectorStore(collection_name=collection_name).search(embedding, limit=limit)


def answer_from_knowledge_base(query: str, limit: int = 5, collection_name: str | None = None):
    results = search_knowledge_base(query, limit=limit, collection_name=collection_name)
    return NvidiaChatClient().answer_with_context(query, results)