Support ollama-style response from LLM, and fix bug cuased by lint fix

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)
 
 

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 `<price> <outflow_account> [<inflow_account>] <payee> [<description>] [#<tag1> [#<tag2>] ...]`, 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")

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")

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')

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