Staging

Pipfile

@@ -6,7 +6,7 @@ name = "pypi"
 [packages]
 fastapi = "~=0.116.1"
 uvicorn = "~=0.35.0"
-litellm = "~=1.76.0"
+litellm = "~=1.76.3"
 falkordb = "~=1.2.0"
 psycopg2-binary = "~=2.9.9"
 pymysql = "~=1.1.0"
@@ -20,10 +20,10 @@ fastapi-mcp = "~=0.4.0"
 graphiti-core = {ref = "staging", git = "git+https://github.com/FalkorDB/graphiti.git"}
 
 [dev-packages]
-pytest = "~=8.4.1"
+pytest = "~=8.4.2"
 pylint = "~=3.3.4"
-playwright = "~=1.54.0"
-pytest-playwright = "~=0.7.0"
+playwright = "~=1.55.0"
+pytest-playwright = "~=0.7.1"
 pytest-asyncio = "~=1.1.0"
 
 [requires]

api/core/__init__.py

@@ -0,0 +1,20 @@
+
+"""
+Core module for QueryWeaver.
+
+This module provides the core functionality for QueryWeaver including
+error handling, database schema loading, and text-to-SQL processing.
+"""
+
+from .errors import InternalError, GraphNotFoundError, InvalidArgumentError
+from .schema_loader import load_database, list_databases
+from .text2sql import MESSAGE_DELIMITER
+
+__all__ = [
+    "InternalError",
+    "GraphNotFoundError",
+    "InvalidArgumentError",
+    "load_database",
+    "list_databases",
+    "MESSAGE_DELIMITER",
+]

api/core/errors.py

@@ -0,0 +1,13 @@
+"""Custom exceptions for the text2sql API."""
+
+# Interal Error Exception
+class InternalError(Exception):
+    """Custom exception for internal errors."""
+
+# Graph not found Exception
+class GraphNotFoundError(Exception):
+    """Custom exception for graph not found errors."""
+
+# Wrong argument Exception
+class InvalidArgumentError(Exception):
+    """Custom exception for invalid argument errors."""

api/core/schema_loader.py

@@ -0,0 +1,161 @@
+"""Database connection routes for the text2sql API."""
+
+import logging
+import json
+import time
+from typing import AsyncGenerator
+
+from pydantic import BaseModel
+
+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
+
+# Use the same delimiter as in the JavaScript frontend for streaming chunks
+MESSAGE_DELIMITER = "|||FALKORDB_MESSAGE_BOUNDARY|||"
+
+
+class DatabaseConnectionRequest(BaseModel):
+    """Database connection request model.
+
+    Args:
+        BaseModel (_type_): _description_
+    """
+
+    url: str
+
+def _step_start(steps_counter: int) -> dict[str, str]:
+    """Yield the starting step message."""
+    return {
+        "type": "reasoning_step",
+        "message": f"Step {steps_counter}: Starting database connection",
+    }
+
+def _step_detect_db_type(steps_counter: int, url: str) -> tuple[type[BaseLoader], dict[str, str]]:
+    """Yield the database type detection step message."""
+    db_type = None
+    loader: type[BaseLoader] = BaseLoader  # type: ignore
+    if url.startswith("postgres://") or url.startswith("postgresql://"):
+        db_type = "postgresql"
+        loader = PostgresLoader
+    elif url.startswith("mysql://"):
+        db_type = "mysql"
+        loader = MySQLLoader
+    else:
+        raise InvalidArgumentError("Invalid database URL format")
+
+    return loader, {
+        "type": "reasoning_step",
+        "message": f"Step {steps_counter}: Detected database type: {db_type}. "
+        "Attempting to load schema...",
+    }
+
+
+async def _step_attempt_load(
+    steps_counter: int, loader: type[BaseLoader], user_id: str, url: str
+) -> AsyncGenerator[dict[str, str | bool], None]:
+    """Yield the attempt to load schema step message."""
+    success, result = [False, ""]
+    try:
+        load_start = time.perf_counter()
+        async for progress in loader.load(user_id, url):
+            success, result = progress
+            if success:
+                steps_counter += 1
+                yield {
+                    "type": "reasoning_step",
+                    "message": f"Step {steps_counter}: {result}",
+                }
+            else:
+                break
+
+        load_elapsed = time.perf_counter() - load_start
+        logging.info("Database load attempt finished in %.2f seconds", load_elapsed)
+
+        if success:
+            yield {
+                "type": "final_result",
+                "success": True,
+                "message": "Database connected and schema loaded successfully",
+            }
+        else:
+            # Don't stream the full internal result; give higher-level error
+            logging.error("Database loader failed: %s", str(result))  # nosemgrep
+            yield {"type": "error", "message": "Failed to load database schema"}
+    except Exception as e:  # pylint: disable=broad-exception-caught
+        logging.exception("Error while loading database schema: %s", str(e))
+        yield {"type": "error", "message": "Error connecting to database"}
+
+
+def _step_result(result) -> str:
+    """Yield the final result message."""
+    return json.dumps(result) + MESSAGE_DELIMITER
+
+
+async def load_database(url: str, user_id: str):
+    """
+    Accepts a JSON payload with a database URL and attempts to connect.
+    Supports both PostgreSQL and MySQL databases.
+    Streams progress steps as a sequence of JSON messages separated by MESSAGE_DELIMITER.
+    """
+
+    # Validate URL format
+    if len(url.strip()) == 0:
+        raise InvalidArgumentError("Invalid URL format")
+
+    async def generate():
+        overall_start = time.perf_counter()
+        steps_counter = 0
+        try:
+            # Step 1: Start
+            steps_counter += 1
+            result = _step_start(steps_counter)
+            yield _step_result(result)
+
+            # Step 2: Determine type
+            steps_counter += 1
+            loader, result = _step_detect_db_type(steps_counter, url)
+            yield _step_result(result)
+
+            # Step 3: Attempt to load schema using the loader
+            async for progress in _step_attempt_load(
+                steps_counter, loader, user_id, url
+            ):
+                yield _step_result(progress)
+
+        except InvalidArgumentError as ia:
+            yield _step_result({"type": "error", "message": str(ia)})
+        except Exception as e:  # pylint: disable=broad-exception-caught
+            logging.exception("Unexpected error in connect_database stream: %s", str(e))
+            yield _step_result({"type": "error", "message": "Internal server error"})
+        finally:
+            overall_elapsed = time.perf_counter() - overall_start
+            logging.info(
+                "connect_database processing completed - Total time: %.2f seconds",
+                overall_elapsed,
+            )
+
+    return generate()
+
+
+async def list_databases(user_id: str, general_prefix: str, db) -> list[str]:
+    """
+    This route is used to list all the graphs (databases names) that are available in the database.
+    """
+    user_graphs = await db.list_graphs()
+
+    # Only include graphs that start with user_id + '_', and strip the prefix
+    filtered_graphs = [
+        graph[len(f"{user_id}_") :]
+        for graph in user_graphs
+        if graph.startswith(f"{user_id}_")
+    ]
+
+    if general_prefix:
+        demo_graphs = [
+            graph for graph in user_graphs if graph.startswith(general_prefix)
+        ]
+        filtered_graphs = filtered_graphs + demo_graphs
+
+    return filtered_graphs