satif-ai 0.2.8__py3-none-any.whl → 0.2.10__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.
@@ -0,0 +1,535 @@
1
+ import json
2
+ import logging
3
+ import shutil
4
+ import sqlite3
5
+ from datetime import datetime
6
+ from pathlib import Path
7
+ from typing import Any, Dict, List
8
+
9
+ from satif_core.types import SDIFPath
10
+ from sdif_db.database import (
11
+ SDIFDatabase, # Assuming this is the conventional import path
12
+ )
13
+
14
+ log = logging.getLogger(__name__)
15
+
16
+
17
+ class _SDIFMerger:
18
+ def __init__(self, target_sdif_path: Path):
19
+ self.target_db = SDIFDatabase(target_sdif_path, overwrite=True)
20
+ # Mappings per source_db_idx:
21
+ self.source_id_map: Dict[int, Dict[int, int]] = {}
22
+ self.table_name_map: Dict[int, Dict[str, str]] = {}
23
+ self.object_name_map: Dict[int, Dict[str, str]] = {}
24
+ self.media_name_map: Dict[int, Dict[str, str]] = {}
25
+
26
+ def _get_new_source_id(self, source_db_idx: int, old_source_id: int) -> int:
27
+ return self.source_id_map[source_db_idx][old_source_id]
28
+
29
+ def _get_new_table_name(self, source_db_idx: int, old_table_name: str) -> str:
30
+ return self.table_name_map[source_db_idx].get(old_table_name, old_table_name)
31
+
32
+ def _get_new_object_name(self, source_db_idx: int, old_object_name: str) -> str:
33
+ return self.object_name_map[source_db_idx].get(old_object_name, old_object_name)
34
+
35
+ def _get_new_media_name(self, source_db_idx: int, old_media_name: str) -> str:
36
+ return self.media_name_map[source_db_idx].get(old_media_name, old_media_name)
37
+
38
+ def _generate_unique_name_in_target(self, base_name: str, list_func) -> str:
39
+ """Generates a unique name for the target DB by appending a suffix if base_name exists."""
40
+ if base_name not in list_func():
41
+ return base_name
42
+ i = 1
43
+ while True:
44
+ new_name = f"{base_name}_{i}"
45
+ if new_name not in list_func():
46
+ return new_name
47
+ i += 1
48
+ if i > 1000: # Safety break
49
+ raise RuntimeError(
50
+ f"Could not generate a unique name for base '{base_name}' after 1000 attempts."
51
+ )
52
+
53
+ def _merge_properties(self, source_db: SDIFDatabase, source_db_idx: int):
54
+ source_props = source_db.get_properties()
55
+ if not source_props:
56
+ log.warning(
57
+ f"Source database {source_db.path} has no properties. Skipping properties merge for this source."
58
+ )
59
+ return
60
+
61
+ if source_props.get("sdif_version") != "1.0":
62
+ # Or allow a configurable expected version
63
+ raise ValueError(
64
+ f"Source database {source_db.path} has unsupported SDIF version: {source_props.get('sdif_version')}. Expected '1.0'."
65
+ )
66
+
67
+ if source_db_idx == 0: # First database sets the version for the target
68
+ try:
69
+ self.target_db.conn.execute(
70
+ "UPDATE sdif_properties SET sdif_version = ?",
71
+ (
72
+ source_props.get("sdif_version", "1.0"),
73
+ ), # Default to 1.0 if somehow missing
74
+ )
75
+ self.target_db.conn.commit()
76
+ except sqlite3.Error as e:
77
+ log.error(
78
+ f"Failed to set sdif_version in target DB from {source_db.path}: {e}"
79
+ )
80
+ raise
81
+ # creation_timestamp will be set at the end of the entire merge process.
82
+
83
+ def _merge_sources(self, source_db: SDIFDatabase, source_db_idx: int):
84
+ self.source_id_map[source_db_idx] = {}
85
+ source_sources = source_db.list_sources()
86
+ for old_source_entry in source_sources:
87
+ old_source_id = old_source_entry["source_id"]
88
+ new_source_id = self.target_db.add_source(
89
+ file_name=old_source_entry["original_file_name"],
90
+ file_type=old_source_entry["original_file_type"],
91
+ description=old_source_entry.get("source_description"),
92
+ )
93
+ # original processing_timestamp is not directly carried over, new one is set by add_source
94
+ self.source_id_map[source_db_idx][old_source_id] = new_source_id
95
+
96
+ def _merge_tables(self, source_db: SDIFDatabase, source_db_idx: int):
97
+ self.table_name_map[source_db_idx] = {}
98
+ source_schema = source_db.get_schema()
99
+ source_tables_schema = source_schema.get("tables", {})
100
+
101
+ # Pass 1: Determine new table names for all tables from this source DB
102
+ # This is to ensure FKs can be remapped correctly to tables from the *same* source db.
103
+ temp_name_map_for_this_source = {}
104
+ for old_table_name in source_tables_schema.keys():
105
+ # Use create_table with if_exists='add' to get a unique name, but only for name generation.
106
+ # This is a bit of a workaround. A dedicated _generate_unique_target_table_name might be cleaner.
107
+ # The SDIFDatabase.create_table(if_exists='add') will actually create metadata entries.
108
+ # This might be acceptable if we're careful.
109
+ # Let's use the simpler approach of generating unique name first.
110
+ effective_new_name = self._generate_unique_name_in_target(
111
+ old_table_name, self.target_db.list_tables
112
+ )
113
+ temp_name_map_for_this_source[old_table_name] = effective_new_name
114
+ self.table_name_map[source_db_idx] = temp_name_map_for_this_source
115
+
116
+ # Pass 2: Create tables with remapped FKs and copy data
117
+ for old_table_name, table_detail_from_schema in source_tables_schema.items():
118
+ new_table_name = self.table_name_map[source_db_idx][old_table_name]
119
+
120
+ columns_for_create: Dict[str, Dict[str, Any]] = {}
121
+ original_columns_detail = table_detail_from_schema.get("columns", [])
122
+
123
+ for col_detail in original_columns_detail:
124
+ col_name = col_detail["name"]
125
+ col_props = {
126
+ "type": col_detail["sqlite_type"],
127
+ "not_null": col_detail["not_null"],
128
+ "primary_key": col_detail[
129
+ "primary_key"
130
+ ], # Assumes single col PK flag
131
+ "description": col_detail.get("description"),
132
+ "original_column_name": col_detail.get("original_column_name"),
133
+ # 'unique' constraint not in get_schema output, assumed not used or handled by primary_key
134
+ }
135
+
136
+ # Remap foreign keys defined for this column
137
+ table_fks_detail = table_detail_from_schema.get("foreign_keys", [])
138
+ for fk_info in table_fks_detail:
139
+ if fk_info["from_column"] == col_name:
140
+ original_fk_target_table = fk_info["target_table"]
141
+ # FKs are assumed to target tables within the same source SDIF file.
142
+ remapped_fk_target_table = self.table_name_map[
143
+ source_db_idx
144
+ ].get(original_fk_target_table)
145
+ if not remapped_fk_target_table:
146
+ log.warning(
147
+ f"Could not remap FK target table '{original_fk_target_table}' for column '{col_name}' in table '{old_table_name}'. FK might be dropped or invalid."
148
+ )
149
+ # Decide: skip FK, or raise error, or create FK pointing to original name (which might conflict or be wrong)
150
+ # For now, we'll proceed without this FK if target not found in map (shouldn't happen if all tables from source are processed)
151
+ continue
152
+
153
+ col_props["foreign_key"] = {
154
+ "table": remapped_fk_target_table,
155
+ "column": fk_info["target_column"],
156
+ "on_delete": fk_info[
157
+ "on_delete"
158
+ ].upper(), # Ensure standard casing
159
+ "on_update": fk_info[
160
+ "on_update"
161
+ ].upper(), # Ensure standard casing
162
+ }
163
+ break # Assuming one FK per 'from_column' for this col_props structure
164
+ columns_for_create[col_name] = col_props
165
+
166
+ source_table_metadata = table_detail_from_schema.get("metadata", {})
167
+ old_source_id_for_table = source_table_metadata.get("source_id")
168
+ if old_source_id_for_table is None:
169
+ raise ValueError(
170
+ f"Table '{old_table_name}' from {source_db.path} is missing source_id in its metadata."
171
+ )
172
+
173
+ new_source_id_for_table = self._get_new_source_id(
174
+ source_db_idx, old_source_id_for_table
175
+ )
176
+
177
+ # Create the table structure in the target database
178
+ # Using if_exists="fail" because new_table_name should already be unique.
179
+ # SDIFDatabase.create_table handles complex PKs via table_constraints reconstruction.
180
+ actual_created_name = self.target_db.create_table(
181
+ table_name=new_table_name,
182
+ columns=columns_for_create,
183
+ source_id=new_source_id_for_table,
184
+ description=source_table_metadata.get("description"),
185
+ original_identifier=source_table_metadata.get("original_identifier"),
186
+ if_exists="fail",
187
+ )
188
+ if actual_created_name != new_table_name:
189
+ # This case should ideally not happen if _generate_unique_target_table_name was correct
190
+ # and create_table used if_exists='fail'. If create_table internally changes name even with 'fail',
191
+ # this is an issue. For now, assume 'fail' means it uses the name or errors.
192
+ log.warning(
193
+ f"Table name discrepancy: expected {new_table_name}, created as {actual_created_name}. Using created name."
194
+ )
195
+ self.table_name_map[source_db_idx][old_table_name] = (
196
+ actual_created_name # Update map
197
+ )
198
+
199
+ # Copy data
200
+ try:
201
+ data_df = source_db.read_table(old_table_name)
202
+ if not data_df.empty:
203
+ # SDIFDatabase.insert_data expects List[Dict].
204
+ # SDIFDatabase.write_dataframe is higher level but might re-create table.
205
+ # Let's use insert_data.
206
+
207
+ # Handle data type conversions that pandas might do, to align with SQLite expectations
208
+ # For example, pandas bools to int 0/1, datetimes to ISO strings.
209
+ # The SDIFDatabase.write_dataframe has logic for this.
210
+ # We can replicate parts or simplify if read_table and insert_data are robust.
211
+ # For now, assume read_table gives compatible data for insert_data
212
+ # or insert_data can handle common pandas types.
213
+ # A quick check: SDIFDatabase.insert_data does not do type conversion.
214
+ # SDIFDatabase.write_dataframe does. So it's safer to go df -> records -> insert
215
+ # after manual conversion like in write_dataframe.
216
+
217
+ df_copy = data_df.copy()
218
+ for col_name_str in df_copy.columns:
219
+ col_name = str(col_name_str) # Ensure string
220
+ if pd.api.types.is_bool_dtype(df_copy[col_name].dtype):
221
+ df_copy[col_name] = df_copy[col_name].astype(int)
222
+ elif pd.api.types.is_datetime64_any_dtype(
223
+ df_copy[col_name].dtype
224
+ ):
225
+ df_copy[col_name] = df_copy[col_name].apply(
226
+ lambda x: x.isoformat() if pd.notnull(x) else None
227
+ )
228
+ elif pd.api.types.is_timedelta64_dtype(df_copy[col_name].dtype):
229
+ df_copy[col_name] = df_copy[col_name].astype(str)
230
+ # Handle potential np.nan to None for JSON compatibility if objects were stored as text
231
+ if df_copy[col_name].dtype == object:
232
+ df_copy[col_name] = df_copy[col_name].replace(
233
+ {np.nan: None}
234
+ )
235
+
236
+ data_records = df_copy.to_dict("records")
237
+ if data_records: # Ensure there are records to insert
238
+ self.target_db.insert_data(actual_created_name, data_records)
239
+ except Exception as e:
240
+ log.error(
241
+ f"Failed to copy data for table {old_table_name} to {actual_created_name}: {e}"
242
+ )
243
+ # Decide: continue with other tables or raise? For robustness, log and continue.
244
+ # Or add a strict mode flag. For now, log and continue.
245
+
246
+ def _merge_objects(self, source_db: SDIFDatabase, source_db_idx: int):
247
+ self.object_name_map[source_db_idx] = {}
248
+ for old_object_name in source_db.list_objects():
249
+ obj_data = source_db.get_object(
250
+ old_object_name, parse_json=False
251
+ ) # Get raw JSON strings
252
+ if not obj_data:
253
+ log.warning(
254
+ f"Could not retrieve object '{old_object_name}' from {source_db.path}. Skipping."
255
+ )
256
+ continue
257
+
258
+ new_object_name = self._generate_unique_name_in_target(
259
+ old_object_name, self.target_db.list_objects
260
+ )
261
+ self.object_name_map[source_db_idx][old_object_name] = new_object_name
262
+
263
+ new_source_id = self._get_new_source_id(
264
+ source_db_idx, obj_data["source_id"]
265
+ )
266
+
267
+ # Data is already string from parse_json=False. Schema hint also string.
268
+ # SDIFDatabase.add_object expects data to be Any (serializable) and schema_hint Dict.
269
+ # So we need to parse them back if they are strings.
270
+ parsed_json_data = json.loads(obj_data["json_data"])
271
+ parsed_schema_hint = (
272
+ json.loads(obj_data["schema_hint"])
273
+ if obj_data.get("schema_hint")
274
+ else None
275
+ )
276
+
277
+ self.target_db.add_object(
278
+ object_name=new_object_name,
279
+ json_data=parsed_json_data,
280
+ source_id=new_source_id,
281
+ description=obj_data.get("description"),
282
+ schema_hint=parsed_schema_hint,
283
+ )
284
+
285
+ def _merge_media(self, source_db: SDIFDatabase, source_db_idx: int):
286
+ self.media_name_map[source_db_idx] = {}
287
+ for old_media_name in source_db.list_media():
288
+ media_entry = source_db.get_media(
289
+ old_media_name, parse_json=False
290
+ ) # Get raw JSON for tech_metadata
291
+ if not media_entry:
292
+ log.warning(
293
+ f"Could not retrieve media '{old_media_name}' from {source_db.path}. Skipping."
294
+ )
295
+ continue
296
+
297
+ new_media_name = self._generate_unique_name_in_target(
298
+ old_media_name, self.target_db.list_media
299
+ )
300
+ self.media_name_map[source_db_idx][old_media_name] = new_media_name
301
+
302
+ new_source_id = self._get_new_source_id(
303
+ source_db_idx, media_entry["source_id"]
304
+ )
305
+
306
+ parsed_tech_metadata = (
307
+ json.loads(media_entry["technical_metadata"])
308
+ if media_entry.get("technical_metadata")
309
+ else None
310
+ )
311
+
312
+ self.target_db.add_media(
313
+ media_name=new_media_name,
314
+ media_data=media_entry["media_data"], # Should be bytes
315
+ media_type=media_entry["media_type"],
316
+ source_id=new_source_id,
317
+ description=media_entry.get("description"),
318
+ original_format=media_entry.get("original_format"),
319
+ technical_metadata=parsed_tech_metadata,
320
+ )
321
+
322
+ def _remap_element_spec(
323
+ self, element_type: str, element_spec_json: str, source_db_idx: int
324
+ ) -> str:
325
+ if not element_spec_json:
326
+ return element_spec_json
327
+
328
+ try:
329
+ spec_dict = json.loads(element_spec_json)
330
+ except json.JSONDecodeError:
331
+ log.warning(
332
+ f"Invalid JSON in element_spec: {element_spec_json}. Returning as is."
333
+ )
334
+ return element_spec_json
335
+
336
+ new_spec_dict = spec_dict.copy()
337
+
338
+ # Remap source_id if present (relevant for 'source' element_type in annotations, not directly in semantic_links spec)
339
+ # Semantic links link to other entities which carry their own source_id.
340
+ # But if spec itself contains a source_id key (e.g. for target_element_type='source' in annotations)
341
+ if "source_id" in new_spec_dict and isinstance(new_spec_dict["source_id"], int):
342
+ new_spec_dict["source_id"] = self._get_new_source_id(
343
+ source_db_idx, new_spec_dict["source_id"]
344
+ )
345
+
346
+ # Remap names based on element_type
347
+ if element_type in ["table", "column"]:
348
+ if "table_name" in new_spec_dict:
349
+ new_spec_dict["table_name"] = self._get_new_table_name(
350
+ source_db_idx, new_spec_dict["table_name"]
351
+ )
352
+ elif element_type == "object": # Direct object reference
353
+ if "object_name" in new_spec_dict:
354
+ new_spec_dict["object_name"] = self._get_new_object_name(
355
+ source_db_idx, new_spec_dict["object_name"]
356
+ )
357
+ elif element_type == "json_path": # JSONPath typically refers to an object
358
+ if (
359
+ "object_name" in new_spec_dict
360
+ ): # If the spec identifies the object container
361
+ new_spec_dict["object_name"] = self._get_new_object_name(
362
+ source_db_idx, new_spec_dict["object_name"]
363
+ )
364
+ elif element_type == "media":
365
+ if "media_name" in new_spec_dict:
366
+ new_spec_dict["media_name"] = self._get_new_media_name(
367
+ source_db_idx, new_spec_dict["media_name"]
368
+ )
369
+ # 'file' type needs no remapping of spec content.
370
+ # 'source' type: primary key is 'source_id', remapped above.
371
+
372
+ return json.dumps(new_spec_dict)
373
+
374
+ def _merge_semantic_links(self, source_db: SDIFDatabase, source_db_idx: int):
375
+ # SDIFDatabase.list_semantic_links default parses JSON spec. We need this.
376
+ source_links = source_db.list_semantic_links(parse_json=True)
377
+
378
+ for link in source_links:
379
+ # The specs are already dicts because parse_json=True was used.
380
+ try:
381
+ from_spec_dict = link["from_element_spec"]
382
+ to_spec_dict = link["to_element_spec"]
383
+
384
+ # Remap the dicts directly
385
+ new_from_spec_dict = self._remap_element_spec_dict(
386
+ link["from_element_type"], from_spec_dict, source_db_idx
387
+ )
388
+ new_to_spec_dict = self._remap_element_spec_dict(
389
+ link["to_element_type"], to_spec_dict, source_db_idx
390
+ )
391
+
392
+ self.target_db.add_semantic_link(
393
+ link_type=link["link_type"],
394
+ from_element_type=link["from_element_type"],
395
+ from_element_spec=new_from_spec_dict, # add_semantic_link takes dict
396
+ to_element_type=link["to_element_type"],
397
+ to_element_spec=new_to_spec_dict, # add_semantic_link takes dict
398
+ description=link.get("description"),
399
+ )
400
+ except Exception as e:
401
+ link_id = link.get("link_id", "Unknown")
402
+ log.error(
403
+ f"Failed to merge semantic link ID {link_id} from {source_db.path}: {e}. Skipping link."
404
+ )
405
+
406
+ def _remap_element_spec_dict(
407
+ self, element_type: str, spec_dict: Dict, source_db_idx: int
408
+ ) -> Dict:
409
+ # Helper for _merge_semantic_links that works with dicts directly
410
+ new_spec_dict = spec_dict.copy()
411
+
412
+ if "source_id" in new_spec_dict and isinstance(new_spec_dict["source_id"], int):
413
+ new_spec_dict["source_id"] = self._get_new_source_id(
414
+ source_db_idx, new_spec_dict["source_id"]
415
+ )
416
+
417
+ if element_type in ["table", "column"]:
418
+ if "table_name" in new_spec_dict:
419
+ new_spec_dict["table_name"] = self._get_new_table_name(
420
+ source_db_idx, new_spec_dict["table_name"]
421
+ )
422
+ elif element_type == "object" or (
423
+ element_type == "json_path" and "object_name" in new_spec_dict
424
+ ):
425
+ if "object_name" in new_spec_dict:
426
+ new_spec_dict["object_name"] = self._get_new_object_name(
427
+ source_db_idx, new_spec_dict["object_name"]
428
+ )
429
+ elif element_type == "media":
430
+ if "media_name" in new_spec_dict:
431
+ new_spec_dict["media_name"] = self._get_new_media_name(
432
+ source_db_idx, new_spec_dict["media_name"]
433
+ )
434
+ return new_spec_dict
435
+
436
+ def merge_all(self, source_sdif_paths: List[SDIFPath]):
437
+ # Import pandas and numpy here to avoid making them a hard dependency of the module if not used.
438
+ # However, SDIFDatabase itself uses them. So they are effectively dependencies.
439
+ global pd, np
440
+ import numpy as np
441
+ import pandas as pd
442
+
443
+ for idx, source_path_item in enumerate(source_sdif_paths):
444
+ source_path = Path(source_path_item) # Ensure Path object
445
+ log.info(
446
+ f"Processing source SDIF ({idx + 1}/{len(source_sdif_paths)}): {source_path}"
447
+ )
448
+ source_db = SDIFDatabase(source_path, read_only=True)
449
+ try: # Ensure source_db is closed
450
+ self._merge_properties(source_db, idx)
451
+ self._merge_sources(source_db, idx)
452
+ self._merge_tables(source_db, idx) # This needs pandas for data reading
453
+ self._merge_objects(source_db, idx)
454
+ self._merge_media(source_db, idx)
455
+ self._merge_semantic_links(source_db, idx)
456
+ # Not merging sdif_annotations in this version.
457
+ finally:
458
+ source_db.close()
459
+
460
+ # Finalize target DB properties
461
+ try:
462
+ current_timestamp_utc_z = datetime.utcnow().strftime("%Y-%m-%dT%H:%M:%SZ")
463
+ self.target_db.conn.execute(
464
+ "UPDATE sdif_properties SET creation_timestamp = ?",
465
+ (current_timestamp_utc_z,),
466
+ )
467
+ self.target_db.conn.commit()
468
+ except sqlite3.Error as e:
469
+ log.error(f"Failed to update final creation_timestamp in target DB: {e}")
470
+ # Non-fatal, proceed.
471
+
472
+ self.target_db.close()
473
+ log.info(
474
+ f"Successfully merged {len(source_sdif_paths)} SDIF files into {self.target_db.path}"
475
+ )
476
+ return self.target_db.path
477
+
478
+
479
+ def merge_sdif_files(sdif_paths: List[SDIFPath], output_path: Path) -> Path:
480
+ """
481
+ Merges multiple SDIF files into a single new SDIF file.
482
+
483
+ Args:
484
+ sdif_paths: A list of paths to the SDIF files to merge.
485
+ output_path: The full path where the merged SDIF file should be saved.
486
+ Its parent directory will be created if it doesn't exist.
487
+ If output_path is an existing file, it will be overwritten.
488
+ If output_path is an existing directory, a ValueError is raised.
489
+
490
+ Returns:
491
+ Path to the newly created merged SDIF file (same as output_path).
492
+
493
+ Raises:
494
+ ValueError: If no SDIF files are provided, or output_path is invalid (e.g., an existing directory).
495
+ FileNotFoundError: If a source SDIF file does not exist.
496
+ sqlite3.Error: For database-related errors during merging.
497
+ RuntimeError: For critical errors like inability to generate unique names.
498
+ """
499
+ if not sdif_paths:
500
+ raise ValueError("No SDIF files provided for merging.")
501
+
502
+ output_path = Path(output_path).resolve()
503
+
504
+ if output_path.is_dir():
505
+ raise ValueError(
506
+ f"Output path '{output_path}' is an existing directory. Please provide a full file path."
507
+ )
508
+
509
+ # Ensure parent directory of output_path exists
510
+ output_path.parent.mkdir(parents=True, exist_ok=True)
511
+
512
+ # Ensure all source paths are Path objects and exist
513
+ processed_sdif_paths: List[Path] = []
514
+ for p in sdif_paths:
515
+ path_obj = Path(p).resolve()
516
+ if not path_obj.exists():
517
+ raise FileNotFoundError(f"Source SDIF file not found: {path_obj}")
518
+ if not path_obj.is_file():
519
+ raise ValueError(f"Source SDIF path is not a file: {path_obj}")
520
+ processed_sdif_paths.append(path_obj)
521
+
522
+ if len(processed_sdif_paths) == 1:
523
+ source_file = processed_sdif_paths[0]
524
+
525
+ # If the source and target are the same file, no copy is needed.
526
+ if source_file == output_path:
527
+ return source_file
528
+
529
+ shutil.copy(source_file, output_path)
530
+ log.info(f"Copied single SDIF file to '{output_path}' as no merge was needed.")
531
+ return output_path
532
+
533
+ # For multiple files, merge them into the output_path
534
+ merger = _SDIFMerger(output_path)
535
+ return merger.merge_all(processed_sdif_paths)
@@ -0,0 +1,97 @@
1
+ import logging
2
+ from typing import Any
3
+
4
+ from agents.mcp.server import CallToolResult, MCPServer, MCPTool
5
+ from fastmcp import FastMCP
6
+
7
+ logger = logging.getLogger(__name__)
8
+
9
+
10
+ class OpenAICompatibleMCP(MCPServer):
11
+ def __init__(self, mcp: FastMCP):
12
+ self.mcp = mcp
13
+ self._is_connected = False # Track connection state
14
+
15
+ async def connect(self):
16
+ """Connect to the server.
17
+ For FastMCP, connection is managed externally when the server is run.
18
+ This method marks the wrapper as connected.
19
+ """
20
+ # Assuming FastMCP instance is already running and configured.
21
+ # No specific connect action required for the FastMCP instance itself here,
22
+ # as its lifecycle (run, stop) is managed outside this wrapper.
23
+ logger.info(
24
+ f"OpenAICompatibleMCP: Simulating connection to FastMCP server '{self.mcp.name}'."
25
+ )
26
+ self._is_connected = True
27
+
28
+ @property
29
+ def name(self) -> str:
30
+ """A readable name for the server."""
31
+ return self.mcp.name
32
+
33
+ async def cleanup(self):
34
+ """Cleanup the server.
35
+ For FastMCP, cleanup is managed externally. This method marks the wrapper as disconnected.
36
+ """
37
+ # Similar to connect, actual server cleanup is external.
38
+ logger.info(
39
+ f"OpenAICompatibleMCP: Simulating cleanup for FastMCP server '{self.mcp.name}'."
40
+ )
41
+ self._is_connected = False
42
+
43
+ async def list_tools(self) -> list[MCPTool]:
44
+ """List the tools available on the server."""
45
+ if not self._is_connected:
46
+ # Or raise an error, depending on desired behavior for disconnected state
47
+ raise RuntimeError(
48
+ "OpenAICompatibleMCP.list_tools called while not connected."
49
+ )
50
+
51
+ # FastMCP's get_tools() returns a dict[str, fastmcp.tools.tool.Tool]
52
+ # Each fastmcp.tools.tool.Tool has a to_mcp_tool(name=key) method
53
+ # MCPTool is an alias for mcp.types.Tool
54
+ try:
55
+ fastmcp_tools = await self.mcp.get_tools()
56
+ mcp_tools_list = [
57
+ tool.to_mcp_tool(name=key) for key, tool in fastmcp_tools.items()
58
+ ]
59
+ return mcp_tools_list
60
+ except Exception as e:
61
+ logger.error(
62
+ f"Error listing tools from FastMCP server '{self.mcp.name}': {e}",
63
+ exc_info=True,
64
+ )
65
+ raise e
66
+
67
+ async def call_tool(
68
+ self, tool_name: str, arguments: dict[str, Any] | None
69
+ ) -> CallToolResult:
70
+ """Invoke a tool on the server."""
71
+ if not self._is_connected:
72
+ logger.warning(
73
+ f"OpenAICompatibleMCP.call_tool '{tool_name}' called while not connected."
74
+ )
75
+ # Return an error CallToolResult
76
+ return CallToolResult(
77
+ content=[{"type": "text", "text": "Server not connected"}], isError=True
78
+ )
79
+
80
+ try:
81
+ # FastMCP's _mcp_call_tool is a protected member, but seems to be what we need.
82
+ # It returns: list[TextContent | ImageContent | EmbeddedResource]
83
+ # This matches the 'content' part of CallToolResult.
84
+ # We need to handle potential errors and wrap the result.
85
+ content = await self.mcp._mcp_call_tool(tool_name, arguments or {})
86
+ return CallToolResult(content=content, isError=False)
87
+ except Exception as e:
88
+ logger.error(
89
+ f"Error calling tool '{tool_name}' on FastMCP server '{self.mcp.name}': {e}",
90
+ exc_info=True,
91
+ )
92
+ error_message = f"Error calling tool '{tool_name}': {type(e).__name__}: {e}"
93
+ # Ensure content is a list of valid MCP content items, even for errors.
94
+ # A TextContent is a safe choice.
95
+ return CallToolResult(
96
+ content=[{"type": "text", "text": error_message}], isError=True
97
+ )