ci: Add mypy function type check (#130)

* ci: Add mypy function type check

* fix typing

* fix rebase conflict

* more typing fix

* fix type

* moreeee fix

---------

Co-authored-by: Averi Kitsch <akitsch@google.com>

Author: duwenxin99

pyproject.toml

@@ -62,6 +62,8 @@ profile = "black"
 [tool.mypy]
 python_version = 3.8
 warn_unused_configs = true
+disallow_incomplete_defs = true
+
 exclude = [
     'docs/*',
     'noxfile.py'

src/langchain_google_cloud_sql_pg/chat_message_history.py

@@ -44,7 +44,7 @@ class PostgresChatMessageHistory(BaseChatMessageHistory):
 
     def __init__(
         self,
-        key,
+        key: object,
         engine: PostgresEngine,
         session_id: str,
         table_name: str,
@@ -77,7 +77,7 @@ async def create(
         engine: PostgresEngine,
         session_id: str,
         table_name: str,
-    ):
+    ) -> PostgresChatMessageHistory:
         """Create a new PostgresChatMessageHistory instance.
 
         Args:
@@ -117,7 +117,7 @@ def create_sync(
         engine: PostgresEngine,
         session_id: str,
         table_name: str,
-    ):
+    ) -> PostgresChatMessageHistory:
         """Create a new PostgresChatMessageHistory instance.
 
         Args:

src/langchain_google_cloud_sql_pg/engine.py

@@ -17,13 +17,23 @@
 import asyncio
 from dataclasses import dataclass
 from threading import Thread
-from typing import TYPE_CHECKING, Awaitable, Dict, List, Optional, TypeVar, Union
+from typing import (
+    TYPE_CHECKING,
+    Awaitable,
+    Dict,
+    List,
+    Optional,
+    Sequence,
+    TypeVar,
+    Union,
+)
 
 import aiohttp
 import google.auth  # type: ignore
 import google.auth.transport.requests  # type: ignore
 from google.cloud.sql.connector import Connector, IPTypes, RefreshStrategy
 from sqlalchemy import MetaData, Table, text
+from sqlalchemy.engine.row import RowMapping
 from sqlalchemy.exc import InvalidRequestError
 from sqlalchemy.ext.asyncio import AsyncEngine, create_async_engine
 
@@ -305,19 +315,21 @@ def from_engine(cls, engine: AsyncEngine) -> PostgresEngine:
         """Create an PostgresEngine instance from an AsyncEngine."""
         return cls(cls.__create_key, engine, None, None)
 
-    async def _aexecute(self, query: str, params: Optional[dict] = None):
+    async def _aexecute(self, query: str, params: Optional[dict] = None) -> None:
         """Execute a SQL query."""
         async with self._engine.connect() as conn:
             await conn.execute(text(query), params)
             await conn.commit()
 
-    async def _aexecute_outside_tx(self, query: str):
+    async def _aexecute_outside_tx(self, query: str) -> None:
         """Execute a SQL query."""
         async with self._engine.connect() as conn:
             await conn.execute(text("COMMIT"))
             await conn.execute(text(query))
 
-    async def _afetch(self, query: str, params: Optional[dict] = None):
+    async def _afetch(
+        self, query: str, params: Optional[dict] = None
+    ) -> Sequence[RowMapping]:
         """Fetch results from a SQL query."""
         async with self._engine.connect() as conn:
             result = await conn.execute(text(query), params)
@@ -326,11 +338,11 @@ async def _afetch(self, query: str, params: Optional[dict] = None):
 
         return result_fetch
 
-    def _execute(self, query: str, params: Optional[dict] = None):
+    def _execute(self, query: str, params: Optional[dict] = None) -> None:
         """Execute a SQL query."""
         return self._run_as_sync(self._aexecute(query, params))
 
-    def _fetch(self, query: str, params: Optional[dict] = None):
+    def _fetch(self, query: str, params: Optional[dict] = None) -> Sequence[RowMapping]:
         """Fetch results from a SQL query."""
         return self._run_as_sync(self._afetch(query, params))
 
@@ -439,7 +451,7 @@ def init_vectorstore_table(
             )
         )
 
-    async def ainit_chat_history_table(self, table_name) -> None:
+    async def ainit_chat_history_table(self, table_name: str) -> None:
         """Create a Cloud SQL table to store chat history.
 
         Args:
@@ -456,7 +468,7 @@ async def ainit_chat_history_table(self, table_name) -> None:
         );"""
         await self._aexecute(create_table_query)
 
-    def init_chat_history_table(self, table_name) -> None:
+    def init_chat_history_table(self, table_name: str) -> None:
         """Create a Cloud SQL table to store chat history.
 
         Args:

src/langchain_google_cloud_sql_pg/loader.py

@@ -36,24 +36,24 @@
 DEFAULT_METADATA_COL = "langchain_metadata"
 
 
-def text_formatter(row, content_columns) -> str:
+def text_formatter(row: dict, content_columns: List[str]) -> str:
     """txt document formatter."""
     return " ".join(str(row[column]) for column in content_columns if column in row)
 
 
-def csv_formatter(row, content_columns) -> str:
+def csv_formatter(row: dict, content_columns: List[str]) -> str:
     """CSV document formatter."""
     return ", ".join(str(row[column]) for column in content_columns if column in row)
 
 
-def yaml_formatter(row, content_columns) -> str:
+def yaml_formatter(row: dict, content_columns: List[str]) -> str:
     """YAML document formatter."""
     return "\n".join(
         f"{column}: {str(row[column])}" for column in content_columns if column in row
     )
 
 
-def json_formatter(row, content_columns) -> str:
+def json_formatter(row: dict, content_columns: List[str]) -> str:
     """JSON document formatter."""
     dictionary = {}
     for column in content_columns:
@@ -116,7 +116,7 @@ class PostgresLoader(BaseLoader):
 
     def __init__(
         self,
-        key,
+        key: object,
         engine: PostgresEngine,
         query: str,
         content_columns: List[str],
@@ -162,7 +162,7 @@ async def create(
         metadata_json_column: Optional[str] = None,
         format: Optional[str] = None,
         formatter: Optional[Callable] = None,
-    ):
+    ) -> PostgresLoader:
         """Create a new PostgresLoader instance.
 
         Args:
@@ -255,7 +255,7 @@ def create_sync(
         metadata_json_column: Optional[str] = None,
         format: Optional[str] = None,
         formatter: Optional[Callable] = None,
-    ):
+    ) -> PostgresLoader:
         """Create a new PostgresLoader instance.
 
         Args:
@@ -340,7 +340,7 @@ class PostgresDocumentSaver:
 
     def __init__(
         self,
-        key,
+        key: object,
         engine: PostgresEngine,
         table_name: str,
         content_column: str,
@@ -378,7 +378,7 @@ async def create(
         content_column: str = DEFAULT_CONTENT_COL,
         metadata_columns: List[str] = [],
         metadata_json_column: Optional[str] = DEFAULT_METADATA_COL,
-    ):
+    ) -> PostgresDocumentSaver:
         """Create an PostgresDocumentSaver instance.
 
         Args:
@@ -435,7 +435,7 @@ def create_sync(
         content_column: str = DEFAULT_CONTENT_COL,
         metadata_columns: List[str] = [],
         metadata_json_column: str = DEFAULT_METADATA_COL,
-    ):
+    ) -> PostgresDocumentSaver:
         """Create an PostgresDocumentSaver instance.
 
         Args:

src/langchain_google_cloud_sql_pg/vectorstore.py

@@ -17,12 +17,13 @@
 
 import json
 import uuid
-from typing import Any, Callable, Iterable, List, Optional, Tuple, Type, Union
+from typing import Any, Callable, Iterable, List, Optional, Sequence, Tuple, Type, Union
 
 import numpy as np
 from langchain_core.documents import Document
 from langchain_core.embeddings import Embeddings
 from langchain_core.vectorstores import VectorStore
+from sqlalchemy.engine.row import RowMapping
 
 from .engine import PostgresEngine
 from .indexes import (
@@ -42,7 +43,7 @@ class PostgresVectorStore(VectorStore):
 
     def __init__(
         self,
-        key,
+        key: object,
         engine: PostgresEngine,
         embedding_service: Embeddings,
         table_name: str,
@@ -114,7 +115,7 @@ async def create(
         fetch_k: int = 20,
         lambda_mult: float = 0.5,
         index_query_options: Optional[QueryOptions] = None,
-    ):
+    ) -> PostgresVectorStore:
         """Create a new PostgresVectorStore instance.
 
         Args:
@@ -218,7 +219,7 @@ def create_sync(
         fetch_k: int = 20,
         lambda_mult: float = 0.5,
         index_query_options: Optional[QueryOptions] = None,
-    ):
+    ) -> PostgresVectorStore:
         """Create a new PostgresVectorStore instance.
 
         Args:
@@ -496,7 +497,7 @@ def from_texts(  # type: ignore[override]
         id_column: str = "langchain_id",
         metadata_json_column: str = "langchain_metadata",
         **kwargs: Any,
-    ):
+    ) -> PostgresVectorStore:
         """Create an PostgresVectorStore instance from texts.
         Args:
             texts (List[str]): Texts to add to the vector store.
@@ -589,7 +590,7 @@ async def __query_collection(
         k: Optional[int] = None,
         filter: Optional[str] = None,
         **kwargs: Any,
-    ) -> List[Any]:
+    ) -> Sequence[RowMapping]:
         """Perform similarity search query on the vector store table."""
         k = k if k else self.k
         operator = self.distance_strategy.operator

tests/test_postgresql_chatmessagehistory.py

@@ -79,14 +79,14 @@ def test_chat_message_history(memory_engine: PostgresEngine) -> None:
     assert len(history.messages) == 0
 
 
-def test_chat_table(memory_engine: Any):
+def test_chat_table(memory_engine: Any) -> None:
     with pytest.raises(ValueError):
         PostgresChatMessageHistory.create_sync(
             engine=memory_engine, session_id="test", table_name="doesnotexist"
         )
 
 
-def test_chat_schema(memory_engine: Any):
+def test_chat_schema(memory_engine: Any) -> None:
     doc_table_name = "test_table" + str(uuid.uuid4())
     memory_engine.init_document_table(table_name=doc_table_name)
     with pytest.raises(IndexError):