|
| 1 | +#!/usr/bin/env python3 |
| 2 | +""" |
| 3 | +Sample Search Index Chunking Customer Function |
| 4 | +
|
| 5 | +This function demonstrates the new signature-based invocation with Pydantic models: |
| 6 | +- Uses SearchIndexChunkingV1Request/Response (Pydantic models) |
| 7 | +- Requires Runtime parameter (for agentic capabilities) |
| 8 | +- Type-safe with direct field access (no wrappers) |
| 9 | +- Automatic validation and conversion |
| 10 | +""" |
| 11 | + |
| 12 | +import logging |
| 13 | + |
| 14 | +from datacustomcode.function.feature_types.chunking import ( |
| 15 | + ChunkType, |
| 16 | + SearchIndexChunkingV1Output, |
| 17 | + SearchIndexChunkingV1Request, |
| 18 | + SearchIndexChunkingV1Response, |
| 19 | +) |
| 20 | +from datacustomcode.function.runtime import Runtime |
| 21 | +from datacustomcode.llm_gateway.types.generate_text_request_builder import ( |
| 22 | + GenerateTextRequestBuilder, |
| 23 | +) |
| 24 | + |
| 25 | +logger = logging.getLogger(__name__) |
| 26 | +logging.basicConfig(level=logging.INFO) |
| 27 | + |
| 28 | + |
| 29 | +def _load_prompt_template(runtime: Runtime) -> str: |
| 30 | + """Load the chunking prompt template from file.""" |
| 31 | + prompt_file = runtime.file.find_file_path("chunking_prompt.txt") |
| 32 | + with open(prompt_file, "r") as f: |
| 33 | + _prompt_template_cache = f.read() |
| 34 | + logger.info(f"Loaded prompt template from {prompt_file}") |
| 35 | + return _prompt_template_cache |
| 36 | + |
| 37 | + |
| 38 | +def function( |
| 39 | + request: SearchIndexChunkingV1Request, runtime: Runtime |
| 40 | +) -> SearchIndexChunkingV1Response: |
| 41 | + """ |
| 42 | + Chunk documents for Search Index. |
| 43 | +
|
| 44 | + Args: |
| 45 | + request: SearchIndexChunkingV1Request with input documents |
| 46 | + runtime: Runtime instance for agentic capabilities (future use) |
| 47 | +
|
| 48 | + Returns: |
| 49 | + SearchIndexChunkingV1Response with chunked output |
| 50 | + """ |
| 51 | + logger.info(f"Received {len(request.input)} documents to chunk") |
| 52 | + |
| 53 | + # Load prompt template (cached after first call) |
| 54 | + prompt_template = _load_prompt_template(runtime) |
| 55 | + |
| 56 | + chunks = [] |
| 57 | + chunk_id = 1 |
| 58 | + |
| 59 | + # Process each document |
| 60 | + for doc_idx, doc in enumerate(request.input): |
| 61 | + # Direct field access - no wrappers! |
| 62 | + text = doc.text |
| 63 | + |
| 64 | + # Use LLM to intelligently chunk the document |
| 65 | + # This creates semantic chunks that preserve context and meaning |
| 66 | + prompt = prompt_template.format(text=text) |
| 67 | + |
| 68 | + builder = GenerateTextRequestBuilder() |
| 69 | + llm_request = ( |
| 70 | + builder.set_model("sfdc_ai__DefaultGPT4Turbo").set_prompt(prompt).build() |
| 71 | + ) |
| 72 | + response = runtime.llm_gateway.generate_text(llm_request) |
| 73 | + |
| 74 | + if response.is_success: |
| 75 | + # Parse LLM response to extract chunks |
| 76 | + llm_chunks = response.text.split("---CHUNK---") |
| 77 | + llm_chunks = [chunk.strip() for chunk in llm_chunks if chunk.strip()] |
| 78 | + |
| 79 | + # Create chunk outputs |
| 80 | + for chunk_text in llm_chunks: |
| 81 | + chunk = SearchIndexChunkingV1Output( |
| 82 | + text=chunk_text, |
| 83 | + seq_no=chunk_id, |
| 84 | + chunk_type=ChunkType.TEXT, |
| 85 | + citations={}, |
| 86 | + ) |
| 87 | + chunks.append(chunk) |
| 88 | + chunk_id += 1 |
| 89 | + |
| 90 | + else: |
| 91 | + # LLM chunking failed - log error and raise exception |
| 92 | + error_msg = ( |
| 93 | + f"LLM chunking failed for document {doc_idx + 1}: {response.error_code}" |
| 94 | + ) |
| 95 | + logger.error(error_msg) |
| 96 | + raise RuntimeError(error_msg) |
| 97 | + |
| 98 | + # Return Pydantic response |
| 99 | + return SearchIndexChunkingV1Response(output=chunks) |
0 commit comments