knowledge2 0.4.0__py3-none-any.whl

sdk/examples/quickstart.py

@@ -0,0 +1,36 @@
+"""Minimal example: create client, project, corpus, upload document, build indexes, search."""
+
+from __future__ import annotations
+
+import os
+
+from sdk import Knowledge2, Knowledge2Error
+
+try:
+    api_key = os.environ.get("K2_API_KEY")
+    if not api_key:
+        raise SystemExit("K2_API_KEY is required")
+
+    client = Knowledge2(
+        api_host=os.environ.get("K2_BASE_URL", "https://api.knowledge2.ai"),
+        api_key=api_key,
+    )
+
+    project = client.create_project("quickstart-demo")
+    corpus = client.create_corpus(project["id"], "quickstart-corpus")
+
+    client.upload_document(
+        corpus["id"],
+        raw_text="Knowledge2 enables semantic search over your documents.",
+        source_uri="doc://intro",
+    )
+
+    client.build_indexes(corpus["id"], dense=True, sparse=True, wait=True)
+
+    results = client.search(corpus["id"], "What does Knowledge2 do?", top_k=5)
+    for r in results.get("results", []):
+        print((r.get("text") or "")[:80])
+
+except Knowledge2Error as e:
+    print(f"API error: {e}")
+    raise

sdk/examples/request_options.py

@@ -0,0 +1,44 @@
+"""Per-call request options — override timeout and retry for specific calls.
+
+Usage:
+    K2_API_KEY=... python -m sdk.examples.request_options
+"""
+
+import os
+
+from sdk import ClientTimeouts, Knowledge2, RequestOptions
+
+
+def main() -> None:
+    client = Knowledge2(
+        api_host=os.environ.get("K2_API_HOST", "https://api.knowledge2.ai"),
+        api_key=os.environ["K2_API_KEY"],
+    )
+
+    # Longer timeout for a known-slow bulk operation
+    slow_opts = RequestOptions(
+        timeout=ClientTimeouts(read=300),
+        max_retries=5,
+    )
+
+    # Passthrough tracing headers
+    traced_opts = RequestOptions(
+        passthrough_headers={
+            "X-Request-ID": "example-123",
+            "X-Correlation-ID": "trace-abc",
+        },
+    )
+
+    # Use per-call options
+    corpora = client.list_corpora(request_options=traced_opts)
+    print(f"Found {corpora.total} corpora (with tracing headers)")
+
+    # Raw response access
+    raw = client.with_raw_response.list_corpora()
+    print(f"Status: {raw.status_code}")
+    print(f"Headers: {dict(list(raw.headers.items())[:3])}...")
+    print(f"Parsed: {type(raw.parsed).__name__} with {len(raw.parsed)} items")
+
+
+if __name__ == "__main__":
+    main()

sdk/examples/search.py

@@ -0,0 +1,64 @@
+"""Search examples: basic search, hybrid config, filters, search_generate (RAG)."""
+
+from __future__ import annotations
+
+import os
+
+from sdk import Knowledge2, Knowledge2Error
+from sdk.types.search import SearchHybridConfig
+
+try:
+    api_key = os.environ.get("K2_API_KEY")
+    if not api_key:
+        raise SystemExit("K2_API_KEY is required")
+
+    client = Knowledge2(api_key=api_key)
+    corpus_id = os.environ.get("K2_CORPUS_ID", "corpus-123")
+
+    # Basic search
+    basic = client.search(corpus_id, "machine learning basics", top_k=5)
+    for r in basic.get("results", []):
+        print((r.get("text") or "")[:80])
+
+    # Hybrid search with TypedDict config
+    hybrid_config: SearchHybridConfig = {
+        "enabled": True,
+        "fusion_mode": "rrf",
+        "rrf_k": 60,
+        "dense_weight": 0.6,
+        "sparse_weight": 0.4,
+    }
+    hybrid = client.search(
+        corpus_id,
+        "neural networks",
+        top_k=5,
+        hybrid=hybrid_config,
+        return_config={"include_text": True, "include_scores": True},
+    )
+    for r in hybrid.get("results", []):
+        print(r.get("score"), (r.get("text") or "")[:60])
+
+    # Search with metadata filters
+    filtered = client.search(
+        corpus_id,
+        "API documentation",
+        top_k=5,
+        filters={"topic": "docs", "org": "acme"},
+    )
+    for r in filtered.get("results", []):
+        print(r.get("chunk_id"), (r.get("text") or "")[:60])
+
+    # RAG: search_generate returns answer + retrieval results
+    rag = client.search_generate(
+        corpus_id,
+        "Summarize best practices for hybrid search.",
+        top_k=5,
+        generation={"temperature": 0.3, "max_tokens": 500},
+    )
+    print("Answer:", (rag.get("answer") or "")[:200])
+    for r in rag.get("results", []):
+        print("Source:", (r.get("text") or "")[:80])
+
+except Knowledge2Error as e:
+    print(f"API error: {e}")
+    raise

sdk/integrations/__init__.py

@@ -0,0 +1,57 @@
+"""Framework integrations for the Knowledge2 SDK.
+
+These imports are lazy so users can install only one framework extra.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+__all__ = [
+    "K2LangChainRetriever",
+    "K2LlamaIndexRetriever",
+    "K2LlamaIndexVectorStore",
+    "create_k2_langchain_tools",
+    "create_k2_llamaindex_tools",
+]
+
+if TYPE_CHECKING:  # pragma: no cover
+    # These are resolved lazily at runtime via __getattr__ so optional dependencies
+    # (langchain-core / llama-index-core) are not required for base installs.
+    from sdk.integrations.langchain import K2LangChainRetriever, create_k2_langchain_tools
+    from sdk.integrations.llamaindex import (
+        K2LlamaIndexRetriever,
+        K2LlamaIndexVectorStore,
+        create_k2_llamaindex_tools,
+    )
+
+
+def __getattr__(name: str):
+    if name in {"K2LangChainRetriever", "create_k2_langchain_tools"}:
+        from sdk.integrations.langchain import K2LangChainRetriever, create_k2_langchain_tools
+
+        mapping = {
+            "K2LangChainRetriever": K2LangChainRetriever,
+            "create_k2_langchain_tools": create_k2_langchain_tools,
+        }
+        return mapping[name]
+
+    if name in {
+        "K2LlamaIndexRetriever",
+        "K2LlamaIndexVectorStore",
+        "create_k2_llamaindex_tools",
+    }:
+        from sdk.integrations.llamaindex import (
+            K2LlamaIndexRetriever,
+            K2LlamaIndexVectorStore,
+            create_k2_llamaindex_tools,
+        )
+
+        mapping = {
+            "K2LlamaIndexRetriever": K2LlamaIndexRetriever,
+            "K2LlamaIndexVectorStore": K2LlamaIndexVectorStore,
+            "create_k2_llamaindex_tools": create_k2_llamaindex_tools,
+        }
+        return mapping[name]
+
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")

sdk/integrations/_client.py

@@ -0,0 +1,101 @@
+from __future__ import annotations
+
+import os
+
+from sdk import AsyncKnowledge2, Knowledge2
+
+DEFAULT_K2_API_HOST = "https://api.knowledge2.ai"
+
+
+def resolve_async_client(
+    *,
+    sync_client: Knowledge2 | None = None,
+    api_key: str | None = None,
+    api_host: str | None = None,
+) -> AsyncKnowledge2:
+    """Return a configured AsyncKnowledge2 client.
+
+    When a *sync_client* is provided, the async client is created with
+    the same credentials (api_key, base_url, admin_token).
+
+    Preference order:
+      1) credentials from sync_client
+      2) api_key/api_host parameters
+      3) environment variables K2_API_KEY/K2_BASE_URL
+    """
+    if sync_client is not None:
+        # NOTE: validate_responses is intentionally NOT propagated here.
+        # Framework integrations (LangChain, LlamaIndex) call .get() on
+        # results expecting raw dicts; Pydantic models would break them.
+        return AsyncKnowledge2(
+            api_key=sync_client.api_key,
+            api_host=sync_client.base_url,
+            admin_token=sync_client.admin_token,
+            bearer_token=sync_client.bearer_token,
+            bearer_token_factory=sync_client._bearer_token_factory,
+            token_cache_ttl=sync_client._token_cache_ttl,
+            org_id=sync_client.org_id,
+            headers=dict(sync_client._default_headers) if sync_client._default_headers else None,
+            max_retries=sync_client._max_retries,
+        )
+
+    resolved_api_key = api_key or os.getenv("K2_API_KEY")
+    if not resolved_api_key:
+        raise ValueError("K2 API key required: pass api_key/client or set K2_API_KEY")
+
+    resolved_api_host = api_host or os.getenv("K2_BASE_URL") or DEFAULT_K2_API_HOST
+    return AsyncKnowledge2(api_key=resolved_api_key, api_host=resolved_api_host)
+
+
+def resolve_client(
+    *,
+    client: Knowledge2 | None,
+    api_key: str | None,
+    api_host: str | None,
+) -> Knowledge2:
+    """Return a configured Knowledge2 client.
+
+    Preference order:
+      1) explicit client
+      2) api_key/api_host parameters
+      3) environment variables K2_API_KEY/K2_BASE_URL
+    """
+    if client is not None:
+        return client
+
+    resolved_api_key = api_key or os.getenv("K2_API_KEY")
+    if not resolved_api_key:
+        raise ValueError("K2 API key required: pass api_key/client or set K2_API_KEY")
+
+    # NOTE: typeshed's os.getenv() typing keeps this as Optional[str] even with a default.
+    resolved_api_host = api_host or os.getenv("K2_BASE_URL") or DEFAULT_K2_API_HOST
+    return Knowledge2(api_key=resolved_api_key, api_host=resolved_api_host)
+
+
+def resolve_corpus_id(corpus_id: str | None) -> str:
+    """Resolve corpus_id from argument or environment."""
+    resolved_corpus_id = corpus_id or os.getenv("K2_CORPUS_ID")
+    if not resolved_corpus_id:
+        raise ValueError("K2 corpus_id required: pass corpus_id or set K2_CORPUS_ID")
+    return resolved_corpus_id
+
+
+def merge_return_config(
+    *,
+    base: dict[str, object] | None,
+    override: dict[str, object] | None,
+    include_text: bool = True,
+    include_scores: bool = True,
+    include_provenance: bool = True,
+) -> dict[str, object]:
+    """Merge return configs with safe defaults for framework adapters."""
+    merged: dict[str, object] = {
+        "include_text": include_text,
+        "include_scores": include_scores,
+        "include_provenance": include_provenance,
+    }
+    if base:
+        merged.update(base)
+    if override:
+        merged.update(override)
+    return merged

sdk/integrations/langchain/__init__.py

@@ -0,0 +1,6 @@
+"""LangChain integrations for Knowledge2."""
+
+from sdk.integrations.langchain.retriever import K2LangChainRetriever
+from sdk.integrations.langchain.tools import create_k2_langchain_tools
+
+__all__ = ["K2LangChainRetriever", "create_k2_langchain_tools"]

sdk/integrations/langchain/retriever.py

@@ -0,0 +1,166 @@
+from __future__ import annotations
+
+from typing import Any, ClassVar
+
+from pydantic import ConfigDict, Field, PrivateAttr
+
+from sdk import AsyncKnowledge2, Knowledge2
+from sdk.integrations._client import (
+    merge_return_config,
+    resolve_async_client,
+    resolve_client,
+    resolve_corpus_id,
+)
+from sdk.types import SearchResult
+
+try:
+    from langchain_core.callbacks import (
+        AsyncCallbackManagerForRetrieverRun,
+        CallbackManagerForRetrieverRun,
+    )
+    from langchain_core.documents import Document
+    from langchain_core.retrievers import BaseRetriever
+except ImportError as exc:  # pragma: no cover - import-time dependency guard
+    raise ImportError(
+        "LangChain integration requires langchain-core. Install with `pip install .[langchain]`."
+    ) from exc
+
+
+class K2LangChainRetriever(BaseRetriever):
+    """LangChain retriever backed by Knowledge2 search APIs."""
+
+    model_config: ClassVar[ConfigDict] = ConfigDict(arbitrary_types_allowed=True)
+
+    client: Any | None = Field(default=None, exclude=True)
+    api_key: str | None = Field(default=None, exclude=True, repr=False)
+    api_host: str | None = None
+    corpus_id: str | None = None
+
+    top_k: int = 10
+    filters: dict[str, Any] | None = None
+    hybrid: dict[str, Any] | None = None
+    rerank: dict[str, Any] | None = None
+    return_config: dict[str, Any] | None = None
+
+    _client: Knowledge2 | Any = PrivateAttr()
+    _async_client: AsyncKnowledge2 | None = PrivateAttr(default=None)
+    _corpus_id: str = PrivateAttr()
+
+    def model_post_init(self, __context: Any) -> None:
+        self._client = resolve_client(
+            client=self.client, api_key=self.api_key, api_host=self.api_host
+        )
+        self._corpus_id = resolve_corpus_id(self.corpus_id)
+
+    def _ensure_async_client(self) -> AsyncKnowledge2:
+        """Lazily create an AsyncKnowledge2 sharing the sync client's credentials."""
+        if self._async_client is None:
+            self._async_client = resolve_async_client(sync_client=self._client)
+        return self._async_client
+
+    async def aclose(self) -> None:
+        """Close the lazily-created async client, releasing connections."""
+        if self._async_client is not None:
+            await self._async_client.close()
+            self._async_client = None
+
+    @staticmethod
+    def _result_to_document(result: SearchResult, corpus_id: str) -> Document:
+        custom_meta = result.get("custom_metadata") or {}
+        system_meta = result.get("system_metadata") or {}
+        if not custom_meta and not system_meta:
+            legacy = result.get("metadata")
+            if isinstance(legacy, dict):
+                custom_meta = legacy
+        if not isinstance(custom_meta, dict):
+            custom_meta = {}
+        if not isinstance(system_meta, dict):
+            system_meta = {}
+        result_metadata = {**system_meta, **custom_meta}
+
+        metadata: dict[str, Any] = {
+            "source": "knowledge2",
+            "corpus_id": corpus_id,
+            "chunk_id": result.get("chunk_id"),
+            "score": result.get("score"),
+            "raw_score": result.get("raw_score"),
+            "offset_start": result.get("offset_start"),
+            "offset_end": result.get("offset_end"),
+            "page_start": result.get("page_start"),
+            "page_end": result.get("page_end"),
+        }
+        metadata.update(result_metadata)
+
+        return Document(page_content=result.get("text") or "", metadata=metadata)
+
+    def _search(
+        self,
+        query: str,
+        *,
+        top_k: int | None = None,
+        filters: dict[str, Any] | None = None,
+        hybrid: dict[str, Any] | None = None,
+        rerank: dict[str, Any] | None = None,
+        return_config: dict[str, Any] | None = None,
+    ) -> list[SearchResult]:
+        payload_return_config = merge_return_config(
+            base=self.return_config,
+            override=return_config,
+            include_text=True,
+            include_scores=True,
+            include_provenance=True,
+        )
+        response = self._client.search(
+            self._corpus_id,
+            query,
+            top_k=top_k if top_k is not None else self.top_k,
+            filters=filters if filters is not None else self.filters,
+            hybrid=hybrid if hybrid is not None else self.hybrid,
+            rerank=rerank if rerank is not None else self.rerank,
+            return_config=payload_return_config,
+        )
+        return response.get("results", [])
+
+    def _get_relevant_documents(
+        self,
+        query: str,
+        *,
+        run_manager: CallbackManagerForRetrieverRun,
+        **kwargs: Any,
+    ) -> list[Document]:
+        results = self._search(
+            query,
+            top_k=kwargs.get("top_k"),
+            filters=kwargs.get("filters"),
+            hybrid=kwargs.get("hybrid"),
+            rerank=kwargs.get("rerank"),
+            return_config=kwargs.get("return_config"),
+        )
+        return [self._result_to_document(result, self._corpus_id) for result in results]
+
+    async def _aget_relevant_documents(
+        self,
+        query: str,
+        *,
+        run_manager: AsyncCallbackManagerForRetrieverRun,
+        **kwargs: Any,
+    ) -> list[Document]:
+        async_client = self._ensure_async_client()
+        payload_return_config = merge_return_config(
+            base=self.return_config,
+            override=kwargs.get("return_config"),
+            include_text=True,
+            include_scores=True,
+            include_provenance=True,
+        )
+        response = await async_client.search(
+            self._corpus_id,
+            query,
+            top_k=kwargs.get("top_k") if kwargs.get("top_k") is not None else self.top_k,  # type: ignore[arg-type]
+            filters=kwargs.get("filters") if kwargs.get("filters") is not None else self.filters,
+            hybrid=kwargs.get("hybrid") if kwargs.get("hybrid") is not None else self.hybrid,
+            rerank=kwargs.get("rerank") if kwargs.get("rerank") is not None else self.rerank,
+            return_config=payload_return_config,
+        )
+        results = response.get("results", [])
+        return [self._result_to_document(result, self._corpus_id) for result in results]

sdk/integrations/langchain/tools.py

@@ -0,0 +1,108 @@
+from __future__ import annotations
+
+from typing import Any, cast
+
+from sdk import Knowledge2
+from sdk.integrations._client import merge_return_config, resolve_client, resolve_corpus_id
+
+try:
+    from langchain_core.tools import BaseTool, tool
+except ImportError as exc:  # pragma: no cover - import-time dependency guard
+    raise ImportError(
+        "LangChain integration requires langchain-core. Install with `pip install .[langchain]`."
+    ) from exc
+
+
+def create_k2_langchain_tools(
+    *,
+    corpus_id: str | None = None,
+    client: Knowledge2 | None = None,
+    api_key: str | None = None,
+    api_host: str | None = None,
+    default_top_k: int = 10,
+    default_hybrid: dict[str, Any] | None = None,
+    default_generation: dict[str, Any] | None = None,
+) -> list[BaseTool]:
+    """Create LangChain tools for K2 search and ingestion flows."""
+    resolved_client = resolve_client(client=client, api_key=api_key, api_host=api_host)
+    resolved_corpus_id = resolve_corpus_id(corpus_id)
+
+    @tool("k2_search")
+    def k2_search(
+        query: str,
+        top_k: int = default_top_k,
+        filters: dict[str, Any] | None = None,
+    ) -> dict[str, Any]:
+        """Search a K2 corpus and return scored chunks."""
+        return cast(
+            "dict[str, Any]",
+            resolved_client.search(
+                resolved_corpus_id,
+                query,
+                top_k=top_k,
+                filters=filters,
+                hybrid=default_hybrid,
+                return_config=merge_return_config(base=None, override=None),
+            ),
+        )
+
+    @tool("k2_ingest_text")
+    def k2_ingest_text(
+        raw_text: str,
+        source_uri: str | None = None,
+        metadata: dict[str, Any] | None = None,
+        auto_index: bool = False,
+    ) -> dict[str, Any]:
+        """Ingest a text document into K2."""
+        return cast(
+            "dict[str, Any]",
+            resolved_client.upload_document(
+                resolved_corpus_id,
+                raw_text=raw_text,
+                source_uri=source_uri,
+                metadata=metadata,
+                auto_index=auto_index,
+            ),
+        )
+
+    @tool("k2_build_indexes")
+    def k2_build_indexes(
+        dense: bool = True,
+        sparse: bool = True,
+        mode: str = "incremental",
+        wait: bool = True,
+    ) -> dict[str, Any]:
+        """Trigger K2 index build for the current corpus."""
+        return cast(
+            "dict[str, Any]",
+            resolved_client.build_indexes(
+                resolved_corpus_id,
+                dense=dense,
+                sparse=sparse,
+                mode=mode,
+                wait=wait,
+            ),
+        )
+
+    @tool("k2_generate_answer")
+    def k2_generate_answer(
+        query: str,
+        top_k: int = default_top_k,
+        filters: dict[str, Any] | None = None,
+        generation: dict[str, Any] | None = None,
+    ) -> dict[str, Any]:
+        """Generate a grounded answer using K2 retrieval + server-side LLM generation."""
+        return cast(
+            "dict[str, Any]",
+            resolved_client.search_generate(
+                resolved_corpus_id,
+                query,
+                top_k=top_k,
+                filters=filters,
+                hybrid=default_hybrid,
+                generation=generation if generation is not None else default_generation,
+                return_config=merge_return_config(base=None, override=None),
+            ),
+        )
+
+    return [k2_search, k2_ingest_text, k2_build_indexes, k2_generate_answer]

sdk/integrations/llamaindex/__init__.py

@@ -0,0 +1,11 @@
+"""LlamaIndex integrations for Knowledge2."""
+
+from sdk.integrations.llamaindex.retriever import K2LlamaIndexRetriever
+from sdk.integrations.llamaindex.tools import create_k2_llamaindex_tools
+from sdk.integrations.llamaindex.vector_store import K2LlamaIndexVectorStore
+
+__all__ = [
+    "K2LlamaIndexRetriever",
+    "K2LlamaIndexVectorStore",
+    "create_k2_llamaindex_tools",
+]

sdk/integrations/llamaindex/filters.py

@@ -0,0 +1,78 @@
+from __future__ import annotations
+
+from typing import Any
+
+try:
+    from llama_index.core.vector_stores.types import MetadataFilter, MetadataFilters
+except ImportError as exc:  # pragma: no cover - import-time dependency guard
+    raise ImportError(
+        "LlamaIndex integration requires llama-index-core. Install with `pip install .[llamaindex]`."
+    ) from exc
+
+
+_OPERATOR_MAP = {
+    "eq": "==",
+    "==": "==",
+    "ne": "!=",
+    "!=": "!=",
+    "gt": ">",
+    ">": ">",
+    "gte": ">=",
+    ">=": ">=",
+    "lt": "<",
+    "<": "<",
+    "lte": "<=",
+    "<=": "<=",
+    "in": "in",
+    "nin": "not_in",
+    "not_in": "not_in",
+    "contains": "contains",
+    "text_match": "text_match",
+}
+
+
+def _normalize_enum_name(value: Any, *, default: str = "") -> str:
+    if value is None:
+        return default
+    if hasattr(value, "value"):
+        raw = str(value.value)
+    else:
+        raw = str(value)
+    return raw.split(".")[-1].strip().lower()
+
+
+def llama_filters_to_k2(filters: MetadataFilters | None) -> dict[str, Any] | None:
+    """Convert LlamaIndex MetadataFilters to K2 structured filter format.
+
+    Supports all K2 filter operators and both AND/OR conditions.
+    Returns the structured format: {"filters": [...], "condition": "and|or"}
+    """
+    if filters is None:
+        return None
+
+    condition = _normalize_enum_name(getattr(filters, "condition", None), default="and")
+
+    converted: list[dict[str, Any]] = []
+    for item in getattr(filters, "filters", []):
+        if not isinstance(item, MetadataFilter):
+            raise ValueError(f"Unsupported metadata filter node type: {type(item).__name__}")
+
+        key = getattr(item, "key", None)
+        if not key:
+            raise ValueError("MetadataFilter key must be set")
+
+        operator_name = _normalize_enum_name(getattr(item, "operator", None), default="eq")
+        k2_op = _OPERATOR_MAP.get(operator_name)
+        if k2_op is None:
+            raise ValueError(
+                f"Unsupported LlamaIndex operator: {operator_name!r}. "
+                f"Supported: {', '.join(sorted(_OPERATOR_MAP))}"
+            )
+
+        value = getattr(item, "value", None)
+        converted.append({"key": key, "op": k2_op, "value": value})
+
+    if not converted:
+        return None
+
+    return {"filters": converted, "condition": condition}