feat: Add QueryWeaver Python SDK for serverless Text2SQL

.github/workflows/tests.yml

@@ -63,8 +63,80 @@ jobs:
 
     - name: Run unit tests
       run: |
-        pipenv run pytest tests/ -k "not e2e" --verbose
+        pipenv run pytest tests/ -k "not e2e and not test_sdk" --verbose
 
     - name: Run lint
       run: |
         make lint
+
+  sdk-tests:
+    runs-on: ubuntu-latest
+
+    services:
+      falkordb:
+        image: falkordb/falkordb:latest
+        ports:
+          - 6379:6379
+        options: >-
+          --health-cmd "redis-cli ping"
+          --health-interval 10s
+          --health-timeout 5s
+          --health-retries 5
+
+      postgres:
+        image: postgres:15
+        env:
+          POSTGRES_USER: postgres
+          POSTGRES_PASSWORD: postgres
+          POSTGRES_DB: testdb
+        ports:
+          - 5432:5432
+        options: >-
+          --health-cmd pg_isready
+          --health-interval 10s
+          --health-timeout 5s
+          --health-retries 5
+
+      mysql:
+        image: mysql:8
+        env:
+          MYSQL_ROOT_PASSWORD: root
+          MYSQL_DATABASE: testdb
+        ports:
+          - 3306:3306
+        options: >-
+          --health-cmd "mysqladmin ping"
+          --health-interval 10s
+          --health-timeout 5s
+          --health-retries 5
+
+    steps:
+    - uses: actions/checkout@v4
+
+    - name: Set up Python
+      uses: actions/setup-python@v5
+      with:
+        python-version: '3.12'
+
+    - name: Install pipenv
+      run: |
+        python -m pip install --upgrade pip
+        pip install pipenv
+
+    - name: Install dependencies
+      run: |
+        pipenv sync --dev
+
+    - name: Create test environment file
+      run: |
+        cp .env.example .env
+        echo "FASTAPI_SECRET_KEY=test-secret-key" >> .env
+
+    - name: Run SDK tests
+      env:
+        FALKORDB_URL: redis://localhost:6379
+        TEST_POSTGRES_URL: postgresql://postgres:postgres@localhost:5432/testdb
+        TEST_MYSQL_URL: mysql://root:root@localhost:3306/testdb
+        OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }}
+      run: |
+        pipenv run pytest tests/test_sdk/ -v

Makefile

@@ -1,19 +1,37 @@
-.PHONY: help install test test-unit test-e2e test-e2e-headed lint format clean setup-dev build lint-frontend
+.PHONY: help install test test-unit test-e2e test-e2e-headed lint format clean setup-dev build lint-frontend test-sdk docker-test-services docker-test-stop build-package
 
 help: ## Show this help message
 	@echo 'Usage: make [target]'
 	@echo ''
 	@echo 'Targets:'
-	@awk 'BEGIN {FS = ":.*?## "} /^[a-zA-Z_-]+:.*?## / {printf "  %-15s %s\n", $$1, $$2}' $(MAKEFILE_LIST)
-
-install: ## Install dependencies
+	@awk 'BEGIN {FS = ":.*?## "} /^[a-zA-Z_-]+:.*?## / {printf "  %-20s %s\n", $$1, $$2}' $(MAKEFILE_LIST)
+
+# Check if uv is available, fallback to pipenv
+UV := $(shell command -v uv 2> /dev/null)
+ifdef UV
+	PKG_MANAGER = uv
+	PIP_CMD = uv pip
+	SYNC_CMD = uv sync
+	RUN_CMD = uv run
+else
+	PKG_MANAGER = pipenv
+	PIP_CMD = pipenv run pip
+	SYNC_CMD = pipenv sync --dev
+	RUN_CMD = pipenv run
+endif
+
+install: ## Install dependencies (uses uv if available, else pipenv)
+ifdef UV
+	uv sync --all-extras
+else
 	pipenv sync --dev
+endif
 	npm install --prefix ./app
 
 
 setup-dev: install ## Set up development environment
-	pipenv run playwright install chromium
-	pipenv run playwright install-deps
+	$(RUN_CMD) playwright install chromium
+	$(RUN_CMD) playwright install-deps
 	@echo "Development environment setup complete!"
 	@echo "Don't forget to copy .env.example to .env and configure your settings"
 
@@ -23,26 +41,34 @@ build-dev:
 build-prod:
 	npm --prefix ./app run build
 
+build-package: ## Build distributable package (wheel + sdist)
+ifdef UV
+	uv build
+else
+	pip install build && python -m build
+endif
+	@echo "Built packages in dist/"
+
 test: build-dev test-unit test-e2e ## Run all tests
 
 test-unit: ## Run unit tests only
-	pipenv run python -m pytest tests/ -k "not e2e" --verbose
+	$(RUN_CMD) python -m pytest tests/ -k "not e2e" --verbose
 
 
 test-e2e: build-dev ## Run E2E tests headless
-	pipenv run python -m pytest tests/e2e/ --browser chromium --video=on --screenshot=on
+	$(RUN_CMD) python -m pytest tests/e2e/ --browser chromium --video=on --screenshot=on
 
 
 test-e2e-headed: build-dev ## Run E2E tests with browser visible
-	pipenv run python -m pytest tests/e2e/ --browser chromium --headed
+	$(RUN_CMD) python -m pytest tests/e2e/ --browser chromium --headed
 
 
 test-e2e-debug: build-dev ## Run E2E tests with debugging enabled
-	pipenv run python -m pytest tests/e2e/ --browser chromium --slowmo=1000
+	$(RUN_CMD) python -m pytest tests/e2e/ --browser chromium --slowmo=1000
 
 lint: ## Run linting (backend + frontend)
 	@echo "Running backend lint (pylint)"
-	pipenv run pylint $(shell git ls-files '*.py') || true
+	$(RUN_CMD) pylint $(shell git ls-files '*.py') || true
 	@echo "Running frontend lint (eslint)"
 	make lint-frontend
 
@@ -57,18 +83,37 @@ clean: ## Clean up test artifacts
 	rm -rf playwright-report/
 	rm -rf tests/e2e/screenshots/
 	rm -rf __pycache__/
+	rm -rf dist/
+	rm -rf *.egg-info/
 	find . -name "*.pyc" -delete
 	find . -name "*.pyo" -delete
 
 run-dev: build-dev ## Run development server
-	pipenv run uvicorn api.index:app --host $${HOST:-127.0.0.1} --port $${PORT:-5000} --reload
+	$(RUN_CMD) uvicorn api.index:app --host $${HOST:-127.0.0.1} --port $${PORT:-5000} --reload
 
 run-prod: build-prod ## Run production server
-	pipenv run uvicorn api.index:app --host $${HOST:-0.0.0.0} --port $${PORT:-5000}
+	$(RUN_CMD) uvicorn api.index:app --host $${HOST:-0.0.0.0} --port $${PORT:-5000}
 
 docker-falkordb: ## Start FalkorDB in Docker for testing
 	docker run -d --name falkordb-test -p 6379:6379 falkordb/falkordb:latest
 
 docker-stop: ## Stop test containers
 	docker stop falkordb-test || true
 	docker rm falkordb-test || true
+
+# SDK Testing
+docker-test-services: ## Start all test services (FalkorDB + PostgreSQL + MySQL)
+	docker compose -f docker-compose.test.yml up -d
+	@echo "Waiting for services to be ready..."
+	@sleep 10
+
+docker-test-stop: ## Stop all test services
+	docker compose -f docker-compose.test.yml down -v
+
+test-sdk: ## Run SDK integration tests (requires docker-test-services)
+	$(RUN_CMD) pytest tests/test_sdk/ -v
+
+test-sdk-quick: ## Run SDK tests without LLM (models and connection only)
+	$(RUN_CMD) pytest tests/test_sdk/test_queryweaver.py::TestModels tests/test_sdk/test_queryweaver.py::TestQueryWeaverInit -v
+
+test-all: test-unit test-sdk test-e2e ## Run all tests

api/core/schema_loader.py

@@ -6,13 +6,15 @@
 from typing import AsyncGenerator, Optional
 
 from pydantic import BaseModel
+from redis import RedisError
 
 from api.extensions import db
 
 from api.core.errors import InvalidArgumentError
 from api.loaders.base_loader import BaseLoader
 from api.loaders.postgres_loader import PostgresLoader
 from api.loaders.mysql_loader import MySQLLoader
+from queryweaver_sdk.models import DatabaseConnection
 
 # Use the same delimiter as in the JavaScript frontend for streaming chunks
 MESSAGE_DELIMITER = "|||FALKORDB_MESSAGE_BOUNDARY|||"
@@ -162,3 +164,74 @@ async def list_databases(user_id: str, general_prefix: Optional[str] = None) ->
         filtered_graphs = filtered_graphs + demo_graphs
 
     return filtered_graphs
+
+
+# =============================================================================
+# SDK Non-Streaming Functions
+# =============================================================================
+
+async def load_database_sync(url: str, user_id: str):
+    """
+    Load a database schema and return structured result (non-streaming).
+
+    SDK-friendly version that returns DatabaseConnection instead of streaming.
+
+    Args:
+        url: Database connection URL (PostgreSQL or MySQL).
+        user_id: User identifier for namespacing.
+
+    Returns:
+        DatabaseConnection with connection status.
+    """
+    # Validate URL format
+    if not url or len(url.strip()) == 0:
+        raise InvalidArgumentError("Invalid URL format")
+
+    # Determine database type and loader
+    loader: type[BaseLoader] = BaseLoader
+    if url.startswith("postgres://") or url.startswith("postgresql://"):
+        loader = PostgresLoader
+    elif url.startswith("mysql://"):
+        loader = MySQLLoader
+    else:
+        raise InvalidArgumentError("Invalid database URL format. Must be PostgreSQL or MySQL.")
+
+    tables_loaded = 0
+    last_message = ""
+    success = False
+
+    try:
+        async for progress_success, progress_message in loader.load(user_id, url):
+            success = progress_success
+            last_message = progress_message
+            if success and "table" in progress_message.lower():
+                # Try to extract table count from message
+                tables_loaded += 1
+
+        if success:
+            # Extract database name from the message or URL
+            # The loader typically returns the graph_id in the final message
+            db_name = url.split("/")[-1].split("?")[0]  # Extract DB name from URL
+
+            return DatabaseConnection(
+                database_id=db_name,
+                success=True,
+                tables_loaded=tables_loaded,
+                message="Database connected and schema loaded successfully",
+            )
+
+        return DatabaseConnection(
+            database_id="",
+            success=False,
+            tables_loaded=0,
+            message=last_message or "Failed to load database schema",
+        )
+
+    except (RedisError, ConnectionError, OSError) as e:
+        logging.exception("Error loading database: %s", str(e))
+        return DatabaseConnection(
+            database_id="",
+            success=False,
+            tables_loaded=0,
+            message=f"Error connecting to database: {str(e)}",
+        )

api/core/text2sql.py

@@ -8,7 +8,7 @@
 import time
 
 from pydantic import BaseModel
-from redis import ResponseError
+from redis import ResponseError, RedisError
 
 from api.core.errors import GraphNotFoundError, InternalError, InvalidArgumentError
 from api.core.schema_loader import load_database
@@ -936,6 +936,6 @@ async def delete_database(user_id: str, graph_id: str):
         return {"success": True, "graph": graph_id}
     except ResponseError as re:
         raise GraphNotFoundError("Failed to delete graph, Graph not found") from re
-    except Exception as e:  # pylint: disable=broad-exception-caught
+    except (RedisError, ConnectionError) as e:
         logging.exception("Failed to delete graph %s: %s", sanitize_log_input(namespaced), e)
         raise InternalError("Failed to delete graph") from e