Small improvements to tool calling and json handling

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "qabot"
-version = "0.7.1"
+version = "0.7.2"
 description = "Query local or remote data files with natural language queries powered by OpenAI and DuckDB."
 authors = [{name = "Brian Thorne", email = "[email protected]" }]
 readme = "README.md"

qabot/agent.py

@@ -194,7 +194,7 @@ def llm_step(self, forced_function_call=None):
         choice = chat_response.choices[0]
         message = choice.message
         self.messages.append(message)
-        if "tool_calls" == choice.finish_reason:
+        if hasattr(message, "tool_calls") and message.tool_calls:
             for tool_call in message.tool_calls:
                 function_name = tool_call.function.name
                 call_id = tool_call.id

qabot/cli.py

@@ -184,7 +184,7 @@ def terminate_session(message: str):
                 #progress.start()
                 continue
 
-            print(format_rocket("Sending query to LLM"))
+            print(format_rocket(f"Sending query to LLM ({settings.agent_model.default_model_name})"))
             print(format_user(query))
 
             t = progress.add_task(description="Processing query...", total=None)

qabot/functions/data_loader.py

@@ -54,6 +54,13 @@ def import_into_duckdb_from_files(
             db_type = "(TYPE postgres, READ_ONLY)" if not dangerously_allow_write_access else "(TYPE postgres)"
             duckdb_connection.execute(f"ATTACH '{file_path}' as postgres_db {db_type};")
 
+            _set_search_path(duckdb_connection)
+        elif file_path.endswith('.json'):
+            # use the filename as the table name
+            table_name, _ = os.path.splitext(os.path.basename(file_path))
+            # remove any non-alphanumeric characters
+            new_table_name = "".join([c for c in table_name if c.isalnum()])
+            duckdb_connection.execute(f"CREATE TABLE {new_table_name} AS select * from read_json('{file_path}');")
             _set_search_path(duckdb_connection)
         elif file_path.endswith('.xlsx'):
             try:

qabot/llm.py

@@ -1,5 +1,5 @@
 import openai
-from openai.types.chat import ChatCompletionToolParam
+from openai.types.chat import ChatCompletionToolParam, ChatCompletionNamedToolChoiceParam
 from openai import RateLimitError, AuthenticationError, OpenAI
 from rich import print
 from tenacity import retry, wait_random_exponential, stop_after_attempt, retry_if_not_exception_type
@@ -22,7 +22,11 @@ def chat_completion_request(
         ]
 
     if function_call is not None:
-        call_data["function_call"] = function_call
+        print(f"Calling function {function_call}")
+        call_data["tool_choice"] = ChatCompletionNamedToolChoiceParam(
+            function=function_call,
+            type='function',
+        )
     try:
         return openai_client.chat.completions.create(**call_data)
     except Exception as e:

qabot/prompts/system.py

@@ -42,19 +42,21 @@
 SELECT count(*) FROM t1 WHERE ST_Within(geom, ST_MakeEnvelope(45, 45, 65, 65));
 ```
 
+Remember JSON is well supported by DuckDB queries and is similar to Postgresql. 
+
 If the question does not seem related to the currently loaded data, Qabot considers other sources for 
 and presents options to the user. Any queries that qabot considers malicious simply returns 
 "I can't help with that" as the answer.
 
 All interactions with the user are carried out through tool calls, Qabot's non-tool call replies are
-treated as Qabot's internal monologue and should not be used as a response to the user. Instead this can be
+treated as Qabot's internal monologue and should not be used as a response to the user. Instead this is 
 used for summarizing failed attempts and planning out future steps.
 
 Qabot prefers to give factual answers backed by data, calculations are to be computed by executing SQL. 
 """
 
 
-research_prompt = """
+research_prompt = f"""
 You are the research and planning component of Qabot, a large language model based application that helpfully answers
 questions involving data and code.
 
@@ -65,5 +67,11 @@
 The main Qabot agent is a smaller LLM that is requesting your help formulating a plan. Ideally provide code
 samples using DuckDB/Postgresql SQL. 
 
+The system prompt for the weaker model delimited by --- --- --- is:
+
+{system_prompt}
+
+--- --- ---
+
 The last messages between the user and qabot:
 """