feat: adding search functions and tests (#56)

Author: totoleon

integration.cloudbuild.yaml

@@ -15,8 +15,12 @@
 steps:
   - id: Install dependencies
     name: python:${_VERSION}
-    entrypoint: pip
-    args: ["install", "--user", "-r", "requirements.txt"]
+    entrypoint: /bin/bash
+    args:
+      - -c
+      - |
+        if [[ $_VERSION == "3.8" ]]; then version="-3.8"; fi
+        pip install --user -r requirements${version}.txt
 
   - id: Install module (and test requirements)
     name: python:${_VERSION}

pyproject.toml

@@ -11,6 +11,7 @@ authors = [
 dependencies = [
     "langchain-core>=0.1.1, <1.0.0",
     "langchain-community>=0.0.18, <1.0.0",
+    "numpy>=1.24.4, <2.0.0",
     "SQLAlchemy>=2.0.7, <3.0.0",
     "cloud-sql-python-connector[pymysql]>=1.7.0, <2.0.0"
 ]

requirements-3.8.txt

@@ -0,0 +1,5 @@
+langchain==0.1.12
+langchain-community==0.0.28
+numpy==1.24.4
+SQLAlchemy==2.0.28
+cloud-sql-python-connector[pymysql]==1.8.0

requirements.txt

@@ -1,5 +1,6 @@
 langchain==0.1.12
 langchain-community==0.0.28
+numpy==1.26.4
 SQLAlchemy==2.0.28
 cloud-sql-python-connector[pymysql]==1.8.0
 

src/langchain_google_cloud_sql_mysql/engine.py

@@ -241,6 +241,13 @@ def _fetch(self, query: str, params: Optional[dict] = None):
             result_fetch = result_map.fetchall()
             return result_fetch
 
+    def _fetch_rows(self, query: str, params: Optional[dict] = None):
+        """Fetch results from a SQL query as rows."""
+        with self.engine.connect() as conn:
+            result = conn.execute(sqlalchemy.text(query), params)
+            result_fetch = result.fetchall()  # Directly fetch rows
+            return result_fetch
+
     def init_chat_history_table(self, table_name: str) -> None:
         """Create table with schema required for MySQLChatMessageHistory class.
 

src/langchain_google_cloud_sql_mysql/indexes.py

@@ -29,18 +29,33 @@ class SearchType(Enum):
     ANN = "ANN"
 
 
+class DistanceMeasure(Enum):
+    """Enumerates the types of distance measures that can be used in searches.
+
+    Attributes:
+        COSINE: Cosine similarity measure.
+        L2_SQUARED: Squared L2 norm (Euclidean) distance.
+        DOT_PRODUCT: Dot product similarity.
+    """
+
+    COSINE = "cosine"
+    L2_SQUARED = "l2_squared"
+    DOT_PRODUCT = "dot_product"
+
+
 @dataclass
 class QueryOptions:
     """Holds configuration options for executing a search query.
 
     Attributes:
         num_partitions (Optional[int]): The number of partitions to divide the search space into. None means default partitioning.
-        num_neighbors (Optional[int]): The number of nearest neighbors to retrieve. None means use the default.
+        num_neighbors (int): The number of nearest neighbors to retrieve. Default to 10.
         search_type (SearchType): The type of search algorithm to use. Defaults to KNN.
     """
 
     num_partitions: Optional[int] = None
-    num_neighbors: Optional[int] = None
+    num_neighbors: int = 10
+    distance_measure: DistanceMeasure = DistanceMeasure.L2_SQUARED
     search_type: SearchType = SearchType.KNN
 
 
@@ -61,20 +76,6 @@ class IndexType(Enum):
     TREE_SQ = "TREE_SQ"
 
 
-class DistanceMeasure(Enum):
-    """Enumerates the types of distance measures that can be used in searches.
-
-    Attributes:
-        COSINE: Cosine similarity measure.
-        SQUARED_L2: Squared L2 norm (Euclidean) distance.
-        DOT_PRODUCT: Dot product similarity.
-    """
-
-    COSINE = "cosine"
-    SQUARED_L2 = "squared_l2"
-    DOT_PRODUCT = "dot_product"
-
-
 class VectorIndex:
     """Represents a vector index for storing and querying vectors.
 

src/langchain_google_cloud_sql_mysql/loader.py

@@ -29,7 +29,7 @@ def _parse_doc_from_row(
     content_columns: Iterable[str],
     metadata_columns: Iterable[str],
     row: Dict,
-    metadata_json_column: str = DEFAULT_METADATA_COL,
+    metadata_json_column: Optional[str] = DEFAULT_METADATA_COL,
 ) -> Document:
     page_content = " ".join(
         str(row[column]) for column in content_columns if column in row