hydra-db-python 0.1.0__py3-none-any.whl

hydra_db/embeddings/client.py

@@ -0,0 +1,486 @@
+# This file was auto-generated by Fern from our API Definition.
+
+import typing
+
+from ..core.client_wrapper import AsyncClientWrapper, SyncClientWrapper
+from ..core.request_options import RequestOptions
+from ..types.delete_result import DeleteResult
+from ..types.insert_result import InsertResult
+from ..types.raw_embedding_document import RawEmbeddingDocument
+from ..types.raw_embedding_search_result import RawEmbeddingSearchResult
+from .raw_client import AsyncRawEmbeddingsClient, RawEmbeddingsClient
+
+# this is used as the default value for optional parameters
+OMIT = typing.cast(typing.Any, ...)
+
+
+class EmbeddingsClient:
+    def __init__(self, *, client_wrapper: SyncClientWrapper):
+        self._raw_client = RawEmbeddingsClient(client_wrapper=client_wrapper)
+
+    @property
+    def with_raw_response(self) -> RawEmbeddingsClient:
+        """
+        Retrieves a raw implementation of this client that returns raw responses.
+
+        Returns
+        -------
+        RawEmbeddingsClient
+        """
+        return self._raw_client
+
+    def insert(
+        self,
+        *,
+        tenant_id: str,
+        embeddings: typing.Sequence[RawEmbeddingDocument],
+        sub_tenant_id: typing.Optional[str] = OMIT,
+        upsert: typing.Optional[bool] = OMIT,
+        request_options: typing.Optional[RequestOptions] = None,
+    ) -> InsertResult:
+        """
+        Parameters
+        ----------
+        tenant_id : str
+            Unique identifier for the tenant/organization
+
+        embeddings : typing.Sequence[RawEmbeddingDocument]
+            List of raw embedding documents to insert
+
+        sub_tenant_id : typing.Optional[str]
+            Optional sub-tenant identifier used to organize data within a tenant. If omitted, the default sub-tenant created during tenant setup will be used.
+
+        upsert : typing.Optional[bool]
+            If True, update existing embeddings; if False, insert only
+
+        request_options : typing.Optional[RequestOptions]
+            Request-specific configuration.
+
+        Returns
+        -------
+        InsertResult
+            Successful Response
+
+        Examples
+        --------
+        from hydra-db import HydraDB, RawEmbeddingDocument, RawEmbeddingVector
+
+        client = HydraDB(token="YOUR_TOKEN", )
+        client.embeddings.insert(tenant_id='tenant_1234', embeddings=[RawEmbeddingDocument(source_id='<source_id>', embeddings=[RawEmbeddingVector(chunk_id='<chunk_id>', embedding=[1.1], ), RawEmbeddingVector(chunk_id='<chunk_id>', embedding=[1.1], )], ), RawEmbeddingDocument(source_id='<source_id>', embeddings=[RawEmbeddingVector(chunk_id='<chunk_id>', embedding=[1.1], ), RawEmbeddingVector(chunk_id='<chunk_id>', embedding=[1.1], )], )], )
+        """
+        _response = self._raw_client.insert(
+            tenant_id=tenant_id,
+            embeddings=embeddings,
+            sub_tenant_id=sub_tenant_id,
+            upsert=upsert,
+            request_options=request_options,
+        )
+        return _response.data
+
+    def search(
+        self,
+        *,
+        tenant_id: str,
+        sub_tenant_id: str,
+        query_embedding: typing.Sequence[float],
+        limit: typing.Optional[int] = OMIT,
+        filter_expr: typing.Optional[str] = OMIT,
+        output_fields: typing.Optional[typing.Sequence[str]] = OMIT,
+        request_options: typing.Optional[RequestOptions] = None,
+    ) -> typing.List[RawEmbeddingSearchResult]:
+        """
+        Parameters
+        ----------
+        tenant_id : str
+            Unique identifier for the tenant/organization
+
+        sub_tenant_id : str
+            Optional sub-tenant identifier used to organize data within a tenant. If omitted, the default sub-tenant created during tenant setup will be used.
+
+        query_embedding : typing.Sequence[float]
+            Query embedding vector to search for
+
+        limit : typing.Optional[int]
+            Maximum number of results to return
+
+        filter_expr : typing.Optional[str]
+            Optional Milvus filter expression for additional filtering
+
+        output_fields : typing.Optional[typing.Sequence[str]]
+            Optional list of fields to return in results (default: chunk_id, source_id, metadata)
+
+        request_options : typing.Optional[RequestOptions]
+            Request-specific configuration.
+
+        Returns
+        -------
+        typing.List[RawEmbeddingSearchResult]
+            Successful Response
+
+        Examples
+        --------
+        from hydra-db import HydraDB
+
+        client = HydraDB(token="YOUR_TOKEN", )
+        client.embeddings.search(tenant_id='tenant_1234', sub_tenant_id='sub_tenant_4567', query_embedding=[1.1], )
+        """
+        _response = self._raw_client.search(
+            tenant_id=tenant_id,
+            sub_tenant_id=sub_tenant_id,
+            query_embedding=query_embedding,
+            limit=limit,
+            filter_expr=filter_expr,
+            output_fields=output_fields,
+            request_options=request_options,
+        )
+        return _response.data
+
+    def filter(
+        self,
+        *,
+        tenant_id: str,
+        sub_tenant_id: str,
+        source_id: typing.Optional[str] = OMIT,
+        chunk_ids: typing.Optional[typing.Sequence[str]] = OMIT,
+        output_fields: typing.Optional[typing.Sequence[str]] = OMIT,
+        limit: typing.Optional[int] = OMIT,
+        request_options: typing.Optional[RequestOptions] = None,
+    ) -> typing.List[RawEmbeddingSearchResult]:
+        """
+        Parameters
+        ----------
+        tenant_id : str
+            Unique identifier for the tenant/organization
+
+        sub_tenant_id : str
+            Optional sub-tenant identifier used to organize data within a tenant. If omitted, the default sub-tenant created during tenant setup will be used.
+
+        source_id : typing.Optional[str]
+            Optional source ID to filter by (mutually exclusive with chunk_ids)
+
+        chunk_ids : typing.Optional[typing.Sequence[str]]
+            Optional list of chunk IDs to filter by (mutually exclusive with source_id)
+
+        output_fields : typing.Optional[typing.Sequence[str]]
+            Optional list of fields to return in results (default: chunk_id, source_id, metadata)
+
+        limit : typing.Optional[int]
+            Maximum number of results to return
+
+        request_options : typing.Optional[RequestOptions]
+            Request-specific configuration.
+
+        Returns
+        -------
+        typing.List[RawEmbeddingSearchResult]
+            Successful Response
+
+        Examples
+        --------
+        from hydra-db import HydraDB
+
+        client = HydraDB(token="YOUR_TOKEN", )
+        client.embeddings.filter(tenant_id='tenant_1234', sub_tenant_id='sub_tenant_4567', )
+        """
+        _response = self._raw_client.filter(
+            tenant_id=tenant_id,
+            sub_tenant_id=sub_tenant_id,
+            source_id=source_id,
+            chunk_ids=chunk_ids,
+            output_fields=output_fields,
+            limit=limit,
+            request_options=request_options,
+        )
+        return _response.data
+
+    def delete(
+        self,
+        *,
+        tenant_id: str,
+        sub_tenant_id: typing.Optional[str] = None,
+        source_id: typing.Optional[str] = None,
+        chunk_ids: typing.Optional[typing.Union[str, typing.Sequence[str]]] = None,
+        request_options: typing.Optional[RequestOptions] = None,
+    ) -> DeleteResult:
+        """
+        Parameters
+        ----------
+        tenant_id : str
+            Unique identifier for the tenant/organization
+
+        sub_tenant_id : typing.Optional[str]
+            Optional sub-tenant ID for scoping deletion
+
+        source_id : typing.Optional[str]
+            Optional source ID to delete by (mutually exclusive with chunk_ids)
+
+        chunk_ids : typing.Optional[typing.Union[str, typing.Sequence[str]]]
+            Optional list of chunk IDs to delete (mutually exclusive with source_id)
+
+        request_options : typing.Optional[RequestOptions]
+            Request-specific configuration.
+
+        Returns
+        -------
+        DeleteResult
+            Successful Response
+
+        Examples
+        --------
+        from hydra-db import HydraDB
+
+        client = HydraDB(token="YOUR_TOKEN", )
+        client.embeddings.delete(tenant_id='tenant_1234', )
+        """
+        _response = self._raw_client.delete(
+            tenant_id=tenant_id,
+            sub_tenant_id=sub_tenant_id,
+            source_id=source_id,
+            chunk_ids=chunk_ids,
+            request_options=request_options,
+        )
+        return _response.data
+
+
+class AsyncEmbeddingsClient:
+    def __init__(self, *, client_wrapper: AsyncClientWrapper):
+        self._raw_client = AsyncRawEmbeddingsClient(client_wrapper=client_wrapper)
+
+    @property
+    def with_raw_response(self) -> AsyncRawEmbeddingsClient:
+        """
+        Retrieves a raw implementation of this client that returns raw responses.
+
+        Returns
+        -------
+        AsyncRawEmbeddingsClient
+        """
+        return self._raw_client
+
+    async def insert(
+        self,
+        *,
+        tenant_id: str,
+        embeddings: typing.Sequence[RawEmbeddingDocument],
+        sub_tenant_id: typing.Optional[str] = OMIT,
+        upsert: typing.Optional[bool] = OMIT,
+        request_options: typing.Optional[RequestOptions] = None,
+    ) -> InsertResult:
+        """
+        Parameters
+        ----------
+        tenant_id : str
+            Unique identifier for the tenant/organization
+
+        embeddings : typing.Sequence[RawEmbeddingDocument]
+            List of raw embedding documents to insert
+
+        sub_tenant_id : typing.Optional[str]
+            Optional sub-tenant identifier used to organize data within a tenant. If omitted, the default sub-tenant created during tenant setup will be used.
+
+        upsert : typing.Optional[bool]
+            If True, update existing embeddings; if False, insert only
+
+        request_options : typing.Optional[RequestOptions]
+            Request-specific configuration.
+
+        Returns
+        -------
+        InsertResult
+            Successful Response
+
+        Examples
+        --------
+        import asyncio
+
+        from hydra-db import AsyncHydraDB, RawEmbeddingDocument, RawEmbeddingVector
+
+        client = AsyncHydraDB(token="YOUR_TOKEN", )
+        async def main() -> None:
+            await client.embeddings.insert(tenant_id='tenant_1234', embeddings=[RawEmbeddingDocument(source_id='<source_id>', embeddings=[RawEmbeddingVector(chunk_id='<chunk_id>', embedding=[1.1], ), RawEmbeddingVector(chunk_id='<chunk_id>', embedding=[1.1], )], ), RawEmbeddingDocument(source_id='<source_id>', embeddings=[RawEmbeddingVector(chunk_id='<chunk_id>', embedding=[1.1], ), RawEmbeddingVector(chunk_id='<chunk_id>', embedding=[1.1], )], )], )
+        asyncio.run(main())
+        """
+        _response = await self._raw_client.insert(
+            tenant_id=tenant_id,
+            embeddings=embeddings,
+            sub_tenant_id=sub_tenant_id,
+            upsert=upsert,
+            request_options=request_options,
+        )
+        return _response.data
+
+    async def search(
+        self,
+        *,
+        tenant_id: str,
+        sub_tenant_id: str,
+        query_embedding: typing.Sequence[float],
+        limit: typing.Optional[int] = OMIT,
+        filter_expr: typing.Optional[str] = OMIT,
+        output_fields: typing.Optional[typing.Sequence[str]] = OMIT,
+        request_options: typing.Optional[RequestOptions] = None,
+    ) -> typing.List[RawEmbeddingSearchResult]:
+        """
+        Parameters
+        ----------
+        tenant_id : str
+            Unique identifier for the tenant/organization
+
+        sub_tenant_id : str
+            Optional sub-tenant identifier used to organize data within a tenant. If omitted, the default sub-tenant created during tenant setup will be used.
+
+        query_embedding : typing.Sequence[float]
+            Query embedding vector to search for
+
+        limit : typing.Optional[int]
+            Maximum number of results to return
+
+        filter_expr : typing.Optional[str]
+            Optional Milvus filter expression for additional filtering
+
+        output_fields : typing.Optional[typing.Sequence[str]]
+            Optional list of fields to return in results (default: chunk_id, source_id, metadata)
+
+        request_options : typing.Optional[RequestOptions]
+            Request-specific configuration.
+
+        Returns
+        -------
+        typing.List[RawEmbeddingSearchResult]
+            Successful Response
+
+        Examples
+        --------
+        import asyncio
+
+        from hydra-db import AsyncHydraDB
+
+        client = AsyncHydraDB(token="YOUR_TOKEN", )
+        async def main() -> None:
+            await client.embeddings.search(tenant_id='tenant_1234', sub_tenant_id='sub_tenant_4567', query_embedding=[1.1], )
+        asyncio.run(main())
+        """
+        _response = await self._raw_client.search(
+            tenant_id=tenant_id,
+            sub_tenant_id=sub_tenant_id,
+            query_embedding=query_embedding,
+            limit=limit,
+            filter_expr=filter_expr,
+            output_fields=output_fields,
+            request_options=request_options,
+        )
+        return _response.data
+
+    async def filter(
+        self,
+        *,
+        tenant_id: str,
+        sub_tenant_id: str,
+        source_id: typing.Optional[str] = OMIT,
+        chunk_ids: typing.Optional[typing.Sequence[str]] = OMIT,
+        output_fields: typing.Optional[typing.Sequence[str]] = OMIT,
+        limit: typing.Optional[int] = OMIT,
+        request_options: typing.Optional[RequestOptions] = None,
+    ) -> typing.List[RawEmbeddingSearchResult]:
+        """
+        Parameters
+        ----------
+        tenant_id : str
+            Unique identifier for the tenant/organization
+
+        sub_tenant_id : str
+            Optional sub-tenant identifier used to organize data within a tenant. If omitted, the default sub-tenant created during tenant setup will be used.
+
+        source_id : typing.Optional[str]
+            Optional source ID to filter by (mutually exclusive with chunk_ids)
+
+        chunk_ids : typing.Optional[typing.Sequence[str]]
+            Optional list of chunk IDs to filter by (mutually exclusive with source_id)
+
+        output_fields : typing.Optional[typing.Sequence[str]]
+            Optional list of fields to return in results (default: chunk_id, source_id, metadata)
+
+        limit : typing.Optional[int]
+            Maximum number of results to return
+
+        request_options : typing.Optional[RequestOptions]
+            Request-specific configuration.
+
+        Returns
+        -------
+        typing.List[RawEmbeddingSearchResult]
+            Successful Response
+
+        Examples
+        --------
+        import asyncio
+
+        from hydra-db import AsyncHydraDB
+
+        client = AsyncHydraDB(token="YOUR_TOKEN", )
+        async def main() -> None:
+            await client.embeddings.filter(tenant_id='tenant_1234', sub_tenant_id='sub_tenant_4567', )
+        asyncio.run(main())
+        """
+        _response = await self._raw_client.filter(
+            tenant_id=tenant_id,
+            sub_tenant_id=sub_tenant_id,
+            source_id=source_id,
+            chunk_ids=chunk_ids,
+            output_fields=output_fields,
+            limit=limit,
+            request_options=request_options,
+        )
+        return _response.data
+
+    async def delete(
+        self,
+        *,
+        tenant_id: str,
+        sub_tenant_id: typing.Optional[str] = None,
+        source_id: typing.Optional[str] = None,
+        chunk_ids: typing.Optional[typing.Union[str, typing.Sequence[str]]] = None,
+        request_options: typing.Optional[RequestOptions] = None,
+    ) -> DeleteResult:
+        """
+        Parameters
+        ----------
+        tenant_id : str
+            Unique identifier for the tenant/organization
+
+        sub_tenant_id : typing.Optional[str]
+            Optional sub-tenant ID for scoping deletion
+
+        source_id : typing.Optional[str]
+            Optional source ID to delete by (mutually exclusive with chunk_ids)
+
+        chunk_ids : typing.Optional[typing.Union[str, typing.Sequence[str]]]
+            Optional list of chunk IDs to delete (mutually exclusive with source_id)
+
+        request_options : typing.Optional[RequestOptions]
+            Request-specific configuration.
+
+        Returns
+        -------
+        DeleteResult
+            Successful Response
+
+        Examples
+        --------
+        import asyncio
+
+        from hydra-db import AsyncHydraDB
+
+        client = AsyncHydraDB(token="YOUR_TOKEN", )
+        async def main() -> None:
+            await client.embeddings.delete(tenant_id='tenant_1234', )
+        asyncio.run(main())
+        """
+        _response = await self._raw_client.delete(
+            tenant_id=tenant_id,
+            sub_tenant_id=sub_tenant_id,
+            source_id=source_id,
+            chunk_ids=chunk_ids,
+            request_options=request_options,
+        )
+        return _response.data