lollms-client 0.17.2__py3-none-any.whl → 0.19.1__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.

Potentially problematic release.


This version of lollms-client might be problematic. Click here for more details.

@@ -0,0 +1,311 @@
1
+ # lollms_client/mcp_bindings/local_mcp/__init__.py
2
+ import json
3
+ import importlib.util
4
+ from pathlib import Path
5
+ from typing import Optional, List, Dict, Any, Union
6
+
7
+ from lollms_client.lollms_mcp_binding import LollmsMCPBinding
8
+ from ascii_colors import ASCIIColors, trace_exception
9
+
10
+ # This variable is used by the LollmsMCPBindingManager to identify the binding class.
11
+ BindingName = "LocalMCPBinding"
12
+
13
+ class LocalMCPBinding(LollmsMCPBinding):
14
+ """
15
+ Local Model Context Protocol (MCP) Binding.
16
+
17
+ This binding discovers and executes tools defined locally in a specified folder.
18
+ Each tool is expected to have:
19
+ 1. A Python file (`<tool_name>.py`) containing an `execute(params: Dict[str, Any]) -> Dict[str, Any]` function.
20
+ 2. A JSON file (`<tool_name>.mcp.json`) defining the tool's MCP metadata (name, description, input_schema, output_schema).
21
+ """
22
+
23
+ def __init__(self,
24
+ tools_folder_path: str|Path|None = None):
25
+ """
26
+ Initialize the LocalMCPBinding.
27
+
28
+ Args:
29
+ binding_name (str): The name of this binding.
30
+ tools_folder_path (str|Path) a folder where to find tools
31
+ """
32
+ super().__init__(binding_name="LocalMCP")
33
+ if tools_folder_path:
34
+ try:
35
+ self.tools_folder_path: Optional[Path] = Path(tools_folder_path)
36
+ except:
37
+ self.tools_folder_path = None
38
+ else:
39
+ self.tools_folder_path = None
40
+ self.discovered_tools: List[Dict[str, Any]] = []
41
+ if not self.tools_folder_path:
42
+ self.tools_folder_path = Path(__file__).parent/"default_tools"
43
+ self._discover_local_tools()
44
+
45
+ def _discover_local_tools(self):
46
+ """Scans the tools_folder_path for valid tool definitions."""
47
+ if not self.tools_folder_path or not self.tools_folder_path.is_dir():
48
+ return
49
+
50
+ self.discovered_tools = []
51
+ ASCIIColors.info(f"Discovering local MCP tools in: {self.tools_folder_path}")
52
+
53
+ for item in self.tools_folder_path.iterdir():
54
+ if item.is_dir(): # Each tool in its own subdirectory
55
+ tool_name = item.name
56
+ mcp_json_file = item / f"{tool_name}.mcp.json"
57
+ python_file = item / f"{tool_name}.py"
58
+
59
+ if mcp_json_file.exists() and python_file.exists():
60
+ try:
61
+ with open(mcp_json_file, 'r', encoding='utf-8') as f:
62
+ tool_def = json.load(f)
63
+
64
+ # Basic validation of MCP definition
65
+ if not all(k in tool_def for k in ["name", "description", "input_schema"]):
66
+ ASCIIColors.warning(f"Tool '{tool_name}' MCP definition is missing required fields (name, description, input_schema). Skipping.")
67
+ continue
68
+ if tool_def["name"] != tool_name:
69
+ ASCIIColors.warning(f"Tool name in MCP JSON ('{tool_def['name']}') does not match folder/file name ('{tool_name}'). Using folder name. Consider aligning them.")
70
+ tool_def["name"] = tool_name # Standardize to folder name
71
+
72
+ # Store the full definition and path to python file for execution
73
+ tool_def['_python_file_path'] = str(python_file.resolve())
74
+ self.discovered_tools.append(tool_def)
75
+ ASCIIColors.green(f"Discovered local tool: {tool_name}")
76
+ except json.JSONDecodeError:
77
+ ASCIIColors.warning(f"Could not parse MCP JSON for tool '{tool_name}'. Skipping.")
78
+ except Exception as e:
79
+ ASCIIColors.warning(f"Error loading tool '{tool_name}': {e}")
80
+ trace_exception(e)
81
+ else:
82
+ if not mcp_json_file.exists():
83
+ ASCIIColors.debug(f"Tool '{tool_name}' missing MCP JSON definition ({mcp_json_file.name}). Skipping.")
84
+ if not python_file.exists():
85
+ ASCIIColors.debug(f"Tool '{tool_name}' missing Python implementation ({python_file.name}). Skipping.")
86
+ ASCIIColors.info(f"Discovery complete. Found {len(self.discovered_tools)} local tools.")
87
+
88
+
89
+ def discover_tools(self, specific_tool_names: Optional[List[str]] = None, **kwargs) -> List[Dict[str, Any]]:
90
+ """
91
+ Discover available local tools.
92
+
93
+ Args:
94
+ specific_tool_names (Optional[List[str]]): If provided, filter discovery
95
+ to only these tool names.
96
+ **kwargs: Ignored by this binding.
97
+
98
+ Returns:
99
+ List[Dict[str, Any]]: A list of discovered tool definitions.
100
+ """
101
+ if not self.tools_folder_path:
102
+ return []
103
+
104
+ # Re-scan if needed, or if discovery hasn't happened
105
+ if not self.discovered_tools and self.tools_folder_path:
106
+ self._discover_local_tools()
107
+
108
+ if specific_tool_names:
109
+ return [tool for tool in self.discovered_tools if tool.get("name") in specific_tool_names]
110
+ return self.discovered_tools
111
+
112
+ def execute_tool(self,
113
+ tool_name: str,
114
+ params: Dict[str, Any],
115
+ lollms_client_instance: Optional[Any] = None, # Added lollms_client_instance
116
+ **kwargs) -> Dict[str, Any]:
117
+ """
118
+ Execute a locally defined Python tool.
119
+
120
+ Args:
121
+ tool_name (str): The name of the tool to execute.
122
+ params (Dict[str, Any]): Parameters for the tool.
123
+ lollms_client_instance (Optional[Any]): The LollmsClient instance, if available.
124
+ **kwargs: Ignored by this binding.
125
+
126
+ Returns:
127
+ Dict[str, Any]: The result from the tool's execute function, or an error dictionary.
128
+ """
129
+ tool_def = next((t for t in self.discovered_tools if t.get("name") == tool_name), None)
130
+
131
+ if not tool_def:
132
+ return {"error": f"Local tool '{tool_name}' not found or not discovered.", "status_code": 404}
133
+
134
+ python_file_path_str = tool_def.get('_python_file_path')
135
+ if not python_file_path_str:
136
+ return {"error": f"Python implementation path missing for tool '{tool_name}'.", "status_code": 500}
137
+
138
+ python_file_path = Path(python_file_path_str)
139
+
140
+ try:
141
+ module_name = f"lollms_client.mcp_bindings.local_mcp.tools.{tool_name}"
142
+ spec = importlib.util.spec_from_file_location(module_name, str(python_file_path))
143
+
144
+ if not spec or not spec.loader:
145
+ return {"error": f"Could not create module spec for tool '{tool_name}'.", "status_code": 500}
146
+
147
+ tool_module = importlib.util.module_from_spec(spec)
148
+ spec.loader.exec_module(tool_module)
149
+
150
+ if not hasattr(tool_module, 'execute'):
151
+ return {"error": f"Tool '{tool_name}' Python file does not have an 'execute' function.", "status_code": 500}
152
+
153
+ execute_function = getattr(tool_module, 'execute')
154
+
155
+ # Inspect the execute function's signature
156
+ import inspect
157
+ sig = inspect.signature(execute_function)
158
+
159
+ exec_params = {}
160
+ if 'params' in sig.parameters: # Always pass params if expected
161
+ exec_params['params'] = params
162
+
163
+ # Conditionally pass lollms_client_instance
164
+ if 'lollms_client_instance' in sig.parameters and lollms_client_instance is not None:
165
+ exec_params['lollms_client_instance'] = lollms_client_instance
166
+ elif 'lollms_client_instance' in sig.parameters and lollms_client_instance is None:
167
+ ASCIIColors.warning(f"Tool '{tool_name}' expects 'lollms_client_instance', but it was not provided to execute_tool.")
168
+
169
+ ASCIIColors.info(f"Executing local tool '{tool_name}' with effective params for its execute(): {exec_params.keys()}")
170
+ result = execute_function(**exec_params) # Pass parameters accordingly
171
+
172
+ return {"output": result, "status_code": 200}
173
+
174
+ except Exception as e:
175
+ trace_exception(e)
176
+ return {"error": f"Error executing tool '{tool_name}': {str(e)}", "status_code": 500}
177
+
178
+
179
+ # --- Example Usage (for testing within this file) ---
180
+ if __name__ == '__main__':
181
+ ASCIIColors.magenta("--- LocalMCPBinding Test ---")
182
+
183
+ # Create a temporary tools directory for testing
184
+ test_tools_base_dir = Path(__file__).parent.parent.parent / "temp_mcp_tools_for_test" # Place it outside the package
185
+ test_tools_base_dir.mkdir(parents=True, exist_ok=True)
186
+
187
+ # Define a sample tool: get_weather
188
+ tool1_dir = test_tools_base_dir / "get_weather"
189
+ tool1_dir.mkdir(exist_ok=True)
190
+
191
+ # MCP JSON for get_weather
192
+ get_weather_mcp = {
193
+ "name": "get_weather",
194
+ "description": "Fetches the current weather for a given city.",
195
+ "input_schema": {
196
+ "type": "object",
197
+ "properties": {
198
+ "city": {"type": "string", "description": "The city name."},
199
+ "unit": {"type": "string", "enum": ["celsius", "fahrenheit"], "default": "celsius"}
200
+ },
201
+ "required": ["city"]
202
+ },
203
+ "output_schema": { # Optional, but good practice
204
+ "type": "object",
205
+ "properties": {
206
+ "temperature": {"type": "number"},
207
+ "condition": {"type": "string"},
208
+ "unit": {"type": "string"}
209
+ }
210
+ }
211
+ }
212
+ with open(tool1_dir / "get_weather.mcp.json", "w") as f:
213
+ json.dump(get_weather_mcp, f, indent=2)
214
+
215
+ # Python code for get_weather
216
+ get_weather_py = """
217
+ import random
218
+ def execute(params: dict) -> dict:
219
+ city = params.get("city")
220
+ unit = params.get("unit", "celsius")
221
+
222
+ if not city:
223
+ return {"error": "City not provided"}
224
+
225
+ # Simulate weather fetching
226
+ conditions = ["sunny", "cloudy", "rainy", "snowy"]
227
+ temp = random.randint(-10 if unit == "celsius" else 14, 35 if unit == "celsius" else 95)
228
+
229
+ return {
230
+ "temperature": temp,
231
+ "condition": random.choice(conditions),
232
+ "unit": unit
233
+ }
234
+ """
235
+ with open(tool1_dir / "get_weather.py", "w") as f:
236
+ f.write(get_weather_py)
237
+
238
+ # Define another sample tool: sum_numbers
239
+ tool2_dir = test_tools_base_dir / "sum_numbers"
240
+ tool2_dir.mkdir(exist_ok=True)
241
+ sum_numbers_mcp = {
242
+ "name": "sum_numbers",
243
+ "description": "Calculates the sum of a list of numbers.",
244
+ "input_schema": {
245
+ "type": "object",
246
+ "properties": {"numbers": {"type": "array", "items": {"type": "number"}}},
247
+ "required": ["numbers"]
248
+ },
249
+ "output_schema": {"type": "object", "properties": {"sum": {"type": "number"}}}
250
+ }
251
+ with open(tool2_dir / "sum_numbers.mcp.json", "w") as f:
252
+ json.dump(sum_numbers_mcp, f, indent=2)
253
+ sum_numbers_py = """
254
+ def execute(params: dict) -> dict:
255
+ numbers = params.get("numbers", [])
256
+ if not isinstance(numbers, list) or not all(isinstance(n, (int, float)) for n in numbers):
257
+ return {"error": "Invalid input: 'numbers' must be a list of numbers."}
258
+ return {"sum": sum(numbers)}
259
+ """
260
+ with open(tool2_dir / "sum_numbers.py", "w") as f:
261
+ f.write(sum_numbers_py)
262
+
263
+
264
+ local_mcp_binding = LocalMCPBinding(binding_name="local_mcp_test", tools_folder_path=test_tools_base_dir)
265
+
266
+ ASCIIColors.cyan("\n1. Discovering all tools...")
267
+ all_tools = local_mcp_binding.discover_tools()
268
+ if all_tools:
269
+ ASCIIColors.green(f"Discovered {len(all_tools)} tools:")
270
+ for tool in all_tools:
271
+ print(f" - Name: {tool.get('name')}, Description: {tool.get('description')}")
272
+ assert "_python_file_path" in tool # Internal check
273
+ else:
274
+ ASCIIColors.warning("No tools discovered.")
275
+
276
+ ASCIIColors.cyan("\n2. Executing 'get_weather' tool...")
277
+ weather_params = {"city": "London", "unit": "celsius"}
278
+ weather_result = local_mcp_binding.execute_tool("get_weather", weather_params)
279
+ ASCIIColors.green(f"Weather result: {weather_result}")
280
+ assert "error" not in weather_result.get("output", {}), f"Weather tool execution failed: {weather_result}"
281
+ assert weather_result.get("status_code") == 200
282
+
283
+ ASCIIColors.cyan("\n3. Executing 'sum_numbers' tool...")
284
+ sum_params = {"numbers": [10, 2.5, 7]}
285
+ sum_result = local_mcp_binding.execute_tool("sum_numbers", sum_params)
286
+ ASCIIColors.green(f"Sum result: {sum_result}")
287
+ assert sum_result.get("output", {}).get("sum") == 19.5, f"Sum tool execution incorrect: {sum_result}"
288
+ assert sum_result.get("status_code") == 200
289
+
290
+ ASCIIColors.cyan("\n4. Executing non-existent tool...")
291
+ non_existent_result = local_mcp_binding.execute_tool("do_magic", {"spell": "abracadabra"})
292
+ ASCIIColors.warning(f"Non-existent tool result: {non_existent_result}")
293
+ assert non_existent_result.get("status_code") == 404
294
+
295
+ ASCIIColors.cyan("\n5. Discovering a specific tool ('sum_numbers')...")
296
+ specific_tools = local_mcp_binding.discover_tools(specific_tool_names=["sum_numbers"])
297
+ if specific_tools and len(specific_tools) == 1 and specific_tools[0].get("name") == "sum_numbers":
298
+ ASCIIColors.green("Successfully discovered specific tool 'sum_numbers'.")
299
+ else:
300
+ ASCIIColors.error(f"Failed to discover specific tool. Found: {specific_tools}")
301
+
302
+
303
+ # Cleanup: Remove the temporary tools directory
304
+ import shutil
305
+ try:
306
+ shutil.rmtree(test_tools_base_dir)
307
+ ASCIIColors.info(f"Cleaned up temporary tools directory: {test_tools_base_dir}")
308
+ except Exception as e_clean:
309
+ ASCIIColors.error(f"Could not clean up temporary tools directory: {e_clean}")
310
+
311
+ ASCIIColors.magenta("\n--- LocalMCPBinding Test Finished ---")
@@ -0,0 +1,74 @@
1
+ from pathlib import Path
2
+ from typing import Dict, Any
3
+
4
+ def execute(params: Dict[str, Any]) -> Dict[str, Any]:
5
+ """
6
+ Writes or appends text content to a specified file.
7
+ """
8
+ file_path_str = params.get("file_path")
9
+ content = params.get("content")
10
+ mode = params.get("mode", "overwrite") # Default to overwrite
11
+
12
+ if not file_path_str:
13
+ return {"status": "error", "message": "file_path parameter is required."}
14
+ if content is None: # Allow empty string, but not None
15
+ return {"status": "error", "message": "content parameter is required."}
16
+ if mode not in ["overwrite", "append"]:
17
+ return {"status": "error", "message": "Invalid mode. Must be 'overwrite' or 'append'."}
18
+
19
+ try:
20
+ # Security consideration: Restrict file paths if necessary in a real application.
21
+ # For this local tool, we'll assume the user/AI provides paths responsibly.
22
+ # However, avoid absolute paths starting from root unless explicitly intended.
23
+ # For simplicity here, we allow relative and 'safe' absolute paths.
24
+
25
+ # Make path relative to a 'workspace' if not absolute, to prevent writing anywhere.
26
+ # For this example, let's assume a 'tool_workspace' directory next to the tools_folder.
27
+ # This needs to be configurable or defined by the LocalMCPBinding.
28
+ # For now, we'll just resolve the path. A real implementation would need sandboxing.
29
+
30
+ target_file = Path(file_path_str)
31
+
32
+ # Create parent directories if they don't exist
33
+ target_file.parent.mkdir(parents=True, exist_ok=True)
34
+
35
+ write_mode = "w" if mode == "overwrite" else "a"
36
+
37
+ with open(target_file, write_mode, encoding='utf-8') as f:
38
+ f.write(content)
39
+
40
+ return {
41
+ "status": "success",
42
+ "message": f"Content successfully {'written to' if mode == 'overwrite' else 'appended to'} file.",
43
+ "file_path": str(target_file.resolve())
44
+ }
45
+ except Exception as e:
46
+ return {
47
+ "status": "error",
48
+ "message": f"Failed to write to file: {str(e)}",
49
+ "file_path": str(Path(file_path_str).resolve() if file_path_str else "N/A")
50
+ }
51
+
52
+ if __name__ == '__main__':
53
+ # Example test
54
+ test_params_overwrite = {
55
+ "file_path": "test_output/example.txt",
56
+ "content": "Hello from the file_writer tool!\nThis is the first line.",
57
+ "mode": "overwrite"
58
+ }
59
+ result_overwrite = execute(test_params_overwrite)
60
+ print(f"Overwrite Test Result: {result_overwrite}")
61
+
62
+ test_params_append = {
63
+ "file_path": "test_output/example.txt",
64
+ "content": "\nThis is an appended line.",
65
+ "mode": "append"
66
+ }
67
+ result_append = execute(test_params_append)
68
+ print(f"Append Test Result: {result_append}")
69
+
70
+ test_params_error = {
71
+ "content": "Missing file path"
72
+ }
73
+ result_error = execute(test_params_error)
74
+ print(f"Error Test Result: {result_error}")
@@ -0,0 +1,195 @@
1
+ from pathlib import Path
2
+ from typing import Dict, Any
3
+ import uuid
4
+ import io
5
+ # We expect LollmsClient instance to be passed if the tool needs it.
6
+ # from lollms_client import LollmsClient # Not imported directly, but type hint is useful
7
+
8
+ def execute(params: Dict[str, Any], lollms_client_instance: Any = None) -> Dict[str, Any]:
9
+ """
10
+ Generates an image using the LollmsClient's TTI binding.
11
+ The lollms_client_instance is expected to be passed by the LocalMCPBinding if available.
12
+ """
13
+ prompt = params.get("prompt")
14
+ negative_prompt = params.get("negative_prompt", "")
15
+ width = params.get("width") # Will be None if not provided
16
+ height = params.get("height") # Will be None if not provided
17
+ output_filename_suggestion = params.get("output_filename_suggestion", "generated_image.png")
18
+
19
+ if not prompt:
20
+ return {"status": "error", "message": "'prompt' parameter is required."}
21
+
22
+ if not lollms_client_instance or not hasattr(lollms_client_instance, 'tti') or not lollms_client_instance.tti:
23
+ return {
24
+ "status": "tti_not_available",
25
+ "message": "LollmsClient instance with an active TTI binding was not provided or TTI is not configured.",
26
+ "image_path": None,
27
+ "image_url": None
28
+ }
29
+
30
+ try:
31
+ tti_binding = lollms_client_instance.tti
32
+
33
+ # Prepare arguments for the TTI binding's generate_image method
34
+ tti_kwargs = {}
35
+ if width is not None:
36
+ tti_kwargs['width'] = width
37
+ if height is not None:
38
+ tti_kwargs['height'] = height
39
+ # Add other common TTI params if desired, e.g. seed, steps, cfg_scale
40
+ # These would need to be added to the MCP JSON input_schema as well.
41
+
42
+ image_bytes = tti_binding.generate_image(
43
+ prompt=prompt,
44
+ negative_prompt=negative_prompt,
45
+ **tti_kwargs
46
+ )
47
+
48
+ if not image_bytes:
49
+ return {
50
+ "status": "error",
51
+ "message": "TTI binding returned no image data.",
52
+ "image_path": None,
53
+ "image_url": None
54
+ }
55
+
56
+ # Define where to save the image.
57
+ # This should ideally be a secure, configurable workspace.
58
+ # For this example, we save it into a 'mcp_generated_images' subdirectory
59
+ # of the current working directory OR a path derived from lollms_paths if available.
60
+
61
+ # Prefer using a path from lollms_client_instance if available (e.g., an output or data path)
62
+ # This part needs careful consideration for a real application.
63
+ save_dir_base = Path.cwd()
64
+ if hasattr(lollms_client_instance, 'lollms_paths_config') and lollms_client_instance.lollms_paths_config.get('personal_outputs_path'):
65
+ save_dir_base = Path(lollms_client_instance.lollms_paths_config['personal_outputs_path'])
66
+ elif hasattr(lollms_client_instance, 'lollms_paths_config') and lollms_client_instance.lollms_paths_config.get('shared_outputs_path'): # type: ignore
67
+ save_dir_base = Path(lollms_client_instance.lollms_paths_config['shared_outputs_path']) # type: ignore
68
+
69
+ save_dir = save_dir_base / "mcp_generated_images"
70
+ save_dir.mkdir(parents=True, exist_ok=True)
71
+
72
+ # Sanitize filename and make it unique
73
+ base, ext = Path(output_filename_suggestion).stem, Path(output_filename_suggestion).suffix
74
+ if not ext: ext = ".png" # Default to png
75
+ safe_base = "".join(c if c.isalnum() or c in ['_', '-'] else '_' for c in base)
76
+ unique_filename = f"{safe_base}_{uuid.uuid4().hex[:8]}{ext}"
77
+
78
+ image_save_path = save_dir / unique_filename
79
+
80
+ with open(image_save_path, "wb") as f:
81
+ f.write(image_bytes)
82
+
83
+ # TODO: How to best represent the URL? For local files, it might be a file:// URL
84
+ # or a relative path that the client application understands.
85
+ # For now, returning a relative path from a conceptual 'outputs' root.
86
+ # A more robust solution would involve the LollmsClient serving these images
87
+ # or providing data URLs.
88
+
89
+ # Create a relative path for client display if possible
90
+ # This assumes the client knows how to interpret "mcp_generated_images/..."
91
+ relative_image_path = f"mcp_generated_images/{unique_filename}"
92
+ image_url = f"file:///{image_save_path.resolve()}" # Example data URL or file path
93
+
94
+ return {
95
+ "status": "success",
96
+ "message": f"Image generated and saved successfully.",
97
+ "image_path": relative_image_path, # More of a hint
98
+ "image_url": image_url # More concrete path
99
+ }
100
+
101
+ except Exception as e:
102
+ # from ascii_colors import trace_exception # If you need full trace
103
+ # trace_exception(e)
104
+ return {
105
+ "status": "error",
106
+ "message": f"Failed to generate image via TTI binding: {str(e)}",
107
+ "image_path": None,
108
+ "image_url": None
109
+ }
110
+
111
+ if __name__ == '__main__':
112
+ import json
113
+ from PIL import Image as PILImage # To avoid conflict with module-level Image if any
114
+ print("--- Generate Image (via LollmsClient TTI) Tool Test ---")
115
+
116
+ # This test requires a LollmsClient instance with a configured TTI binding.
117
+ # We'll mock it for a standalone test of the execute function's logic.
118
+
119
+ class MockTTIBinding:
120
+ def __init__(self, works=True):
121
+ self.works = works
122
+ self.config = {"default_width": 512, "default_height": 512} # Mock config
123
+
124
+ def generate_image(self, prompt, negative_prompt="", width=None, height=None, **kwargs):
125
+ if not self.works:
126
+ # return None # Simulate TTI returning no data
127
+ raise ValueError("Simulated TTI error in generate_image")
128
+
129
+ print(f"MockTTI: Generating image for prompt: '{prompt}', neg: '{negative_prompt}', W:{width}, H:{height}")
130
+ # Create a dummy PIL image and return its bytes
131
+ img = PILImage.new('RGB', (width or 512, height or 512), color = 'skyblue' if "sky" in prompt else "lightcoral")
132
+ from PIL import ImageDraw
133
+ draw = ImageDraw.Draw(img)
134
+ draw.text((10,10), prompt[:30], fill=(0,0,0))
135
+ byte_arr = io.BytesIO()
136
+ img.save(byte_arr, format='PNG')
137
+ return byte_arr.getvalue()
138
+
139
+ class MockLollmsClient:
140
+ def __init__(self, tti_works=True):
141
+ self.tti = MockTTIBinding(works=tti_works)
142
+ self.lollms_paths_config = {"personal_outputs_path": Path("./test_lollms_client_outputs")} # Mock path
143
+ Path(self.lollms_paths_config["personal_outputs_path"]).mkdir(exist_ok=True)
144
+
145
+
146
+ mock_lc_success = MockLollmsClient(tti_works=True)
147
+ mock_lc_tti_fail = MockLollmsClient(tti_works=False)
148
+ mock_lc_no_tti = MockLollmsClient()
149
+ mock_lc_no_tti.tti = None # Simulate TTI not configured
150
+
151
+ # Test 1: Successful generation
152
+ params1 = {"prompt": "A beautiful sunset over mountains", "width": 256, "height": 256}
153
+ result1 = execute(params1, lollms_client_instance=mock_lc_success)
154
+ print(f"\nTest 1 Result (Success):\n{json.dumps(result1, indent=2)}")
155
+ assert result1["status"] == "success"
156
+ assert result1["image_path"] is not None
157
+ if result1["image_path"]:
158
+ # Verify file was created (adjust path based on where it's saved by tool)
159
+ # For this test, it's saved in ./test_lollms_client_outputs/mcp_generated_images/
160
+ full_saved_path = Path(mock_lc_success.lollms_paths_config["personal_outputs_path"]) / result1["image_path"]
161
+ print(f"Checking for image at: {full_saved_path}")
162
+ assert full_saved_path.exists()
163
+ # Optional: Clean up created image after test
164
+ # full_saved_path.unlink(missing_ok=True)
165
+
166
+
167
+ # Test 2: TTI binding itself fails to generate
168
+ params2 = {"prompt": "A futuristic city"}
169
+ result2 = execute(params2, lollms_client_instance=mock_lc_tti_fail)
170
+ print(f"\nTest 2 Result (TTI Fails):\n{json.dumps(result2, indent=2)}")
171
+ assert result2["status"] == "error"
172
+ assert "Simulated TTI error" in result2["message"]
173
+
174
+ # Test 3: No TTI binding available in LollmsClient
175
+ params3 = {"prompt": "A cat wearing a hat"}
176
+ result3 = execute(params3, lollms_client_instance=mock_lc_no_tti)
177
+ print(f"\nTest 3 Result (No TTI Binding):\n{json.dumps(result3, indent=2)}")
178
+ assert result3["status"] == "tti_not_available"
179
+
180
+ # Test 4: Missing prompt
181
+ params4 = {}
182
+ result4 = execute(params4, lollms_client_instance=mock_lc_success)
183
+ print(f"\nTest 4 Result (Missing Prompt):\n{json.dumps(result4, indent=2)}")
184
+ assert result4["status"] == "error"
185
+ assert "'prompt' parameter is required" in result4["message"]
186
+
187
+ # Test 5: No LollmsClient instance passed (should also result in tti_not_available)
188
+ params5 = {"prompt": "A dog"}
189
+ result5 = execute(params5, lollms_client_instance=None)
190
+ print(f"\nTest 5 Result (No LollmsClient passed):\n{json.dumps(result5, indent=2)}")
191
+ assert result5["status"] == "tti_not_available"
192
+
193
+
194
+ print("\n--- Tests Finished ---")
195
+ print(f"Generated test images (if any) are in subdirectories of: {mock_lc_success.lollms_paths_config['personal_outputs_path']}")