querychat 0.2.1__tar.gz → 0.2.2__tar.gz

{querychat-0.2.1 → querychat-0.2.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: querychat
-Version: 0.2.1
+Version: 0.2.2
 Summary: Chat with your data using natural language
 Project-URL: Homepage, https://github.com/posit-dev/querychat
 Project-URL: Repository, https://github.com/posit-dev/querychat

{querychat-0.2.1 → querychat-0.2.2}/pkg-py/src/querychat/datasource.py

@@ -3,12 +3,13 @@ from __future__ import annotations
 from typing import TYPE_CHECKING, ClassVar, Protocol
 
 import duckdb
-import narwhals as nw
+import narwhals.stable.v1 as nw
 import pandas as pd
 from sqlalchemy import inspect, text
 from sqlalchemy.sql import sqltypes
 
 if TYPE_CHECKING:
+    from narwhals.stable.v1.typing import IntoFrame
     from sqlalchemy.engine import Connection, Engine
 
 
@@ -58,8 +59,9 @@ class DataFrameSource:
     """A DataSource implementation that wraps a pandas DataFrame using DuckDB."""
 
     db_engine: ClassVar[str] = "DuckDB"
+    _df: nw.DataFrame | nw.LazyFrame
 
-    def __init__(self, df: pd.DataFrame, table_name: str):
+    def __init__(self, df: IntoFrame, table_name: str):
         """
         Initialize with a pandas DataFrame.
 
@@ -69,9 +71,10 @@ class DataFrameSource:
 
         """
         self._conn = duckdb.connect(database=":memory:")
-        self._df = df
+        self._df = nw.from_native(df)
         self._table_name = table_name
-        self._conn.register(table_name, df)
+        # TODO(@gadenbuie): If the data frame is already SQL-backed, maybe we shouldn't be making a new copy here.
+        self._conn.register(table_name, self._df.lazy().collect().to_pandas())
 
     def get_schema(self, *, categorical_threshold: int) -> str:
         """
@@ -86,10 +89,15 @@ class DataFrameSource:
             String describing the schema
 
         """
-        ndf = nw.from_native(self._df)
-
         schema = [f"Table: {self._table_name}", "Columns:"]
 
+        # Ensure we're working with a DataFrame, not a LazyFrame
+        ndf = (
+            self._df.head(10).collect()
+            if isinstance(self._df, nw.LazyFrame)
+            else self._df
+        )
+
         for column in ndf.columns:
             # Map pandas dtypes to SQL-like types
             dtype = ndf[column].dtype
@@ -149,7 +157,8 @@ class DataFrameSource:
             The complete dataset as a pandas DataFrame
 
         """
-        return self._df.copy()
+        # TODO(@gadenbuie): This should just return `self._df` and not a pandas DataFrame
+        return self._df.lazy().collect().to_pandas()
 
 
 class SQLAlchemySource:

{querychat-0.2.1 → querychat-0.2.2}/pkg-py/src/querychat/querychat.py

@@ -414,15 +414,12 @@ def init(
     data_source_obj: DataSource
     if isinstance(data_source, sqlalchemy.Engine):
         data_source_obj = SQLAlchemySource(data_source, table_name)
-    elif isinstance(data_source, (nw.DataFrame, nw.LazyFrame)):
+    else:
         data_source_obj = DataFrameSource(
-            nw.to_native(data_source),
+            data_source,
             table_name,
         )
-    else:
-        raise TypeError(
-            "`data_source` must be a Narwhals DataFrame or LazyFrame, or a SQLAlchemy Engine",
-        )
+
     # Process greeting
     if greeting is None:
         print(

{querychat-0.2.1 → querychat-0.2.2}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "querychat"
-version = "0.2.1"
+version = "0.2.2"
 description = "Chat with your data using natural language"
 readme = "pkg-py/README.md"
 requires-python = ">=3.9"
@@ -124,6 +124,7 @@ extend-ignore = [
     "D107", # Missing docstring in __init__
     "D205", # 1 blank line required between summary line and description
     "UP045", # Use `X | NULL` for type annotations, not `Optional[X]`
+    "TD003", # TODO doesn't need to have an issue link
 ]
 extend-select = [
     # "C90", # C90; mccabe: https://docs.astral.sh/ruff/rules/complex-structure/