stores 0.1.6__py3-none-any.whl → 0.1.7.dev2__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.
- stores/indexes/base_index.py +177 -44
- stores/indexes/index.py +11 -1
- stores/indexes/local_index.py +20 -4
- stores/indexes/remote_index.py +7 -1
- stores/indexes/venv_utils.py +170 -51
- {stores-0.1.6.dist-info → stores-0.1.7.dev2.dist-info}/METADATA +1 -1
- stores-0.1.7.dev2.dist-info/RECORD +15 -0
- stores-0.1.6.dist-info/RECORD +0 -15
- {stores-0.1.6.dist-info → stores-0.1.7.dev2.dist-info}/WHEEL +0 -0
- {stores-0.1.6.dist-info → stores-0.1.7.dev2.dist-info}/licenses/LICENSE +0 -0
stores/indexes/base_index.py
CHANGED
@@ -128,6 +128,29 @@ def _undo_non_string_literal(annotation: type, value: Any, literal_map: dict):
|
|
128
128
|
return value
|
129
129
|
|
130
130
|
|
131
|
+
def _preprocess_args(
|
132
|
+
original_signature: inspect.Signature, literal_maps: dict, args: list, kwargs: dict
|
133
|
+
):
|
134
|
+
# Inject default values within wrapper
|
135
|
+
bound_args = original_signature.bind(*args, **kwargs)
|
136
|
+
bound_args.apply_defaults()
|
137
|
+
for k, v in bound_args.arguments.items():
|
138
|
+
if (
|
139
|
+
v is None
|
140
|
+
and original_signature.parameters[k].default is not Parameter.empty
|
141
|
+
):
|
142
|
+
bound_args.arguments[k] = original_signature.parameters[k].default
|
143
|
+
_cast_bound_args(bound_args)
|
144
|
+
# Inject correct Literals
|
145
|
+
for k, v in bound_args.arguments.items():
|
146
|
+
if k in literal_maps:
|
147
|
+
param = original_signature.parameters[k]
|
148
|
+
bound_args.arguments[k] = _undo_non_string_literal(
|
149
|
+
param.annotation, v, literal_maps[k]
|
150
|
+
)
|
151
|
+
return bound_args
|
152
|
+
|
153
|
+
|
131
154
|
def wrap_tool(tool: Callable):
|
132
155
|
"""
|
133
156
|
Wrap tool to make it compatible with LLM libraries
|
@@ -189,42 +212,37 @@ def wrap_tool(tool: Callable):
|
|
189
212
|
new_args.append(new_arg)
|
190
213
|
new_sig = original_signature.replace(parameters=new_args)
|
191
214
|
|
192
|
-
if inspect.
|
215
|
+
if inspect.isasyncgenfunction(tool):
|
216
|
+
|
217
|
+
async def wrapper(*args, **kwargs):
|
218
|
+
bound_args = _preprocess_args(
|
219
|
+
original_signature, literal_maps, args, kwargs
|
220
|
+
)
|
221
|
+
async for value in tool(*bound_args.args, **bound_args.kwargs):
|
222
|
+
yield value
|
223
|
+
|
224
|
+
elif inspect.isgeneratorfunction(tool):
|
225
|
+
|
226
|
+
def wrapper(*args, **kwargs):
|
227
|
+
bound_args = _preprocess_args(
|
228
|
+
original_signature, literal_maps, args, kwargs
|
229
|
+
)
|
230
|
+
for value in tool(*bound_args.args, **bound_args.kwargs):
|
231
|
+
yield value
|
232
|
+
|
233
|
+
elif inspect.iscoroutinefunction(tool):
|
193
234
|
|
194
235
|
async def wrapper(*args, **kwargs):
|
195
|
-
|
196
|
-
|
197
|
-
|
198
|
-
_cast_bound_args(bound_args)
|
199
|
-
# Inject correct Literals
|
200
|
-
for k, v in bound_args.arguments.items():
|
201
|
-
if k in literal_maps:
|
202
|
-
param = original_signature.parameters[k]
|
203
|
-
bound_args.arguments[k] = _undo_non_string_literal(
|
204
|
-
param.annotation, v, literal_maps[k]
|
205
|
-
)
|
236
|
+
bound_args = _preprocess_args(
|
237
|
+
original_signature, literal_maps, args, kwargs
|
238
|
+
)
|
206
239
|
return await tool(*bound_args.args, **bound_args.kwargs)
|
207
240
|
else:
|
208
241
|
|
209
242
|
def wrapper(*args, **kwargs):
|
210
|
-
|
211
|
-
|
212
|
-
|
213
|
-
# Inject correct Literals
|
214
|
-
for k, v in bound_args.arguments.items():
|
215
|
-
if (
|
216
|
-
v is None
|
217
|
-
and original_signature.parameters[k].default is not Parameter.empty
|
218
|
-
):
|
219
|
-
bound_args.arguments[k] = original_signature.parameters[k].default
|
220
|
-
|
221
|
-
_cast_bound_args(bound_args)
|
222
|
-
for k, v in bound_args.arguments.items():
|
223
|
-
if k in literal_maps:
|
224
|
-
param = original_signature.parameters[k]
|
225
|
-
bound_args.arguments[k] = _undo_non_string_literal(
|
226
|
-
param.annotation, v, literal_maps[k]
|
227
|
-
)
|
243
|
+
bound_args = _preprocess_args(
|
244
|
+
original_signature, literal_maps, args, kwargs
|
245
|
+
)
|
228
246
|
return tool(*bound_args.args, **bound_args.kwargs)
|
229
247
|
|
230
248
|
wrapped = create_function(
|
@@ -249,14 +267,7 @@ class BaseIndex:
|
|
249
267
|
def tools_dict(self):
|
250
268
|
return {tool.__name__: tool for tool in self.tools}
|
251
269
|
|
252
|
-
def
|
253
|
-
loop = asyncio.new_event_loop()
|
254
|
-
asyncio.set_event_loop(loop)
|
255
|
-
return loop.run_until_complete(self.async_execute(toolname, kwargs))
|
256
|
-
|
257
|
-
async def async_execute(self, toolname: str, kwargs: dict | None = None):
|
258
|
-
kwargs = kwargs or {}
|
259
|
-
|
270
|
+
def _get_tool(self, toolname: str):
|
260
271
|
# Use regex since we need to match cases where we perform
|
261
272
|
# substitutions such as replace(".", "-")
|
262
273
|
pattern = re.compile(":?" + re.sub("-|\\.", "(-|\\.)", toolname) + "$")
|
@@ -272,11 +283,124 @@ class BaseIndex:
|
|
272
283
|
else:
|
273
284
|
toolname = matching_tools[0]
|
274
285
|
|
275
|
-
|
276
|
-
|
277
|
-
|
286
|
+
return self.tools_dict[toolname]
|
287
|
+
|
288
|
+
def execute(self, toolname: str, kwargs: dict | None = None, collect_results=False):
|
289
|
+
tool_fn = self._get_tool(toolname)
|
290
|
+
kwargs = kwargs or {}
|
291
|
+
if inspect.isasyncgenfunction(tool_fn):
|
292
|
+
# Handle async generator
|
293
|
+
|
294
|
+
async def collect():
|
295
|
+
results = []
|
296
|
+
async for value in tool_fn(**kwargs):
|
297
|
+
results.append(value)
|
298
|
+
if collect_results:
|
299
|
+
return results
|
300
|
+
else:
|
301
|
+
return results[-1]
|
302
|
+
|
303
|
+
loop = asyncio.new_event_loop()
|
304
|
+
asyncio.set_event_loop(loop)
|
305
|
+
return loop.run_until_complete(collect())
|
306
|
+
elif inspect.isgeneratorfunction(tool_fn):
|
307
|
+
# Handle sync generator
|
308
|
+
results = []
|
309
|
+
for value in tool_fn(**kwargs):
|
310
|
+
results.append(value)
|
311
|
+
if collect_results:
|
312
|
+
return results
|
313
|
+
else:
|
314
|
+
return results[-1]
|
315
|
+
elif inspect.iscoroutinefunction(tool_fn):
|
316
|
+
# Handle async
|
317
|
+
loop = asyncio.new_event_loop()
|
318
|
+
asyncio.set_event_loop(loop)
|
319
|
+
return loop.run_until_complete(tool_fn(**kwargs))
|
320
|
+
else:
|
321
|
+
# Handle sync
|
322
|
+
return tool_fn(**kwargs)
|
323
|
+
|
324
|
+
async def aexecute(
|
325
|
+
self, toolname: str, kwargs: dict | None = None, collect_results=False
|
326
|
+
):
|
327
|
+
tool_fn = self._get_tool(toolname)
|
328
|
+
kwargs = kwargs or {}
|
329
|
+
if inspect.isasyncgenfunction(tool_fn):
|
330
|
+
# Handle async generator
|
331
|
+
results = []
|
332
|
+
async for value in tool_fn(**kwargs):
|
333
|
+
results.append(value)
|
334
|
+
if collect_results:
|
335
|
+
return results
|
336
|
+
else:
|
337
|
+
return results[-1]
|
338
|
+
elif inspect.isgeneratorfunction(tool_fn):
|
339
|
+
# Handle sync generator
|
340
|
+
results = []
|
341
|
+
for value in tool_fn(**kwargs):
|
342
|
+
results.append(value)
|
343
|
+
if collect_results:
|
344
|
+
return results
|
345
|
+
else:
|
346
|
+
return results[-1]
|
347
|
+
elif inspect.iscoroutinefunction(tool_fn):
|
348
|
+
# Handle async
|
349
|
+
return await tool_fn(**kwargs)
|
350
|
+
else:
|
351
|
+
# Handle sync
|
352
|
+
return tool_fn(**kwargs)
|
353
|
+
|
354
|
+
def stream_execute(self, toolname: str, kwargs: dict | None = None):
|
355
|
+
tool_fn = self._get_tool(toolname)
|
356
|
+
kwargs = kwargs or {}
|
357
|
+
if inspect.isasyncgenfunction(tool_fn):
|
358
|
+
# Handle async generator
|
359
|
+
|
360
|
+
async def collect():
|
361
|
+
async for value in tool_fn(**kwargs):
|
362
|
+
yield value
|
363
|
+
|
364
|
+
loop = asyncio.new_event_loop()
|
365
|
+
asyncio.set_event_loop(loop)
|
366
|
+
agen = collect()
|
367
|
+
try:
|
368
|
+
while True:
|
369
|
+
yield loop.run_until_complete(agen.__anext__())
|
370
|
+
except StopAsyncIteration:
|
371
|
+
pass
|
372
|
+
finally:
|
373
|
+
loop.close()
|
374
|
+
elif inspect.isgeneratorfunction(tool_fn):
|
375
|
+
# Handle sync generator
|
376
|
+
for value in tool_fn(**kwargs):
|
377
|
+
yield value
|
378
|
+
elif inspect.iscoroutinefunction(tool_fn):
|
379
|
+
# Handle async
|
380
|
+
loop = asyncio.new_event_loop()
|
381
|
+
asyncio.set_event_loop(loop)
|
382
|
+
yield loop.run_until_complete(tool_fn(**kwargs))
|
383
|
+
else:
|
384
|
+
# Handle sync
|
385
|
+
yield tool_fn(**kwargs)
|
386
|
+
|
387
|
+
async def astream_execute(self, toolname: str, kwargs: dict | None = None):
|
388
|
+
tool_fn = self._get_tool(toolname)
|
389
|
+
kwargs = kwargs or {}
|
390
|
+
if inspect.isasyncgenfunction(tool_fn):
|
391
|
+
# Handle async generator
|
392
|
+
async for value in tool_fn(**kwargs):
|
393
|
+
yield value
|
394
|
+
elif inspect.isgeneratorfunction(tool_fn):
|
395
|
+
# Handle sync generator
|
396
|
+
for value in tool_fn(**kwargs):
|
397
|
+
yield value
|
398
|
+
elif inspect.iscoroutinefunction(tool_fn):
|
399
|
+
# Handle async
|
400
|
+
yield await tool_fn(**kwargs)
|
278
401
|
else:
|
279
|
-
|
402
|
+
# Handle sync
|
403
|
+
yield tool_fn(**kwargs)
|
280
404
|
|
281
405
|
def parse_and_execute(self, msg: str):
|
282
406
|
toolcall = llm_parse_json(msg, keys=["toolname", "kwargs"])
|
@@ -284,9 +408,18 @@ class BaseIndex:
|
|
284
408
|
|
285
409
|
async def async_parse_and_execute(self, msg: str):
|
286
410
|
toolcall = llm_parse_json(msg, keys=["toolname", "kwargs"])
|
287
|
-
return await self.
|
411
|
+
return await self.aexecute(toolcall.get("toolname"), toolcall.get("kwargs"))
|
412
|
+
|
413
|
+
def stream_parse_and_execute(self, msg: str):
|
414
|
+
toolcall = llm_parse_json(msg, keys=["toolname", "kwargs"])
|
415
|
+
return self.stream_execute(toolcall.get("toolname"), toolcall.get("kwargs"))
|
416
|
+
|
417
|
+
async def astream_parse_and_execute(self, msg: str):
|
418
|
+
toolcall = llm_parse_json(msg, keys=["toolname", "kwargs"])
|
419
|
+
async for value in self.astream_parse_and_execute(
|
288
420
|
toolcall.get("toolname"), toolcall.get("kwargs")
|
289
|
-
)
|
421
|
+
):
|
422
|
+
yield value
|
290
423
|
|
291
424
|
def format_tools(self, provider: ProviderFormat):
|
292
425
|
return format_tools(self.tools, provider)
|
stores/indexes/index.py
CHANGED
@@ -17,11 +17,15 @@ class Index(BaseIndex):
|
|
17
17
|
self,
|
18
18
|
tools: list[Callable, os.PathLike] | None = None,
|
19
19
|
env_var: dict[str, dict] | None = None,
|
20
|
+
include: dict[str, list[str]] | None = None,
|
21
|
+
exclude: dict[str, list[str]] | None = None,
|
20
22
|
cache_dir: Optional[os.PathLike] = None,
|
21
23
|
reset_cache=False,
|
22
24
|
sys_executable: str | None = None,
|
23
25
|
):
|
24
26
|
self.env_var = env_var or {}
|
27
|
+
include = include or {}
|
28
|
+
exclude = exclude or {}
|
25
29
|
tools = tools or []
|
26
30
|
|
27
31
|
_tools = []
|
@@ -32,7 +36,11 @@ class Index(BaseIndex):
|
|
32
36
|
if Path(index_name).exists():
|
33
37
|
# Load LocalIndex
|
34
38
|
try:
|
35
|
-
loaded_index = LocalIndex(
|
39
|
+
loaded_index = LocalIndex(
|
40
|
+
index_name,
|
41
|
+
include=include.get(index_name),
|
42
|
+
exclude=exclude.get(index_name),
|
43
|
+
)
|
36
44
|
except Exception:
|
37
45
|
logger.warning(
|
38
46
|
f'Unable to load index "{index_name}"', exc_info=True
|
@@ -43,6 +51,8 @@ class Index(BaseIndex):
|
|
43
51
|
loaded_index = RemoteIndex(
|
44
52
|
index_name,
|
45
53
|
env_var=self.env_var.get(index_name),
|
54
|
+
include=include.get(index_name),
|
55
|
+
exclude=exclude.get(index_name),
|
46
56
|
cache_dir=cache_dir,
|
47
57
|
reset_cache=reset_cache,
|
48
58
|
sys_executable=sys_executable,
|
stores/indexes/local_index.py
CHANGED
@@ -26,10 +26,14 @@ class LocalIndex(BaseIndex):
|
|
26
26
|
index_folder: os.PathLike,
|
27
27
|
create_venv: bool = False,
|
28
28
|
env_var: dict | None = None,
|
29
|
+
include: list[str] | None = None,
|
30
|
+
exclude: list[str] | None = None,
|
29
31
|
sys_executable: str | None = None,
|
30
32
|
):
|
31
33
|
self.index_folder = Path(index_folder)
|
32
34
|
self.env_var = env_var or {}
|
35
|
+
include = include or []
|
36
|
+
exclude = exclude or []
|
33
37
|
|
34
38
|
if not self.index_folder.exists():
|
35
39
|
raise ValueError(
|
@@ -50,21 +54,31 @@ class LocalIndex(BaseIndex):
|
|
50
54
|
)
|
51
55
|
install_venv_deps(self.index_folder)
|
52
56
|
# Initialize tools
|
53
|
-
tools = init_venv_tools(
|
57
|
+
tools = init_venv_tools(
|
58
|
+
self.index_folder,
|
59
|
+
env_var=self.env_var,
|
60
|
+
include=include,
|
61
|
+
exclude=exclude,
|
62
|
+
)
|
54
63
|
else:
|
55
64
|
if self.env_var:
|
56
65
|
raise ValueError(
|
57
66
|
"Environment variables will only be restricted if create_venv=True when initializing LocalIndex"
|
58
67
|
)
|
59
|
-
tools = self._init_tools()
|
68
|
+
tools = self._init_tools(include=include, exclude=exclude)
|
60
69
|
super().__init__(tools)
|
61
70
|
|
62
|
-
def _init_tools(
|
71
|
+
def _init_tools(
|
72
|
+
self, include: list[str] | None = None, exclude: list[str] | None = None
|
73
|
+
):
|
63
74
|
"""
|
64
75
|
Load local tools.toml file and import tool functions
|
65
76
|
|
66
77
|
NOTE: Can we just add index_folder to sys.path and import the functions?
|
67
78
|
"""
|
79
|
+
include = include or []
|
80
|
+
exclude = exclude or []
|
81
|
+
|
68
82
|
index_manifest = self.index_folder / TOOLS_CONFIG_FILENAME
|
69
83
|
if not index_manifest.exists():
|
70
84
|
raise ValueError(f"Unable to load index - {index_manifest} does not exist")
|
@@ -73,7 +87,9 @@ class LocalIndex(BaseIndex):
|
|
73
87
|
manifest = tomllib.load(file)["index"]
|
74
88
|
|
75
89
|
tools = []
|
76
|
-
for tool_id in manifest.get("tools", []):
|
90
|
+
for tool_id in include or manifest.get("tools", []):
|
91
|
+
if tool_id in exclude:
|
92
|
+
continue
|
77
93
|
module_name = ".".join(tool_id.split(".")[:-1])
|
78
94
|
tool_name = tool_id.split(".")[-1]
|
79
95
|
|
stores/indexes/remote_index.py
CHANGED
@@ -51,6 +51,8 @@ class RemoteIndex(BaseIndex):
|
|
51
51
|
self,
|
52
52
|
index_id: str,
|
53
53
|
env_var: dict | None = None,
|
54
|
+
include: list[str] | None = None,
|
55
|
+
exclude: list[str] | None = None,
|
54
56
|
cache_dir: Optional[PathLike] = None,
|
55
57
|
reset_cache=False,
|
56
58
|
sys_executable: str | None = None,
|
@@ -64,6 +66,8 @@ class RemoteIndex(BaseIndex):
|
|
64
66
|
shutil.rmtree(cache_dir)
|
65
67
|
self.index_folder = cache_dir / self.index_id
|
66
68
|
self.env_var = env_var or {}
|
69
|
+
include = include or []
|
70
|
+
exclude = exclude or []
|
67
71
|
if not self.index_folder.exists():
|
68
72
|
logger.info(f"Installing {index_id}...")
|
69
73
|
commit_like = None
|
@@ -99,5 +103,7 @@ class RemoteIndex(BaseIndex):
|
|
99
103
|
venv.create(self.venv, symlinks=True, with_pip=True, upgrade_deps=True)
|
100
104
|
install_venv_deps(self.index_folder)
|
101
105
|
# Initialize tools
|
102
|
-
tools = init_venv_tools(
|
106
|
+
tools = init_venv_tools(
|
107
|
+
self.index_folder, env_var=self.env_var, include=include, exclude=exclude
|
108
|
+
)
|
103
109
|
super().__init__(tools)
|
stores/indexes/venv_utils.py
CHANGED
@@ -1,3 +1,4 @@
|
|
1
|
+
import asyncio
|
1
2
|
import hashlib
|
2
3
|
import inspect
|
3
4
|
import json
|
@@ -104,16 +105,25 @@ def install_venv_deps(index_folder: os.PathLike):
|
|
104
105
|
return message
|
105
106
|
|
106
107
|
|
107
|
-
def init_venv_tools(
|
108
|
+
def init_venv_tools(
|
109
|
+
index_folder: os.PathLike,
|
110
|
+
env_var: dict | None = None,
|
111
|
+
include: list[str] | None = None,
|
112
|
+
exclude: list[str] | None = None,
|
113
|
+
):
|
108
114
|
index_folder = Path(index_folder)
|
109
115
|
env_var = env_var or {}
|
116
|
+
include = include or []
|
117
|
+
exclude = exclude or []
|
110
118
|
|
111
119
|
index_manifest = index_folder / TOOLS_CONFIG_FILENAME
|
112
120
|
with open(index_manifest, "rb") as file:
|
113
121
|
manifest = tomllib.load(file)["index"]
|
114
122
|
|
115
123
|
tools = []
|
116
|
-
for tool_id in manifest.get("tools", []):
|
124
|
+
for tool_id in include or manifest.get("tools", []):
|
125
|
+
if tool_id in exclude:
|
126
|
+
continue
|
117
127
|
tool_sig = get_tool_signature(
|
118
128
|
tool_id=tool_id,
|
119
129
|
index_folder=index_folder,
|
@@ -215,7 +225,9 @@ try:
|
|
215
225
|
"tool_id": "{tool_id}",
|
216
226
|
"params": params,
|
217
227
|
"return": return_info,
|
218
|
-
"
|
228
|
+
"iscoroutinefunction": inspect.iscoroutinefunction({tool_name}),
|
229
|
+
"isgeneratorfunction": inspect.isgeneratorfunction({tool_name}),
|
230
|
+
"isasyncgenfunction": inspect.isasyncgenfunction({tool_name}),
|
219
231
|
"doc": inspect.getdoc({tool_name}),
|
220
232
|
}},
|
221
233
|
}},
|
@@ -293,25 +305,56 @@ def parse_tool_signature(
|
|
293
305
|
"""
|
294
306
|
env_var = env_var or {}
|
295
307
|
|
296
|
-
|
297
|
-
|
298
|
-
|
299
|
-
|
300
|
-
|
301
|
-
|
302
|
-
|
303
|
-
|
304
|
-
|
308
|
+
if signature_dict.get("isasyncgenfunction"):
|
309
|
+
|
310
|
+
async def func_handler(*args, **kwargs):
|
311
|
+
# TODO: Make this truly async
|
312
|
+
async for value in run_remote_tool(
|
313
|
+
tool_id=signature_dict["tool_id"],
|
314
|
+
index_folder=index_folder,
|
315
|
+
args=args,
|
316
|
+
kwargs=kwargs,
|
317
|
+
venv=venv,
|
318
|
+
env_var=env_var,
|
319
|
+
stream=True,
|
320
|
+
):
|
321
|
+
yield value
|
322
|
+
elif signature_dict.get("isgeneratorfunction"):
|
323
|
+
|
324
|
+
def func_handler(*args, **kwargs):
|
325
|
+
for value in run_remote_tool(
|
326
|
+
tool_id=signature_dict["tool_id"],
|
327
|
+
index_folder=index_folder,
|
328
|
+
args=args,
|
329
|
+
kwargs=kwargs,
|
330
|
+
venv=venv,
|
331
|
+
env_var=env_var,
|
332
|
+
stream=True,
|
333
|
+
):
|
334
|
+
yield value
|
335
|
+
elif signature_dict.get("iscoroutinefunction"):
|
336
|
+
|
337
|
+
async def func_handler(*args, **kwargs):
|
338
|
+
# TODO: Make this truly async
|
339
|
+
return run_remote_tool(
|
340
|
+
tool_id=signature_dict["tool_id"],
|
341
|
+
index_folder=index_folder,
|
342
|
+
args=args,
|
343
|
+
kwargs=kwargs,
|
344
|
+
venv=venv,
|
345
|
+
env_var=env_var,
|
346
|
+
)
|
347
|
+
else:
|
305
348
|
|
306
|
-
|
307
|
-
|
308
|
-
|
309
|
-
|
310
|
-
|
311
|
-
|
312
|
-
|
313
|
-
|
314
|
-
|
349
|
+
def func_handler(*args, **kwargs):
|
350
|
+
return run_remote_tool(
|
351
|
+
tool_id=signature_dict["tool_id"],
|
352
|
+
index_folder=index_folder,
|
353
|
+
args=args,
|
354
|
+
kwargs=kwargs,
|
355
|
+
venv=venv,
|
356
|
+
env_var=env_var,
|
357
|
+
)
|
315
358
|
|
316
359
|
# Reconstruct signature from list of args
|
317
360
|
params = []
|
@@ -329,7 +372,7 @@ def parse_tool_signature(
|
|
329
372
|
signature = inspect.Signature(params, return_annotation=return_type)
|
330
373
|
func = create_function(
|
331
374
|
signature,
|
332
|
-
|
375
|
+
func_handler,
|
333
376
|
qualname=signature_dict["tool_id"],
|
334
377
|
doc=signature_dict.get("doc"),
|
335
378
|
)
|
@@ -345,6 +388,7 @@ def run_remote_tool(
|
|
345
388
|
kwargs: dict | None = None,
|
346
389
|
venv: str = VENV_NAME,
|
347
390
|
env_var: dict | None = None,
|
391
|
+
stream: bool = False,
|
348
392
|
):
|
349
393
|
args = args or []
|
350
394
|
kwargs = kwargs or {}
|
@@ -365,56 +409,131 @@ def run_remote_tool(
|
|
365
409
|
listener.listen(1)
|
366
410
|
_, port = listener.getsockname()
|
367
411
|
|
368
|
-
|
412
|
+
result_data = {}
|
413
|
+
|
414
|
+
def handle_connection_sync():
|
369
415
|
conn, _ = listener.accept()
|
370
416
|
with conn:
|
371
|
-
|
417
|
+
buffer = ""
|
372
418
|
while True:
|
373
|
-
chunk = conn.recv(4096)
|
419
|
+
chunk = conn.recv(4096).decode("utf-8")
|
374
420
|
if not chunk:
|
375
421
|
break
|
376
|
-
|
377
|
-
|
378
|
-
|
379
|
-
|
380
|
-
|
381
|
-
|
382
|
-
|
383
|
-
|
384
|
-
|
422
|
+
buffer += chunk
|
423
|
+
while "\n" in buffer:
|
424
|
+
line, buffer = buffer.split("\n", 1)
|
425
|
+
if not line.strip():
|
426
|
+
continue
|
427
|
+
msg = json.loads(line)
|
428
|
+
if msg.get("ok") and "stream" in msg:
|
429
|
+
result_data.setdefault("stream", []).append(msg["stream"])
|
430
|
+
elif msg.get("ok") and "result" in msg:
|
431
|
+
result_data["result"] = msg["result"]
|
432
|
+
elif "error" in msg:
|
433
|
+
result_data["error"] = msg["error"]
|
434
|
+
elif msg.get("done"):
|
435
|
+
return
|
436
|
+
|
437
|
+
async def handle_connection_async():
|
438
|
+
loop = asyncio.get_running_loop()
|
439
|
+
conn, _ = await loop.sock_accept(listener)
|
440
|
+
conn.setblocking(False)
|
441
|
+
buffer = ""
|
442
|
+
try:
|
443
|
+
while True:
|
444
|
+
chunk = await loop.sock_recv(conn, 4096)
|
445
|
+
if not chunk:
|
446
|
+
break
|
447
|
+
buffer += chunk.decode("utf-8")
|
448
|
+
while "\n" in buffer:
|
449
|
+
line, buffer = buffer.split("\n", 1)
|
450
|
+
if not line.strip():
|
451
|
+
continue
|
452
|
+
msg = json.loads(line)
|
453
|
+
if msg.get("ok") and "stream" in msg:
|
454
|
+
yield msg["stream"]
|
455
|
+
elif msg.get("ok") and "result" in msg:
|
456
|
+
yield msg["result"]
|
457
|
+
elif "error" in msg:
|
458
|
+
raise RuntimeError(f"Subprocess error:\n{msg['error']}")
|
459
|
+
elif msg.get("done"):
|
460
|
+
return
|
461
|
+
finally:
|
462
|
+
conn.close()
|
463
|
+
|
464
|
+
if not stream:
|
465
|
+
thread = threading.Thread(target=lambda: list(handle_connection_sync()))
|
466
|
+
thread.start()
|
385
467
|
|
386
468
|
runner = f"""
|
387
469
|
import asyncio, inspect, json, socket, sys, traceback
|
388
470
|
sys.path.insert(0, "{index_folder}")
|
471
|
+
|
472
|
+
def send(sock, payload):
|
473
|
+
sock.sendall((json.dumps(payload) + "\\n").encode("utf-8"))
|
474
|
+
|
475
|
+
sock = socket.create_connection(("localhost", {port}))
|
476
|
+
|
389
477
|
try:
|
390
478
|
from {module_name} import {tool_name}
|
391
479
|
params = json.load(sys.stdin)
|
392
480
|
args = params.get("args", [])
|
393
481
|
kwargs = params.get("kwargs", {{}})
|
394
|
-
|
482
|
+
|
483
|
+
func = {tool_name}
|
484
|
+
|
485
|
+
if inspect.isasyncgenfunction(func):
|
486
|
+
async def run():
|
487
|
+
async for value in func(*args, **kwargs):
|
488
|
+
send(sock, {{"ok": True, "stream": value}})
|
489
|
+
loop = asyncio.new_event_loop()
|
490
|
+
asyncio.set_event_loop(loop)
|
491
|
+
loop.run_until_complete(run())
|
492
|
+
elif inspect.isgeneratorfunction(func):
|
493
|
+
for value in func(*args, **kwargs):
|
494
|
+
send(sock, {{"ok": True, "stream": value}})
|
495
|
+
elif inspect.iscoroutinefunction(func):
|
395
496
|
loop = asyncio.new_event_loop()
|
396
497
|
asyncio.set_event_loop(loop)
|
397
|
-
result = loop.run_until_complete(
|
498
|
+
result = loop.run_until_complete(func(*args, **kwargs))
|
499
|
+
send(sock, {{"ok": True, "result": result}})
|
398
500
|
else:
|
399
|
-
result =
|
400
|
-
|
501
|
+
result = func(*args, **kwargs)
|
502
|
+
send(sock, {{"ok": True, "result": result}})
|
503
|
+
send(sock, {{"done": True}})
|
401
504
|
except Exception as e:
|
402
505
|
err = traceback.format_exc()
|
403
|
-
|
404
|
-
sock
|
405
|
-
|
406
|
-
|
506
|
+
try:
|
507
|
+
send(sock, {{"ok": False, "error": err}})
|
508
|
+
except:
|
509
|
+
pass
|
510
|
+
finally:
|
511
|
+
try:
|
512
|
+
sock.close()
|
513
|
+
except:
|
514
|
+
pass
|
407
515
|
"""
|
408
|
-
|
516
|
+
|
517
|
+
proc = subprocess.Popen(
|
409
518
|
[get_python_command(Path(index_folder) / venv), "-c", runner],
|
410
|
-
|
411
|
-
|
519
|
+
stdin=subprocess.PIPE,
|
520
|
+
stdout=subprocess.DEVNULL,
|
521
|
+
stderr=subprocess.PIPE,
|
412
522
|
env=env_var or None,
|
413
523
|
)
|
414
|
-
|
415
|
-
|
416
|
-
|
417
|
-
if
|
418
|
-
|
524
|
+
proc.stdin.write(payload)
|
525
|
+
proc.stdin.close()
|
526
|
+
|
527
|
+
if not stream:
|
528
|
+
thread.join()
|
529
|
+
|
530
|
+
if "error" in result_data:
|
531
|
+
raise RuntimeError(f"Subprocess failed with error:\n{result_data['error']}")
|
532
|
+
elif "result" in result_data:
|
533
|
+
return result_data["result"]
|
534
|
+
elif "stream" in result_data:
|
535
|
+
return result_data["stream"]
|
536
|
+
else:
|
537
|
+
raise RuntimeError("Subprocess completed without returning data.")
|
419
538
|
else:
|
420
|
-
|
539
|
+
return handle_connection_async()
|
@@ -0,0 +1,15 @@
|
|
1
|
+
stores/__init__.py,sha256=KYpKkNrMLx6ssVUbxHnn9wFBq5F5KnaFchcimIfDf9g,186
|
2
|
+
stores/constants.py,sha256=7WqFmoGCtmUKHA5WHxOJvvK7g-yYu_KGoqnuVFADNao,57
|
3
|
+
stores/format.py,sha256=LduYBVDiUDB1J1HDyu9jHrRG1V97pw6C5g76OirpJJY,7792
|
4
|
+
stores/parse.py,sha256=HYPNPzQod2vpu1Cln7yQ8aVkZT1Mw2IN0sZ2A1DIaqE,4967
|
5
|
+
stores/utils.py,sha256=GPWT6lCoGobwP3PlEOHyJfKyd0dobamjyErcR7lgm7M,242
|
6
|
+
stores/indexes/__init__.py,sha256=s-RNqml8uGREQhxwSdDoxcbcxeD8soB9BcL5dBKsQfI,215
|
7
|
+
stores/indexes/base_index.py,sha256=YrEwETZ5eXj3rXK5qxOllRXqFifQoteYdzPAasbvEyg,15536
|
8
|
+
stores/indexes/index.py,sha256=Cub5mtnYGipHfPR8BexJYRSKfuJmcGPp0B3ou2bGNqs,2901
|
9
|
+
stores/indexes/local_index.py,sha256=Gg9LkEo1L0_NZZYPItsF_-1y6nFP369C3QlPvPyvPx8,3708
|
10
|
+
stores/indexes/remote_index.py,sha256=-GV4l2c7GL6_bcVOQnSK5rzYru237bwC5L6eiO-QFzM,3438
|
11
|
+
stores/indexes/venv_utils.py,sha256=tS2ZYt2Cf7rqHxBXwn5EoEdUFjyaAIXVRiy8iXjCMb4,17397
|
12
|
+
stores-0.1.7.dev2.dist-info/METADATA,sha256=AciMTIC_8qRvud2y3guHn78X-F_n5j-rlCOUNeFLBAc,3081
|
13
|
+
stores-0.1.7.dev2.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
|
14
|
+
stores-0.1.7.dev2.dist-info/licenses/LICENSE,sha256=VTidYE7_Dam0Dwyq095EhhDIqi47g03oVpLAHQgKws0,1066
|
15
|
+
stores-0.1.7.dev2.dist-info/RECORD,,
|
stores-0.1.6.dist-info/RECORD
DELETED
@@ -1,15 +0,0 @@
|
|
1
|
-
stores/__init__.py,sha256=KYpKkNrMLx6ssVUbxHnn9wFBq5F5KnaFchcimIfDf9g,186
|
2
|
-
stores/constants.py,sha256=7WqFmoGCtmUKHA5WHxOJvvK7g-yYu_KGoqnuVFADNao,57
|
3
|
-
stores/format.py,sha256=LduYBVDiUDB1J1HDyu9jHrRG1V97pw6C5g76OirpJJY,7792
|
4
|
-
stores/parse.py,sha256=HYPNPzQod2vpu1Cln7yQ8aVkZT1Mw2IN0sZ2A1DIaqE,4967
|
5
|
-
stores/utils.py,sha256=GPWT6lCoGobwP3PlEOHyJfKyd0dobamjyErcR7lgm7M,242
|
6
|
-
stores/indexes/__init__.py,sha256=s-RNqml8uGREQhxwSdDoxcbcxeD8soB9BcL5dBKsQfI,215
|
7
|
-
stores/indexes/base_index.py,sha256=U12ZuvyuHSZJGdVdTMakdz6yHJQ1No0wOY45W5FoN0s,11073
|
8
|
-
stores/indexes/index.py,sha256=C3i5JAwYoKEK7jIJRAgnLrQxjeewtInKcMinep6yvyA,2432
|
9
|
-
stores/indexes/local_index.py,sha256=TVk7W8BXF5KEqjEK0LYqq3CtLMfHaF-Tk_5cZ6bjsBU,3189
|
10
|
-
stores/indexes/remote_index.py,sha256=CHa9kQVDQlgI8qw3XQXchrE1cjXJk9fNY2j2HWL2gio,3226
|
11
|
-
stores/indexes/venv_utils.py,sha256=2D0VasFfjwzEu_-LEyXREDN-SkWHGzx5TImvMpuKbjU,13313
|
12
|
-
stores-0.1.6.dist-info/METADATA,sha256=u3ytaMMFLUQCgij7JRgstw-RRyAqlrsp4wXF2-55i8g,3076
|
13
|
-
stores-0.1.6.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
|
14
|
-
stores-0.1.6.dist-info/licenses/LICENSE,sha256=VTidYE7_Dam0Dwyq095EhhDIqi47g03oVpLAHQgKws0,1066
|
15
|
-
stores-0.1.6.dist-info/RECORD,,
|
File without changes
|
File without changes
|