Source code for taimoe.platform.resources.knowledge_bases

"""Knowledge bases resource for Taimoe Platform SDK."""

from __future__ import annotations

from typing import TYPE_CHECKING

from taimoe.platform.types.knowledge_base import KnowledgeBase, QueryResult

if TYPE_CHECKING:
    from taimoe.platform.client.client import TaimoeClient




class KnowledgeBasesResource:
    """Resource for interacting with knowledge base endpoints."""

    def __init__(self, client: TaimoeClient) -> None:
        self._client = client



    def list(self, team: str, org: str | None = None) -> list[KnowledgeBase]:
        """List knowledge bases for a team. Backend returns a bare array."""
        org_name = org or self._client.org
        path = f"/api/v1/organizations/{org_name}/teams/{team}/knowledge-bases"
        items = self._client._request("GET", path)
        return [KnowledgeBase.model_validate(item) for item in items]




    async def list_async(self, team: str, org: str | None = None) -> list[KnowledgeBase]:
        """List knowledge bases for a team asynchronously."""
        org_name = org or self._client.org
        path = f"/api/v1/organizations/{org_name}/teams/{team}/knowledge-bases"
        items = await self._client._request_async("GET", path)
        return [KnowledgeBase.model_validate(item) for item in items]




    def get(self, slug: str, org: str | None = None) -> KnowledgeBase:
        """Get a specific knowledge base by slug.

        Backend only exposes GET-by-UUID, so we list the team and filter.
        """
        if "/" not in slug:
            raise ValueError(f"Invalid slug '{slug}'. Expected format 'team-slug/kb-slug'.")
        team, kb_slug = slug.split("/", 1)
        
        kbs = self.list(team=team, org=org)
        for kb in kbs:
            if kb.slug == kb_slug:
                return kb
        raise ValueError(f"Knowledge base '{slug}' not found.")




    async def get_async(self, slug: str, org: str | None = None) -> KnowledgeBase:
        """Get a specific knowledge base asynchronously by slug."""
        if "/" not in slug:
            raise ValueError(f"Invalid slug '{slug}'. Expected format 'team-slug/kb-slug'.")
        team, kb_slug = slug.split("/", 1)
        
        kbs = await self.list_async(team=team, org=org)
        for kb in kbs:
            if kb.slug == kb_slug:
                return kb
        raise ValueError(f"Knowledge base '{slug}' not found.")




    def test_query(self, slug: str, query: str, org: str | None = None) -> QueryResult:
        """Test a query against a knowledge base using its slug."""
        kb = self.get(slug=slug, org=org)
        team = slug.split("/", 1)[0]
        org_name = org or self._client.org
        path = f"/api/v1/organizations/{org_name}/teams/{team}/knowledge-bases/{kb.id}/test-query"
        payload = {"query": query}
        data = self._client._request("POST", path, json=payload)
        return QueryResult.model_validate(data)




    async def test_query_async(self, slug: str, query: str, org: str | None = None) -> QueryResult:
        """Test a query against a knowledge base asynchronously using its slug."""
        kb = await self.get_async(slug=slug, org=org)
        team = slug.split("/", 1)[0]
        org_name = org or self._client.org
        path = f"/api/v1/organizations/{org_name}/teams/{team}/knowledge-bases/{kb.id}/test-query"
        payload = {"query": query}
        data = await self._client._request_async("POST", path, json=payload)
        return QueryResult.model_validate(data)