satif-ai 0.2.9__py3-none-any.whl → 0.2.11__py3-none-any.whl

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.
satif_ai/adapters/tidy.py CHANGED
@@ -9,7 +9,7 @@ from pathlib import Path
9
9
  from typing import Optional, Union
10
10
 
11
11
  from agents import Agent, Runner, function_tool
12
- from agents.mcp.server import MCPServerStdio
12
+ from agents.mcp.server import MCPServer
13
13
  from mcp import ClientSession
14
14
  from satif_core.adapters.base import Adapter
15
15
  from satif_core.types import Datasource, SDIFPath
@@ -224,7 +224,7 @@ class TidyAdapter(Adapter):
224
224
 
225
225
  def __init__(
226
226
  self,
227
- mcp_server: MCPServerStdio,
227
+ mcp_server: MCPServer,
228
228
  mcp_session: ClientSession,
229
229
  llm_model: str = "o4-mini",
230
230
  max_iterations: int = 5,
@@ -233,7 +233,7 @@ class TidyAdapter(Adapter):
233
233
  Initialize the TidyAdapter.
234
234
 
235
235
  Args:
236
- mcp_server: An instance of MCPServerStdio for agent communication.
236
+ mcp_server: An instance of MCPServer for agent communication.
237
237
  mcp_session: An instance of ClientSession for resource/prompt fetching.
238
238
  llm_model: Name of the language model to use for the agent.
239
239
  max_iterations: Maximum number of attempts the agent gets to refine the code.
@@ -349,7 +349,7 @@ class TidyAdapter(Adapter):
349
349
  if isinstance(sdif, SDIFDatabase):
350
350
  input_path = Path(sdif.path)
351
351
  else:
352
- input_path = sdif
352
+ input_path = Path(sdif)
353
353
  if not input_path.exists():
354
354
  raise FileNotFoundError(f"Input SDIF file not found: {input_path}")
355
355
 
@@ -11,6 +11,7 @@ from satif_core.standardizers.base import AsyncStandardizer
11
11
  from satif_core.types import Datasource, FilePath, SDIFPath, StandardizationResult
12
12
 
13
13
  from satif_ai.adapters.tidy import TidyAdapter
14
+ from satif_ai.standardizers.ai_xlsx import AIXLSXStandardizer
14
15
  from satif_ai.utils.merge_sdif import merge_sdif_files
15
16
  from satif_ai.utils.zip import extract_zip_archive_async
16
17
 
@@ -43,8 +44,9 @@ class AIStandardizer(AsyncStandardizer):
43
44
 
44
45
  self.ai_standardizer_map: Dict[str, Type[AsyncStandardizer]] = {
45
46
  ".csv": AICSVStandardizer,
46
- # Future standardizers:
47
- # ".xlsx": AIXLSXStandardizer,
47
+ ".xlsx": AIXLSXStandardizer,
48
+ ".xls": AIXLSXStandardizer,
49
+ ".xlsm": AIXLSXStandardizer,
48
50
  # ".pdf": AIPDFStandardizer,
49
51
  # ".json": AIJSONStandardizer,
50
52
  # ".xml": AIXMLStandardizer,
@@ -332,11 +334,9 @@ class AIStandardizer(AsyncStandardizer):
332
334
  logger.info(
333
335
  f"Merging {len(intermediate_sdif_files)} intermediate SDIF SQLite files into '{final_sdif_file_target}'."
334
336
  )
335
- # merge_sdif_files must accept a list of source SQLite file paths and a target SQLite file path.
336
- merged_target_path = await merge_sdif_files(
337
+ merged_target_path = merge_sdif_files(
337
338
  intermediate_sdif_files,
338
339
  final_sdif_file_target,
339
- overwrite=False, # We handled overwrite for final_sdif_file_target
340
340
  )
341
341
  final_sdif_path_str = str(merged_target_path)
342
342
 
@@ -414,13 +414,11 @@ class AIStandardizer(AsyncStandardizer):
414
414
  file_processing_temp_dir.mkdir(parents=True, exist_ok=True)
415
415
 
416
416
  try:
417
- # 1. Resolve input datasource to a list of processable file paths
418
417
  resolved_files = await self._resolve_input_files(
419
418
  datasource, file_processing_temp_dir
420
419
  )
421
420
  logger.info(f"Resolved {len(resolved_files)} file(s) for standardization.")
422
421
 
423
- # 2. Group files by the AI standardizer responsible for them
424
422
  grouped_by_std, unsupported = self._group_files_by_standardizer(
425
423
  resolved_files
426
424
  )
@@ -438,7 +436,6 @@ class AIStandardizer(AsyncStandardizer):
438
436
  f"File groups for standardization: { {cls.__name__: [f.name for f in paths] for cls, paths in grouped_by_std.items()} }"
439
437
  )
440
438
 
441
- # 3. Process each group of files, generating intermediate SDIF SQLite files
442
439
  (
443
440
  intermediate_sdif_files,
444
441
  aggregated_file_configs,
@@ -454,7 +451,6 @@ class AIStandardizer(AsyncStandardizer):
454
451
  f"Successfully generated {len(intermediate_sdif_files)} intermediate SDIF SQLite file(s)."
455
452
  )
456
453
 
457
- # 4. Consolidate intermediate SDIF files into the final target file
458
454
  final_result = await self._consolidate_results(
459
455
  intermediate_sdif_files,
460
456
  aggregated_file_configs,
@@ -12,6 +12,7 @@ from agents import Agent, Runner, function_tool
12
12
  from agents.mcp.server import MCPServerStdio
13
13
  from charset_normalizer import detect
14
14
  from mcp import ClientSession
15
+ from satif_core import AsyncStandardizer
15
16
  from satif_core.types import Datasource, SDIFPath, StandardizationResult
16
17
  from satif_sdk.standardizers.csv import (
17
18
  CSVStandardizer,
@@ -274,7 +275,9 @@ async def read_raw_lines(
274
275
 
275
276
 
276
277
  # --- AICSVStandardizer Class ---
277
- class AICSVStandardizer(CSVStandardizer): # Inherits from the enhanced CSVStandardizer
278
+ class AICSVStandardizer(
279
+ CSVStandardizer, AsyncStandardizer
280
+ ): # Inherits from the enhanced CSVStandardizer
278
281
  def __init__(
279
282
  self,
280
283
  mcp_server: Optional[MCPServerStdio] = None,
@@ -0,0 +1,372 @@
1
+ import logging
2
+ import shutil
3
+ import tempfile
4
+ import uuid
5
+ from pathlib import Path
6
+ from typing import Any, Dict, List, Optional, Tuple
7
+
8
+ try:
9
+ from xlsx_to_sdif.graph import graph as xlsx_graph
10
+ from xlsx_to_sdif.state import State as XLSXState
11
+ except ImportError:
12
+ xlsx_graph = None # type: ignore
13
+ XLSXState = None # type: ignore
14
+ logging.getLogger(__name__).warning(
15
+ "Failed to import xlsx_to_sdif. AIXLSXStandardizer will not be functional."
16
+ )
17
+
18
+
19
+ from satif_core.standardizers.base import AsyncStandardizer
20
+ from satif_core.types import Datasource, SDIFPath, StandardizationResult
21
+
22
+ from satif_ai.utils.merge_sdif import merge_sdif_files
23
+
24
+ logger = logging.getLogger(__name__)
25
+
26
+
27
+ class AIXLSXStandardizer(AsyncStandardizer):
28
+ """
29
+ An asynchronous standardizer for XLSX files that leverages the `xlsx-to-sdif` library.
30
+
31
+ This standardizer processes one or more XLSX files, converts each to an
32
+ intermediate SDIF (Standardized Data Interchange Format) file using the
33
+ `xlsx-to-sdif` processing graph, and then consolidates these intermediate
34
+ files into a single final SDIF file.
35
+ """
36
+
37
+ def __init__(self, *args: Any, **kwargs: Any):
38
+ """
39
+ Initializes the AIXLSXStandardizer.
40
+
41
+ Args:
42
+ ...
43
+ """
44
+
45
+ async def _invoke_xlsx_graph(
46
+ self, input_file_path: Path, graph_config: Dict[str, Any]
47
+ ) -> Path:
48
+ """
49
+ Invokes the `xlsx-to-sdif` graph for a single XLSX file.
50
+
51
+ Args:
52
+ input_file_path: Path to the input XLSX file.
53
+ graph_config: Configuration for the `xlsx-to-sdif` graph invocation,
54
+ including a unique `thread_id`.
55
+
56
+ Returns:
57
+ Path to the SDIF file produced by the graph.
58
+
59
+ Raises:
60
+ RuntimeError: If the `xlsx-to-sdif` graph is not available, fails to
61
+ return a final state, or does not produce an output path.
62
+ FileNotFoundError: If the graph reports an output file that doesn't exist.
63
+ """
64
+ if not xlsx_graph or not XLSXState:
65
+ raise RuntimeError(
66
+ "xlsx_to_sdif is not available. "
67
+ "Please ensure 'xlsx-to-sdif' library is installed correctly."
68
+ )
69
+
70
+ initial_state: XLSXState = {"spreadsheet_path": str(input_file_path)} # type: ignore
71
+
72
+ thread_id = graph_config.get("configurable", {}).get(
73
+ "thread_id", "unknown_thread"
74
+ )
75
+ logger.info(
76
+ f"Invoking xlsx_to_sdif graph for: {input_file_path.name} with thread_id: {thread_id}"
77
+ )
78
+
79
+ # Stream events for logging or potential progress updates
80
+ async for event in xlsx_graph.astream_events(
81
+ initial_state, graph_config, version="v1"
82
+ ):
83
+ event_type = event["event"]
84
+ event_name = event.get("name", "")
85
+ if event_type in ["on_tool_start", "on_chain_start"]:
86
+ logger.debug(
87
+ f"Graph event for {input_file_path.name} (Thread: {thread_id}): {event_type} - {event_name}"
88
+ )
89
+ elif event_type in ["on_tool_error", "on_chain_error", "on_llm_error"]:
90
+ logger.warning(
91
+ f"Graph error event for {input_file_path.name} (Thread: {thread_id}): {event_type} - {event_name}. Data: {event.get('data')}"
92
+ )
93
+
94
+ final_snapshot = await xlsx_graph.aget_state(graph_config)
95
+ if not final_snapshot or not final_snapshot.values:
96
+ raise RuntimeError(
97
+ f"xlsx_to_sdif graph did not return a final state for {input_file_path.name} (Thread: {thread_id})."
98
+ )
99
+
100
+ output_sdif_path_str = final_snapshot.values.get("output_sdif_path")
101
+ if not output_sdif_path_str:
102
+ raise RuntimeError(
103
+ f"xlsx_to_sdif graph for {input_file_path.name} (Thread: {thread_id}) "
104
+ f"did not produce an 'output_sdif_path' in its final state. State: {final_snapshot.values}"
105
+ )
106
+
107
+ output_sdif_path = Path(output_sdif_path_str)
108
+ if not output_sdif_path.is_file():
109
+ raise FileNotFoundError(
110
+ f"xlsx_to_sdif graph for {input_file_path.name} (Thread: {thread_id}) "
111
+ f"reported output file '{output_sdif_path}', but it does not exist or is not a file."
112
+ )
113
+
114
+ logger.info(
115
+ f"xlsx_to_sdif graph successfully processed {input_file_path.name} (Thread: {thread_id}). Output at {output_sdif_path}"
116
+ )
117
+ return output_sdif_path
118
+
119
+ async def _resolve_and_filter_input_files(
120
+ self, datasource: Datasource
121
+ ) -> List[Path]:
122
+ """Resolves and validates datasource, returning a list of XLSX file paths."""
123
+ input_files: List[Path]
124
+ if isinstance(datasource, (str, Path)):
125
+ input_files = [Path(datasource)]
126
+ elif isinstance(datasource, list) and all(
127
+ isinstance(p, (str, Path)) for p in datasource
128
+ ):
129
+ input_files = [Path(p) for p in datasource]
130
+ else:
131
+ raise ValueError(
132
+ "Datasource must be a file path (str or Path) or a list of such paths."
133
+ )
134
+
135
+ if not input_files:
136
+ raise ValueError("No input XLSX files provided in the datasource.")
137
+
138
+ xlsx_input_files = []
139
+ for f_path in input_files:
140
+ if not f_path.is_file():
141
+ raise FileNotFoundError(f"Input file not found: {f_path}")
142
+ if f_path.suffix.lower() not in (
143
+ ".xlsx",
144
+ ".xlsm",
145
+ ".xlsb",
146
+ ".xls",
147
+ ): # Common Excel extensions
148
+ logger.warning(
149
+ f"File {f_path.name} is not a typical XLSX file extension, but will be attempted."
150
+ )
151
+ xlsx_input_files.append(f_path)
152
+
153
+ if not xlsx_input_files:
154
+ raise ValueError(
155
+ "No processable XLSX files found in the datasource after filtering."
156
+ )
157
+ return xlsx_input_files
158
+
159
+ def _prepare_final_output_path(
160
+ self, output_path: SDIFPath, overwrite: bool
161
+ ) -> Path:
162
+ """Prepares the final output path, handling overwrites and directory creation."""
163
+ final_output_path = Path(output_path)
164
+ if final_output_path.exists() and not overwrite:
165
+ raise FileExistsError(
166
+ f"Output file {final_output_path} already exists and overwrite is False."
167
+ )
168
+ elif final_output_path.exists() and overwrite:
169
+ logger.info(
170
+ f"Overwrite active: Deleting existing output file {final_output_path}"
171
+ )
172
+ try:
173
+ if (
174
+ final_output_path.is_dir()
175
+ ): # Should not happen if SDIFPath is file path
176
+ raise IsADirectoryError(
177
+ f"Output path {final_output_path} is a directory."
178
+ )
179
+ final_output_path.unlink()
180
+ except OSError as e:
181
+ raise RuntimeError(
182
+ f"Failed to delete existing output file {final_output_path}: {e}"
183
+ ) from e
184
+
185
+ final_output_path.parent.mkdir(parents=True, exist_ok=True)
186
+ return final_output_path
187
+
188
+ def _setup_temp_directories(self) -> Tuple[Path, Path, Path]:
189
+ """Creates and returns paths for temporary working directories."""
190
+ run_temp_dir = Path(tempfile.mkdtemp(prefix="satif_aixlsx_run_"))
191
+ intermediate_sdif_dir = run_temp_dir / "intermediate_sdifs"
192
+ intermediate_sdif_dir.mkdir()
193
+ temp_input_copies_dir = (
194
+ run_temp_dir / "temp_input_copies"
195
+ ) # Directory for temporary input copies
196
+ temp_input_copies_dir.mkdir()
197
+ return run_temp_dir, intermediate_sdif_dir, temp_input_copies_dir
198
+
199
+ async def _process_single_file_to_intermediate_sdif(
200
+ self,
201
+ input_xlsx_file: Path,
202
+ final_output_path_stem: str,
203
+ temp_input_copies_dir: Path,
204
+ intermediate_sdif_dir: Path,
205
+ ) -> Path:
206
+ """Processes a single XLSX file to an intermediate SDIF in a controlled location."""
207
+ logger.info(f"Processing file: {input_xlsx_file.name}")
208
+ graph_thread_id = f"satif_aixlsx_{final_output_path_stem}_{input_xlsx_file.stem}_{uuid.uuid4().hex[:8]}"
209
+
210
+ temp_input_file_for_graph = (
211
+ temp_input_copies_dir
212
+ / f"{input_xlsx_file.stem}_{graph_thread_id}{input_xlsx_file.suffix}"
213
+ )
214
+ shutil.copy2(input_xlsx_file, temp_input_file_for_graph)
215
+ logger.debug(
216
+ f"Created temporary copy of {input_xlsx_file.name} at {temp_input_file_for_graph}"
217
+ )
218
+
219
+ graph_config_for_file = {
220
+ "configurable": {"thread_id": graph_thread_id},
221
+ "recursion_limit": 50, # Default, make configurable if needed
222
+ }
223
+
224
+ try:
225
+ graph_produced_sdif_path = await self._invoke_xlsx_graph(
226
+ temp_input_file_for_graph, graph_config_for_file
227
+ )
228
+
229
+ target_intermediate_sdif_path = (
230
+ intermediate_sdif_dir
231
+ / f"intermediate_{input_xlsx_file.stem}_{graph_thread_id}.sdif"
232
+ )
233
+ shutil.move(
234
+ str(graph_produced_sdif_path),
235
+ str(target_intermediate_sdif_path),
236
+ )
237
+ logger.info(
238
+ f"Moved graph output for {input_xlsx_file.name} to {target_intermediate_sdif_path}"
239
+ )
240
+ return target_intermediate_sdif_path
241
+ except Exception as e:
242
+ error_msg = f"Failed to process file {input_xlsx_file.name} (using copy {temp_input_file_for_graph.name}) with xlsx-to-sdif graph: {e}"
243
+ logger.error(error_msg, exc_info=True)
244
+ # Re-raise to be caught by the main standardize method's loop or error handling
245
+ raise RuntimeError(
246
+ f"Error processing {input_xlsx_file.name}. Halting batch."
247
+ ) from e
248
+
249
+ async def _consolidate_intermediate_sdifs(
250
+ self, intermediate_sdif_paths: List[Path], final_output_path: Path
251
+ ) -> None:
252
+ """Consolidates intermediate SDIF files into the final output path."""
253
+ if not intermediate_sdif_paths:
254
+ # This case should ideally be handled before calling, but as a safeguard:
255
+ raise RuntimeError(
256
+ "No intermediate SDIF files were provided for consolidation."
257
+ )
258
+
259
+ if len(intermediate_sdif_paths) == 1:
260
+ logger.info(
261
+ f"Only one intermediate SDIF generated. Moving {intermediate_sdif_paths[0]} to {final_output_path}"
262
+ )
263
+ shutil.move(str(intermediate_sdif_paths[0]), str(final_output_path))
264
+ else:
265
+ logger.info(
266
+ f"Merging {len(intermediate_sdif_paths)} intermediate SDIF files into {final_output_path}"
267
+ )
268
+ merge_sdif_files(
269
+ source_db_paths=intermediate_sdif_paths,
270
+ target_db_path=final_output_path,
271
+ )
272
+
273
+ async def standardize(
274
+ self,
275
+ datasource: Datasource,
276
+ output_path: SDIFPath,
277
+ *,
278
+ overwrite: bool = False,
279
+ config: Optional[Dict[str, Any]] = None,
280
+ **kwargs: Any,
281
+ ) -> StandardizationResult:
282
+ """
283
+ Standardizes one or more XLSX files into a single SDIF file.
284
+
285
+ Args:
286
+ datasource: A single file path (str or Path) or a list of file paths
287
+ to XLSX files.
288
+ output_path: The path where the final consolidated SDIF file will be saved.
289
+ overwrite: If True, overwrite the output_path if it already exists.
290
+ Defaults to False.
291
+ config: General configuration options (currently not used by this standardizer
292
+ for graph interaction but preserved for API consistency).
293
+ **kwargs: Additional keyword arguments (currently ignored).
294
+
295
+ Returns:
296
+ A StandardizationResult object containing the path to the final SDIF file.
297
+
298
+ Raises:
299
+ ValueError: If the datasource is invalid or no XLSX files are found.
300
+ RuntimeError: If critical errors occur during processing, such as the
301
+ `xlsx-to-sdif` graph not being available or failing.
302
+ FileNotFoundError: If input files are not found or graph outputs are invalid.
303
+ FileExistsError: If output_path exists and overwrite is False.
304
+ """
305
+ if not xlsx_graph or not XLSXState:
306
+ raise RuntimeError(
307
+ "AIXLSXStandardizer cannot operate because `xlsx_to_sdif.graph` or `xlsx_to_sdif.state` is not available. "
308
+ "Please ensure the 'xlsx-to-sdif' library is installed and accessible."
309
+ )
310
+
311
+ xlsx_input_files = await self._resolve_and_filter_input_files(datasource)
312
+ final_output_path = self._prepare_final_output_path(output_path, overwrite)
313
+ run_temp_dir, intermediate_sdif_dir, temp_input_copies_dir = (
314
+ self._setup_temp_directories()
315
+ )
316
+
317
+ intermediate_sdif_paths: List[Path] = []
318
+ processing_errors: List[str] = []
319
+
320
+ try:
321
+ # Process each file sequentially. Consider asyncio.gather for parallel if graph supports it well for many files.
322
+ for i, input_xlsx_file in enumerate(xlsx_input_files):
323
+ try:
324
+ logger.info(
325
+ f"Starting processing for file {i + 1}/{len(xlsx_input_files)}: {input_xlsx_file.name}"
326
+ )
327
+ intermediate_sdif_path = (
328
+ await self._process_single_file_to_intermediate_sdif(
329
+ input_xlsx_file,
330
+ final_output_path.stem, # Pass stem for unique naming
331
+ temp_input_copies_dir,
332
+ intermediate_sdif_dir,
333
+ )
334
+ )
335
+ intermediate_sdif_paths.append(intermediate_sdif_path)
336
+ except Exception:
337
+ logger.error(
338
+ f"Halting standardization due to error processing {input_xlsx_file.name}."
339
+ )
340
+ raise # Re-raise the exception to be caught by the outer try/finally
341
+
342
+ if not intermediate_sdif_paths:
343
+ # This condition might be redundant if _process_single_file_to_intermediate_sdif always raises on failure
344
+ # and we re-raise immediately.
345
+ if processing_errors: # This list would be empty if we fail fast
346
+ raise RuntimeError(
347
+ f"No XLSX files were successfully processed. Errors: {'; '.join(processing_errors)}"
348
+ )
349
+ else:
350
+ raise RuntimeError(
351
+ "No intermediate SDIF files were generated, though no specific errors were caught."
352
+ )
353
+
354
+ await self._consolidate_intermediate_sdifs(
355
+ intermediate_sdif_paths, final_output_path
356
+ )
357
+
358
+ logger.info(f"Successfully created final SDIF: {final_output_path}")
359
+ return StandardizationResult(
360
+ output_path=final_output_path, file_configs=None
361
+ ) # file_configs not available from this process
362
+
363
+ finally:
364
+ if run_temp_dir.exists():
365
+ try:
366
+ shutil.rmtree(run_temp_dir)
367
+ logger.debug(f"Cleaned up temporary directory: {run_temp_dir}")
368
+ except Exception as e_clean:
369
+ logger.error(
370
+ f"Error cleaning up temporary directory {run_temp_dir}: {e_clean}",
371
+ exc_info=True,
372
+ )
satif_ai/transform.py CHANGED
@@ -1,7 +1,7 @@
1
1
  from pathlib import Path
2
2
  from typing import Any, Dict, List, Optional
3
3
 
4
- from fastmcp import FastMCP
4
+ from fastmcp import Client, FastMCP
5
5
  from fastmcp.client.transports import FastMCPTransport
6
6
  from satif_core.code_executors.base import CodeExecutor
7
7
  from satif_core.transformation_builders.base import AsyncTransformationBuilder
@@ -28,7 +28,7 @@ async def atransform(
28
28
  transformation_builder: Optional[AsyncTransformationBuilder] = None,
29
29
  code_executor: Optional[CodeExecutor] = None,
30
30
  mcp_server: Optional[FastMCP] = None,
31
- mcp_transport: Optional[FastMCPTransport] = None,
31
+ mcp_client: Optional[Client] = None,
32
32
  llm_model: str = "o4-mini",
33
33
  schema_only: bool = False,
34
34
  representer_kwargs: Optional[Dict[str, Any]] = None,
@@ -38,7 +38,7 @@ async def atransform(
38
38
  an AI-generated or provided transformation code.
39
39
 
40
40
  This function orchestrates the process of:
41
- 1. Optionally generating transformation code using an AI model via a `CodeBuilder`
41
+ 1. Optionally generating transformation code using an AI model via a `TransformationBuilder`
42
42
  if `transformation_code` is not provided.
43
43
  explicitly passed.
44
44
  2. Executing the transformation code using a `CodeTransformer` and a `CodeExecutor`.
@@ -53,18 +53,21 @@ async def atransform(
53
53
  output filenames if the transformation result keys match.
54
54
  instructions: Optional. Natural language instructions for the AI to generate
55
55
  the transformation code. Used if `transformation_code` is None.
56
+ output_path: Path to the directory where transformation outputs will be saved.
56
57
  transformation_code: Optional. Pre-existing Python code for the transformation.
57
58
  If None, code will be generated by the `transformation_builder`.
58
59
  transformation_builder: Optional. An `AsyncTransformationBuilder` instance responsible for generating
59
60
  the transformation code if `transformation_code` is not provided.
60
- If None, a `TransformationAsyncCodeBuilder` is instantiated.
61
+ If None, a `SyncpulseTransformationBuilder` is instantiated.
61
62
  code_executor: Optional. A `CodeExecutor` instance for running the transformation
62
63
  code. If None, a `LocalCodeExecutor` is used.
63
64
  mcp_server: Optional. A `FastMCP` server instance for the AI code builder.
64
- Defaults to the global `mcp` instance if `transformation_builder` is None.
65
- mcp_transport: Optional. A `FastMCPTransport` instance for communication with
66
- the `mcp_server`. Defaults to a new transport using `mcp_server`
67
- if `transformation_builder` is None.
65
+ Defaults to the global `mcp` instance if `transformation_builder` is None and
66
+ a new `SyncpulseTransformationBuilder` is being created.
67
+ mcp_client: Optional. A user-provided `Client` instance. If provided when
68
+ `transformation_builder` is None, it will be used by the internally
69
+ created `SyncpulseTransformationBuilder`. The caller is responsible for
70
+ managing the lifecycle of a provided client.
68
71
  llm_model: The language model to use for code generation (e.g., "o4-mini").
69
72
  Used if `transformation_builder` is None.
70
73
  schema_only: If True, the transformation aims to match only the schema (headers)
@@ -78,44 +81,75 @@ async def atransform(
78
81
  A `TransformationResult` object containing the path to the output
79
82
  and the transformation code used.
80
83
  """
81
- if transformation_builder is None:
82
- if mcp_server is None:
83
- mcp_server = mcp
84
-
85
- if mcp_transport is None:
86
- mcp_transport = FastMCPTransport(mcp=mcp_server)
87
-
88
- openai_compatible_mcp = OpenAICompatibleMCP(mcp=mcp_server)
89
- await openai_compatible_mcp.connect()
90
-
91
- transformation_builder = SyncpulseTransformationBuilder(
92
- mcp_server=openai_compatible_mcp,
93
- mcp_session=mcp_transport,
94
- llm_model=llm_model,
95
- )
96
-
97
- if transformation_code is None:
98
- function_code = await transformation_builder.build(
99
- sdif=sdif,
100
- output_target_files=output_target_files,
101
- instructions=instructions,
102
- schema_only=schema_only,
103
- representer_kwargs=representer_kwargs,
104
- )
105
- else:
106
- function_code = transformation_code
107
-
108
- if code_executor is None:
109
- code_executor = LocalCodeExecutor()
84
+ current_transformation_code: Optional[str] = transformation_code
85
+ active_builder: Optional[AsyncTransformationBuilder] = transformation_builder
86
+
87
+ _openai_mcp_instance: Optional[OpenAICompatibleMCP] = None
88
+ openai_mcp_managed_locally = False
89
+
90
+ # If code isn't provided, we need a builder. If a builder isn't provided, we create one.
91
+ if current_transformation_code is None:
92
+ if active_builder is None:
93
+ # Create SyncpulseTransformationBuilder
94
+ _effective_mcp_server = mcp_server if mcp_server is not None else mcp
95
+
96
+ _openai_mcp_instance = OpenAICompatibleMCP(mcp=_effective_mcp_server)
97
+ await _openai_mcp_instance.connect()
98
+ openai_mcp_managed_locally = True
99
+
100
+ if mcp_client is None: # No user-provided client, create and manage one
101
+ mcp_transport = FastMCPTransport(mcp=_effective_mcp_server)
102
+ async with Client(mcp_transport) as new_client:
103
+ active_builder = SyncpulseTransformationBuilder(
104
+ mcp_server=_openai_mcp_instance,
105
+ mcp_session=new_client.session,
106
+ llm_model=llm_model,
107
+ )
108
+ current_transformation_code = await active_builder.build(
109
+ sdif=sdif,
110
+ output_target_files=output_target_files,
111
+ instructions=instructions,
112
+ schema_only=schema_only,
113
+ representer_kwargs=representer_kwargs,
114
+ )
115
+ else:
116
+ active_builder = SyncpulseTransformationBuilder(
117
+ mcp_server=_openai_mcp_instance,
118
+ mcp_session=mcp_client, # Use the provided client
119
+ llm_model=llm_model,
120
+ )
121
+ current_transformation_code = await active_builder.build(
122
+ sdif=sdif,
123
+ output_target_files=output_target_files,
124
+ instructions=instructions,
125
+ schema_only=schema_only,
126
+ representer_kwargs=representer_kwargs,
127
+ )
128
+
129
+ # Disconnect OpenAICompatibleMCP if it was created and connected locally
130
+ if (
131
+ openai_mcp_managed_locally
132
+ and _openai_mcp_instance
133
+ and _openai_mcp_instance._is_connected
134
+ ):
135
+ await _openai_mcp_instance.cleanup()
136
+
137
+ if current_transformation_code is None:
138
+ raise ValueError("Transformation code could not be obtained or generated.")
139
+
140
+ # Code Executor and Transformation
141
+ _code_executor = code_executor if code_executor is not None else LocalCodeExecutor()
110
142
 
111
143
  transformer = CodeTransformer(
112
- function=function_code,
113
- code_executor=code_executor,
144
+ function=current_transformation_code,
145
+ code_executor=_code_executor,
114
146
  )
115
147
 
116
- output_path = transformer.export(
148
+ exported_artifact_path = transformer.export(
117
149
  sdif=sdif,
118
150
  output_path=output_path,
119
151
  )
120
152
 
121
- return TransformationResult(output_path=output_path, function_code=function_code)
153
+ return TransformationResult(
154
+ output_path=exported_artifact_path, function_code=current_transformation_code
155
+ )