allow lmql.F with positionals

src/lmql/algorithms/cache.py

@@ -3,7 +3,7 @@
 import inspect
 
 from lmql.runtime.lmql_runtime import LMQLQueryFunction
-from lmql import LMQLResult
+from lmql import LMQLResult, F
 
 # cache query results by query code and arguments
 global cache_file
@@ -46,9 +46,13 @@ def persist_cache():
  with open(cache_file, "wb") as f:
  pickle.dump(cache, f)
 
-async def apply(q, *args):
+async def apply(q, *args, **kwargs):
  global cache
 
+ if type(q) is str:
+ where = kwargs.pop("where", None)
+ q = F(q, constraints=where, is_async=True)
+
  # handle non-LMQL queries
  if type(q) is not LMQLQueryFunction:
  if inspect.iscoroutinefunction(q):
@@ -74,7 +78,7 @@ async def apply(q, *args):
  return cache[key]
  else:
  try:
- result = await q(*args)
+ result = await q(*args, **kwargs)
  if len(result) == 1:
  result = result[0]
  if type(result) is LMQLResult:

src/lmql/algorithms/functools.py

@@ -16,7 +16,7 @@ async def map(q, items, chunksize=None, progress=False, **kwargs):
  chunks = tqdm.tqdm(chunks, file=sys.stdout)
 
  for chunk in chunks:
- results = await asyncio.gather(*[apply(q, x) for x in chunk])
+ results = await asyncio.gather(*[apply(q, x, **kwargs) for x in chunk])
  total_results += results
 
  return total_results

src/lmql/runtime/lmql_runtime.py

@@ -96,6 +96,23 @@ def output_keys(self) -> List[str]:
  def force_model(self, model):
  self.model = model
 
+ def try_bind_positional_to_kwargs(self, signature, *args, **query_kwargs):
+ """
+ Best-effort attempt to bind positional arguments to keyword arguments in order of self.args. 
+
+ Only enabled for lmql.F for now, may have unexpected effects depending on the order query
+ arguments as determined by the compiler.
+ """
+ # only bind if kwargs are empty and no signature is provided (lmql.F or lmql.run)
+ if len(signature.parameters) != 0 or len(self.args) != len(args):
+ return
+ kwargs = {**{k:v for k,v in zip(self.args, args)}, **query_kwargs}
+
+ return inspect.BoundArguments(
+ signature=inspect.Signature(parameters=[inspect.Parameter(name=k, kind=inspect.Parameter.POSITIONAL_OR_KEYWORD) for k in self.args]),
+ arguments=kwargs
+ )
+
  def make_kwargs(self, *args, **kwargs):
  """
  Binds args and kwargs to the function signature and returns a dict of all user-defined kwargs.
@@ -122,11 +139,19 @@ def make_kwargs(self, *args, **kwargs):
  try:
  signature = signature.bind(*args, **query_kwargs)
  except TypeError as e:
- if len(e.args) == 1 and e.args[0].startswith("missing "):
- e.args = (f"Call to @lmql.query function is " + e.args[0] + "." + f" Expecting {signature}, but got positional args {args} and {kwargs}.",)
- elif len(e.args) == 1:
- e.args = (e.args[0] + "." + f" Expecting {signature}, but got positional args {args} and {kwargs}.",)
- raise e
+ if "too many positional arguments" in str(e):
+ # this is different from Python behavior, but we allow it for lmql.F and lmql.run
+ pos_as_kw = self.try_bind_positional_to_kwargs(signature, *args, **kwargs)
+ if pos_as_kw is not None:
+ signature = pos_as_kw
+ else:
+ raise e
+ else: 
+ if len(e.args) == 1 and e.args[0].startswith("missing "):
+ e.args = (f"Call to @lmql.query function is " + e.args[0] + "." + f" Expecting {signature}, but got positional args {args} and {kwargs}.",)
+ elif len(e.args) == 1:
+ e.args = (e.args[0] + "." + f" Expecting {signature}, but got positional args {args} and {kwargs}.",)
+ raise e
 
  # special case, if signature is empty (no input variables provided)
  if len(signature.arguments) == 0:

src/lmql/tests/test_f.py

@@ -0,0 +1,53 @@
+import lmql
+import lmql.algorithms as la
+
+from lmql.tests.expr_test_utils import run_all_tests
+
+def test_f_functions():
+ a = lmql.F("Summarize this text: {text}: [SUMMARY]", "len(TOKENS(SUMMARY)) < 10", model=lmql.model("random", seed=123))
+ assert type(a(text="This is a test.")) is str
+
+def test_f_positional_args():
+ a = lmql.F("Summarize this text: {text}: [SUMMARY]", "len(TOKENS(SUMMARY)) < 10", model=lmql.model("random", seed=123))
+ assert type(a("This is a test.")) is str
+
+def test_f_multiple_positional_args():
+ a = lmql.F("Summarize this text: {text1}{text2}: [SUMMARY]", "len(TOKENS(SUMMARY)) < 10", model=lmql.model("random", seed=123))
+ b = lmql.F("Summarize this text: {text}: [SUMMARY]", "len(TOKENS(SUMMARY)) < 10", model=lmql.model("random", seed=123))
+ r1 = type(a("This is ", "a test."))
+ r2 = type(b("This is a test."))
+
+ assert r1 is str, "output should be a string"
+ assert r1 == r2, "using positional and keyword args should result in the same output"
+
+
+async def test_async_f_functions():
+ a = lmql.F("Summarize this text: {text}: [SUMMARY]", "len(TOKENS(SUMMARY)) < 10", model=lmql.model("random", seed=123), is_async=True)
+ assert type(await a(text="This is a test.")) is str
+
+async def test_map_async_f_functions():
+ # a = lmql.F("Summarize this text: {text}: [SUMMARY]", "len(TOKENS(SUMMARY)) < 10", model=lmql.model("random", seed=123), is_async=True)
+ data = [
+ "A dog walks into a bar.",
+ "A cat goes for a walk.",
+ ]
+ q = lmql.F("Summarize this text: {text}: [SUMMARY]", "len(TOKENS(SUMMARY)) < 10", model=lmql.model("random", seed=123), is_async=True)
+ r = await la.map(q, data)
+ assert len(r) == 2
+
+ r = await la.map("Summarize this text: {text}: [SUMMARY]", data, where="len(TOKENS(SUMMARY)) < 10", model=lmql.model("random", seed=123))
+ assert len(r) == 2
+
+def test_run_pos_args():
+ source = '''lmql
+ "Summarize this text: {t1}{t2}: [SUMMARY]" where len(TOKENS(SUMMARY)) < 4
+ return SUMMARY
+ '''
+
+ r1 = lmql.run_sync(source, "This is ", "a test.", model=lmql.model("random", seed=123))
+ r2 = lmql.run_sync(source, t1="This is ", t2="a test.", model=lmql.model("random", seed=123))
+
+ assert r1 == r2, "using positional and keyword args should result in the same output"
+
+if __name__ == "__main__":
+ run_all_tests(globals())