quillsql 2.2.8__py3-none-any.whl → 2.2.10__py3-none-any.whl

quillsql/core.py

@@ -1,5 +1,8 @@
 import os
 import codecs
+import uuid
+from datetime import date, datetime
+from decimal import Decimal
 from dotenv import load_dotenv
 
 import requests
@@ -393,6 +396,8 @@ class Quill:
         tenants,
         metadata,
         flags=None,
+        filters=None,
+        admin_enabled=None,
     ):
         if not tenants:
             raise ValueError("You may not pass an empty tenants array.")
@@ -409,51 +414,26 @@ class Quill:
         try:
             # Set tenant IDs in the connection
             self.target_connection.tenant_ids = extract_tenant_ids(tenants)
+            if task in ("chat", "agent"):
+                for event in self._agentic_chat_loop(
+                    tenants,
+                    metadata,
+                    flags,
+                    filters,
+                    admin_enabled,
+                ):
+                    yield event
+                return
 
-            # Handle tenant flags synthesis
-            tenant_flags = None
-            if tenants[0] == SINGLE_TENANT and flags:
-                if flags and isinstance(flags[0], dict):
-                    tenant_flags = [{'tenantField': SINGLE_TENANT, 'flags': flags}]
-                else:
-                    tenant_flags = flags
-
-            payload = { 
-                **metadata,
-                "tenants": tenants,
-                "flags": tenant_flags,
-            }
-            # Custom JSON Encoder to handle Enums
-            class EnumEncoder(json.JSONEncoder):
-                def default(self, obj):
-                    if isinstance(obj, Enum):
-                        return obj.value  # Convert enum to its value (string in this case)
-                    return super().default(obj)
-            url = f"{self.baseUrl}/sdk/{task}"
-            headers = {"Authorization": f"Bearer {self.private_key}", "Content-Type": "application/json","Accept": "text/event-stream"}
-            encoded = json.dumps(payload, cls=EnumEncoder)
-
-            resp = requests.post(url, data=encoded, headers=headers, stream=True)
-            decoder = codecs.getincrementaldecoder('utf-8')()
-            buf = ""
-            for chunk in resp.iter_content(chunk_size=4096):
-                buf += decoder.decode(chunk)
-                while "\n\n" in buf:
-                    raw_event, buf = buf.split("\n\n", 1)
-                    data_lines = []
-                    for line in raw_event.splitlines():
-                        if line.startswith("data:"):
-                            data_lines.append(line[len("data:"):].strip())
-                    if not data_lines:
-                        continue
-                    payload = "\n".join(data_lines)
-                    if payload == "[DONE]":
-                        break
-                    yield json.loads(payload)
-
-            # flush any partial code points at the end
-            buf += decoder.decode(b"", final=True)
-            yield buf
+            for event in self._stream_sse(
+                task,
+                tenants,
+                metadata,
+                flags,
+                filters,
+                admin_enabled,
+            ):
+                yield event
             return
         except Exception as err:
             yield {
@@ -462,6 +442,228 @@ class Quill:
             }
             return
 
+    def _normalize_tenant_flags(self, tenants, flags):
+        tenant_flags = None
+        if tenants and tenants[0] == SINGLE_TENANT and flags:
+            if flags and isinstance(flags[0], dict):
+                tenant_flags = [{"tenantField": SINGLE_TENANT, "flags": flags}]
+            else:
+                tenant_flags = flags
+        return tenant_flags
+
+    def _agentic_chat_loop(self, tenants, metadata, flags, filters, admin_enabled):
+        messages = list(metadata.get("messages") or [])
+        max_iterations = 10
+
+        for _ in range(max_iterations):
+            payload = {
+                **metadata,
+                "messages": messages,
+            }
+
+            has_tool_calls = False
+            assistant_text = ""
+            tool_results = []
+            tool_calls = []
+
+            for event in self._stream_sse(
+                "agent",
+                tenants,
+                payload,
+                flags,
+                filters,
+                admin_enabled,
+            ):
+                yield event
+
+                if event.get("type") == "text-delta":
+                    assistant_text += event.get("delta", "")
+
+                if event.get("type") == "tool-input-available":
+                    tool_name = event.get("toolName")
+                    tool_call_id = event.get("toolCallId")
+                    tool_input = event.get("input") or {}
+
+                    if tool_call_id is None:
+                        yield {
+                            "type": "error",
+                            "errorText": "Missing toolCallId for tool-input-available event.",
+                        }
+                        continue
+
+                    has_tool_calls = True
+                    yield {
+                        "type": "tool-executing",
+                        "toolCallId": tool_call_id,
+                        "toolName": tool_name,
+                    }
+
+                    result = self._execute_tool_locally(tool_name, tool_input)
+                    tool_results.append(
+                        {
+                            "toolCallId": tool_call_id,
+                            "toolName": tool_name,
+                            "input": tool_input,
+                            "result": result,
+                        }
+                    )
+                    tool_calls.append(
+                        {
+                            "id": tool_call_id,
+                            "type": "function",
+                            "function": {
+                                "name": tool_name,
+                                "arguments": json.dumps(tool_input or {}),
+                            },
+                        }
+                    )
+
+                    yield {
+                        "type": "tool-result",
+                        "toolCallId": tool_call_id,
+                        "result": result,
+                    }
+
+                if event.get("type") in ("finish", "error"):
+                    break
+
+            if not has_tool_calls:
+                break
+
+            def _build_tool_result_for_history(tool_result):
+                result = tool_result.get("result") or {}
+                has_rows = isinstance(result.get("rows"), list)
+                has_fields = isinstance(result.get("fields"), list)
+                is_query_result = has_rows or has_fields or result.get("dbMismatched")
+                if not is_query_result:
+                    return result
+                tool_input = tool_result.get("input") or {}
+                error = result.get("error") or tool_input.get("error")
+                status = "error" if error or result.get("dbMismatched") else "success"
+                payload = {"status": status}
+                if tool_input.get("sql"):
+                    payload["sql"] = tool_input.get("sql")
+                if error:
+                    payload["error"] = error
+                if result.get("dbMismatched"):
+                    payload["meta"] = {
+                        "dbMismatched": True,
+                        "backendDatabaseType": result.get("backendDatabaseType"),
+                    }
+                elif not error:
+                    payload["meta"] = {
+                        "rowsFetchedSuccessfully": True,
+                        "rowCount": len(result.get("rows") or []),
+                    }
+                return payload
+
+            messages.append(
+                {
+                    "role": "assistant",
+                    "content": assistant_text or None,
+                    "tool_calls": tool_calls,
+                }
+            )
+            for tool_result in tool_results:
+                messages.append(
+                    {
+                        "role": "tool",
+                        "tool_call_id": tool_result["toolCallId"],
+                        "content": json.dumps(_build_tool_result_for_history(tool_result)),
+                    }
+                )
+
+        yield {"type": "done"}
+
+    def _execute_tool_locally(self, tool_name, tool_input):
+        if tool_name == "generateReport":
+            if tool_input.get("error"):
+                return {"error": tool_input.get("error")}
+            sql = tool_input.get("sql")
+            if not sql:
+                return {"error": "No SQL provided"}
+            results = self.run_queries(
+                [sql],
+                self.target_connection.database_type,
+            )
+            if results.get("dbMismatched"):
+                return results
+            query_results = results.get("queryResults") or []
+            if query_results and isinstance(query_results[0], dict):
+                if query_results[0].get("error"):
+                    return query_results[0]
+                return {
+                    "rows": query_results[0].get("rows", []),
+                    "fields": query_results[0].get("fields", []),
+                }
+            return {"rows": [], "fields": []}
+
+        if tool_name == "createChart":
+            return {"chartConfig": tool_input}
+
+        return {"error": f"Unknown tool: {tool_name}"}
+
+    def _json_default(self, obj):
+        if isinstance(obj, Enum):
+            return obj.value
+        if isinstance(obj, uuid.UUID):
+            return str(obj)
+        if isinstance(obj, (datetime, date)):
+            return obj.isoformat()
+        if isinstance(obj, Decimal):
+            return float(obj)
+        raise TypeError(f"Object of type {type(obj).__name__} is not JSON serializable")
+
+    def _stream_sse(self, endpoint, tenants, payload, flags, filters, admin_enabled):
+        tenant_flags = self._normalize_tenant_flags(tenants, flags)
+        request_payload = {
+            **payload,
+            "tenants": tenants,
+            "flags": tenant_flags,
+        }
+        if filters:
+            request_payload["sdkFilters"] = [convert_custom_filter(f) for f in filters]
+        if admin_enabled is not None:
+            request_payload["adminEnabled"] = admin_enabled
+
+        url = f"{self.baseUrl}/sdk/{endpoint}"
+        headers = {
+            "Authorization": f"Bearer {self.private_key}",
+            "Content-Type": "application/json",
+            "Accept": "text/event-stream",
+        }
+        encoded = json.dumps(request_payload, default=self._json_default)
+
+        resp = requests.post(url, data=encoded, headers=headers, stream=True)
+        decoder = codecs.getincrementaldecoder("utf-8")()
+        buf = ""
+        for chunk in resp.iter_content(chunk_size=4096):
+            buf += decoder.decode(chunk)
+            while "\n\n" in buf:
+                raw_event, buf = buf.split("\n\n", 1)
+                data_lines = []
+                for line in raw_event.splitlines():
+                    if line.startswith("data:"):
+                        data_lines.append(line[len("data:"):].strip())
+                if not data_lines:
+                    continue
+                payload = "\n".join(data_lines)
+                if payload == "[DONE]":
+                    return
+                try:
+                    parsed = json.loads(payload)
+                    if isinstance(parsed, str):
+                        yield {"type": "text-delta", "id": "0", "delta": parsed}
+                    else:
+                        yield parsed
+                except json.JSONDecodeError:
+                    continue
+
+        # flush any partial code points at the end
+        buf += decoder.decode(b"", final=True)
+        yield buf
+        return
+
     def apply_limit(self, query, limit):
         # Simple logic: if query already has a limit, don't add another
         if getattr(self.target_connection, 'database_type', '').lower() == 'mssql':
@@ -633,16 +835,9 @@ class Quill:
         return results
     
     def post_quill(self, path, payload):
-        # Custom JSON Encoder to handle Enums
-        class EnumEncoder(json.JSONEncoder):
-            def default(self, obj):
-                if isinstance(obj, Enum):
-                    return obj.value  # Convert enum to its value (string in this case)
-                return super().default(obj)
-        
         url = f"{self.baseUrl}/sdk/{path}"
         # Set content type to application/json
         headers = {"Authorization": f"Bearer {self.private_key}", "Content-Type": "application/json"}
-        encoded = json.dumps(payload, cls=EnumEncoder)
+        encoded = json.dumps(payload, default=self._json_default)
         response = requests.post(url, data=encoded, headers=headers)
         return response.json()

{quillsql-2.2.8.dist-info → quillsql-2.2.10.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: quillsql
-Version: 2.2.8
+Version: 2.2.10
 Summary: Quill SDK for Python.
 Home-page: https://github.com/quill-sql/quill-python
 Author: Quill

{quillsql-2.2.8.dist-info → quillsql-2.2.10.dist-info}/RECORD

@@ -1,5 +1,5 @@
 quillsql/__init__.py,sha256=wjJfszle5vheUbgUfJMHQqtqhx2W3UaDN4ndcRIfmkQ,236
-quillsql/core.py,sha256=yLDGW_BZTcY5FBxJVefCxzPBgzChy_vLHfJNylz6wuo,26737
+quillsql/core.py,sha256=cxuAfOgQ3AMt26hQgNTFejGv7L_fEAVXtV7dsNdiT7k,33592
 quillsql/error.py,sha256=n9VKHw4FAgg7ZEAz2YQ8L_8FdRG_1shwGngf2iWhUSM,175
 quillsql/assets/__init__.py,sha256=oXQ2ZS5XDXkXTYjADxNfGt55cIn_rqfgWL2EDqjTyoI,45
 quillsql/assets/pgtypes.py,sha256=-B_2wUaoAsdX7_HnJhUlx4ptZQ6x-cXwuST9ACgGFdE,33820
@@ -15,7 +15,7 @@ quillsql/utils/post_quill_executor.py,sha256=DB1RHNfqHPYarMM10vSv--UjpCZqe4qYTjq
 quillsql/utils/run_query_processes.py,sha256=QwnMr5UwXdtO_W88lv5nBaf6pJ_h5oWQnYd8K9oHQ5s,1030
 quillsql/utils/schema_conversion.py,sha256=TFfMibN9nOsxNRhHw5YIFl3jGTvipG81bxX4LFDulUY,314
 quillsql/utils/tenants.py,sha256=ZD2FuKz0gjBVSsThHDv1P8PU6EL8E009NWihE5hAH-Q,2022
-quillsql-2.2.8.dist-info/METADATA,sha256=qfMIakgsSk045GZzk7sNtLu-2gVIQ82toNiAp3lwif0,3052
-quillsql-2.2.8.dist-info/WHEEL,sha256=qELbo2s1Yzl39ZmrAibXA2jjPLUYfnVhUNTlyF1rq0Y,92
-quillsql-2.2.8.dist-info/top_level.txt,sha256=eU2vHnVqwpYQJ3ADl1Q-DIBzbYejZRUhcMdN_4zMCz8,9
-quillsql-2.2.8.dist-info/RECORD,,
+quillsql-2.2.10.dist-info/METADATA,sha256=LiIaD8DgC3gUuUTkW5Opgk2Sh5alIdVh-2UzhIJxAHQ,3053
+quillsql-2.2.10.dist-info/WHEEL,sha256=wUyA8OaulRlbfwMtmQsvNngGrxQHAvkKcvRmdizlJi0,92
+quillsql-2.2.10.dist-info/top_level.txt,sha256=eU2vHnVqwpYQJ3ADl1Q-DIBzbYejZRUhcMdN_4zMCz8,9
+quillsql-2.2.10.dist-info/RECORD,,

{quillsql-2.2.8.dist-info → quillsql-2.2.10.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (80.10.1)
+Generator: setuptools (80.10.2)
 Root-Is-Purelib: true
 Tag: py3-none-any