|
| 1 | +#!/usr/bin/env python3 |
| 2 | +# Copyright (c) 2025, Salesforce, Inc. |
| 3 | +# SPDX-License-Identifier: Apache-2 |
| 4 | + |
| 5 | +""" |
| 6 | +Document Chunking with a Gemini Named Credential Callout |
| 7 | +
|
| 8 | +Splits each input document into paragraph-sized chunks and classifies every |
| 9 | +chunk via Google's Gemini ``generateContent`` API, reached through a Named |
| 10 | +Credential (``callout:gemini``) so the endpoint URL and API key are resolved |
| 11 | +outside this code. The classification is attached to each chunk as citations. |
| 12 | +""" |
| 13 | + |
| 14 | +import json |
| 15 | +import logging |
| 16 | + |
| 17 | +from datacustomcode.function import Runtime |
| 18 | +from datacustomcode.function.feature_types.chunking import ( |
| 19 | + ChunkType, |
| 20 | + SearchIndexChunkingV1Output, |
| 21 | + SearchIndexChunkingV1Request, |
| 22 | + SearchIndexChunkingV1Response, |
| 23 | +) |
| 24 | +from datacustomcode.named_credential.types.http_method import HTTPMethod |
| 25 | +from datacustomcode.named_credential.types.http_request_builder import ( |
| 26 | + HTTPRequestBuilder, |
| 27 | +) |
| 28 | + |
| 29 | +logger = logging.getLogger(__name__) |
| 30 | +logging.basicConfig(level=logging.INFO) |
| 31 | + |
| 32 | +CALLOUT_URL = "callout:gemini" |
| 33 | + |
| 34 | +_ANALYSIS_FIELDS = ("summary", "category", "sentiment") |
| 35 | + |
| 36 | +_PROMPT = ( |
| 37 | + "Analyze the following document chunk and classify it. Respond with its " |
| 38 | + "one-sentence summary, a single-word category, overall sentiment " |
| 39 | + "(positive, negative, or neutral), and up to five key topics.\n\nChunk:\n" |
| 40 | +) |
| 41 | + |
| 42 | +# Force Gemini to return the classification as JSON in a fixed shape. |
| 43 | +_GENERATION_CONFIG = { |
| 44 | + "responseMimeType": "application/json", |
| 45 | + "responseSchema": { |
| 46 | + "type": "object", |
| 47 | + "properties": { |
| 48 | + "summary": {"type": "string"}, |
| 49 | + "category": {"type": "string"}, |
| 50 | + "sentiment": {"type": "string"}, |
| 51 | + "topics": {"type": "array", "items": {"type": "string"}}, |
| 52 | + }, |
| 53 | + "required": ["summary", "category", "sentiment", "topics"], |
| 54 | + }, |
| 55 | +} |
| 56 | + |
| 57 | + |
| 58 | +def _chunk_text(text: str, max_words: int = 80) -> list[str]: |
| 59 | + """Split text into paragraph-aligned chunks of at most ``max_words`` words.""" |
| 60 | + paragraphs = [p.strip() for p in text.split("\n\n") if p.strip()] |
| 61 | + |
| 62 | + chunks: list[str] = [] |
| 63 | + current: list[str] = [] |
| 64 | + current_words = 0 |
| 65 | + |
| 66 | + for paragraph in paragraphs: |
| 67 | + paragraph_words = len(paragraph.split()) |
| 68 | + if current and current_words + paragraph_words > max_words: |
| 69 | + chunks.append("\n\n".join(current)) |
| 70 | + current = [] |
| 71 | + current_words = 0 |
| 72 | + current.append(paragraph) |
| 73 | + current_words += paragraph_words |
| 74 | + |
| 75 | + if current: |
| 76 | + chunks.append("\n\n".join(current)) |
| 77 | + |
| 78 | + return chunks |
| 79 | + |
| 80 | + |
| 81 | +def _extract_model_json(body: str) -> dict: |
| 82 | + """Decode the model's JSON classification from a Gemini response. |
| 83 | +
|
| 84 | + The generated text sits at ``candidates[0].content.parts[0].text`` and is |
| 85 | + itself a JSON string, so decode twice. Any malformed layer yields ``{}``. |
| 86 | + """ |
| 87 | + try: |
| 88 | + envelope = json.loads(body) if body else {} |
| 89 | + except json.JSONDecodeError: |
| 90 | + return {} |
| 91 | + |
| 92 | + try: |
| 93 | + text = envelope["candidates"][0]["content"]["parts"][0]["text"] |
| 94 | + except (KeyError, IndexError, TypeError): |
| 95 | + return {} |
| 96 | + |
| 97 | + try: |
| 98 | + payload = json.loads(text) |
| 99 | + except json.JSONDecodeError: |
| 100 | + return {} |
| 101 | + return payload if isinstance(payload, dict) else {} |
| 102 | + |
| 103 | + |
| 104 | +def _analyze_chunk(chunk_text: str, runtime: Runtime) -> dict[str, str]: |
| 105 | + """Classify one chunk via the Gemini callout and return it as citations.""" |
| 106 | + request = ( |
| 107 | + HTTPRequestBuilder() |
| 108 | + .set_url(CALLOUT_URL) |
| 109 | + .set_method(HTTPMethod.POST) |
| 110 | + .set_headers({"Content-Type": "application/json", "Accept": "application/json"}) |
| 111 | + .set_response_timeout_seconds(60) |
| 112 | + .build() |
| 113 | + ) |
| 114 | + |
| 115 | + payload = { |
| 116 | + "contents": [{"parts": [{"text": _PROMPT + chunk_text}]}], |
| 117 | + "generationConfig": _GENERATION_CONFIG, |
| 118 | + } |
| 119 | + response = runtime.named_credential.request(request, json.dumps(payload)) |
| 120 | + |
| 121 | + # Don't raise: a single failed callout shouldn't abort the whole job. |
| 122 | + if not response.is_success: |
| 123 | + logger.error(f"Gemini callout failed with status {response.status_code}") |
| 124 | + return {"analysis_status": "failed", "http_status": str(response.status_code)} |
| 125 | + |
| 126 | + data = _extract_model_json(response.body) |
| 127 | + citations = {"analysis_status": "success"} |
| 128 | + for field in _ANALYSIS_FIELDS: |
| 129 | + value = data.get(field) |
| 130 | + citations[field] = str(value) if value is not None else "unavailable" |
| 131 | + |
| 132 | + topics = data.get("topics") |
| 133 | + if isinstance(topics, list): |
| 134 | + citations["topics"] = ", ".join(str(topic) for topic in topics) |
| 135 | + |
| 136 | + return citations |
| 137 | + |
| 138 | + |
| 139 | +def function( |
| 140 | + request: SearchIndexChunkingV1Request, runtime: Runtime |
| 141 | +) -> SearchIndexChunkingV1Response: |
| 142 | + """Chunk each input document and classify every chunk via the Gemini API.""" |
| 143 | + logger.info(f"Received {len(request.input)} documents to chunk") |
| 144 | + |
| 145 | + chunks = [] |
| 146 | + chunk_id = 1 |
| 147 | + |
| 148 | + for doc in request.input: |
| 149 | + for chunk_text in _chunk_text(doc.text): |
| 150 | + citations = _analyze_chunk(chunk_text, runtime) |
| 151 | + |
| 152 | + chunk = SearchIndexChunkingV1Output( |
| 153 | + text=chunk_text, |
| 154 | + seq_no=chunk_id, |
| 155 | + chunk_type=ChunkType.TEXT, |
| 156 | + citations=citations, |
| 157 | + ) |
| 158 | + chunks.append(chunk) |
| 159 | + chunk_id += 1 |
| 160 | + |
| 161 | + logger.info(f"Produced {len(chunks)} classified chunks") |
| 162 | + return SearchIndexChunkingV1Response(output=chunks) |
0 commit comments