From f16bb6b8cfda319589bc42ee3469ddc6119a3b6e Mon Sep 17 00:00:00 2001
From: Guillem Borrell <guillemborrell@gmail.com>
Date: Mon, 29 Jul 2024 17:20:56 +0200
Subject: [PATCH] Working on SQL tools

---
 .gitignore                 | 3 +--
 src/hellocomputer/tools.py | 9 ++++++---
 2 files changed, 7 insertions(+), 5 deletions(-)

diff --git a/.gitignore b/.gitignore
index bbffe28..75ff67b 100644
--- a/.gitignore
+++ b/.gitignore
@@ -165,5 +165,4 @@ cython_debug/
 test/data/output/*
 
 .pytest_cache
-.ruff_cache
-*.ipynb
\ No newline at end of file
+.ruff_cache
\ No newline at end of file
diff --git a/src/hellocomputer/tools.py b/src/hellocomputer/tools.py
index b031d7f..e654c99 100644
--- a/src/hellocomputer/tools.py
+++ b/src/hellocomputer/tools.py
@@ -15,16 +15,19 @@ class DuckdbQueryInput(BaseModel):
 class DuckdbQueryTool(BaseTool):
     name: str = "sql_query"
     description: str = "Run a SQL query in the database containing all the datasets "
-    "and provide a summary of the results"
+    "and provide a summary of the results, and the name of the table with them if the "
+    "volume of the results is large"
     args_schema: Type[BaseModel] = DuckdbQueryInput
 
     def _run(self, query: str, session_id: str) -> str:
         """Run the query"""
-        db = SessionDB(settings, session_id)
+        session = SessionDB(settings, session_id)
+        session.db.sql(query)
 
     async def _arun(self, query: str, session_id: str) -> str:
         """Use the tool asynchronously."""
-        db = SessionDB(settings, session_id)
+        session = SessionDB(settings, session_id)
+        session.db.sql(query)
         return "Table"