guillem/hellocomputer
guillem/hellocomputer
1
0
Fork 0
Code Issues Pull requests Packages Projects Releases 1 Wiki Activity

Working on SQL tools
All checks were successful
ci/woodpecker/push/woodpecker Pipeline was successful
Details

Browse source
This commit is contained in:
Guillem Borrell 2024-07-29 17:20:56 +02:00
parent 9e1da05276
commit f16bb6b8cf
2 changed files with 7 additions and 5 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

1
.gitignore vendored
View file

@ -166,4 +166,3 @@ test/data/output/*
.pytest_cache .pytest_cache
.ruff_cache .ruff_cache
*.ipynb

9
src/hellocomputer/tools.py
View file

@ -15,16 +15,19 @@ class DuckdbQueryInput(BaseModel):
class DuckdbQueryTool(BaseTool): class DuckdbQueryTool(BaseTool):
name: str = "sql_query" name: str = "sql_query"
description: str = "Run a SQL query in the database containing all the datasets " 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 args_schema: Type[BaseModel] = DuckdbQueryInput
def _run(self, query: str, session_id: str) -> str: def _run(self, query: str, session_id: str) -> str:
"""Run the query""" """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: async def _arun(self, query: str, session_id: str) -> str:
"""Use the tool asynchronously.""" """Use the tool asynchronously."""
db = SessionDB(settings, session_id) session = SessionDB(settings, session_id)
session.db.sql(query)
return "Table" return "Table"