Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add Sqlite plugin #154

Merged
merged 4 commits into from
Jan 10, 2025
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions app/plugins/loader.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@
from app.plugins.airtable.handler import Airtable
from app.plugins.website.handler import Website
from app.plugins.document.handler import Document
from app.plugins.sqlite.handler import Sqlite
from loguru import logger

class DSLoader:
Expand All @@ -18,6 +19,7 @@ def load_ds(self):
"airtable": Airtable,
"website": Website,
"document" : Document,
"sqlite" : Sqlite,
}
db_type = self.config.get("type")
connection_params = self.config.get("params")
Expand Down
150 changes: 150 additions & 0 deletions app/plugins/sqlite/__init__.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,150 @@
from app.models.prompt import Prompt
from collections import OrderedDict
from app.models.request import ConnectionArgument

# Plugin Metadata
__version__ = '1.0.0'
__plugin_name__ = 'sqlite'
__display_name__ = "SQLite DB"
__description__ = 'SQLite integration for handling SQLite database operations.'
__icon__ = '/assets/plugins/logos/sqlite.svg'
__category__ = 2


# Connection arguments
__connection_args__ = OrderedDict(
db_name = ConnectionArgument(
type = 1,
generic_name= 'Sqlite Database Name',
description = 'Database name',
order= 1,
required = True,
value = None,
slug = "db_name"
),
)

# Prompt
__prompt__ = Prompt(**{
"base_prompt": "{system_prompt}{user_prompt}",
"system_prompt": {
"template": """
You are an Sqlite expert. Your job is to answer questions about a Sqlite database using only the provided schema details and rules.

go through the schema details given below
-- start db schema section--
{schema}
-- end db schema section--

A brief description about the schema is given below

-- start db context section--
{context}
-- end db context section--

Sample sql queries with their questions are given below

-- start query samples section--
$suggestions
-- end query samples section--

Adhere to the given rules without failure

-- start rules section --
- Use Table Aliases always to prevent ambiguity . For example, `SELECT table1.col1, table2.col1 FROM table1 JOIN table2 ON table1.id = table2.id`.
- use LIKE operator with LOWER function for string comparison or equality
- Always use alias/table name for fields in WHERE condition
- Do not use non existing tables or fields
- id columns are mandatory for all operations
- Do not use JSON_BUILD_OBJECT operation
- Do not use unwanted joins
- Do not return incomplete queries
- Adher to sqlite query syntax
-- end rules section --
"""
},
"user_prompt":{
"template": """
Follow these steps to generate query to solve the question `$question`

1. Deliberately go through schema, context, rules deliberately
2. Understand the question and check whether it's doable with the given context
3. Do only the task asked, Don't hallucinate and overdo the task
4. Strictly return all the fields in the schema during listing operations
5. Strictly return at least 1 text fields and an id field during aggregation/group by operations
6. Generate a query to solve the problem using the schema, context, and strictly follow the rules
7. output in the given json format, extra explanation is strictly prohibited

{
"explanation": "Explain how you finalized the sql query using the schemas and rules provided",
"query" : "sqlite query",
"operation_kind" : "aggregation|list",
"schema": "used schema details separated by comma",
"confidence" : "confidence in 100",
"visualisation": {
"type": "chart type (bar chart, line chart, pie chart) or 'table' for tabular format; 'none' if operation_kind is 'list'",
"x-axis": ["fields that can be used as x axis"],
"y-axis": ["fields that can be used as y axis"],
"title": "layout title name"
},
"general_message": "a general message describing the answers like 'here is your list of incidents' or 'look what i found'",
"main_entity" : "main entity for the query",
"next_questions" : [Produce 3 related questions(maximum 8 words) aligned with the current question, db context and which can be answered with only two table . While creating questions strictly prohibit questions which tells to specify for a specific item]

}
"""
},
"regeneration_prompt": {
"template": """
You were trying to answer the following user question by writing SQL query to answer the question given in `[question][/question]`
[question]
$question
[/question]

You generated this query given in `[query][/query]`
[query]
{query_generated}
[/query]

But upon execution you encountered some error , error traceback is given in [query_error][/query_error]
[query_error]
{exception_log}
[/query_error]

Follow these steps to generate the query

1. Deliberately go through schema, context, rules deliberately
2. Understand the question and check whether it's doable with the given context
3. Use survey answers if available and include it in query for filtering values
4. Do only the task asked, Don't hallucinate and overdo the task
5. Strictly return all the fields in the schema during listing operations
6. Strictly return at least 1 text fields and an id field during aggregation/group by operations
7. Generate a query to solve the problem using the schema, context and the rules and based on the previous query try to rectify the query error
8. output in the given json format, extra explanation is strictly prohibited

{
"explanation": "Explain how you finalized the sql query using the schemas and rules provided",
"query" : "sqlite query",
"operation_kind" : "aggregation|list",
"visualisation": {
"type": "chart type (bar chart, line chart, pie chart) or 'table' for tabular format; 'none' if operation_kind is 'list'",
"value_field": "fields in which values are stored",
"x-axis": "field that can be used as x axis",
"y-axis": "field that can be used as y axis",
"title": "layout title name"
},
"schema": "used schema details separated by comma",
"confidence" : "confidence in 100",
"general_message": "a general message describing the answers like 'here is your list of incidents' or 'look what i found'",
"main_entity" : "main entity for the query",
"next_questions" : [Produce 3 related questions(maximum 8 words) aligned with the current question, db context and which can be answered with only two table . While creating questions strictly prohibit questions which tells to specify for a specific item]
}
"""
}
})



__all__ = [
__version__, __plugin_name__, __display_name__ , __description__, __icon__, __category__, __prompt__
]
96 changes: 96 additions & 0 deletions app/plugins/sqlite/formatter.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,96 @@
from typing import Any
from loguru import logger


class Formatter:
def format(self, data: Any,input) -> (dict):
"""
Main entry point for formatting the data based on the input parameters.
Handles different formatting strategies based on operation kind.

:param data: The data to format.
:param input_params: Dictionary containing operation and formatting details.
:return: A dictionary containing the formatted response.
"""
response = {}
self.main_entity = input.get("main_entity")
self.kind = input.get("operation_kind", "").lower()
self.general_message = input.get("general_message")
self.empty_message = input.get("empty_message")

logger.info("Formatting output using inference for sqlite")

if self.kind == "list":
response = self.basic_formatter(data, input)
elif self.kind == "aggregation":
response = self.aggregation_formatter(data, input)
else:
response["data"] = data
response["kind"] = "list"


response.update({
"main_entity": self.main_entity,
"main_format": self.kind,
"role": "assistant",
"content": self.general_message,
"empty_message": self.empty_message,
})

return response

def basic_formatter(self, data: Any, input:Any) -> dict :
"""
Formats data as a list, handling cases for none, single, and multiple entries.

:param data: The data to format.
:return: A dictionary containing the formatted list response.
"""
logger.info("Formatting data as a list")

if data is None:
response = {"data": [], "kind": "none"}
elif len(data) == 1:
response = {"data": data, "kind": "single"}
else:
response = {"data": data, "kind": "list"}

return response


def aggregation_formatter(self, data:Any, input:Any) -> dict :
"""
Formats data for aggregation visualisation, supporting table and chart formats.

:param data: The data to format.
:param visualisation: Dictionary containing visualisation details (e.g., x-axis, y-axis, chart type).
:return: A dictionary containing the formatted aggregation response.
"""

logger.info("Formatting data as aggregation")

visualisation = input.get("visualisation", {})
response = {}


if data is None or len(data) == 0:
response = {"data": [], "kind": "none"}
elif len(data) == 1:
response = {"data": data, "kind": "table"}
else:
value_fields = visualisation.get("y-axis", [])
key_fields = visualisation.get("x-axis", [])
title = visualisation.get("title", "")

visualisaton_kind = visualisation["type"].replace(" ", "_") if visualisation["type"] is not None else "table"

if visualisaton_kind.lower() in ["bar_chart", "line_chart", "pie_chart"] and len(value_fields) > 0 and len(key_fields) > 0:
response["kind"] = visualisaton_kind
response["data"] = data
response["x"] = key_fields
response["y"] = value_fields
response["title"] = title
else:
response = {"kind": "table", "data": data}

return response
Loading