salesforce-data-customcode 6.0.2__tar.gz → 6.0.3.dev1__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (101) hide show
  1. {salesforce_data_customcode-6.0.2 → salesforce_data_customcode-6.0.3.dev1}/PKG-INFO +47 -1
  2. {salesforce_data_customcode-6.0.2 → salesforce_data_customcode-6.0.3.dev1}/README.md +46 -0
  3. {salesforce_data_customcode-6.0.2 → salesforce_data_customcode-6.0.3.dev1}/pyproject.toml +1 -1
  4. {salesforce_data_customcode-6.0.2 → salesforce_data_customcode-6.0.3.dev1}/src/datacustomcode/__init__.py +15 -0
  5. {salesforce_data_customcode-6.0.2 → salesforce_data_customcode-6.0.3.dev1}/src/datacustomcode/cli.py +3 -6
  6. {salesforce_data_customcode-6.0.2 → salesforce_data_customcode-6.0.3.dev1}/src/datacustomcode/client.py +92 -1
  7. {salesforce_data_customcode-6.0.2 → salesforce_data_customcode-6.0.3.dev1}/src/datacustomcode/config.yaml +3 -0
  8. {salesforce_data_customcode-6.0.2 → salesforce_data_customcode-6.0.3.dev1}/src/datacustomcode/einstein_platform_config.py +8 -4
  9. {salesforce_data_customcode-6.0.2 → salesforce_data_customcode-6.0.3.dev1}/src/datacustomcode/einstein_predictions_config.py +1 -1
  10. {salesforce_data_customcode-6.0.2 → salesforce_data_customcode-6.0.3.dev1}/src/datacustomcode/function/feature_types/chunking.py +1 -0
  11. {salesforce_data_customcode-6.0.2 → salesforce_data_customcode-6.0.3.dev1}/src/datacustomcode/io/writer/base.py +1 -11
  12. {salesforce_data_customcode-6.0.2 → salesforce_data_customcode-6.0.3.dev1}/src/datacustomcode/llm_gateway/__init__.py +6 -0
  13. salesforce_data_customcode-6.0.2/src/datacustomcode/proxy/base.py → salesforce_data_customcode-6.0.3.dev1/src/datacustomcode/llm_gateway/errors.py +17 -5
  14. salesforce_data_customcode-6.0.3.dev1/src/datacustomcode/llm_gateway/spark_base.py +53 -0
  15. salesforce_data_customcode-6.0.3.dev1/src/datacustomcode/llm_gateway/spark_default.py +125 -0
  16. {salesforce_data_customcode-6.0.2 → salesforce_data_customcode-6.0.3.dev1}/src/datacustomcode/llm_gateway_config.py +43 -2
  17. {salesforce_data_customcode-6.0.2 → salesforce_data_customcode-6.0.3.dev1}/src/datacustomcode/templates/function/.devcontainer/devcontainer.json +1 -1
  18. {salesforce_data_customcode-6.0.2 → salesforce_data_customcode-6.0.3.dev1}/src/datacustomcode/templates/function/example/chunking_with_llm/entrypoint.py +3 -0
  19. {salesforce_data_customcode-6.0.2 → salesforce_data_customcode-6.0.3.dev1}/src/datacustomcode/templates/function/example/chunking_with_prediction/entrypoint.py +3 -0
  20. {salesforce_data_customcode-6.0.2 → salesforce_data_customcode-6.0.3.dev1}/src/datacustomcode/templates/function/payload/entrypoint.py +11 -12
  21. {salesforce_data_customcode-6.0.2 → salesforce_data_customcode-6.0.3.dev1}/src/datacustomcode/templates/script/.devcontainer/devcontainer.json +1 -1
  22. salesforce_data_customcode-6.0.3.dev1/src/datacustomcode/templates/script/payload/entrypoint.py +51 -0
  23. {salesforce_data_customcode-6.0.2 → salesforce_data_customcode-6.0.3.dev1}/src/datacustomcode/token_provider.py +36 -70
  24. salesforce_data_customcode-6.0.2/src/datacustomcode/proxy/__init__.py +0 -14
  25. salesforce_data_customcode-6.0.2/src/datacustomcode/proxy/client/__init__.py +0 -14
  26. salesforce_data_customcode-6.0.2/src/datacustomcode/proxy/client/base.py +0 -32
  27. salesforce_data_customcode-6.0.2/src/datacustomcode/templates/script/payload/entrypoint.py +0 -25
  28. {salesforce_data_customcode-6.0.2 → salesforce_data_customcode-6.0.3.dev1}/LICENSE.txt +0 -0
  29. {salesforce_data_customcode-6.0.2 → salesforce_data_customcode-6.0.3.dev1}/src/datacustomcode/auth.py +0 -0
  30. {salesforce_data_customcode-6.0.2 → salesforce_data_customcode-6.0.3.dev1}/src/datacustomcode/cmd.py +0 -0
  31. {salesforce_data_customcode-6.0.2 → salesforce_data_customcode-6.0.3.dev1}/src/datacustomcode/common_config.py +0 -0
  32. {salesforce_data_customcode-6.0.2 → salesforce_data_customcode-6.0.3.dev1}/src/datacustomcode/config.py +0 -0
  33. {salesforce_data_customcode-6.0.2 → salesforce_data_customcode-6.0.3.dev1}/src/datacustomcode/constants.py +0 -0
  34. {salesforce_data_customcode-6.0.2 → salesforce_data_customcode-6.0.3.dev1}/src/datacustomcode/credentials.py +0 -0
  35. {salesforce_data_customcode-6.0.2 → salesforce_data_customcode-6.0.3.dev1}/src/datacustomcode/deploy.py +0 -0
  36. {salesforce_data_customcode-6.0.2 → salesforce_data_customcode-6.0.3.dev1}/src/datacustomcode/einstein_platform_client.py +0 -0
  37. {salesforce_data_customcode-6.0.2 → salesforce_data_customcode-6.0.3.dev1}/src/datacustomcode/einstein_predictions/__init__.py +0 -0
  38. {salesforce_data_customcode-6.0.2 → salesforce_data_customcode-6.0.3.dev1}/src/datacustomcode/einstein_predictions/base.py +0 -0
  39. {salesforce_data_customcode-6.0.2 → salesforce_data_customcode-6.0.3.dev1}/src/datacustomcode/einstein_predictions/impl/default.py +0 -0
  40. {salesforce_data_customcode-6.0.2 → salesforce_data_customcode-6.0.3.dev1}/src/datacustomcode/einstein_predictions/types.py +0 -0
  41. {salesforce_data_customcode-6.0.2 → salesforce_data_customcode-6.0.3.dev1}/src/datacustomcode/file/__init__.py +0 -0
  42. {salesforce_data_customcode-6.0.2 → salesforce_data_customcode-6.0.3.dev1}/src/datacustomcode/file/base.py +0 -0
  43. {salesforce_data_customcode-6.0.2 → salesforce_data_customcode-6.0.3.dev1}/src/datacustomcode/file/path/__init__.py +0 -0
  44. {salesforce_data_customcode-6.0.2 → salesforce_data_customcode-6.0.3.dev1}/src/datacustomcode/file/path/default.py +0 -0
  45. {salesforce_data_customcode-6.0.2 → salesforce_data_customcode-6.0.3.dev1}/src/datacustomcode/function/__init__.py +0 -0
  46. {salesforce_data_customcode-6.0.2 → salesforce_data_customcode-6.0.3.dev1}/src/datacustomcode/function/base.py +0 -0
  47. {salesforce_data_customcode-6.0.2 → salesforce_data_customcode-6.0.3.dev1}/src/datacustomcode/function/feature_types/__init__.py +0 -0
  48. {salesforce_data_customcode-6.0.2 → salesforce_data_customcode-6.0.3.dev1}/src/datacustomcode/function/runtime.py +0 -0
  49. {salesforce_data_customcode-6.0.2 → salesforce_data_customcode-6.0.3.dev1}/src/datacustomcode/function_utils.py +0 -0
  50. {salesforce_data_customcode-6.0.2 → salesforce_data_customcode-6.0.3.dev1}/src/datacustomcode/io/__init__.py +0 -0
  51. {salesforce_data_customcode-6.0.2 → salesforce_data_customcode-6.0.3.dev1}/src/datacustomcode/io/base.py +0 -0
  52. {salesforce_data_customcode-6.0.2 → salesforce_data_customcode-6.0.3.dev1}/src/datacustomcode/io/reader/__init__.py +0 -0
  53. {salesforce_data_customcode-6.0.2 → salesforce_data_customcode-6.0.3.dev1}/src/datacustomcode/io/reader/base.py +0 -0
  54. {salesforce_data_customcode-6.0.2 → salesforce_data_customcode-6.0.3.dev1}/src/datacustomcode/io/reader/query_api.py +0 -0
  55. {salesforce_data_customcode-6.0.2 → salesforce_data_customcode-6.0.3.dev1}/src/datacustomcode/io/reader/sf_cli.py +0 -0
  56. {salesforce_data_customcode-6.0.2 → salesforce_data_customcode-6.0.3.dev1}/src/datacustomcode/io/reader/utils.py +0 -0
  57. {salesforce_data_customcode-6.0.2 → salesforce_data_customcode-6.0.3.dev1}/src/datacustomcode/io/writer/__init__.py +0 -0
  58. {salesforce_data_customcode-6.0.2 → salesforce_data_customcode-6.0.3.dev1}/src/datacustomcode/io/writer/csv.py +0 -0
  59. {salesforce_data_customcode-6.0.2 → salesforce_data_customcode-6.0.3.dev1}/src/datacustomcode/io/writer/print.py +0 -0
  60. {salesforce_data_customcode-6.0.2 → salesforce_data_customcode-6.0.3.dev1}/src/datacustomcode/llm_gateway/base.py +0 -0
  61. {salesforce_data_customcode-6.0.2 → salesforce_data_customcode-6.0.3.dev1}/src/datacustomcode/llm_gateway/default.py +0 -0
  62. {salesforce_data_customcode-6.0.2 → salesforce_data_customcode-6.0.3.dev1}/src/datacustomcode/llm_gateway/types/__init__.py +0 -0
  63. {salesforce_data_customcode-6.0.2 → salesforce_data_customcode-6.0.3.dev1}/src/datacustomcode/llm_gateway/types/generate_text_request.py +0 -0
  64. {salesforce_data_customcode-6.0.2 → salesforce_data_customcode-6.0.3.dev1}/src/datacustomcode/llm_gateway/types/generate_text_request_builder.py +0 -0
  65. {salesforce_data_customcode-6.0.2 → salesforce_data_customcode-6.0.3.dev1}/src/datacustomcode/llm_gateway/types/generate_text_response.py +0 -0
  66. {salesforce_data_customcode-6.0.2 → salesforce_data_customcode-6.0.3.dev1}/src/datacustomcode/llm_gateway/types/generate_text_response_builder.py +0 -0
  67. {salesforce_data_customcode-6.0.2 → salesforce_data_customcode-6.0.3.dev1}/src/datacustomcode/mixin.py +0 -0
  68. {salesforce_data_customcode-6.0.2 → salesforce_data_customcode-6.0.3.dev1}/src/datacustomcode/py.typed +0 -0
  69. {salesforce_data_customcode-6.0.2 → salesforce_data_customcode-6.0.3.dev1}/src/datacustomcode/run.py +0 -0
  70. {salesforce_data_customcode-6.0.2 → salesforce_data_customcode-6.0.3.dev1}/src/datacustomcode/scan.py +0 -0
  71. {salesforce_data_customcode-6.0.2 → salesforce_data_customcode-6.0.3.dev1}/src/datacustomcode/spark/__init__.py +0 -0
  72. {salesforce_data_customcode-6.0.2 → salesforce_data_customcode-6.0.3.dev1}/src/datacustomcode/spark/base.py +0 -0
  73. {salesforce_data_customcode-6.0.2 → salesforce_data_customcode-6.0.3.dev1}/src/datacustomcode/spark/default.py +0 -0
  74. {salesforce_data_customcode-6.0.2 → salesforce_data_customcode-6.0.3.dev1}/src/datacustomcode/template.py +0 -0
  75. {salesforce_data_customcode-6.0.2 → salesforce_data_customcode-6.0.3.dev1}/src/datacustomcode/templates/__init__.py +0 -0
  76. {salesforce_data_customcode-6.0.2 → salesforce_data_customcode-6.0.3.dev1}/src/datacustomcode/templates/function/Dockerfile.dependencies +0 -0
  77. {salesforce_data_customcode-6.0.2 → salesforce_data_customcode-6.0.3.dev1}/src/datacustomcode/templates/function/README.md +0 -0
  78. {salesforce_data_customcode-6.0.2 → salesforce_data_customcode-6.0.3.dev1}/src/datacustomcode/templates/function/build_native_dependencies.sh +0 -0
  79. {salesforce_data_customcode-6.0.2 → salesforce_data_customcode-6.0.3.dev1}/src/datacustomcode/templates/function/chunking/payload/config.json +0 -0
  80. {salesforce_data_customcode-6.0.2 → salesforce_data_customcode-6.0.3.dev1}/src/datacustomcode/templates/function/chunking/payload/entrypoint.py +0 -0
  81. {salesforce_data_customcode-6.0.2 → salesforce_data_customcode-6.0.3.dev1}/src/datacustomcode/templates/function/chunking/requirements.txt +0 -0
  82. {salesforce_data_customcode-6.0.2 → salesforce_data_customcode-6.0.3.dev1}/src/datacustomcode/templates/function/example/chunking_with_llm/config.json +0 -0
  83. {salesforce_data_customcode-6.0.2 → salesforce_data_customcode-6.0.3.dev1}/src/datacustomcode/templates/function/example/chunking_with_llm/files/chunking_prompt.txt +0 -0
  84. {salesforce_data_customcode-6.0.2 → salesforce_data_customcode-6.0.3.dev1}/src/datacustomcode/templates/function/example/chunking_with_llm/tests/test.json +0 -0
  85. {salesforce_data_customcode-6.0.2 → salesforce_data_customcode-6.0.3.dev1}/src/datacustomcode/templates/function/example/chunking_with_prediction/config.json +0 -0
  86. {salesforce_data_customcode-6.0.2 → salesforce_data_customcode-6.0.3.dev1}/src/datacustomcode/templates/function/example/chunking_with_prediction/tests/test.json +0 -0
  87. {salesforce_data_customcode-6.0.2 → salesforce_data_customcode-6.0.3.dev1}/src/datacustomcode/templates/function/payload/config.json +0 -0
  88. {salesforce_data_customcode-6.0.2 → salesforce_data_customcode-6.0.3.dev1}/src/datacustomcode/templates/function/requirements-dev.txt +0 -0
  89. {salesforce_data_customcode-6.0.2 → salesforce_data_customcode-6.0.3.dev1}/src/datacustomcode/templates/function/requirements.txt +0 -0
  90. {salesforce_data_customcode-6.0.2 → salesforce_data_customcode-6.0.3.dev1}/src/datacustomcode/templates/script/Dockerfile +0 -0
  91. {salesforce_data_customcode-6.0.2 → salesforce_data_customcode-6.0.3.dev1}/src/datacustomcode/templates/script/Dockerfile.dependencies +0 -0
  92. {salesforce_data_customcode-6.0.2 → salesforce_data_customcode-6.0.3.dev1}/src/datacustomcode/templates/script/README.md +0 -0
  93. {salesforce_data_customcode-6.0.2 → salesforce_data_customcode-6.0.3.dev1}/src/datacustomcode/templates/script/account.ipynb +0 -0
  94. {salesforce_data_customcode-6.0.2 → salesforce_data_customcode-6.0.3.dev1}/src/datacustomcode/templates/script/build_native_dependencies.sh +0 -0
  95. {salesforce_data_customcode-6.0.2 → salesforce_data_customcode-6.0.3.dev1}/src/datacustomcode/templates/script/examples/employee_hierarchy/employee_data.csv +0 -0
  96. {salesforce_data_customcode-6.0.2 → salesforce_data_customcode-6.0.3.dev1}/src/datacustomcode/templates/script/examples/employee_hierarchy/entrypoint.py +0 -0
  97. {salesforce_data_customcode-6.0.2 → salesforce_data_customcode-6.0.3.dev1}/src/datacustomcode/templates/script/jupyterlab.sh +0 -0
  98. {salesforce_data_customcode-6.0.2 → salesforce_data_customcode-6.0.3.dev1}/src/datacustomcode/templates/script/payload/config.json +0 -0
  99. {salesforce_data_customcode-6.0.2 → salesforce_data_customcode-6.0.3.dev1}/src/datacustomcode/templates/script/requirements-dev.txt +0 -0
  100. {salesforce_data_customcode-6.0.2 → salesforce_data_customcode-6.0.3.dev1}/src/datacustomcode/templates/script/requirements.txt +0 -0
  101. {salesforce_data_customcode-6.0.2 → salesforce_data_customcode-6.0.3.dev1}/src/datacustomcode/version.py +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: salesforce-data-customcode
3
- Version: 6.0.2
3
+ Version: 6.0.3.dev1
4
4
  Summary: Data Cloud Custom Code SDK
5
5
  License-Expression: Apache-2.0
6
6
  License-File: LICENSE.txt
@@ -330,6 +330,52 @@ Options:
330
330
  - `--function-invoke-opt TEXT`: Currently we support only `UnstructuredChunking` for functions.
331
331
 
332
332
 
333
+ ## Testing LLM Gateway
334
+
335
+ You can use AI models configured in Salesforce to generate responses while transforming your data. Below is a sample code example:
336
+
337
+ ```
338
+ from datacustomcode.client import Client, llm_gateway_generate_text_col
339
+
340
+
341
+ def main():
342
+ client = Client()
343
+ df = client.read_dlo("Input__dll")
344
+ df_generated = df.withColumn(
345
+ "greeting__c",
346
+ llm_gateway_generate_text_col(
347
+ "In one sentence, greet {name} from {city}.",
348
+ {"name": col("name__c"), "city": col("homecity__c")},
349
+ model_id="sfdc_ai__DefaultGPT4Omni", # An AI model in your org
350
+ ),
351
+ )
352
+
353
+ dlo_name = "Output_dll"
354
+ client.write_to_dlo(dlo_name, df_upper1, write_mode=WriteMode.APPEND)
355
+
356
+ greeting = client.llm_gateway_generate_text("In one sentence, generate a greeting message", "sfdc_ai__DefaultGPT52")
357
+
358
+ if __name__ == "__main__":
359
+ main()
360
+ ```
361
+
362
+ In order to test this code on your local machine before deploying it to Data Cloud, you must first set up an External Client App that allows access to the Agent API. Follow this guide to create the ECA https://developer.salesforce.com/docs/ai/agentforce/guide/agent-api-get-started.html#create-a-salesforce-app. You must use `http://localhost:1717/OauthRedirect` as the callback URL.
363
+
364
+ Once the ECA is set up, log in to your org using this ECA
365
+ ```
366
+ sf org login web \
367
+ --alias myorg \
368
+ --instance-url https://{MY_DOMAIN_URL} \
369
+ --client-id {CONSUMER_KEY} \
370
+ --scopes "sfap_api api"
371
+ ```
372
+
373
+ then you can test your code using `myorg` alias
374
+ ```
375
+ datacustomcode run ./payload/entrypoint.py --sf-cli-org myorg
376
+ ```
377
+
378
+
333
379
  ## Docker usage
334
380
 
335
381
  The SDK provides Docker-based development options that allow you to test your code in an environment that closely resembles Data Cloud's execution environment.
@@ -305,6 +305,52 @@ Options:
305
305
  - `--function-invoke-opt TEXT`: Currently we support only `UnstructuredChunking` for functions.
306
306
 
307
307
 
308
+ ## Testing LLM Gateway
309
+
310
+ You can use AI models configured in Salesforce to generate responses while transforming your data. Below is a sample code example:
311
+
312
+ ```
313
+ from datacustomcode.client import Client, llm_gateway_generate_text_col
314
+
315
+
316
+ def main():
317
+ client = Client()
318
+ df = client.read_dlo("Input__dll")
319
+ df_generated = df.withColumn(
320
+ "greeting__c",
321
+ llm_gateway_generate_text_col(
322
+ "In one sentence, greet {name} from {city}.",
323
+ {"name": col("name__c"), "city": col("homecity__c")},
324
+ model_id="sfdc_ai__DefaultGPT4Omni", # An AI model in your org
325
+ ),
326
+ )
327
+
328
+ dlo_name = "Output_dll"
329
+ client.write_to_dlo(dlo_name, df_upper1, write_mode=WriteMode.APPEND)
330
+
331
+ greeting = client.llm_gateway_generate_text("In one sentence, generate a greeting message", "sfdc_ai__DefaultGPT52")
332
+
333
+ if __name__ == "__main__":
334
+ main()
335
+ ```
336
+
337
+ In order to test this code on your local machine before deploying it to Data Cloud, you must first set up an External Client App that allows access to the Agent API. Follow this guide to create the ECA https://developer.salesforce.com/docs/ai/agentforce/guide/agent-api-get-started.html#create-a-salesforce-app. You must use `http://localhost:1717/OauthRedirect` as the callback URL.
338
+
339
+ Once the ECA is set up, log in to your org using this ECA
340
+ ```
341
+ sf org login web \
342
+ --alias myorg \
343
+ --instance-url https://{MY_DOMAIN_URL} \
344
+ --client-id {CONSUMER_KEY} \
345
+ --scopes "sfap_api api"
346
+ ```
347
+
348
+ then you can test your code using `myorg` alias
349
+ ```
350
+ datacustomcode run ./payload/entrypoint.py --sf-cli-org myorg
351
+ ```
352
+
353
+
308
354
  ## Docker usage
309
355
 
310
356
  The SDK provides Docker-based development options that allow you to test your code in an environment that closely resembles Data Cloud's execution environment.
@@ -18,7 +18,7 @@ license = "Apache-2.0"
18
18
  name = "salesforce-data-customcode"
19
19
  readme = "README.md"
20
20
  requires-python = ">=3.10,<3.12"
21
- version = "6.0.2"
21
+ version = "6.0.3.dev1"
22
22
 
23
23
  [tool.black]
24
24
  exclude = '''
@@ -17,8 +17,11 @@ __all__ = [
17
17
  "AuthType",
18
18
  "Client",
19
19
  "Credentials",
20
+ "DefaultSparkLLMGateway",
20
21
  "PrintDataCloudWriter",
21
22
  "QueryAPIDataCloudReader",
23
+ "SparkLLMGateway",
24
+ "llm_gateway_generate_text_col",
22
25
  ]
23
26
 
24
27
 
@@ -44,4 +47,16 @@ def __getattr__(name: str):
44
47
  from datacustomcode.io.reader.query_api import QueryAPIDataCloudReader
45
48
 
46
49
  return QueryAPIDataCloudReader
50
+ elif name == "SparkLLMGateway":
51
+ from datacustomcode.llm_gateway import SparkLLMGateway
52
+
53
+ return SparkLLMGateway
54
+ elif name == "DefaultSparkLLMGateway":
55
+ from datacustomcode.llm_gateway import DefaultSparkLLMGateway
56
+
57
+ return DefaultSparkLLMGateway
58
+ elif name == "llm_gateway_generate_text_col":
59
+ from datacustomcode.client import llm_gateway_generate_text_col
60
+
61
+ return llm_gateway_generate_text_col
47
62
  raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
@@ -198,7 +198,6 @@ def zip(path: str, network: str):
198
198
  default=None,
199
199
  help="SF CLI org alias or username. Fetches credentials via `sf org display`.",
200
200
  )
201
- @click.option("--use-in-feature", default=None, hidden=True, deprecated=True)
202
201
  def deploy(
203
202
  path: str,
204
203
  name: str,
@@ -208,7 +207,6 @@ def deploy(
208
207
  profile: str,
209
208
  network: str,
210
209
  sf_cli_org: Optional[str],
211
- use_in_feature: Optional[str],
212
210
  ):
213
211
  from datacustomcode.constants import USE_IN_FEATURE_MAPPING_FOR_CONNECT_API
214
212
  from datacustomcode.deploy import (
@@ -251,14 +249,13 @@ def deploy(
251
249
  logger.info(f"Inferred use_in_feature: {use_in_feature}")
252
250
  else:
253
251
  click.secho(
254
- "Error: Function signature does not match a supported type. "
255
- "Use SearchIndexChunkingV1Request and "
256
- "SearchIndexChunkingV1Response in function signature.",
252
+ "Error: Could not infer function invoke options. "
253
+ "Please provide --use-in-feature",
257
254
  fg="red",
258
255
  )
259
256
  raise click.Abort()
260
257
 
261
- # Map feature names to Connect API names
258
+ # Map user-provided feature names to API names
262
259
  mapped_feature = USE_IN_FEATURE_MAPPING_FOR_CONNECT_API.get(
263
260
  use_in_feature, use_in_feature
264
261
  )
@@ -18,24 +18,75 @@ from enum import Enum
18
18
  from typing import (
19
19
  TYPE_CHECKING,
20
20
  ClassVar,
21
+ Dict,
21
22
  Optional,
23
+ Union,
22
24
  )
23
25
 
24
26
  from datacustomcode.config import config
25
27
  from datacustomcode.file.path.default import DefaultFindFilePath
26
28
  from datacustomcode.io.reader.base import BaseDataCloudReader
29
+ from datacustomcode.llm_gateway_config import spark_llm_gateway_config
27
30
  from datacustomcode.spark.default import DefaultSparkSessionProvider
28
31
 
29
32
  if TYPE_CHECKING:
30
33
  from pathlib import Path
31
34
 
32
- from pyspark.sql import DataFrame as PySparkDataFrame
35
+ from pyspark.sql import Column, DataFrame as PySparkDataFrame
33
36
 
34
37
  from datacustomcode.io.reader.base import BaseDataCloudReader
35
38
  from datacustomcode.io.writer.base import BaseDataCloudWriter, WriteMode
39
+ from datacustomcode.llm_gateway.spark_base import SparkLLMGateway
36
40
  from datacustomcode.spark.base import BaseSparkSessionProvider
37
41
 
38
42
 
43
+ def _build_spark_llm_gateway() -> "SparkLLMGateway":
44
+ """Instantiate the SDK-configured :class:`SparkLLMGateway`.
45
+
46
+ Raises:
47
+ RuntimeError: If no ``spark_llm_gateway_config`` has been loaded.
48
+ """
49
+ cfg = spark_llm_gateway_config.spark_llm_gateway_config
50
+ if cfg is None:
51
+ raise RuntimeError(
52
+ "spark_llm_gateway_config is not configured. Add a "
53
+ "'spark_llm_gateway_config' section to config.yaml."
54
+ )
55
+ return cfg.to_object()
56
+
57
+
58
+ def llm_gateway_generate_text_col(
59
+ template: str,
60
+ values: Union[Dict[str, "Column"], "Column"],
61
+ model_id: Optional[str] = None,
62
+ ) -> "Column":
63
+ """Build a Spark Column that runs the LLM Gateway per row.
64
+
65
+ Example:
66
+
67
+ >>> df.withColumn(
68
+ ... "greeting__c",
69
+ ... llm_gateway_generate_text_col(
70
+ ... "In one sentence, greet {name} from {city}.",
71
+ ... {"name": col("name__c"), "city": col("homecity__c")},
72
+ ... model_id="sfdc_ai__DefaultGPT4Omni",
73
+ ... ),
74
+ ... )
75
+
76
+ Args:
77
+ template: The prompt template, with ``{field}`` placeholders matching
78
+ keys in ``values``. Substitution uses ``str.format``.
79
+ values: Either a mapping from placeholder name to Spark ``Column``, or
80
+ a single ``Column`` whose value is already a struct.
81
+ model_id: LLM model id. Defaults to ``sfdc_ai__DefaultGPT4Omni``.
82
+
83
+ Returns:
84
+ A Spark ``Column`` that, when evaluated, produces the generated text.
85
+ """
86
+ gateway = Client()._get_spark_llm_gateway()
87
+ return gateway.llm_gateway_generate_text_col(template, values, model_id=model_id)
88
+
89
+
39
90
  class DataCloudObjectType(Enum):
40
91
  DLO = "dlo"
41
92
  DMO = "dmo"
@@ -94,18 +145,21 @@ class Client:
94
145
  finder: Find a file path
95
146
  reader: A custom reader to use for reading Data Cloud objects.
96
147
  writer: A custom writer to use for writing Data Cloud objects.
148
+ spark_llm_gateway: Optional custom :class:`SparkLLMGateway`.
97
149
 
98
150
  Example:
99
151
  >>> client = Client()
100
152
  >>> file_path = client.find_file_path("data.csv")
101
153
  >>> dlo = client.read_dlo("my_dlo")
102
154
  >>> client.write_to_dmo("my_dmo", dlo)
155
+ >>> answer = client.llm_gateway_generate_text("Generate a greeting message")
103
156
  """
104
157
 
105
158
  _instance: ClassVar[Optional[Client]] = None
106
159
  _reader: BaseDataCloudReader
107
160
  _writer: BaseDataCloudWriter
108
161
  _file: DefaultFindFilePath
162
+ _spark_llm_gateway: Optional[SparkLLMGateway]
109
163
  _data_layer_history: dict[DataCloudObjectType, set[str]]
110
164
  _code_type: str
111
165
 
@@ -114,11 +168,13 @@ class Client:
114
168
  reader: Optional[BaseDataCloudReader] = None,
115
169
  writer: Optional[BaseDataCloudWriter] = None,
116
170
  spark_provider: Optional[BaseSparkSessionProvider] = None,
171
+ spark_llm_gateway: Optional[SparkLLMGateway] = None,
117
172
  code_type: str = "script",
118
173
  ) -> Client:
119
174
 
120
175
  if cls._instance is None:
121
176
  cls._instance = super().__new__(cls)
177
+ cls._instance._spark_llm_gateway = spark_llm_gateway
122
178
  # Initialize Readers and Writers from config
123
179
  # and/or provided reader and writer
124
180
  if reader is None or writer is None:
@@ -225,6 +281,41 @@ class Client:
225
281
 
226
282
  return self._file.find_file_path(file_name) # type: ignore[no-any-return]
227
283
 
284
+ def llm_gateway_generate_text(
285
+ self,
286
+ prompt: str,
287
+ model_id: Optional[str] = None,
288
+ ) -> str:
289
+ """Issue a one-shot LLM Gateway call. This is the scalar counterpart to
290
+ :func:`llm_gateway_generate_text_col`: it runs **once** — not per row.
291
+ Use the column helper method instead when you want to fan a prompt out across
292
+ every row of a DataFrame.
293
+
294
+ Example:
295
+
296
+ >>> response = Client().llm_gateway_generate_text(
297
+ ... "Generate a greeting message"
298
+ ... )
299
+
300
+ Args:
301
+ prompt: The literal prompt to send. Plain text — no
302
+ ``{field}`` substitution is performed on this string.
303
+ model_id: LLM model id to target. Defaults to
304
+ ``sfdc_ai__DefaultGPT4Omni`` when ``None``.
305
+
306
+ Returns:
307
+ The generated text as a plain Python ``str``; empty when the
308
+ gateway response carries no generated text.
309
+ """
310
+ return self._get_spark_llm_gateway().llm_gateway_generate_text(
311
+ prompt, model_id=model_id
312
+ )
313
+
314
+ def _get_spark_llm_gateway(self) -> SparkLLMGateway:
315
+ if self._spark_llm_gateway is None:
316
+ self._spark_llm_gateway = _build_spark_llm_gateway()
317
+ return self._spark_llm_gateway
318
+
228
319
  def _validate_data_layer_history_does_not_contain(
229
320
  self, data_cloud_object_type: DataCloudObjectType
230
321
  ) -> None:
@@ -28,3 +28,6 @@ llm_gateway_config:
28
28
  type_config_name: DefaultLLMGateway
29
29
  options:
30
30
  credentials_profile: default
31
+
32
+ spark_llm_gateway_config:
33
+ type_config_name: DefaultSparkLLMGateway
@@ -15,20 +15,23 @@
15
15
 
16
16
  from typing import (
17
17
  ClassVar,
18
+ Generic,
18
19
  Optional,
19
20
  Type,
20
- cast,
21
+ TypeVar,
21
22
  )
22
23
 
23
24
  from datacustomcode.common_config import BaseObjectConfig
24
25
 
26
+ _T = TypeVar("_T")
25
27
 
26
- class CredentialsObjectConfig(BaseObjectConfig):
28
+
29
+ class CredentialsObjectConfig(BaseObjectConfig, Generic[_T]):
27
30
  type_to_create: ClassVar[Type]
28
31
  credentials_profile: Optional[str] = None
29
32
  sf_cli_org: Optional[str] = None
30
33
 
31
- def to_object(self):
34
+ def to_object(self) -> _T:
32
35
  """Create an object instance, automatically including credentials in options"""
33
36
 
34
37
  options = self.options.copy()
@@ -38,4 +41,5 @@ class CredentialsObjectConfig(BaseObjectConfig):
38
41
  options["sf_cli_org"] = self.sf_cli_org
39
42
 
40
43
  type_ = self.type_to_create.subclass_from_config_name(self.type_config_name)
41
- return cast(type_, type_(**options))
44
+ instance: _T = type_(**options)
45
+ return instance
@@ -28,7 +28,7 @@ from datacustomcode.einstein_predictions.base import EinsteinPredictions
28
28
  _E = TypeVar("_E", bound=EinsteinPredictions)
29
29
 
30
30
 
31
- class EinsteinPredictionsObjectConfig(CredentialsObjectConfig, Generic[_E]):
31
+ class EinsteinPredictionsObjectConfig(CredentialsObjectConfig[_E], Generic[_E]):
32
32
  type_to_create: ClassVar[Type[EinsteinPredictions]] = EinsteinPredictions # type: ignore[type-abstract]
33
33
 
34
34
 
@@ -16,6 +16,7 @@
16
16
  """
17
17
  Pydantic models for Search Index Chunking V1
18
18
  """
19
+
19
20
  from enum import Enum
20
21
  from typing import (
21
22
  Dict,
@@ -27,21 +27,11 @@ if TYPE_CHECKING:
27
27
  class WriteMode(str, Enum):
28
28
  APPEND = "append"
29
29
  OVERWRITE = "overwrite"
30
- OVERWRITE_PARTITIONS = "overwrite_partitions" # Deprecated: raises error if used
30
+ OVERWRITE_PARTITIONS = "overwrite_partitions"
31
31
  MERGE = "merge"
32
32
  MERGE_UPSERT_DELETE = "merge_upsert_delete"
33
33
 
34
34
 
35
- class MergeRecordType(str, Enum):
36
- """Values for the _merge_record_type column used by MERGE_UPSERT_DELETE."""
37
-
38
- UPSERT = "UPSERT"
39
- DELETE = "DELETE"
40
-
41
-
42
- MERGE_RECORD_TYPE_COLUMN = "_merge_record_type"
43
-
44
-
45
35
  class BaseDataCloudWriter(BaseDataAccessLayer):
46
36
  """Base class for Data Cloud writers."""
47
37
 
@@ -15,8 +15,14 @@
15
15
 
16
16
  from datacustomcode.llm_gateway.base import LLMGateway
17
17
  from datacustomcode.llm_gateway.default import DefaultLLMGateway
18
+ from datacustomcode.llm_gateway.errors import LLMGatewayCallError
19
+ from datacustomcode.llm_gateway.spark_base import SparkLLMGateway
20
+ from datacustomcode.llm_gateway.spark_default import DefaultSparkLLMGateway
18
21
 
19
22
  __all__ = [
20
23
  "DefaultLLMGateway",
24
+ "DefaultSparkLLMGateway",
21
25
  "LLMGateway",
26
+ "LLMGatewayCallError",
27
+ "SparkLLMGateway",
22
28
  ]
@@ -12,13 +12,25 @@
12
12
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13
13
  # See the License for the specific language governing permissions and
14
14
  # limitations under the License.
15
+ """Exceptions raised by LLM Gateway implementations."""
16
+
15
17
  from __future__ import annotations
16
18
 
17
- from abc import ABC
19
+ from typing import Optional
18
20
 
19
- from datacustomcode.mixin import UserExtendableNamedConfigMixin
20
21
 
22
+ class LLMGatewayCallError(RuntimeError):
23
+ """Raised when an LLM Gateway call returns an error."""
21
24
 
22
- class BaseProxyAccessLayer(ABC, UserExtendableNamedConfigMixin):
23
- def __init__(self):
24
- pass
25
+ def __init__(
26
+ self,
27
+ message: str,
28
+ *,
29
+ status: Optional[object] = None,
30
+ error_code: Optional[str] = None,
31
+ error_message: Optional[str] = None,
32
+ ) -> None:
33
+ super().__init__(message)
34
+ self.status = status
35
+ self.error_code = error_code
36
+ self.error_message = error_message
@@ -0,0 +1,53 @@
1
+ # Copyright (c) 2025, Salesforce, Inc.
2
+ # SPDX-License-Identifier: Apache-2
3
+ #
4
+ # Licensed under the Apache License, Version 2.0 (the "License");
5
+ # you may not use this file except in compliance with the License.
6
+ # You may obtain a copy of the License at
7
+ #
8
+ # http://www.apache.org/licenses/LICENSE-2.0
9
+ #
10
+ # Unless required by applicable law or agreed to in writing, software
11
+ # distributed under the License is distributed on an "AS IS" BASIS,
12
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13
+ # See the License for the specific language governing permissions and
14
+ # limitations under the License.
15
+ from __future__ import annotations
16
+
17
+ from abc import ABC, abstractmethod
18
+ from typing import (
19
+ TYPE_CHECKING,
20
+ Any,
21
+ Dict,
22
+ Optional,
23
+ Union,
24
+ )
25
+
26
+ from datacustomcode.mixin import UserExtendableNamedConfigMixin
27
+
28
+ if TYPE_CHECKING:
29
+ from pyspark.sql import Column
30
+
31
+
32
+ class SparkLLMGateway(ABC, UserExtendableNamedConfigMixin):
33
+ CONFIG_NAME: str
34
+
35
+ def __init__(self, **kwargs: Any) -> None:
36
+ pass
37
+
38
+ @abstractmethod
39
+ def llm_gateway_generate_text(
40
+ self,
41
+ prompt: str,
42
+ model_id: Optional[str] = None,
43
+ ) -> str:
44
+ """Issue a one-shot LLM Gateway call and return the generated text."""
45
+
46
+ @abstractmethod
47
+ def llm_gateway_generate_text_col(
48
+ self,
49
+ template: str,
50
+ values: Union[Dict[str, "Column"], "Column"],
51
+ model_id: Optional[str] = None,
52
+ ) -> "Column":
53
+ """Build a Spark ``Column`` that invokes the LLM Gateway per row."""
@@ -0,0 +1,125 @@
1
+ # Copyright (c) 2025, Salesforce, Inc.
2
+ # SPDX-License-Identifier: Apache-2
3
+ #
4
+ # Licensed under the Apache License, Version 2.0 (the "License");
5
+ # you may not use this file except in compliance with the License.
6
+ # You may obtain a copy of the License at
7
+ #
8
+ # http://www.apache.org/licenses/LICENSE-2.0
9
+ #
10
+ # Unless required by applicable law or agreed to in writing, software
11
+ # distributed under the License is distributed on an "AS IS" BASIS,
12
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13
+ # See the License for the specific language governing permissions and
14
+ # limitations under the License.
15
+ from __future__ import annotations
16
+
17
+ from typing import (
18
+ TYPE_CHECKING,
19
+ Any,
20
+ Dict,
21
+ Optional,
22
+ Union,
23
+ )
24
+
25
+ from datacustomcode.llm_gateway.spark_base import SparkLLMGateway
26
+
27
+ if TYPE_CHECKING:
28
+ from pyspark.sql import Column
29
+
30
+ from datacustomcode.llm_gateway.base import LLMGateway
31
+
32
+
33
+ _DEFAULT_LLM_MODEL_ID = "sfdc_ai__DefaultGPT4Omni"
34
+
35
+
36
+ class DefaultSparkLLMGateway(SparkLLMGateway):
37
+
38
+ CONFIG_NAME = "DefaultSparkLLMGateway"
39
+
40
+ def __init__(
41
+ self,
42
+ llm_gateway: Optional["LLMGateway"] = None,
43
+ **kwargs: Any,
44
+ ) -> None:
45
+ super().__init__(**kwargs)
46
+ if llm_gateway is None:
47
+ llm_gateway = _build_underlying_gateway()
48
+ self._llm_gateway: "LLMGateway" = llm_gateway
49
+
50
+ def llm_gateway_generate_text(
51
+ self,
52
+ prompt: str,
53
+ model_id: Optional[str] = None,
54
+ ) -> str:
55
+ return _invoke_llm_gateway(self._llm_gateway, prompt, model_id)
56
+
57
+ def llm_gateway_generate_text_col(
58
+ self,
59
+ template: str,
60
+ values: Union[Dict[str, "Column"], "Column"],
61
+ model_id: Optional[str] = None,
62
+ ) -> "Column":
63
+
64
+ from pyspark.sql.functions import struct, udf
65
+ from pyspark.sql.types import StringType
66
+
67
+ if isinstance(values, dict):
68
+ values_col = struct(*[v.alias(k) for k, v in values.items()])
69
+ else:
70
+ values_col = values
71
+
72
+ gateway = self._llm_gateway
73
+
74
+ def _generate(values_row: Any) -> str:
75
+ if values_row is None:
76
+ return ""
77
+ subs = (
78
+ values_row.asDict()
79
+ if hasattr(values_row, "asDict")
80
+ else dict(values_row)
81
+ )
82
+ prompt = template.format(**subs)
83
+ return _invoke_llm_gateway(gateway, prompt, model_id)
84
+
85
+ return udf(_generate, StringType())(values_col)
86
+
87
+
88
+ def _build_underlying_gateway() -> "LLMGateway":
89
+ from datacustomcode.llm_gateway_config import llm_gateway_config
90
+
91
+ cfg = llm_gateway_config.llm_gateway_config
92
+ if cfg is None:
93
+ raise RuntimeError(
94
+ "llm_gateway_config is not configured. Add an 'llm_gateway_config' "
95
+ "section to config.yaml."
96
+ )
97
+ return cfg.to_object()
98
+
99
+
100
+ def _invoke_llm_gateway(
101
+ gateway: "LLMGateway",
102
+ prompt: str,
103
+ model_id: Optional[str],
104
+ ) -> str:
105
+ from datacustomcode.llm_gateway.errors import LLMGatewayCallError
106
+ from datacustomcode.llm_gateway.types.generate_text_request_builder import (
107
+ GenerateTextRequestBuilder,
108
+ )
109
+
110
+ builder = (
111
+ GenerateTextRequestBuilder()
112
+ .set_prompt(prompt)
113
+ .set_model(model_id or _DEFAULT_LLM_MODEL_ID)
114
+ )
115
+ response = gateway.generate_text(builder.build())
116
+ if response.is_error:
117
+ raise LLMGatewayCallError(
118
+ f"LLM Gateway call failed: status_code={response.status_code}, "
119
+ f"error_code={response.error_code!r}, "
120
+ f"message={response.data!r}",
121
+ status=response.status_code,
122
+ error_code=response.error_code or None,
123
+ error_message=str(response.data) if response.data else None,
124
+ )
125
+ return response.text