Skip to content
Draft
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
@@ -1,3 +1,4 @@
import uuid
from pathlib import Path
from typing import Annotated, Any, Dict, List, Optional, Tuple, Union

Expand Down Expand Up @@ -33,11 +34,13 @@
)
from .context_grounding_index import ContextGroundingIndex
from .context_grounding_payloads import (
AttachmentsDataSource,
BucketDataSource,
BucketSourceConfig,
ConfluenceDataSource,
ConfluenceSourceConfig,
CreateIndexPayload,
CreateJitIndexPayload,
DropboxDataSource,
DropboxSourceConfig,
GoogleDriveDataSource,
Expand Down Expand Up @@ -398,6 +401,32 @@ def create_index(

return ContextGroundingIndex.model_validate(response.json())

@resource_override(resource_type="index")
@traced(name="contextgrounding_create_jit_index", run_type="uipath")
def create_jit_index(
self,
usage: str,
attachments: list[uuid.UUID],
folder_key: Optional[str] = None,
folder_path: Optional[str] = None,
) -> ContextGroundingIndex:
"""Create a new context jit grounding index."""
spec = self._create_jit_spec(
usage,
attachments,
folder_path=folder_path,
folder_key=folder_key,
)

response = self.request(
spec.method,
spec.endpoint,
json=spec.json,
headers=spec.headers,
)

return ContextGroundingIndex.model_validate(response.json())

@resource_override(resource_type="index")
@traced(name="contextgrounding_create_index", run_type="uipath")
async def create_index_async(
Expand Down Expand Up @@ -1197,6 +1226,41 @@ def _create_spec(
},
)

def _create_jit_spec(
self,
usage: str,
attachments: list[uuid.UUID] = None,
folder_key: Optional[str] = None,
folder_path: Optional[str] = None,
) -> RequestSpec:
"""Create request spec for index creation."""
folder_key = self._resolve_folder_key(folder_key, folder_path)

data_source_dict = self._build_jit_data_source(attachments)

payload = CreateJitIndexPayload(
usage=usage,
data_source=data_source_dict,
)

return RequestSpec(
method="POST",
endpoint=Endpoint("/ecs_/v2/indexes/createephemeral"),
json=payload.model_dump(by_alias=True, exclude_none=True),
headers={
**header_folder(folder_key, None),
},
)

def _build_jit_data_source(self, attachments: list[uuid.UUID]) -> Dict[str, Any]:
data_source: AttachmentsDataSource
data_source = AttachmentsDataSource(attachments=attachments)
return data_source.model_dump(
by_alias=True,
exclude_none=True,
mode="json",
)

def _build_data_source(self, source: SourceConfig) -> Dict[str, Any]:
"""Build data source configuration from typed source config.

Expand Down
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
"""Payload models for context grounding index creation and configuration."""

import re
import uuid
from typing import Any, Dict, Literal, Optional, Union

from pydantic import BaseModel, ConfigDict, Field, model_validator
Expand Down Expand Up @@ -82,6 +83,10 @@ class ConfluenceDataSource(DataSourceBase):
space_id: str = Field(alias="spaceId", description="Space ID")


class AttachmentsDataSource(BaseModel):
attachments: list[uuid.UUID] = Field(description="List of attachment ids")


class Indexer(BaseModel):
"""Configuration for periodic indexing of data sources."""

Expand Down Expand Up @@ -136,6 +141,17 @@ class CreateIndexPayload(BaseModel):
model_config = ConfigDict(populate_by_name=True)


class CreateJitIndexPayload(BaseModel):
""" """

usage: str = Field(description="Index usage")
data_source: Dict[str, Any] = Field(
alias="dataSource", description="Data source configuration"
)

model_config = ConfigDict(populate_by_name=True)


# user-facing source configuration models
class BaseSourceConfig(BaseModel):
"""Base configuration for all source types."""
Expand Down
Loading