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