oclaw/runtime/tools/experts/productivity/kb_tools.py

from __future__ import annotations

import hashlib
from typing import Any

from oclaw.platform.config.paths import db_path
from oclaw.platform.embeddings.embedding_client import build_default_embedding_client
from oclaw.platform.persistence.sqlite_store import SqliteStore
from oclaw.runtime.tools.base import ToolSpec


def _chunk_id(source: str, text: str) -> str:
    raw = f"{source}\n{text}".encode("utf-8", errors="ignore")
    return hashlib.sha1(raw).hexdigest()


def kb_add_tool() -> ToolSpec:
    def handler(args: dict[str, Any]) -> dict[str, Any]:
        try:
            tenant_id = str(args.get("tenant_id") or "").strip()
            user_id = str(args.get("user_id") or "").strip()
            text = str(args.get("text") or "").strip()
            title = str(args.get("title") or "").strip()
            if not tenant_id or not user_id or not text:
                return {"ok": False, "error": "tenant_id, user_id, text are required"}
            source = f"builtin:tenant:{tenant_id}:kb"
            if title:
                source = f"{source}:{title[:48]}"
            cid = _chunk_id(source, text)
            store = SqliteStore(db_path())
            store.upsert_knowledge_chunk(
                chunk_id=cid,
                source=source,
                content=text,
                metadata={"tenant_id": tenant_id, "user_id": user_id, "title": title, "source": source},
            )
            client = build_default_embedding_client()
            emb = client.embed(text[:8000])
            store.upsert_knowledge_embedding(chunk_id=cid, model=emb.model, vector=emb.vector)
            return {"ok": True, "chunk_id": cid, "source": source, "embedding_model": emb.model}
        except Exception as e:
            return {"ok": False, "error": f"{type(e).__name__}: {e}"}

    return ToolSpec(
        name="kb_add",
        description="Add a knowledge snippet for a tenant into the vector knowledge base.",
        parameters={
            "type": "object",
            "properties": {"tenant_id": {"type": "string"}, "user_id": {"type": "string"}, "title": {"type": "string", "description": "Optional title/label."}, "text": {"type": "string", "description": "Knowledge content to store."}},
            "required": ["tenant_id", "user_id", "text"],
            "additionalProperties": False,
        },
        handler=handler,
        tags=frozenset({"productivity", "rag", "write"}),
    )


def kb_search_tool() -> ToolSpec:
    def handler(args: dict[str, Any]) -> dict[str, Any]:
        try:
            tenant_id = str(args.get("tenant_id") or "").strip()
            query = str(args.get("query") or "").strip()
            limit = int(args.get("limit") or 3)
            if not tenant_id or not query:
                return {"ok": False, "error": "tenant_id and query are required"}
            store = SqliteStore(db_path())
            from oclaw.runtime.orchestration.memory import retrieve_context

            rows = retrieve_context(store, query, limit=max(1, min(limit, 6)))
            filtered = [r for r in rows if str(r.get("source") or "").startswith(f"builtin:tenant:{tenant_id}:")]
            hits = filtered[: max(1, min(limit, 6))]
            if not hits:
                like_rows = store.search_knowledge(query=query, limit=max(1, min(limit, 6)))
                hits = [r for r in like_rows if str(r.get("source") or "").startswith(f"builtin:tenant:{tenant_id}:")][: max(1, min(limit, 6))]
            refs = []
            for h in hits:
                refs.append({"source": str(h.get("source") or ""), "snippet": str(h.get("content") or "")[:240]})
            return {"ok": True, "hits": refs}
        except Exception as e:
            return {"ok": False, "error": f"{type(e).__name__}: {e}"}

    return ToolSpec(
        name="kb_search",
        description="Search tenant knowledge base and return citations/snippets.",
        parameters={
            "type": "object",
            "properties": {"tenant_id": {"type": "string"}, "query": {"type": "string"}, "limit": {"type": "integer", "default": 3}},
            "required": ["tenant_id", "query"],
            "additionalProperties": False,
        },
        handler=handler,
        tags=frozenset({"productivity", "rag"}),
    )


__all__ = ["kb_add_tool", "kb_search_tool"]