quillsql 2.2.7__tar.gz → 2.2.9__tar.gz

{quillsql-2.2.7/quillsql.egg-info → quillsql-2.2.9}/PKG-INFO

@@ -1,6 +1,6 @@
 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.9/README.md

@@ -0,0 +1,87 @@
+# Quill Python SDK
+
+## Quickstart
+
+First, install the quillsql package by running:
+
+```bash
+$ pip install quillsql
+```
+
+Then, add a `/quill` endpoint to your existing python server. For example, if
+you were running a FASTAPI app, you would just add the endpoint like this:
+
+```python
+from quillsql import Quill
+
+quill = Quill(
+    private_key=os.getenv("QULL_PRIVATE_KEY"),
+    database_connection_string=os.getenv("POSTGRES_READ"),
+    database_type="postgresql"
+)
+
+security = HTTPBearer()
+
+async def authenticate_jwt(token: str = Depends(security)):
+    # Your JWT validation logic here
+    # Return user object or raise HTTPException
+    user = validate_jwt_token(token.credentials)
+    return user
+
+@app.post("/quill")
+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")
+
+    result = quill.query(
+        tenants=[{"tenantField": "user_id", "tenantIds": [user_id]}],
+        metadata=metadata
+    )
+    return result
+```
+
+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 → quillsql-2.2.9}/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 → quillsql-2.2.9/quillsql.egg-info}/PKG-INFO

@@ -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 → quillsql-2.2.9}/setup.py

@@ -2,7 +2,7 @@ from setuptools import setup, find_packages
 
 setup(
     name="quillsql",
-    version="2.2.7",
+    version="2.2.9",
     packages=find_packages(),
     install_requires=[
         "psycopg2-binary",

quillsql-2.2.7/README.md

@@ -1,46 +0,0 @@
-# Quill Python SDK
-
-## Quickstart
-
-First, install the quillsql package by running:
-
-```bash
-$ pip install quillsql
-```
-
-Then, add a `/quill` endpoint to your existing python server. For example, if
-you were running a FASTAPI app, you would just add the endpoint like this:
-
-```python
-from quillsql import Quill
-
-quill = Quill(
-    private_key=os.getenv("QULL_PRIVATE_KEY"),
-    database_connection_string=os.getenv("POSTGRES_READ"),
-    database_type="postgresql"
-)
-
-security = HTTPBearer()
-
-async def authenticate_jwt(token: str = Depends(security)):
-    # Your JWT validation logic here
-    # Return user object or raise HTTPException
-    user = validate_jwt_token(token.credentials)
-    return user
-
-@app.post("/quill")
-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")
-
-    result = quill.query(
-        tenants=[{"tenantField": "user_id", "tenantIds": [user_id]}],
-        metadata=metadata
-    )
-    return result
-```
-
-Then you can run your app like normally. Pass in this route to our react library
-on the frontend and you all set!