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.
@@ -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.iscoroutinefunction(tool):
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
- # Inject default values within wrapper
196
- bound_args = original_signature.bind(*args, **kwargs)
197
- bound_args.apply_defaults()
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
- # Inject default values within wrapper
211
- bound_args = original_signature.bind(*args, **kwargs)
212
- bound_args.apply_defaults()
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 execute(self, toolname: str, kwargs: dict | None = None):
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
- tool = self.tools_dict[toolname]
276
- if inspect.iscoroutinefunction(tool):
277
- return await tool(**kwargs)
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
- return tool(**kwargs)
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.async_execute(
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(index_name)
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,
@@ -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(self.index_folder, self.env_var)
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(self):
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
 
@@ -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(self.index_folder, self.env_var)
106
+ tools = init_venv_tools(
107
+ self.index_folder, env_var=self.env_var, include=include, exclude=exclude
108
+ )
103
109
  super().__init__(tools)
@@ -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(index_folder: os.PathLike, env_var: dict | None = None):
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
- "is_async": inspect.iscoroutinefunction({tool_name}),
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
- def func_handler(*args, **kwargs):
297
- return run_remote_tool(
298
- tool_id=signature_dict["tool_id"],
299
- index_folder=index_folder,
300
- args=args,
301
- kwargs=kwargs,
302
- venv=venv,
303
- env_var=env_var,
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
- async def async_func_handler(*args, **kwargs):
307
- return run_remote_tool(
308
- tool_id=signature_dict["tool_id"],
309
- index_folder=index_folder,
310
- args=args,
311
- kwargs=kwargs,
312
- venv=venv,
313
- env_var=env_var,
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
- async_func_handler if signature_dict.get("is_async") else func_handler,
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
- def handle_connection():
412
+ result_data = {}
413
+
414
+ def handle_connection_sync():
369
415
  conn, _ = listener.accept()
370
416
  with conn:
371
- data = b""
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
- data += chunk
377
- listener.close()
378
- return data
379
-
380
- result_data = {}
381
- t = threading.Thread(
382
- target=lambda: result_data.setdefault("data", handle_connection())
383
- )
384
- t.start()
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
- if inspect.iscoroutinefunction({tool_name}):
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({tool_name}(*args, **kwargs))
498
+ result = loop.run_until_complete(func(*args, **kwargs))
499
+ send(sock, {{"ok": True, "result": result}})
398
500
  else:
399
- result = {tool_name}(*args, **kwargs)
400
- response = json.dumps({{"ok": True, "result": result}})
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
- response = json.dumps({{"ok": False, "error": err}})
404
- sock = socket.create_connection(("localhost", {port}))
405
- sock.sendall(response.encode("utf-8"))
406
- sock.close()
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
- result = subprocess.run(
516
+
517
+ proc = subprocess.Popen(
409
518
  [get_python_command(Path(index_folder) / venv), "-c", runner],
410
- input=payload,
411
- capture_output=True,
519
+ stdin=subprocess.PIPE,
520
+ stdout=subprocess.DEVNULL,
521
+ stderr=subprocess.PIPE,
412
522
  env=env_var or None,
413
523
  )
414
-
415
- t.join()
416
- response = json.loads(result_data["data"].decode("utf-8"))
417
- if response.get("ok"):
418
- return response["result"]
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
- raise RuntimeError(f"Subprocess failed with error:\n{response['error']}")
539
+ return handle_connection_async()
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: stores
3
- Version: 0.1.6
3
+ Version: 0.1.7.dev2
4
4
  Summary: Repository of Python functions and tools for LLMs
5
5
  License-File: LICENSE
6
6
  Requires-Python: >=3.10
@@ -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,,
@@ -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,,