response

Author: joroscoSF

src/datacustomcode/client.py

@@ -151,11 +151,9 @@ def __new__(
                         "Proxy config is required when reader is built from config"
                     )
                 proxy_init = config.proxy_config.to_object(spark)
-                print(f"chuy1 reader, proxy_init: {proxy_init}")
 
                 reader_init = config.reader_config.to_object(spark)  # type: ignore
             else:
-                print("chuy2 reader")
                 reader_init = reader
                 if config.proxy_config is None:
                     raise ValueError("Proxy config is required when reader is provided")
@@ -232,8 +230,8 @@ def write_to_dmo(
         self._validate_data_layer_history_does_not_contain(DataCloudObjectType.DLO)
         return self._writer.write_to_dmo(name, dataframe, write_mode, **kwargs)
 
-    def call_llm_gateway(self, LLM_MODEL_ID: str, prompt: str, maxTokens: int):
-        self._proxy.call_llm_gateway(LLM_MODEL_ID, prompt, maxTokens)
+    def call_llm_gateway(self, LLM_MODEL_ID: str, prompt: str, maxTokens: int) -> str:
+        return self._proxy.call_llm_gateway(LLM_MODEL_ID, prompt, maxTokens)
 
     def find_file_path(self, file_name: str) -> Path:
         """Return a file path"""

src/datacustomcode/proxy/client/base.py

@@ -25,4 +25,4 @@ def __init__(self, spark=None, **kwargs):
             super().__init__(spark)
 
     @abstractmethod
-    def call_llm_gateway(self, llmModelId: str, prompt: str, maxTokens: int): ...
+    def call_llm_gateway(self, llmModelId: str, prompt: str, maxTokens: int) -> str: ...

src/datacustomcode/proxy/client/local_proxy_client.py

@@ -22,5 +22,5 @@ class LocalProxyClientProvider(BaseProxyClient):
 
     CONFIG_NAME = "LocalProxyClientProvider"
 
-    def call_llm_gateway(self, llmModelId: str, prompt: str, maxTokens: int):
-        print("Hello, from LocalProxyClientProvider call_llm_gateway() World!")
+    def call_llm_gateway(self, llmModelId: str, prompt: str, maxTokens: int) -> str:
+        return f"Hello, thanks for using {llmModelId}. So many tokens: {maxTokens}"

src/datacustomcode/run.py

@@ -62,10 +62,6 @@ def run_entrypoint(
         )
 
     package_type = get_package_type(entrypoint_dir)
-    print(
-        f"Chuy entrypoint: {entrypoint} directory: {entrypoint_dir} "
-        f"package type: {package_type}"
-    )
 
     try:
         with open(config_json_path, "r") as f: