diff --git a/bean_utils/bean.py b/bean_utils/bean.py index e7795fa..aaabbbe 100644 --- a/bean_utils/bean.py +++ b/bean_utils/bean.py @@ -1,5 +1,5 @@ import contextlib -import shutil +import shlex from pathlib import Path import re from datetime import datetime @@ -41,14 +41,14 @@ def _load(self): # Fill mtime for f in self._options["include"]: - self.mtimes[f] = Path.stat(f).st_mtime + self.mtimes[f] = Path(f).stat().st_mtime def _auto_reload(self, accounts_only=False): # Check and reload for f, mtime in self.mtimes.items(): if accounts_only and ("accounts" not in f): continue - if mtime != Path.stat(f).st_mtime: + if mtime != Path(f).stat().st_mtime: self._load() return @@ -191,7 +191,7 @@ def commit_trx(self, data): fname = self.fname with open(fname, 'a') as f: f.write("\n" + data + "\n") - subprocess.run(["bean-format", "-o", shutil.quote(fname), shutil.quote(fname)], # noqa: S607,S603 + subprocess.run(["bean-format", "-o", shlex.quote(fname), shlex.quote(fname)], # noqa: S607,S603 shell=False) diff --git a/bean_utils/rag.py b/bean_utils/rag.py index 5576ea3..6344bfb 100644 --- a/bean_utils/rag.py +++ b/bean_utils/rag.py @@ -4,6 +4,7 @@ import conf +_TIMEOUT = 30 # flake8: noqa _PROMPT_TEMPLATE = """The user is using Beancount for bookkeeping. For simplicity, there is currently a set of accounting grammar that is converted by a program into complete transaction records. The format of the grammar is ` [] [] [# [#] ...]`, where the inflow and outflow accounts are subject to fuzzy matching. @@ -42,6 +43,7 @@ def complete_rag(args, date, accounts): prompt = _PROMPT_TEMPLATE.format(date=date, reference_records=reference_records, accounts=accounts) payload = { "model": rag_config.model, + "stream": False, "messages": [ { "role": "system", @@ -59,5 +61,10 @@ def complete_rag(args, date, accounts): } response = requests.post(rag_config.api_url, json=payload, headers=headers, timeout=_TIMEOUT) data = response.json() - content = data["choices"][0]["message"]["content"] + if "choices" in data: + # ChatGPT-format response + content = data["choices"][0]["message"]["content"] + else: + # Ollama-format response + content = data["message"]["content"] return content.strip("`\n") diff --git a/bots/telegram_bot.py b/bots/telegram_bot.py index e37daa2..d9078b0 100644 --- a/bots/telegram_bot.py +++ b/bots/telegram_bot.py @@ -130,7 +130,7 @@ async def callback(update, context): if choice == "提交": result_msg = "已提交 ✅" bean_manager.commit_trx(trx) - logging.info("Commit transaction\n", trx) + logging.info("Commit transaction: %s\n", trx) else: result_msg = "已取消 ❌" logging.info("Cancel transaction") diff --git a/main.py b/main.py index ead5f6d..aad2de3 100644 --- a/main.py +++ b/main.py @@ -1,8 +1,13 @@ import argparse import conf +import logging if __name__ == "__main__": + # Init logging + logging.basicConfig(level=logging.INFO) + # logging.getLogger().addHandler(logging.StreamHandler()) + parser = argparse.ArgumentParser(prog='beanbot', description='Bot to translate text into beancount transaction') subparser = parser.add_subparsers(title='sub command', dest='command') diff --git a/vec_db/sqlite_vec_db.py b/vec_db/sqlite_vec_db.py index 5746503..a32c06f 100644 --- a/vec_db/sqlite_vec_db.py +++ b/vec_db/sqlite_vec_db.py @@ -84,5 +84,4 @@ def query_by_embedding(embedding, sentence, candidate_amount): candidates.append(tx_row) candidates.sort(key=itemgetter("score"), reverse=True) - # print(candidates) return candidates