usecortex-ai 0.3.6__py3-none-any.whl → 0.5.0__py3-none-any.whl

usecortex_ai/embeddings/client.py

@@ -4,10 +4,10 @@ import typing
 
 from ..core.client_wrapper import AsyncClientWrapper, SyncClientWrapper
 from ..core.request_options import RequestOptions
-from ..types.embeddings_create_collection_data import EmbeddingsCreateCollectionData
-from ..types.embeddings_delete_data import EmbeddingsDeleteData
-from ..types.embeddings_get_data import EmbeddingsGetData
-from ..types.embeddings_search_data import EmbeddingsSearchData
+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
@@ -29,47 +29,61 @@ class EmbeddingsClient:
         """
         return self._raw_client
 
-    def delete(
+    def insert(
         self,
         *,
-        chunk_ids: typing.Sequence[str],
         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,
-    ) -> EmbeddingsDeleteData:
+    ) -> InsertResult:
         """
-        Delete embedding chunks by chunk ID.
+        Upload pre-computed embeddings for advanced similarity search.
+
+        This endpoint accepts vector embeddings that you’ve generated externally,
+         allowing you to integrate with custom embedding models or existing vector databases.
+          The embeddings represent chunks of your content as numerical vectors.
 
-        Use this to remove specific chunks from your embeddings index when they are no longer valid or should not appear in results.
+        The system stores these embeddings and makes them available for semantic search and similarity matching.
+         Use this when you want to leverage specialized embedding models or have existing vector representations.
+        When upsert=True, existing embeddings with the same chunk_id will be updated.
 
         Parameters
         ----------
-        chunk_ids : typing.Sequence[str]
-            The chunk IDs of the source you want to delete
-
         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
         -------
-        EmbeddingsDeleteData
+        InsertResult
             Successful Response
 
         Examples
         --------
-        from usecortex-ai import CortexAI
+        from usecortex-ai import CortexAI, RawEmbeddingDocument, RawEmbeddingVector
 
         client = CortexAI(token="YOUR_TOKEN", )
-        client.embeddings.delete(chunk_ids=['CortexEmbeddings123_0', 'CortexEmbeddings123_1'], tenant_id='tenant_1234', )
+        client.embeddings.insert(tenant_id='tenant_id', embeddings=[RawEmbeddingDocument(source_id='source_id', embeddings=[RawEmbeddingVector(chunk_id='chunk_id', embedding=[1.1], )], )], )
         """
-        _response = self._raw_client.delete(
-            chunk_ids=chunk_ids, tenant_id=tenant_id, sub_tenant_id=sub_tenant_id, request_options=request_options
+        _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
 
@@ -77,18 +91,19 @@ class EmbeddingsClient:
         self,
         *,
         tenant_id: str,
-        embeddings: typing.Optional[typing.Sequence[float]] = OMIT,
-        sub_tenant_id: typing.Optional[str] = OMIT,
-        max_chunks: typing.Optional[int] = OMIT,
+        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,
-    ) -> EmbeddingsSearchData:
+    ) -> typing.List[RawEmbeddingSearchResult]:
         """
         Find similar chunks using an embedding vector.
 
         Use this to retrieve the most similar chunk IDs to a single query embedding.
 
-
-        Expected outcome
+        Expected outcome:
         - You receive the closest chunk IDs with optional similarity scores.
 
         Parameters
@@ -96,20 +111,27 @@ class EmbeddingsClient:
         tenant_id : str
             Unique identifier for the tenant/organization
 
-        embeddings : typing.Optional[typing.Sequence[float]]
-            The embedding vector for search
-
-        sub_tenant_id : typing.Optional[str]
+        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.
 
-        max_chunks : typing.Optional[int]
+        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
         -------
-        EmbeddingsSearchData
+        typing.List[RawEmbeddingSearchResult]
             Successful Response
 
         Examples
@@ -117,47 +139,60 @@ class EmbeddingsClient:
         from usecortex-ai import CortexAI
 
         client = CortexAI(token="YOUR_TOKEN", )
-        client.embeddings.search(tenant_id='tenant_1234', )
+        client.embeddings.search(tenant_id='tenant_id', sub_tenant_id='sub_tenant_id', query_embedding=[1.1], )
         """
         _response = self._raw_client.search(
             tenant_id=tenant_id,
-            embeddings=embeddings,
             sub_tenant_id=sub_tenant_id,
-            max_chunks=max_chunks,
+            query_embedding=query_embedding,
+            limit=limit,
+            filter_expr=filter_expr,
+            output_fields=output_fields,
             request_options=request_options,
         )
         return _response.data
 
-    def get_by_chunk_ids(
+    def filter(
         self,
         *,
-        chunk_ids: typing.Sequence[str],
         tenant_id: str,
-        sub_tenant_id: typing.Optional[str] = OMIT,
+        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,
-    ) -> EmbeddingsGetData:
+    ) -> typing.List[RawEmbeddingSearchResult]:
         """
         Retrieve embeddings for specific chunk IDs.
-
         Use this when you already know the chunk IDs and need their corresponding embeddings.
 
         Parameters
         ----------
-        chunk_ids : typing.Sequence[str]
-            The chunk IDs of the source you want to get embeddings for
-
         tenant_id : str
             Unique identifier for the tenant/organization
 
-        sub_tenant_id : typing.Optional[str]
+        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
         -------
-        EmbeddingsGetData
+        typing.List[RawEmbeddingSearchResult]
             Successful Response
 
         Examples
@@ -165,32 +200,54 @@ class EmbeddingsClient:
         from usecortex-ai import CortexAI
 
         client = CortexAI(token="YOUR_TOKEN", )
-        client.embeddings.get_by_chunk_ids(chunk_ids=['CortexEmbeddings123_0', 'CortexEmbeddings123_1'], tenant_id='tenant_1234', )
+        client.embeddings.filter(tenant_id='tenant_id', sub_tenant_id='sub_tenant_id', )
         """
-        _response = self._raw_client.get_by_chunk_ids(
-            chunk_ids=chunk_ids, tenant_id=tenant_id, sub_tenant_id=sub_tenant_id, request_options=request_options
+        _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 create_collection(
-        self, *, tenant_id: str, request_options: typing.Optional[RequestOptions] = None
-    ) -> EmbeddingsCreateCollectionData:
+    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:
         """
-        Create an embeddings collection for a tenant.
+        Delete embedding chunks by chunk ID.
 
-        Use this to initialize storage so you can index and query embeddings for the tenant.
+        Use this to remove specific chunks from your embeddings index when they are no longer
+        valid or should not appear in results.
 
         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
         -------
-        EmbeddingsCreateCollectionData
+        DeleteResult
             Successful Response
 
         Examples
@@ -198,9 +255,15 @@ class EmbeddingsClient:
         from usecortex-ai import CortexAI
 
         client = CortexAI(token="YOUR_TOKEN", )
-        client.embeddings.create_collection(tenant_id='tenant_1234', )
+        client.embeddings.delete(tenant_id='tenant_id', )
         """
-        _response = self._raw_client.create_collection(tenant_id=tenant_id, request_options=request_options)
+        _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
 
 
@@ -219,51 +282,66 @@ class AsyncEmbeddingsClient:
         """
         return self._raw_client
 
-    async def delete(
+    async def insert(
         self,
         *,
-        chunk_ids: typing.Sequence[str],
         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,
-    ) -> EmbeddingsDeleteData:
+    ) -> InsertResult:
         """
-        Delete embedding chunks by chunk ID.
+        Upload pre-computed embeddings for advanced similarity search.
+
+        This endpoint accepts vector embeddings that you’ve generated externally,
+         allowing you to integrate with custom embedding models or existing vector databases.
+          The embeddings represent chunks of your content as numerical vectors.
 
-        Use this to remove specific chunks from your embeddings index when they are no longer valid or should not appear in results.
+        The system stores these embeddings and makes them available for semantic search and similarity matching.
+         Use this when you want to leverage specialized embedding models or have existing vector representations.
+        When upsert=True, existing embeddings with the same chunk_id will be updated.
 
         Parameters
         ----------
-        chunk_ids : typing.Sequence[str]
-            The chunk IDs of the source you want to delete
-
         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
         -------
-        EmbeddingsDeleteData
+        InsertResult
             Successful Response
 
         Examples
         --------
         import asyncio
 
-        from usecortex-ai import AsyncCortexAI
+        from usecortex-ai import (AsyncCortexAI, RawEmbeddingDocument,
+                                  RawEmbeddingVector)
 
         client = AsyncCortexAI(token="YOUR_TOKEN", )
         async def main() -> None:
-            await client.embeddings.delete(chunk_ids=['CortexEmbeddings123_0', 'CortexEmbeddings123_1'], tenant_id='tenant_1234', )
+            await client.embeddings.insert(tenant_id='tenant_id', embeddings=[RawEmbeddingDocument(source_id='source_id', embeddings=[RawEmbeddingVector(chunk_id='chunk_id', embedding=[1.1], )], )], )
         asyncio.run(main())
         """
-        _response = await self._raw_client.delete(
-            chunk_ids=chunk_ids, tenant_id=tenant_id, sub_tenant_id=sub_tenant_id, request_options=request_options
+        _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
 
@@ -271,18 +349,19 @@ class AsyncEmbeddingsClient:
         self,
         *,
         tenant_id: str,
-        embeddings: typing.Optional[typing.Sequence[float]] = OMIT,
-        sub_tenant_id: typing.Optional[str] = OMIT,
-        max_chunks: typing.Optional[int] = OMIT,
+        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,
-    ) -> EmbeddingsSearchData:
+    ) -> typing.List[RawEmbeddingSearchResult]:
         """
         Find similar chunks using an embedding vector.
 
         Use this to retrieve the most similar chunk IDs to a single query embedding.
 
-
-        Expected outcome
+        Expected outcome:
         - You receive the closest chunk IDs with optional similarity scores.
 
         Parameters
@@ -290,20 +369,27 @@ class AsyncEmbeddingsClient:
         tenant_id : str
             Unique identifier for the tenant/organization
 
-        embeddings : typing.Optional[typing.Sequence[float]]
-            The embedding vector for search
-
-        sub_tenant_id : typing.Optional[str]
+        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.
 
-        max_chunks : typing.Optional[int]
+        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
         -------
-        EmbeddingsSearchData
+        typing.List[RawEmbeddingSearchResult]
             Successful Response
 
         Examples
@@ -314,48 +400,61 @@ class AsyncEmbeddingsClient:
 
         client = AsyncCortexAI(token="YOUR_TOKEN", )
         async def main() -> None:
-            await client.embeddings.search(tenant_id='tenant_1234', )
+            await client.embeddings.search(tenant_id='tenant_id', sub_tenant_id='sub_tenant_id', query_embedding=[1.1], )
         asyncio.run(main())
         """
         _response = await self._raw_client.search(
             tenant_id=tenant_id,
-            embeddings=embeddings,
             sub_tenant_id=sub_tenant_id,
-            max_chunks=max_chunks,
+            query_embedding=query_embedding,
+            limit=limit,
+            filter_expr=filter_expr,
+            output_fields=output_fields,
             request_options=request_options,
         )
         return _response.data
 
-    async def get_by_chunk_ids(
+    async def filter(
         self,
         *,
-        chunk_ids: typing.Sequence[str],
         tenant_id: str,
-        sub_tenant_id: typing.Optional[str] = OMIT,
+        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,
-    ) -> EmbeddingsGetData:
+    ) -> typing.List[RawEmbeddingSearchResult]:
         """
         Retrieve embeddings for specific chunk IDs.
-
         Use this when you already know the chunk IDs and need their corresponding embeddings.
 
         Parameters
         ----------
-        chunk_ids : typing.Sequence[str]
-            The chunk IDs of the source you want to get embeddings for
-
         tenant_id : str
             Unique identifier for the tenant/organization
 
-        sub_tenant_id : typing.Optional[str]
+        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
         -------
-        EmbeddingsGetData
+        typing.List[RawEmbeddingSearchResult]
             Successful Response
 
         Examples
@@ -366,33 +465,55 @@ class AsyncEmbeddingsClient:
 
         client = AsyncCortexAI(token="YOUR_TOKEN", )
         async def main() -> None:
-            await client.embeddings.get_by_chunk_ids(chunk_ids=['CortexEmbeddings123_0', 'CortexEmbeddings123_1'], tenant_id='tenant_1234', )
+            await client.embeddings.filter(tenant_id='tenant_id', sub_tenant_id='sub_tenant_id', )
         asyncio.run(main())
         """
-        _response = await self._raw_client.get_by_chunk_ids(
-            chunk_ids=chunk_ids, tenant_id=tenant_id, sub_tenant_id=sub_tenant_id, request_options=request_options
+        _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 create_collection(
-        self, *, tenant_id: str, request_options: typing.Optional[RequestOptions] = None
-    ) -> EmbeddingsCreateCollectionData:
+    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:
         """
-        Create an embeddings collection for a tenant.
+        Delete embedding chunks by chunk ID.
 
-        Use this to initialize storage so you can index and query embeddings for the tenant.
+        Use this to remove specific chunks from your embeddings index when they are no longer
+        valid or should not appear in results.
 
         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
         -------
-        EmbeddingsCreateCollectionData
+        DeleteResult
             Successful Response
 
         Examples
@@ -403,8 +524,14 @@ class AsyncEmbeddingsClient:
 
         client = AsyncCortexAI(token="YOUR_TOKEN", )
         async def main() -> None:
-            await client.embeddings.create_collection(tenant_id='tenant_1234', )
+            await client.embeddings.delete(tenant_id='tenant_id', )
         asyncio.run(main())
         """
-        _response = await self._raw_client.create_collection(tenant_id=tenant_id, request_options=request_options)
+        _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