quillsql 2.2.7__py3-none-any.whl → 2.2.9__py3-none-any.whl

quillsql/core.py

@@ -1,4 +1,5 @@
 import os
+import codecs
 from dotenv import load_dotenv
 
 import requests
@@ -29,7 +30,7 @@ load_dotenv()
 
 ENV = os.getenv("PYTHON_ENV")
 DEV_HOST = "http://localhost:8080"
-PROD_HOST = "https://quill-344421.uc.r.appspot.com"
+PROD_HOST = "https://api.quill.co"
 HOST = DEV_HOST if ENV == "development" else PROD_HOST
 
 SINGLE_TENANT = "QUILL_SINGLE_TENANT"
@@ -386,6 +387,275 @@ class Quill:
                 "status": "error",
                 "data": responseMetadata,
             }
+        
+    async def stream(
+        self,
+        tenants,
+        metadata,
+        flags=None,
+        filters=None,
+        admin_enabled=None,
+    ):
+        if not tenants:
+            raise ValueError("You may not pass an empty tenants array.")
+
+        if not metadata:
+            yield {"type": "error", "errorText": "Missing metadata."}
+            return
+
+        task = metadata.get("task")
+        if not task:
+            yield {"type": "error", "errorText": "Missing task."}
+            return
+
+        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
+
+            for event in self._stream_sse(
+                task,
+                tenants,
+                metadata,
+                flags,
+                filters,
+                admin_enabled,
+            ):
+                yield event
+            return
+        except Exception as err:
+            yield {
+                "type": "error",
+                "errorText": str(err).splitlines()[0],
+            }
+            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 _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
+
+        # 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/{endpoint}"
+        headers = {
+            "Authorization": f"Bearer {self.private_key}",
+            "Content-Type": "application/json",
+            "Accept": "text/event-stream",
+        }
+        encoded = json.dumps(request_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]":
+                    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

{quillsql-2.2.7.dist-info → quillsql-2.2.9.dist-info}/METADATA

@@ -1,6 +1,6 @@
-Metadata-Version: 2.2
+Metadata-Version: 2.4
 Name: quillsql
-Version: 2.2.7
+Version: 2.2.9
 Summary: Quill SDK for Python.
 Home-page: https://github.com/quill-sql/quill-python
 Author: Quill
@@ -67,3 +67,44 @@ async def quill_post(data: Request, user: dict = Depends(authenticate_jwt)):
 
 Then you can run your app like normally. Pass in this route to our react library
 on the frontend and you all set!
+
+## Streaming
+
+```python
+from quillsql import Quill
+from fastapi.responses import StreamingResponse
+import asyncio
+
+quill = Quill(
+    private_key=os.getenv("QULL_PRIVATE_KEY"),
+    database_connection_string=os.getenv("POSTGRES_READ"),
+    database_type="postgresql"
+)
+
+@app.post("/quill-stream")
+async def quill_post(data: Request, user: dict = Depends(authenticate_jwt)):
+    # assuming user fetched via auth middleware has an userId
+    user_id = user["user_id"]
+    body = await data.json()
+    metadata = body.get("metadata")
+
+    quill_stream = quill.stream(
+        tenants=[{"tenantField": "user_id", "tenantIds": [user_id]}],
+        metadata=metadata,
+    )
+
+    async def event_generator():
+        # Full event types list: https://ai-sdk.dev/docs/ai-sdk-ui/stream-protocol#data-stream-protocol
+        async for event in quill_stream:
+            if event["type"] == "start":
+                pass
+            elif event["type"] == "text-delta":
+                yield event['delta']
+            elif event["type"] == "finish":
+                return
+            elif event["type"] == "error":
+                yield event['errorText']
+            await asyncio.sleep(0)
+
+    return StreamingResponse(event_generator(), media_type="text/event-stream")
+```

{quillsql-2.2.7.dist-info → quillsql-2.2.9.dist-info}/RECORD

@@ -1,5 +1,5 @@
 quillsql/__init__.py,sha256=wjJfszle5vheUbgUfJMHQqtqhx2W3UaDN4ndcRIfmkQ,236
-quillsql/core.py,sha256=IQZlNdEAfiEuWF4PnA4hI6iK3Yk5OI8QKwKVklwW8GA,23943
+quillsql/core.py,sha256=VzWtu0up8tcSxijHt_x3bSbZnPjbZFIyGKXCwerpPrg,33704
 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.7.dist-info/METADATA,sha256=EE-TzuvH71mqurYprv3i-rzqd4_REL4bB14WEHqPCaA,1786
-quillsql-2.2.7.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-quillsql-2.2.7.dist-info/top_level.txt,sha256=eU2vHnVqwpYQJ3ADl1Q-DIBzbYejZRUhcMdN_4zMCz8,9
-quillsql-2.2.7.dist-info/RECORD,,
+quillsql-2.2.9.dist-info/METADATA,sha256=zVlzKEUZAAIQR30Re6w6tdebu6hFjtTQnOeIXlOwZz4,3052
+quillsql-2.2.9.dist-info/WHEEL,sha256=wUyA8OaulRlbfwMtmQsvNngGrxQHAvkKcvRmdizlJi0,92
+quillsql-2.2.9.dist-info/top_level.txt,sha256=eU2vHnVqwpYQJ3ADl1Q-DIBzbYejZRUhcMdN_4zMCz8,9
+quillsql-2.2.9.dist-info/RECORD,,

{quillsql-2.2.7.dist-info → quillsql-2.2.9.dist-info}/WHEEL

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