kb/mcp/engine.py

"""HTTP client for the kb engine API."""

import httpx

from config import KB_ENGINE_URL, KB_API_KEY


def _auth_headers() -> dict[str, str]:
    h: dict[str, str] = {}
    if KB_API_KEY:
        h["Authorization"] = f"Bearer {KB_API_KEY}"
    return h


def _client() -> httpx.Client:
    return httpx.Client(base_url=KB_ENGINE_URL, headers=_auth_headers(), timeout=60.0)


def search(query: str, top: int = 10, tags: list[str] | None = None,
           doc_type: str | None = None, fts_only: bool = False,
           vec_only: bool = False, threshold: float | None = None,
           explain: bool = False, rerank: bool | None = None) -> dict:
    body: dict = {"query": query, "top": top}
    if tags:
        body["tags"] = tags
    if doc_type:
        body["doc_type"] = doc_type
    if fts_only:
        body["fts_only"] = True
    if vec_only:
        body["vec_only"] = True
    if threshold is not None:
        body["threshold"] = threshold
    if explain:
        body["explain"] = True
    if rerank is not None:
        body["rerank"] = rerank
    with _client() as c:
        r = c.post("/api/v1/search", json=body)
        r.raise_for_status()
        return r.json()


def add_note(text: str, tags: list[str] | None = None,
             title: str | None = None) -> dict:
    fields = {"note": text}
    if tags:
        fields["tags"] = ",".join(tags)
    if title:
        fields["title"] = title
    with _client() as c:
        r = c.post("/api/v1/jobs", data=fields)
        r.raise_for_status()
        return r.json()


def update_note(doc_id: int, text: str) -> dict:
    with _client() as c:
        r = c.patch(f"/api/v1/notes/{doc_id}", json={"text": text})
        r.raise_for_status()
        return r.json()


def get_document(doc_id: int) -> dict:
    with _client() as c:
        r = c.get(f"/api/v1/documents/{doc_id}")
        r.raise_for_status()
        return r.json()


def list_documents(doc_type: str | None = None,
                   tags: str | None = None) -> list[dict]:
    params: dict = {}
    if doc_type:
        params["type"] = doc_type
    if tags:
        params["tags"] = tags
    with _client() as c:
        r = c.get("/api/v1/documents", params=params)
        r.raise_for_status()
        return r.json()


def get_status() -> dict:
    with _client() as c:
        r = c.get("/api/v1/status")
        r.raise_for_status()
        return r.json()


def list_jobs(status: str | None = None) -> list[dict]:
    params: dict = {}
    if status:
        params["status"] = status
    with _client() as c:
        r = c.get("/api/v1/jobs", params=params)
        r.raise_for_status()
        return r.json()


def update_tags(doc_id: int, add: list[str] | None = None,
                remove: list[str] | None = None) -> dict:
    body: dict = {}
    if add:
        body["add"] = add
    if remove:
        body["remove"] = remove
    with _client() as c:
        r = c.put(f"/api/v1/documents/{doc_id}/tags", json=body)
        r.raise_for_status()
        return r.json()


def delete_document(doc_id: int) -> dict:
    with _client() as c:
        r = c.delete(f"/api/v1/documents/{doc_id}")
        r.raise_for_status()
        return r.json()


def _bulk_body(
    document_ids: list[int] | None = None,
    tags: list[str] | None = None,
    doc_type: str | None = None,
    from_id: int | None = None,
    to_id: int | None = None,
    force: bool = False,
    **extra,
) -> dict:
    body: dict = {}
    if document_ids:
        body["document_ids"] = document_ids
    if tags:
        body["tags"] = tags
    if doc_type:
        body["doc_type"] = doc_type
    if from_id is not None:
        body["from_id"] = from_id
    if to_id is not None:
        body["to_id"] = to_id
    if force:
        body["force"] = True
    body.update(extra)
    return body


def bulk_delete(
    document_ids: list[int] | None = None,
    tags: list[str] | None = None,
    doc_type: str | None = None,
    from_id: int | None = None,
    to_id: int | None = None,
    force: bool = False,
) -> dict:
    body = _bulk_body(document_ids, tags, doc_type, from_id, to_id, force)
    with _client() as c:
        r = c.post("/api/v1/bulk/delete", json=body)
        r.raise_for_status()
        return r.json()


def bulk_tags(
    document_ids: list[int] | None = None,
    tags: list[str] | None = None,
    doc_type: str | None = None,
    from_id: int | None = None,
    to_id: int | None = None,
    add: list[str] | None = None,
    remove: list[str] | None = None,
    force: bool = False,
) -> dict:
    extra = {}
    if add:
        extra["add"] = add
    if remove:
        extra["remove"] = remove
    body = _bulk_body(document_ids, tags, doc_type, from_id, to_id, force, **extra)
    with _client() as c:
        r = c.post("/api/v1/bulk/tags", json=body)
        r.raise_for_status()
        return r.json()


def bulk_set_tags(
    document_ids: list[int] | None = None,
    tags: list[str] | None = None,
    doc_type: str | None = None,
    from_id: int | None = None,
    to_id: int | None = None,
    new_tags: list[str] | None = None,
    force: bool = False,
) -> dict:
    extra = {"new_tags": new_tags or []}
    body = _bulk_body(document_ids, tags, doc_type, from_id, to_id, force, **extra)
    with _client() as c:
        r = c.post("/api/v1/bulk/set-tags", json=body)
        r.raise_for_status()
        return r.json()


def upload_file(filename: str, file_bytes: bytes,
                tags: list[str] | None = None) -> dict:
    fields: dict = {}
    if tags:
        fields["tags"] = ",".join(tags)
    with _client() as c:
        r = c.post(
            "/api/v1/jobs",
            data=fields,
            files={"file": (filename, file_bytes)},
        )
        r.raise_for_status()
        return r.json()