diff --git a/samples/index_tuning_sample/README.md b/samples/index_tuning_sample/README.md index 519b5830..5140152f 100644 --- a/samples/index_tuning_sample/README.md +++ b/samples/index_tuning_sample/README.md @@ -187,7 +187,7 @@ class HNSWIndex( index_type: str = "hnsw", # Distance strategy does not affect recall and has minimal little on latency; refer to this guide to learn more https://cloud.google.com/spanner/docs/choose-vector-distance-function distance_strategy: DistanceStrategy = lambda : DistanceStrategy.COSINE_DISTANCE, - partial_indexes: List[str] | None = None, + partial_indexes: list[str] | None = None, m: int = 16, ef_construction: int = 64 ) @@ -235,7 +235,7 @@ class IVFFlatIndex( name: str = DEFAULT_INDEX_NAME, index_type: str = "ivfflat", distance_strategy: DistanceStrategy = lambda : DistanceStrategy.COSINE_DISTANCE, - partial_indexes: List[str] | None = None, + partial_indexes: list[str] | None = None, lists: int = 1 ) diff --git a/samples/langchain_on_vertexai/prebuilt_langchain_agent_template.py b/samples/langchain_on_vertexai/prebuilt_langchain_agent_template.py index ea8ef509..cb2f2cf6 100644 --- a/samples/langchain_on_vertexai/prebuilt_langchain_agent_template.py +++ b/samples/langchain_on_vertexai/prebuilt_langchain_agent_template.py @@ -12,7 +12,6 @@ # See the License for the specific language governing permissions and # limitations under the License. import os -from typing import List import vertexai # type: ignore from config import ( @@ -39,14 +38,14 @@ engine = None # Use global variable to share connection pooling -def similarity_search(query: str) -> List[Document]: +def similarity_search(query: str) -> list[Document]: """Searches and returns movies. Args: query: The user query to search for related items Returns: - List[Document]: A list of Documents + list[Document]: A list of Documents """ global engine if not engine: # Reuse connection pool diff --git a/src/langchain_google_alloydb_pg/async_chat_message_history.py b/src/langchain_google_alloydb_pg/async_chat_message_history.py index 1b668e2b..e2e6e72b 100644 --- a/src/langchain_google_alloydb_pg/async_chat_message_history.py +++ b/src/langchain_google_alloydb_pg/async_chat_message_history.py @@ -15,7 +15,7 @@ from __future__ import annotations import json -from typing import List, Sequence +from typing import Sequence from langchain_core.chat_history import BaseChatMessageHistory from langchain_core.messages import BaseMessage, messages_from_dict @@ -128,7 +128,7 @@ async def aclear(self) -> None: await conn.execute(text(query), {"session_id": self.session_id}) await conn.commit() - async def _aget_messages(self) -> List[BaseMessage]: + async def _aget_messages(self) -> list[BaseMessage]: """Retrieve the messages from AlloyDB.""" query = f"""SELECT data, type FROM "{self.schema_name}"."{self.table_name}" WHERE session_id = :session_id ORDER BY id;""" async with self.pool.connect() as conn: diff --git a/src/langchain_google_alloydb_pg/async_loader.py b/src/langchain_google_alloydb_pg/async_loader.py index cc34ec50..7611faf9 100644 --- a/src/langchain_google_alloydb_pg/async_loader.py +++ b/src/langchain_google_alloydb_pg/async_loader.py @@ -15,7 +15,7 @@ from __future__ import annotations import json -from typing import Any, AsyncIterator, Callable, Dict, Iterable, List, Optional +from typing import Any, AsyncIterator, Callable, Iterable, Optional from langchain_core.document_loaders.base import BaseLoader from langchain_core.documents import Document @@ -28,24 +28,24 @@ DEFAULT_METADATA_COL = "langchain_metadata" -def text_formatter(row: dict, content_columns: List[str]) -> 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: dict, content_columns: List[str]) -> 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: dict, content_columns: List[str]) -> 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: dict, content_columns: List[str]) -> str: +def json_formatter(row: dict, content_columns: list[str]) -> str: """JSON document formatter.""" dictionary = {} for column in content_columns: @@ -63,7 +63,7 @@ def _parse_doc_from_row( ) -> Document: """Parse row into document.""" page_content = formatter(row, content_columns) - metadata: Dict[str, Any] = {} + metadata: dict[str, Any] = {} # unnest metadata from langchain_metadata column if metadata_json_column and row.get(metadata_json_column): for k, v in row[metadata_json_column].items(): @@ -81,10 +81,10 @@ def _parse_row_from_doc( column_names: Iterable[str], content_column: str = DEFAULT_CONTENT_COL, metadata_json_column: Optional[str] = DEFAULT_METADATA_COL, -) -> Dict: +) -> dict: """Parse document into a dictionary of rows.""" doc_metadata = doc.metadata.copy() - row: Dict[str, Any] = {content_column: doc.page_content} + row: dict[str, Any] = {content_column: doc.page_content} for entry in doc.metadata: if entry in column_names: row[entry] = doc_metadata[entry] @@ -111,8 +111,8 @@ def __init__( key: object, pool: AsyncEngine, query: str, - content_columns: List[str], - metadata_columns: List[str], + content_columns: list[str], + metadata_columns: list[str], formatter: Callable, metadata_json_column: Optional[str] = None, ) -> None: @@ -122,8 +122,8 @@ def __init__( key (object): Prevent direct constructor usage. engine (AlloyDBEngine): AsyncEngine with pool connection to the postgres database query (Optional[str], optional): SQL query. Defaults to None. - content_columns (Optional[List[str]], optional): Column that represent a Document's page_content. Defaults to the first column. - metadata_columns (Optional[List[str]], optional): Column(s) that represent a Document's metadata. Defaults to None. + content_columns (Optional[list[str]], optional): Column that represent a Document's page_content. Defaults to the first column. + metadata_columns (Optional[list[str]], optional): Column(s) that represent a Document's metadata. Defaults to None. formatter (Optional[Callable], optional): A function to format page content (OneOf: format, formatter). Defaults to None. metadata_json_column (Optional[str], optional): Column to store metadata as JSON. Defaults to "langchain_metadata". @@ -149,8 +149,8 @@ async def create( query: Optional[str] = None, table_name: Optional[str] = None, schema_name: str = "public", - content_columns: Optional[List[str]] = None, - metadata_columns: Optional[List[str]] = None, + content_columns: Optional[list[str]] = None, + metadata_columns: Optional[list[str]] = None, metadata_json_column: Optional[str] = None, format: Optional[str] = None, formatter: Optional[Callable] = None, @@ -162,8 +162,8 @@ async def create( query (Optional[str], optional): SQL query. Defaults to None. table_name (Optional[str], optional): Name of table to query. Defaults to None. schema_name (str, optional): Name of the schema where table is located. Defaults to "public". - content_columns (Optional[List[str]], optional): Column that represent a Document's page_content. Defaults to the first column. - metadata_columns (Optional[List[str]], optional): Column(s) that represent a Document's metadata. Defaults to None. + content_columns (Optional[list[str]], optional): Column that represent a Document's page_content. Defaults to the first column. + metadata_columns (Optional[list[str]], optional): Column(s) that represent a Document's metadata. Defaults to None. metadata_json_column (Optional[str], optional): Column to store metadata as JSON. Defaults to "langchain_metadata". format (Optional[str], optional): Format of page content (OneOf: text, csv, YAML, JSON). Defaults to 'text'. formatter (Optional[Callable], optional): A function to format page content (OneOf: format, formatter). Defaults to None. @@ -236,7 +236,7 @@ async def create( metadata_json_column, ) - async def aload(self) -> List[Document]: + async def aload(self) -> list[Document]: """Load PostgreSQL data into Document objects.""" return [doc async for doc in self.alazy_load()] @@ -282,7 +282,7 @@ def __init__( table_name: str, content_column: str, schema_name: str = "public", - metadata_columns: List[str] = [], + metadata_columns: list[str] = [], metadata_json_column: Optional[str] = None, ): """AsyncAlloyDBDocumentSaver constructor. @@ -293,7 +293,7 @@ def __init__( table_name (str): Name of table to query. schema_name (str, optional): Name of schema where the table is located. Defaults to "public". content_column (str, optional): Column that represent a Document's page_content. Defaults to "page_content". - metadata_columns (Optional[List[str]], optional): Column(s) that represent a Document's metadata. Defaults to an empty list. + metadata_columns (Optional[list[str]], optional): Column(s) that represent a Document's metadata. Defaults to an empty list. metadata_json_column (str, optional): Column to store metadata as JSON. Defaults to "langchain_metadata". Raises: @@ -317,7 +317,7 @@ async def create( table_name: str, schema_name: str = "public", content_column: str = DEFAULT_CONTENT_COL, - metadata_columns: List[str] = [], + metadata_columns: list[str] = [], metadata_json_column: Optional[str] = DEFAULT_METADATA_COL, ) -> AsyncAlloyDBDocumentSaver: """Create an AsyncAlloyDBDocumentSaver instance. @@ -327,7 +327,7 @@ async def create( table_name (str): Name of table to query. schema_name (str, optional): Name of schema where the table is located. Defaults to "public". content_column (str, optional): Column that represent a Document's page_content. Defaults to "page_content". - metadata_columns (List[str], optional): Column(s) that represent a Document's metadata. Defaults to an empty list. + metadata_columns (list[str], optional): Column(s) that represent a Document's metadata. Defaults to an empty list. metadata_json_column (Optional[str], optional): Column to store metadata as JSON. Defaults to "langchain_metadata". Returns: @@ -370,13 +370,13 @@ async def create( metadata_json_column, ) - async def aadd_documents(self, docs: List[Document]) -> None: + async def aadd_documents(self, docs: list[Document]) -> None: """ Save documents in the DocumentSaver table. Document’s metadata is added to columns if found or stored in langchain_metadata JSON column. Args: - docs (List[langchain_core.documents.Document]): a list of documents to be saved. + docs (list[langchain_core.documents.Document]): a list of documents to be saved. """ for doc in docs: @@ -414,13 +414,13 @@ async def aadd_documents(self, docs: List[Document]) -> None: await conn.execute(text(query), row) await conn.commit() - async def adelete(self, docs: List[Document]) -> None: + async def adelete(self, docs: list[Document]) -> None: """ Delete all instances of a document from the DocumentSaver table by matching the entire Document object. Args: - docs (List[langchain_core.documents.Document]): a list of documents to be deleted. + docs (list[langchain_core.documents.Document]): a list of documents to be deleted. """ for doc in docs: row = _parse_row_from_doc( diff --git a/src/langchain_google_alloydb_pg/async_vectorstore.py b/src/langchain_google_alloydb_pg/async_vectorstore.py index fb3bfc1e..36cd4e2f 100644 --- a/src/langchain_google_alloydb_pg/async_vectorstore.py +++ b/src/langchain_google_alloydb_pg/async_vectorstore.py @@ -19,7 +19,7 @@ import json import re import uuid -from typing import Any, Callable, Iterable, List, Optional, Sequence, Tuple, Type +from typing import Any, Callable, Iterable, Optional, Sequence import numpy as np import requests @@ -57,7 +57,7 @@ def __init__( schema_name: str = "public", content_column: str = "content", embedding_column: str = "embedding", - metadata_columns: List[str] = [], + metadata_columns: list[str] = [], id_column: str = "langchain_id", metadata_json_column: Optional[str] = "langchain_metadata", distance_strategy: DistanceStrategy = DEFAULT_DISTANCE_STRATEGY, @@ -75,7 +75,7 @@ def __init__( schema_name (str, optional): Name of the database schema. Defaults to "public". content_column (str): Column that represent a Document’s page_content. Defaults to "content". embedding_column (str): Column for embedding vectors. The embedding is generated from the document value. Defaults to "embedding". - metadata_columns (List[str]): Column(s) that represent a document's metadata. + metadata_columns (list[str]): Column(s) that represent a document's metadata. id_column (str): Column that represents the Document's id. Defaults to "langchain_id". metadata_json_column (str): Column to store metadata as JSON. Defaults to "langchain_metadata". distance_strategy (DistanceStrategy): Distance strategy to use for vector similarity search. Defaults to COSINE_DISTANCE. @@ -110,15 +110,15 @@ def __init__( @classmethod async def create( - cls: Type[AsyncAlloyDBVectorStore], + cls: type[AsyncAlloyDBVectorStore], engine: AlloyDBEngine, embedding_service: Embeddings, table_name: str, schema_name: str = "public", content_column: str = "content", embedding_column: str = "embedding", - metadata_columns: List[str] = [], - ignore_metadata_columns: Optional[List[str]] = None, + metadata_columns: list[str] = [], + ignore_metadata_columns: Optional[list[str]] = None, id_column: str = "langchain_id", metadata_json_column: Optional[str] = "langchain_metadata", distance_strategy: DistanceStrategy = DEFAULT_DISTANCE_STRATEGY, @@ -136,8 +136,8 @@ async def create( schema_name (str, optional): Name of the database schema. Defaults to "public". content_column (str): Column that represent a Document’s page_content. Defaults to "content". embedding_column (str): Column for embedding vectors. The embedding is generated from the document value. Defaults to "embedding". - metadata_columns (List[str]): Column(s) that represent a document's metadata. - ignore_metadata_columns (List[str]): Column(s) to ignore in pre-existing tables for a document's metadata. Can not be used with metadata_columns. Defaults to None. + metadata_columns (list[str]): Column(s) that represent a document's metadata. + ignore_metadata_columns (list[str]): Column(s) to ignore in pre-existing tables for a document's metadata. Can not be used with metadata_columns. Defaults to None. id_column (str): Column that represents the Document's id. Defaults to "langchain_id". metadata_json_column (str): Column to store metadata as JSON. Defaults to "langchain_metadata". distance_strategy (DistanceStrategy): Distance strategy to use for vector similarity search. Defaults to COSINE_DISTANCE. @@ -225,11 +225,11 @@ def embeddings(self) -> Embeddings: async def aadd_embeddings( self, texts: Iterable[str], - embeddings: List[List[float]], - metadatas: Optional[List[dict]] = None, - ids: Optional[List] = None, + embeddings: list[list[float]], + metadatas: Optional[list[dict]] = None, + ids: Optional[list] = None, **kwargs: Any, - ) -> List[str]: + ) -> list[str]: """Add data along with embeddings to the table. Raises: @@ -282,17 +282,17 @@ async def aadd_embeddings( async def aadd_texts( self, texts: Iterable[str], - metadatas: Optional[List[dict]] = None, - ids: Optional[List] = None, + metadatas: Optional[list[dict]] = None, + ids: Optional[list] = None, **kwargs: Any, - ) -> List[str]: + ) -> list[str]: """Embed texts and add to the table. Raises: :class:`InvalidTextRepresentationError `: if the `ids` data type does not match that of the `id_column`. """ if isinstance(self.embedding_service, AlloyDBEmbeddings): - embeddings: List[List[float]] = [[] for _ in list(texts)] + embeddings: list[list[float]] = [[] for _ in list(texts)] else: embeddings = await self.embedding_service.aembed_documents(list(texts)) @@ -303,10 +303,10 @@ async def aadd_texts( async def aadd_documents( self, - documents: List[Document], - ids: Optional[List] = None, + documents: list[Document], + ids: Optional[list] = None, **kwargs: Any, - ) -> List[str]: + ) -> list[str]: """Embed documents and add to the table. Raises: @@ -338,17 +338,17 @@ def _encode_image(self, uri: str) -> str: async def aadd_images( self, - uris: List[str], - metadatas: Optional[List[dict]] = None, - ids: Optional[List[str]] = None, + uris: list[str], + metadatas: Optional[list[dict]] = None, + ids: Optional[list[str]] = None, **kwargs: Any, - ) -> List[str]: + ) -> list[str]: """Embed images and add to the table. Args: - uris (List[str]): List of local image URIs to add to the table. - metadatas (Optional[List[dict]]): List of metadatas to add to table records. - ids: (Optional[List[str]]): List of IDs to add to table records. + uris (list[str]): List of local image URIs to add to the table. + metadatas (Optional[list[dict]]): List of metadatas to add to table records. + ids: (Optional[list[str]]): List of IDs to add to table records. Returns: List of record IDs added. @@ -369,7 +369,7 @@ async def aadd_images( async def adelete( self, - ids: Optional[List] = None, + ids: Optional[list] = None, **kwargs: Any, ) -> Optional[bool]: """Delete records from the table. @@ -389,18 +389,18 @@ async def adelete( @classmethod async def afrom_texts( # type: ignore[override] - cls: Type[AsyncAlloyDBVectorStore], - texts: List[str], + cls: type[AsyncAlloyDBVectorStore], + texts: list[str], embedding: Embeddings, engine: AlloyDBEngine, table_name: str, schema_name: str = "public", - metadatas: Optional[List[dict]] = None, - ids: Optional[List] = None, + metadatas: Optional[list[dict]] = None, + ids: Optional[list] = None, content_column: str = "content", embedding_column: str = "embedding", - metadata_columns: List[str] = [], - ignore_metadata_columns: Optional[List[str]] = None, + metadata_columns: list[str] = [], + ignore_metadata_columns: Optional[list[str]] = None, id_column: str = "langchain_id", metadata_json_column: str = "langchain_metadata", distance_strategy: DistanceStrategy = DEFAULT_DISTANCE_STRATEGY, @@ -413,16 +413,16 @@ async def afrom_texts( # type: ignore[override] """Create an AsyncAlloyDBVectorStore instance from texts. Args: - texts (List[str]): Texts to add to the vector store. + texts (list[str]): Texts to add to the vector store. embedding (Embeddings): Text embedding model to use. engine (AlloyDBEngine): Connection pool engine for managing connections to AlloyDB database. table_name (str): Name of an existing table. - metadatas (Optional[List[dict]]): List of metadatas to add to table records. - ids: (Optional[List[str]]): List of IDs to add to table records. + metadatas (Optional[list[dict]]): List of metadatas to add to table records. + ids: (Optional[list[str]]): List of IDs to add to table records. content_column (str): Column that represent a Document’s page_content. Defaults to "content". embedding_column (str): Column for embedding vectors. The embedding is generated from the document value. Defaults to "embedding". - metadata_columns (List[str]): Column(s) that represent a document's metadata. - ignore_metadata_columns (List[str]): Column(s) to ignore in pre-existing tables for a document's metadata. Can not be used with metadata_columns. Defaults to None. + metadata_columns (list[str]): Column(s) that represent a document's metadata. + ignore_metadata_columns (list[str]): Column(s) to ignore in pre-existing tables for a document's metadata. Can not be used with metadata_columns. Defaults to None. id_column (str): Column that represents the Document's id. Defaults to "langchain_id". metadata_json_column (str): Column to store metadata as JSON. Defaults to "langchain_metadata". distance_strategy (DistanceStrategy): Distance strategy to use for vector similarity search. Defaults to COSINE_DISTANCE. @@ -459,17 +459,17 @@ async def afrom_texts( # type: ignore[override] @classmethod async def afrom_documents( # type: ignore[override] - cls: Type[AsyncAlloyDBVectorStore], - documents: List[Document], + cls: type[AsyncAlloyDBVectorStore], + documents: list[Document], embedding: Embeddings, engine: AlloyDBEngine, table_name: str, schema_name: str = "public", - ids: Optional[List] = None, + ids: Optional[list] = None, content_column: str = "content", embedding_column: str = "embedding", - metadata_columns: List[str] = [], - ignore_metadata_columns: Optional[List[str]] = None, + metadata_columns: list[str] = [], + ignore_metadata_columns: Optional[list[str]] = None, id_column: str = "langchain_id", metadata_json_column: str = "langchain_metadata", distance_strategy: DistanceStrategy = DEFAULT_DISTANCE_STRATEGY, @@ -482,16 +482,16 @@ async def afrom_documents( # type: ignore[override] """Create an AsyncAlloyDBVectorStore instance from documents. Args: - documents (List[Document]): Documents to add to the vector store. + documents (list[Document]): Documents to add to the vector store. embedding (Embeddings): Text embedding model to use. engine (AlloyDBEngine): Connection pool engine for managing connections to AlloyDB database. table_name (str): Name of an existing table. - metadatas (Optional[List[dict]]): List of metadatas to add to table records. - ids: (Optional[List[str]]): List of IDs to add to table records. + metadatas (Optional[list[dict]]): List of metadatas to add to table records. + ids: (Optional[list[str]]): List of IDs to add to table records. content_column (str): Column that represent a Document’s page_content. Defaults to "content". embedding_column (str): Column for embedding vectors. The embedding is generated from the document value. Defaults to "embedding". - metadata_columns (List[str]): Column(s) that represent a document's metadata. - ignore_metadata_columns (List[str]): Column(s) to ignore in pre-existing tables for a document's metadata. Can not be used with metadata_columns. Defaults to None. + metadata_columns (list[str]): Column(s) that represent a document's metadata. + ignore_metadata_columns (list[str]): Column(s) to ignore in pre-existing tables for a document's metadata. Can not be used with metadata_columns. Defaults to None. id_column (str): Column that represents the Document's id. Defaults to "langchain_id". metadata_json_column (str): Column to store metadata as JSON. Defaults to "langchain_metadata". distance_strategy (DistanceStrategy): Distance strategy to use for vector similarity search. Defaults to COSINE_DISTANCE. @@ -531,7 +531,7 @@ async def afrom_documents( # type: ignore[override] async def __query_collection( self, - embedding: List[float], + embedding: list[float], k: Optional[int] = None, filter: Optional[str] = None, **kwargs: Any, @@ -571,7 +571,7 @@ async def asimilarity_search( k: Optional[int] = None, filter: Optional[str] = None, **kwargs: Any, - ) -> List[Document]: + ) -> list[Document]: """Return docs selected by similarity search on query.""" embedding = ( [] @@ -584,7 +584,7 @@ async def asimilarity_search( embedding=embedding, k=k, filter=filter, **kwargs ) - def _images_embedding_helper(self, image_uris: List[str]) -> List[List[float]]: + def _images_embedding_helper(self, image_uris: list[str]) -> list[list[float]]: # check if either `embed_images()` or `embed_image()` API is supported by the embedding service used if hasattr(self.embedding_service, "embed_images"): try: @@ -612,7 +612,7 @@ async def asimilarity_search_image( k: Optional[int] = None, filter: Optional[str] = None, **kwargs: Any, - ) -> List[Document]: + ) -> list[Document]: """Return docs selected by similarity search on query.""" embedding = self._images_embedding_helper([image_uri])[0] @@ -637,7 +637,7 @@ async def asimilarity_search_with_score( k: Optional[int] = None, filter: Optional[str] = None, **kwargs: Any, - ) -> List[Tuple[Document, float]]: + ) -> list[tuple[Document, float]]: """Return docs and distance scores selected by similarity search on query.""" embedding = ( [] @@ -653,11 +653,11 @@ async def asimilarity_search_with_score( async def asimilarity_search_by_vector( self, - embedding: List[float], + embedding: list[float], k: Optional[int] = None, filter: Optional[str] = None, **kwargs: Any, - ) -> List[Document]: + ) -> list[Document]: """Return docs selected by vector similarity search.""" docs_and_scores = await self.asimilarity_search_with_score_by_vector( embedding=embedding, k=k, filter=filter, **kwargs @@ -667,11 +667,11 @@ async def asimilarity_search_by_vector( async def asimilarity_search_with_score_by_vector( self, - embedding: List[float], + embedding: list[float], k: Optional[int] = None, filter: Optional[str] = None, **kwargs: Any, - ) -> List[Tuple[Document, float]]: + ) -> list[tuple[Document, float]]: """Return docs and distance scores selected by vector similarity search.""" results = await self.__query_collection( embedding=embedding, k=k, filter=filter, **kwargs @@ -706,7 +706,7 @@ async def amax_marginal_relevance_search( lambda_mult: Optional[float] = None, filter: Optional[str] = None, **kwargs: Any, - ) -> List[Document]: + ) -> list[Document]: """Return docs selected using the maximal marginal relevance.""" embedding = await self.embedding_service.aembed_query(text=query) @@ -721,13 +721,13 @@ async def amax_marginal_relevance_search( async def amax_marginal_relevance_search_by_vector( self, - embedding: List[float], + embedding: list[float], k: Optional[int] = None, fetch_k: Optional[int] = None, lambda_mult: Optional[float] = None, filter: Optional[str] = None, **kwargs: Any, - ) -> List[Document]: + ) -> list[Document]: """Return docs selected using the maximal marginal relevance.""" docs_and_scores = ( await self.amax_marginal_relevance_search_with_score_by_vector( @@ -744,13 +744,13 @@ async def amax_marginal_relevance_search_by_vector( async def amax_marginal_relevance_search_with_score_by_vector( self, - embedding: List[float], + embedding: list[float], k: Optional[int] = None, fetch_k: Optional[int] = None, lambda_mult: Optional[float] = None, filter: Optional[str] = None, **kwargs: Any, - ) -> List[Tuple[Document, float]]: + ) -> list[tuple[Document, float]]: """Return docs and distance scores selected using the maximal marginal relevance.""" results = await self.__query_collection( embedding=embedding, k=fetch_k, filter=filter, **kwargs @@ -875,38 +875,38 @@ async def is_valid_index( def add_texts( self, texts: Iterable[str], - metadatas: Optional[List[dict]] = None, - ids: Optional[List] = None, + metadatas: Optional[list[dict]] = None, + ids: Optional[list] = None, **kwargs: Any, - ) -> List[str]: + ) -> list[str]: raise NotImplementedError( "Sync methods are not implemented for AsyncAlloyDBVectorStore. Use AlloyDBVectorStore interface instead." ) def add_documents( self, - documents: List[Document], - ids: Optional[List] = None, + documents: list[Document], + ids: Optional[list] = None, **kwargs: Any, - ) -> List[str]: + ) -> list[str]: raise NotImplementedError( "Sync methods are not implemented for AsyncAlloyDBVectorStore. Use AlloyDBVectorStore interface instead." ) def add_images( self, - uris: List[str], - metadatas: Optional[List[dict]] = None, - ids: Optional[List[str]] = None, + uris: list[str], + metadatas: Optional[list[dict]] = None, + ids: Optional[list[str]] = None, **kwargs: Any, - ) -> List[str]: + ) -> list[str]: raise NotImplementedError( "Sync methods are not implemented for AsyncAlloyDBVectorStore. Use AlloyDBVectorStore interface instead." ) def delete( self, - ids: Optional[List] = None, + ids: Optional[list] = None, **kwargs: Any, ) -> Optional[bool]: raise NotImplementedError( @@ -915,17 +915,17 @@ def delete( @classmethod def from_texts( # type: ignore[override] - cls: Type[AsyncAlloyDBVectorStore], - texts: List[str], + cls: type[AsyncAlloyDBVectorStore], + texts: list[str], embedding: Embeddings, engine: AlloyDBEngine, table_name: str, - metadatas: Optional[List[dict]] = None, - ids: Optional[List] = None, + metadatas: Optional[list[dict]] = None, + ids: Optional[list] = None, content_column: str = "content", embedding_column: str = "embedding", - metadata_columns: List[str] = [], - ignore_metadata_columns: Optional[List[str]] = None, + metadata_columns: list[str] = [], + ignore_metadata_columns: Optional[list[str]] = None, id_column: str = "langchain_id", metadata_json_column: str = "langchain_metadata", **kwargs: Any, @@ -936,16 +936,16 @@ def from_texts( # type: ignore[override] @classmethod def from_documents( # type: ignore[override] - cls: Type[AsyncAlloyDBVectorStore], - documents: List[Document], + cls: type[AsyncAlloyDBVectorStore], + documents: list[Document], embedding: Embeddings, engine: AlloyDBEngine, table_name: str, - ids: Optional[List] = None, + ids: Optional[list] = None, content_column: str = "content", embedding_column: str = "embedding", - metadata_columns: List[str] = [], - ignore_metadata_columns: Optional[List[str]] = None, + metadata_columns: list[str] = [], + ignore_metadata_columns: Optional[list[str]] = None, id_column: str = "langchain_id", metadata_json_column: str = "langchain_metadata", **kwargs: Any, @@ -960,7 +960,7 @@ def similarity_search( k: Optional[int] = None, filter: Optional[str] = None, **kwargs: Any, - ) -> List[Document]: + ) -> list[Document]: raise NotImplementedError( "Sync methods are not implemented for AsyncAlloyDBVectorStore. Use AlloyDBVectorStore interface instead." ) @@ -971,7 +971,7 @@ def similarity_search_image( k: Optional[int] = None, filter: Optional[str] = None, **kwargs: Any, - ) -> List[Document]: + ) -> list[Document]: raise NotImplementedError( "Sync methods are not implemented for AsyncAlloyDBVectorStore. Use AlloyDBVectorStore interface instead." ) @@ -982,29 +982,29 @@ def similarity_search_with_score( k: Optional[int] = None, filter: Optional[str] = None, **kwargs: Any, - ) -> List[Tuple[Document, float]]: + ) -> list[tuple[Document, float]]: raise NotImplementedError( "Sync methods are not implemented for AsyncAlloyDBVectorStore. Use AlloyDBVectorStore interface instead." ) def similarity_search_by_vector( self, - embedding: List[float], + embedding: list[float], k: Optional[int] = None, filter: Optional[str] = None, **kwargs: Any, - ) -> List[Document]: + ) -> list[Document]: raise NotImplementedError( "Sync methods are not implemented for AsyncAlloyDBVectorStore. Use AlloyDBVectorStore interface instead." ) def similarity_search_with_score_by_vector( self, - embedding: List[float], + embedding: list[float], k: Optional[int] = None, filter: Optional[str] = None, **kwargs: Any, - ) -> List[Tuple[Document, float]]: + ) -> list[tuple[Document, float]]: raise NotImplementedError( "Sync methods are not implemented for AsyncAlloyDBVectorStore. Use AlloyDBVectorStore interface instead." ) @@ -1017,33 +1017,33 @@ def max_marginal_relevance_search( lambda_mult: Optional[float] = None, filter: Optional[str] = None, **kwargs: Any, - ) -> List[Document]: + ) -> list[Document]: raise NotImplementedError( "Sync methods are not implemented for AsyncAlloyDBVectorStore. Use AlloyDBVectorStore interface instead." ) def max_marginal_relevance_search_by_vector( self, - embedding: List[float], + embedding: list[float], k: Optional[int] = None, fetch_k: Optional[int] = None, lambda_mult: Optional[float] = None, filter: Optional[str] = None, **kwargs: Any, - ) -> List[Document]: + ) -> list[Document]: raise NotImplementedError( "Sync methods are not implemented for AsyncAlloyDBVectorStore. Use AlloyDBVectorStore interface instead." ) def max_marginal_relevance_search_with_score_by_vector( self, - embedding: List[float], + embedding: list[float], k: Optional[int] = None, fetch_k: Optional[int] = None, lambda_mult: Optional[float] = None, filter: Optional[str] = None, **kwargs: Any, - ) -> List[Tuple[Document, float]]: + ) -> list[tuple[Document, float]]: raise NotImplementedError( "Sync methods are not implemented for AsyncAlloyDBVectorStore. Use AlloyDBVectorStore interface instead." ) diff --git a/src/langchain_google_alloydb_pg/chat_message_history.py b/src/langchain_google_alloydb_pg/chat_message_history.py index 5384a41e..b3c8d07c 100644 --- a/src/langchain_google_alloydb_pg/chat_message_history.py +++ b/src/langchain_google_alloydb_pg/chat_message_history.py @@ -14,7 +14,7 @@ from __future__ import annotations -from typing import List, Sequence +from typing import Sequence from langchain_core.chat_history import BaseChatMessageHistory from langchain_core.messages import BaseMessage @@ -108,7 +108,7 @@ def create_sync( return cls(cls.__create_key, engine, history) @property # type: ignore[override] - def messages(self) -> List[BaseMessage]: + def messages(self) -> list[BaseMessage]: """The abstraction required a property.""" return self._engine._run_as_sync(self.__history._aget_messages()) diff --git a/src/langchain_google_alloydb_pg/embeddings.py b/src/langchain_google_alloydb_pg/embeddings.py index a50799cd..4cb111b4 100644 --- a/src/langchain_google_alloydb_pg/embeddings.py +++ b/src/langchain_google_alloydb_pg/embeddings.py @@ -16,7 +16,6 @@ from __future__ import annotations import json -from typing import List, Type from langchain_core.embeddings import Embeddings from sqlalchemy import text @@ -50,7 +49,7 @@ def __init__(self, key: object, engine: AlloyDBEngine, model_id: str): @classmethod async def create( - cls: Type[AlloyDBEmbeddings], engine: AlloyDBEngine, model_id: str + cls: type[AlloyDBEmbeddings], engine: AlloyDBEngine, model_id: str ) -> AlloyDBEmbeddings: """Create AlloyDBEmbeddings instance. @@ -72,7 +71,7 @@ async def create( @classmethod def create_sync( - cls: Type[AlloyDBEmbeddings], engine: AlloyDBEngine, model_id: str + cls: type[AlloyDBEmbeddings], engine: AlloyDBEngine, model_id: str ) -> AlloyDBEmbeddings: """Create AlloyDBEmbeddings instance. @@ -119,12 +118,12 @@ async def __amodel_exists(self) -> bool: return True return False - def embed_documents(self, texts: List[str]) -> List[List[float]]: + def embed_documents(self, texts: list[str]) -> list[list[float]]: raise NotImplementedError( "Embedding functions are not implemented. Use VertexAIEmbeddings interface instead." ) - async def aembed_documents(self, texts: List[str]) -> List[List[float]]: + async def aembed_documents(self, texts: list[str]) -> list[list[float]]: raise NotImplementedError( "Embedding functions are not implemented. Use VertexAIEmbeddings interface instead." ) @@ -132,37 +131,37 @@ async def aembed_documents(self, texts: List[str]) -> List[List[float]]: def embed_query_inline(self, query: str) -> str: return f"embedding('{self.model_id}', '{query}')::vector" - async def aembed_query(self, text: str) -> List[float]: + async def aembed_query(self, text: str) -> list[float]: """Asynchronous Embed query text. Args: query (str): Text to embed. Returns: - List[float]: Embedding. + list[float]: Embedding. """ embeddings = await self._engine._run_as_async(self.__aembed_query(text)) return embeddings - def embed_query(self, text: str) -> List[float]: + def embed_query(self, text: str) -> list[float]: """Embed query text. Args: query (str): Text to embed. Returns: - List[float]: Embedding. + list[float]: Embedding. """ return self._engine._run_as_sync(self.__aembed_query(text)) - async def __aembed_query(self, query: str) -> List[float]: + async def __aembed_query(self, query: str) -> list[float]: """Coroutine for generating embeddings for a given query. Args: query (str): Text to embed. Returns: - List[float]: Embedding. + list[float]: Embedding. """ query = f" SELECT embedding('{self.model_id}', '{query}')::vector " async with self._engine._pool.connect() as conn: diff --git a/src/langchain_google_alloydb_pg/engine.py b/src/langchain_google_alloydb_pg/engine.py index 0782a2a9..84efc2f0 100644 --- a/src/langchain_google_alloydb_pg/engine.py +++ b/src/langchain_google_alloydb_pg/engine.py @@ -17,17 +17,7 @@ from concurrent.futures import Future from dataclasses import dataclass from threading import Thread -from typing import ( - TYPE_CHECKING, - Any, - Awaitable, - Dict, - List, - Optional, - Type, - TypeVar, - Union, -) +from typing import TYPE_CHECKING, Any, Awaitable, Optional, TypeVar, Union import aiohttp import google.auth # type: ignore @@ -76,7 +66,7 @@ async def _get_iam_principal_email( url = f"https://oauth2.googleapis.com/tokeninfo?access_token={credentials.token}" async with aiohttp.ClientSession() as client: response = await client.get(url, raise_for_status=True) - response_json: Dict = await response.json() + response_json: dict = await response.json() email = response_json.get("email") if email is None: raise ValueError( @@ -179,7 +169,7 @@ def __start_background_loop( @classmethod def from_instance( - cls: Type[AlloyDBEngine], + cls: type[AlloyDBEngine], project_id: str, region: str, cluster: str, @@ -221,7 +211,7 @@ def from_instance( @classmethod async def _create( - cls: Type[AlloyDBEngine], + cls: type[AlloyDBEngine], project_id: str, region: str, cluster: str, @@ -305,7 +295,7 @@ async def getconn() -> asyncpg.Connection: @classmethod async def afrom_instance( - cls: Type[AlloyDBEngine], + cls: type[AlloyDBEngine], project_id: str, region: str, cluster: str, @@ -347,7 +337,7 @@ async def afrom_instance( @classmethod def from_engine( - cls: Type[AlloyDBEngine], + cls: type[AlloyDBEngine], engine: AsyncEngine, loop: Optional[asyncio.AbstractEventLoop] = None, ) -> AlloyDBEngine: @@ -418,7 +408,7 @@ async def _ainit_vectorstore_table( schema_name: str = "public", content_column: str = "content", embedding_column: str = "embedding", - metadata_columns: List[Column] = [], + metadata_columns: list[Column] = [], metadata_json_column: str = "langchain_metadata", id_column: Union[str, Column] = "langchain_id", overwrite_existing: bool = False, @@ -436,7 +426,7 @@ async def _ainit_vectorstore_table( Default: "page_content". embedding_column (str) : Name of the column to store vector embeddings. Default: "embedding". - metadata_columns (List[Column]): A list of Columns to create for custom + metadata_columns (list[Column]): A list of Columns to create for custom metadata. Default: []. Optional. metadata_json_column (str): The column to store extra metadata in JSON format. Default: "langchain_metadata". Optional. @@ -486,7 +476,7 @@ async def ainit_vectorstore_table( schema_name: str = "public", content_column: str = "content", embedding_column: str = "embedding", - metadata_columns: List[Column] = [], + metadata_columns: list[Column] = [], metadata_json_column: str = "langchain_metadata", id_column: Union[str, Column] = "langchain_id", overwrite_existing: bool = False, @@ -504,7 +494,7 @@ async def ainit_vectorstore_table( Default: "page_content". embedding_column (str) : Name of the column to store vector embeddings. Default: "embedding". - metadata_columns (List[Column]): A list of Columns to create for custom + metadata_columns (list[Column]): A list of Columns to create for custom metadata. Default: []. Optional. metadata_json_column (str): The column to store extra metadata in JSON format. Default: "langchain_metadata". Optional. @@ -536,7 +526,7 @@ def init_vectorstore_table( schema_name: str = "public", content_column: str = "content", embedding_column: str = "embedding", - metadata_columns: List[Column] = [], + metadata_columns: list[Column] = [], metadata_json_column: str = "langchain_metadata", id_column: Union[str, Column] = "langchain_id", overwrite_existing: bool = False, @@ -554,7 +544,7 @@ def init_vectorstore_table( Default: "page_content". embedding_column (str) : Name of the column to store vector embeddings. Default: "embedding". - metadata_columns (List[Column]): A list of Columns to create for custom + metadata_columns (list[Column]): A list of Columns to create for custom metadata. Default: []. Optional. metadata_json_column (str): The column to store extra metadata in JSON format. Default: "langchain_metadata". Optional. @@ -643,7 +633,7 @@ async def _ainit_document_table( table_name: str, schema_name: str = "public", content_column: str = "page_content", - metadata_columns: List[Column] = [], + metadata_columns: list[Column] = [], metadata_json_column: str = "langchain_metadata", store_metadata: bool = True, ) -> None: @@ -657,7 +647,7 @@ async def _ainit_document_table( Default: "public". content_column (str): Name of the column to store document content. Default: "page_content". - metadata_columns (List[Column]): A list of Columns + metadata_columns (list[Column]): A list of Columns to create for custom metadata. Optional. metadata_json_column (str): The column to store extra metadata in JSON format. Default: "langchain_metadata". Optional. @@ -684,7 +674,7 @@ async def ainit_document_table( table_name: str, schema_name: str = "public", content_column: str = "page_content", - metadata_columns: List[Column] = [], + metadata_columns: list[Column] = [], metadata_json_column: str = "langchain_metadata", store_metadata: bool = True, ) -> None: @@ -697,7 +687,7 @@ async def ainit_document_table( Default: "public". content_column (str): Name of the column to store document content. Default: "page_content". - metadata_columns (List[sqlalchemy.Column]): A list of SQLAlchemy Columns + metadata_columns (list[sqlalchemy.Column]): A list of SQLAlchemy Columns to create for custom metadata. Optional. metadata_json_column (str): The column to store extra metadata in JSON format. Default: "langchain_metadata". Optional. @@ -723,7 +713,7 @@ def init_document_table( table_name: str, schema_name: str = "public", content_column: str = "page_content", - metadata_columns: List[Column] = [], + metadata_columns: list[Column] = [], metadata_json_column: str = "langchain_metadata", store_metadata: bool = True, ) -> None: @@ -736,7 +726,7 @@ def init_document_table( Default: "public". content_column (str): Name of the column to store document content. Default: "page_content". - metadata_columns (List[sqlalchemy.Column]): A list of SQLAlchemy Columns + metadata_columns (list[sqlalchemy.Column]): A list of SQLAlchemy Columns to create for custom metadata. Optional. metadata_json_column (str): The column to store extra metadata in JSON format. Default: "langchain_metadata". Optional. diff --git a/src/langchain_google_alloydb_pg/indexes.py b/src/langchain_google_alloydb_pg/indexes.py index 43de268c..1c9b46cc 100644 --- a/src/langchain_google_alloydb_pg/indexes.py +++ b/src/langchain_google_alloydb_pg/indexes.py @@ -15,7 +15,7 @@ import enum from abc import ABC, abstractmethod from dataclasses import dataclass, field -from typing import List, Optional, Type, TypeVar +from typing import Optional @dataclass @@ -45,7 +45,7 @@ class BaseIndex(ABC): distance_strategy: DistanceStrategy = field( default_factory=lambda: DistanceStrategy.COSINE_DISTANCE ) - partial_indexes: Optional[List[str]] = None + partial_indexes: Optional[list[str]] = None @abstractmethod def index_options(self) -> str: diff --git a/src/langchain_google_alloydb_pg/loader.py b/src/langchain_google_alloydb_pg/loader.py index 8d97cc37..acb629c5 100644 --- a/src/langchain_google_alloydb_pg/loader.py +++ b/src/langchain_google_alloydb_pg/loader.py @@ -14,7 +14,7 @@ from __future__ import annotations -from typing import Any, AsyncIterator, Callable, Dict, Iterator, List, Optional +from typing import AsyncIterator, Callable, Iterator, Optional from langchain_core.document_loaders.base import BaseLoader from langchain_core.documents import Document @@ -46,8 +46,8 @@ def __init__( key (object): Prevent direct constructor usage. engine (AlloyDBEngine): AsyncEngine with pool connection to the postgres database query (Optional[str], optional): SQL query. Defaults to None. - content_columns (Optional[List[str]], optional): Column that represent a Document's page_content. Defaults to the first column. - metadata_columns (Optional[List[str]], optional): Column(s) that represent a Document's metadata. Defaults to None. + content_columns (Optional[list[str]], optional): Column that represent a Document's page_content. Defaults to the first column. + metadata_columns (Optional[list[str]], optional): Column(s) that represent a Document's metadata. Defaults to None. formatter (Optional[Callable], optional): A function to format page content (OneOf: format, formatter). Defaults to None. metadata_json_column (Optional[str], optional): Column to store metadata as JSON. Defaults to "langchain_metadata". @@ -70,8 +70,8 @@ async def create( query: Optional[str] = None, table_name: Optional[str] = None, schema_name: str = "public", - content_columns: Optional[List[str]] = None, - metadata_columns: Optional[List[str]] = None, + content_columns: Optional[list[str]] = None, + metadata_columns: Optional[list[str]] = None, metadata_json_column: Optional[str] = None, format: Optional[str] = None, formatter: Optional[Callable] = None, @@ -83,8 +83,8 @@ async def create( query (Optional[str], optional): SQL query. Defaults to None. table_name (Optional[str], optional): Name of table to query. Defaults to None. schema_name (str, optional): Name of the schema where table is located. Defaults to "public". - content_columns (Optional[List[str]], optional): Column that represent a Document's page_content. Defaults to the first column. - metadata_columns (Optional[List[str]], optional): Column(s) that represent a Document's metadata. Defaults to None. + content_columns (Optional[list[str]], optional): Column that represent a Document's page_content. Defaults to the first column. + metadata_columns (Optional[list[str]], optional): Column(s) that represent a Document's metadata. Defaults to None. metadata_json_column (Optional[str], optional): Column to store metadata as JSON. Defaults to "langchain_metadata". format (Optional[str], optional): Format of page content (OneOf: text, csv, YAML, JSON). Defaults to 'text'. formatter (Optional[Callable], optional): A function to format page content (OneOf: format, formatter). Defaults to None. @@ -113,8 +113,8 @@ def create_sync( query: Optional[str] = None, table_name: Optional[str] = None, schema_name: str = "public", - content_columns: Optional[List[str]] = None, - metadata_columns: Optional[List[str]] = None, + content_columns: Optional[list[str]] = None, + metadata_columns: Optional[list[str]] = None, metadata_json_column: Optional[str] = None, format: Optional[str] = None, formatter: Optional[Callable] = None, @@ -126,8 +126,8 @@ def create_sync( query (Optional[str], optional): SQL query. Defaults to None. table_name (Optional[str], optional): Name of table to query. Defaults to None. schema_name (str, optional): Name of the schema where table is located. Defaults to "public". - content_columns (Optional[List[str]], optional): Column that represent a Document's page_content. Defaults to the first column. - metadata_columns (Optional[List[str]], optional): Column(s) that represent a Document's metadata. Defaults to None. + content_columns (Optional[list[str]], optional): Column that represent a Document's page_content. Defaults to the first column. + metadata_columns (Optional[list[str]], optional): Column(s) that represent a Document's metadata. Defaults to None. metadata_json_column (Optional[str], optional): Column to store metadata as JSON. Defaults to "langchain_metadata". format (Optional[str], optional): Format of page content (OneOf: text, csv, YAML, JSON). Defaults to 'text'. formatter (Optional[Callable], optional): A function to format page content (OneOf: format, formatter). Defaults to None. @@ -149,11 +149,11 @@ def create_sync( loader = engine._run_as_sync(coro) return cls(cls.__create_key, engine, loader) - def load(self) -> List[Document]: + def load(self) -> list[Document]: """Load PostgreSQL data into Document objects.""" return self._engine._run_as_sync(self.__loader.aload()) - async def aload(self) -> List[Document]: + async def aload(self) -> list[Document]: """Load PostgreSQL data into Document objects.""" return await self._engine._run_as_async(self.__loader.aload()) @@ -197,7 +197,7 @@ def __init__( table_name (str): Name of table to query. content_column (str): Column that represent a Document's page_content. schema_name (str, optional): Name of the schema where table is located. Defaults to "public". - metadata_columns (List[str], optional): Column(s) that represent a Document's metadata. Defaults to empty list. + metadata_columns (list[str], optional): Column(s) that represent a Document's metadata. Defaults to empty list. metadata_json_column (Optional[str], optional): Column to store metadata as JSON. Defaults to None. Raises: @@ -217,7 +217,7 @@ async def create( table_name: str, schema_name: str = "public", content_column: str = DEFAULT_CONTENT_COL, - metadata_columns: List[str] = [], + metadata_columns: list[str] = [], metadata_json_column: Optional[str] = DEFAULT_METADATA_COL, ) -> AlloyDBDocumentSaver: """Create an AlloyDBDocumentSaver instance. @@ -227,7 +227,7 @@ async def create( table_name (str): Name of table to query. schema_name (str, optional): Name of schema where the table is located. Defaults to "public". content_column (str, optional): Column that represent a Document's page_content. Defaults to "page_content". - metadata_columns (List[str], optional): Column(s) that represent a Document's metadata. Defaults to an empty list. + metadata_columns (list[str], optional): Column(s) that represent a Document's metadata. Defaults to an empty list. metadata_json_column (Optional[str], optional): Column to store metadata as JSON. Defaults to "langchain_metadata". Returns: @@ -251,7 +251,7 @@ def create_sync( table_name: str, schema_name: str = "public", content_column: str = DEFAULT_CONTENT_COL, - metadata_columns: List[str] = [], + metadata_columns: list[str] = [], metadata_json_column: str = DEFAULT_METADATA_COL, ) -> AlloyDBDocumentSaver: """Create an AlloyDBDocumentSaver instance. @@ -261,7 +261,7 @@ def create_sync( table_name (str): Name of table to query. schema_name (str, optional): Name of schema where the table is located. Defaults to "public". content_column (str, optional): Column that represent a Document's page_content. Defaults to "page_content". - metadata_columns (List[str], optional): Column(s) that represent a Document's metadata. Defaults to an empty list. + metadata_columns (list[str], optional): Column(s) that represent a Document's metadata. Defaults to an empty list. metadata_json_column (Optional[str], optional): Column to store metadata as JSON. Defaults to "langchain_metadata". Returns: @@ -278,42 +278,42 @@ def create_sync( saver = engine._run_as_sync(coro) return cls(cls.__create_key, engine, saver) - async def aadd_documents(self, docs: List[Document]) -> None: + async def aadd_documents(self, docs: list[Document]) -> None: """ Save documents in the DocumentSaver table. Document’s metadata is added to columns if found or stored in langchain_metadata JSON column. Args: - docs (List[langchain_core.documents.Document]): a list of documents to be saved. + docs (list[langchain_core.documents.Document]): List of documents to be saved. """ await self._engine._run_as_async(self._saver.aadd_documents(docs)) - def add_documents(self, docs: List[Document]) -> None: + def add_documents(self, docs: list[Document]) -> None: """ Save documents in the DocumentSaver table. Document’s metadata is added to columns if found or stored in langchain_metadata JSON column. Args: - docs (List[langchain_core.documents.Document]): a list of documents to be saved. + docs (list[langchain_core.documents.Document]): List of documents to be saved. """ self._engine._run_as_sync(self._saver.aadd_documents(docs)) - async def adelete(self, docs: List[Document]) -> None: + async def adelete(self, docs: list[Document]) -> None: """ Delete all instances of a document from the DocumentSaver table by matching the entire Document object. Args: - docs (List[langchain_core.documents.Document]): a list of documents to be deleted. + docs (list[langchain_core.documents.Document]): List of documents to be deleted. """ await self._engine._run_as_async(self._saver.adelete(docs)) - def delete(self, docs: List[Document]) -> None: + def delete(self, docs: list[Document]) -> None: """ Delete all instances of a document from the DocumentSaver table by matching the entire Document object. Args: - docs (List[langchain_core.documents.Document]): a list of documents to be deleted. + docs (list[langchain_core.documents.Document]): List of documents to be deleted. """ self._engine._run_as_sync(self._saver.adelete(docs)) diff --git a/src/langchain_google_alloydb_pg/model_manager.py b/src/langchain_google_alloydb_pg/model_manager.py index 336a0548..6b253bf2 100644 --- a/src/langchain_google_alloydb_pg/model_manager.py +++ b/src/langchain_google_alloydb_pg/model_manager.py @@ -16,7 +16,7 @@ from __future__ import annotations from dataclasses import dataclass -from typing import List, Optional, Sequence, Type +from typing import Optional, Sequence from sqlalchemy import text from sqlalchemy.engine.row import RowMapping @@ -62,7 +62,7 @@ def __init__( @classmethod async def create( - cls: Type[AlloyDBModelManager], + cls: type[AlloyDBModelManager], engine: AlloyDBEngine, ) -> AlloyDBModelManager: manager = AlloyDBModelManager(cls.__create_key, engine) @@ -72,7 +72,7 @@ async def create( @classmethod def create_sync( - cls: Type[AlloyDBModelManager], + cls: type[AlloyDBModelManager], engine: AlloyDBEngine, ) -> AlloyDBModelManager: manager = AlloyDBModelManager(cls.__create_key, engine) @@ -93,11 +93,11 @@ async def aget_model(self, model_id: str) -> Optional[AlloyDBModel]: result = await self._engine._run_as_async(self.__aget_model(model_id=model_id)) return result - async def alist_models(self) -> List[AlloyDBModel]: + async def alist_models(self) -> list[AlloyDBModel]: """Lists all the models and its details. Returns: - List[`AlloyDBModel`] of all available model.. + list[`AlloyDBModel`] of all available model.. """ results = await self._engine._run_as_async(self.__alist_models()) return results @@ -206,7 +206,7 @@ async def __aget_model(self, model_id: str) -> Optional[AlloyDBModel]: data_class = self.__convert_dict_to_dataclass(result)[0] return data_class - async def __alist_models(self) -> List[AlloyDBModel]: + async def __alist_models(self) -> list[AlloyDBModel]: """Lists all the models and its details.""" query = "SELECT * FROM google_ml.model_info_view;" result = await self.__query_db(query) @@ -294,7 +294,7 @@ async def __fetch_db_flag(self) -> str: def __convert_dict_to_dataclass( self, list_of_rows: Sequence[RowMapping] - ) -> List[AlloyDBModel]: + ) -> list[AlloyDBModel]: """Converts a list of DB rows to list of AlloyDBModel dataclass. Args: diff --git a/src/langchain_google_alloydb_pg/utils/pgvector_migrator.py b/src/langchain_google_alloydb_pg/utils/pgvector_migrator.py index 91977319..34927ee0 100644 --- a/src/langchain_google_alloydb_pg/utils/pgvector_migrator.py +++ b/src/langchain_google_alloydb_pg/utils/pgvector_migrator.py @@ -13,7 +13,7 @@ # limitations under the License. import asyncio import warnings -from typing import Any, AsyncIterator, Iterator, List, Optional, Sequence, TypeVar +from typing import Any, AsyncIterator, Iterator, Optional, Sequence, TypeVar from sqlalchemy import RowMapping, text from sqlalchemy.exc import ProgrammingError, SQLAlchemyError @@ -189,7 +189,7 @@ async def __amigrate_pgvector_collection( async def __alist_pgvector_collection_names( engine: AlloyDBEngine, -) -> List[str]: +) -> list[str]: """Lists all collection names present in PGVector table.""" try: query = f"SELECT name from {COLLECTIONS_TABLE}" @@ -232,7 +232,7 @@ async def aextract_pgvector_collection( async def alist_pgvector_collection_names( engine: AlloyDBEngine, -) -> List[str]: +) -> list[str]: """Lists all collection names present in PGVector table.""" return await engine._run_as_async(__alist_pgvector_collection_names(engine)) @@ -296,7 +296,7 @@ def extract_pgvector_collection( break -def list_pgvector_collection_names(engine: AlloyDBEngine) -> List[str]: +def list_pgvector_collection_names(engine: AlloyDBEngine) -> list[str]: """Lists all collection names present in PGVector table.""" return engine._run_as_sync(__alist_pgvector_collection_names(engine)) diff --git a/src/langchain_google_alloydb_pg/vectorstore.py b/src/langchain_google_alloydb_pg/vectorstore.py index 9f86a93f..b0e078bd 100644 --- a/src/langchain_google_alloydb_pg/vectorstore.py +++ b/src/langchain_google_alloydb_pg/vectorstore.py @@ -15,7 +15,7 @@ # TODO: Remove below import when minimum supported Python version is 3.10 from __future__ import annotations -from typing import Any, Callable, Iterable, List, Optional, Tuple, Type +from typing import Any, Callable, Iterable, Optional from langchain_core.documents import Document from langchain_core.embeddings import Embeddings @@ -57,15 +57,15 @@ def __init__(self, key: object, engine: AlloyDBEngine, vs: AsyncAlloyDBVectorSto @classmethod async def create( - cls: Type[AlloyDBVectorStore], + cls: type[AlloyDBVectorStore], engine: AlloyDBEngine, embedding_service: Embeddings, table_name: str, schema_name: str = "public", content_column: str = "content", embedding_column: str = "embedding", - metadata_columns: List[str] = [], - ignore_metadata_columns: Optional[List[str]] = None, + metadata_columns: list[str] = [], + ignore_metadata_columns: Optional[list[str]] = None, id_column: str = "langchain_id", metadata_json_column: Optional[str] = "langchain_metadata", distance_strategy: DistanceStrategy = DEFAULT_DISTANCE_STRATEGY, @@ -83,8 +83,8 @@ async def create( schema_name (str, optional): Name of the database schema. Defaults to "public". content_column (str): Column that represent a Document’s page_content. Defaults to "content". embedding_column (str): Column for embedding vectors. The embedding is generated from the document value. Defaults to "embedding". - metadata_columns (List[str]): Column(s) that represent a document's metadata. - ignore_metadata_columns (List[str]): Column(s) to ignore in pre-existing tables for a document's metadata. Can not be used with metadata_columns. Defaults to None. + metadata_columns (list[str]): Column(s) that represent a document's metadata. + ignore_metadata_columns (list[str]): Column(s) to ignore in pre-existing tables for a document's metadata. Can not be used with metadata_columns. Defaults to None. id_column (str): Column that represents the Document's id. Defaults to "langchain_id". metadata_json_column (str): Column to store metadata as JSON. Defaults to "langchain_metadata". distance_strategy (DistanceStrategy): Distance strategy to use for vector similarity search. Defaults to COSINE_DISTANCE. @@ -125,8 +125,8 @@ def create_sync( schema_name: str = "public", content_column: str = "content", embedding_column: str = "embedding", - metadata_columns: List[str] = [], - ignore_metadata_columns: Optional[List[str]] = None, + metadata_columns: list[str] = [], + ignore_metadata_columns: Optional[list[str]] = None, id_column: str = "langchain_id", metadata_json_column: str = "langchain_metadata", distance_strategy: DistanceStrategy = DEFAULT_DISTANCE_STRATEGY, @@ -145,8 +145,8 @@ def create_sync( schema_name (str, optional): Name of the database schema. Defaults to "public". content_column (str, optional): Column that represent a Document’s page_content. Defaults to "content". embedding_column (str, optional): Column for embedding vectors. The embedding is generated from the document value. Defaults to "embedding". - metadata_columns (List[str]): Column(s) that represent a document's metadata. Defaults to an empty list. - ignore_metadata_columns (Optional[List[str]]): Column(s) to ignore in pre-existing tables for a document's metadata. Can not be used with metadata_columns. Defaults to None. + metadata_columns (list[str]): Column(s) that represent a document's metadata. Defaults to an empty list. + ignore_metadata_columns (Optional[list[str]]): Column(s) to ignore in pre-existing tables for a document's metadata. Can not be used with metadata_columns. Defaults to None. id_column (str, optional): Column that represents the Document's id. Defaults to "langchain_id". metadata_json_column (str, optional): Column to store metadata as JSON. Defaults to "langchain_metadata". distance_strategy (DistanceStrategy, optional): Distance strategy to use for vector similarity search. Defaults to COSINE_DISTANCE. @@ -185,11 +185,11 @@ def embeddings(self) -> Embeddings: async def aadd_embeddings( self, texts: Iterable[str], - embeddings: List[List[float]], - metadatas: Optional[List[dict]] = None, - ids: Optional[List[str]] = None, + embeddings: list[list[float]], + metadatas: Optional[list[dict]] = None, + ids: Optional[list[str]] = None, **kwargs: Any, - ) -> List[str]: + ) -> list[str]: """Add data along with embeddings to the table.""" return await self._engine._run_as_async( self.__vs.aadd_embeddings(texts, embeddings, metadatas, ids, **kwargs) @@ -198,10 +198,10 @@ async def aadd_embeddings( async def aadd_texts( self, texts: Iterable[str], - metadatas: Optional[List[dict]] = None, - ids: Optional[List] = None, + metadatas: Optional[list[dict]] = None, + ids: Optional[list] = None, **kwargs: Any, - ) -> List[str]: + ) -> list[str]: """Embed texts and add to the table. Raises: @@ -213,10 +213,10 @@ async def aadd_texts( async def aadd_documents( self, - documents: List[Document], - ids: Optional[List] = None, + documents: list[Document], + ids: Optional[list] = None, **kwargs: Any, - ) -> List[str]: + ) -> list[str]: """Embed documents and add to the table. Raises: @@ -228,11 +228,11 @@ async def aadd_documents( async def aadd_images( self, - uris: List[str], - metadatas: Optional[List[dict]] = None, - ids: Optional[List[str]] = None, + uris: list[str], + metadatas: Optional[list[dict]] = None, + ids: Optional[list[str]] = None, **kwargs: Any, - ) -> List[str]: + ) -> list[str]: """Embed images and add to the table.""" return await self._engine._run_as_async( self.__vs.aadd_images(uris, metadatas, ids, **kwargs) @@ -241,11 +241,11 @@ async def aadd_images( def add_embeddings( self, texts: Iterable[str], - embeddings: List[List[float]], - metadatas: Optional[List[dict]] = None, - ids: Optional[List[str]] = None, + embeddings: list[list[float]], + metadatas: Optional[list[dict]] = None, + ids: Optional[list[str]] = None, **kwargs: Any, - ) -> List[str]: + ) -> list[str]: """Add data along with embeddings to the table.""" return self._engine._run_as_sync( self.__vs.aadd_embeddings(texts, embeddings, metadatas, ids, **kwargs) @@ -254,10 +254,10 @@ def add_embeddings( def add_texts( self, texts: Iterable[str], - metadatas: Optional[List[dict]] = None, - ids: Optional[List] = None, + metadatas: Optional[list[dict]] = None, + ids: Optional[list] = None, **kwargs: Any, - ) -> List[str]: + ) -> list[str]: """Embed texts and add to the table. Raises: @@ -269,10 +269,10 @@ def add_texts( def add_documents( self, - documents: List[Document], - ids: Optional[List] = None, + documents: list[Document], + ids: Optional[list] = None, **kwargs: Any, - ) -> List[str]: + ) -> list[str]: """Embed documents and add to the table. Raises: @@ -284,11 +284,11 @@ def add_documents( def add_images( self, - uris: List[str], - metadatas: Optional[List[dict]] = None, - ids: Optional[List[str]] = None, + uris: list[str], + metadatas: Optional[list[dict]] = None, + ids: Optional[list[str]] = None, **kwargs: Any, - ) -> List[str]: + ) -> list[str]: """Embed images and add to the table.""" return self._engine._run_as_sync( self.__vs.aadd_images(uris, metadatas, ids, **kwargs) @@ -296,7 +296,7 @@ def add_images( async def adelete( self, - ids: Optional[List] = None, + ids: Optional[list] = None, **kwargs: Any, ) -> Optional[bool]: """Delete records from the table. @@ -308,7 +308,7 @@ async def adelete( def delete( self, - ids: Optional[List] = None, + ids: Optional[list] = None, **kwargs: Any, ) -> Optional[bool]: """Delete records from the table. @@ -320,18 +320,18 @@ def delete( @classmethod async def afrom_texts( # type: ignore[override] - cls: Type[AlloyDBVectorStore], - texts: List[str], + cls: type[AlloyDBVectorStore], + texts: list[str], embedding: Embeddings, engine: AlloyDBEngine, table_name: str, schema_name: str = "public", - metadatas: Optional[List[dict]] = None, - ids: Optional[List] = None, + metadatas: Optional[list[dict]] = None, + ids: Optional[list] = None, content_column: str = "content", embedding_column: str = "embedding", - metadata_columns: List[str] = [], - ignore_metadata_columns: Optional[List[str]] = None, + metadata_columns: list[str] = [], + ignore_metadata_columns: Optional[list[str]] = None, id_column: str = "langchain_id", metadata_json_column: str = "langchain_metadata", distance_strategy: DistanceStrategy = DEFAULT_DISTANCE_STRATEGY, @@ -344,17 +344,17 @@ async def afrom_texts( # type: ignore[override] """Create an AlloyDBVectorStore instance from texts. Args: - texts (List[str]): Texts to add to the vector store. + texts (list[str]): Texts to add to the vector store. embedding (Embeddings): Text embedding model to use. engine (AlloyDBEngine): Connection pool engine for managing connections to AlloyDB database. table_name (str): Name of an existing table. schema_name (str, optional): Name of the database schema. Defaults to "public". - metadatas (Optional[List[dict]], optional): List of metadatas to add to table records. Defaults to None. - ids: (Optional[List]): List of IDs to add to table records. Defaults to None. + metadatas (Optional[list[dict]], optional): List of metadatas to add to table records. Defaults to None. + ids: (Optional[list]): List of IDs to add to table records. Defaults to None. content_column (str, optional): Column that represent a Document’s page_content. Defaults to "content". embedding_column (str, optional): Column for embedding vectors. The embedding is generated from the document value. Defaults to "embedding". - metadata_columns (List[str], optional): Column(s) that represent a document's metadata. Defaults to an empty list. - ignore_metadata_columns (Optional[List[str]], optional): Column(s) to ignore in pre-existing tables for a document's metadata. Can not be used with metadata_columns. Defaults to None. + metadata_columns (list[str], optional): Column(s) that represent a document's metadata. Defaults to an empty list. + ignore_metadata_columns (Optional[list[str]], optional): Column(s) to ignore in pre-existing tables for a document's metadata. Can not be used with metadata_columns. Defaults to None. id_column (str, optional): Column that represents the Document's id. Defaults to "langchain_id". metadata_json_column (str, optional): Column to store metadata as JSON. Defaults to "langchain_metadata". distance_strategy (DistanceStrategy): Distance strategy to use for vector similarity search. Defaults to COSINE_DISTANCE. @@ -391,17 +391,17 @@ async def afrom_texts( # type: ignore[override] @classmethod async def afrom_documents( # type: ignore[override] - cls: Type[AlloyDBVectorStore], - documents: List[Document], + cls: type[AlloyDBVectorStore], + documents: list[Document], embedding: Embeddings, engine: AlloyDBEngine, table_name: str, schema_name: str = "public", - ids: Optional[List] = None, + ids: Optional[list] = None, content_column: str = "content", embedding_column: str = "embedding", - metadata_columns: List[str] = [], - ignore_metadata_columns: Optional[List[str]] = None, + metadata_columns: list[str] = [], + ignore_metadata_columns: Optional[list[str]] = None, id_column: str = "langchain_id", metadata_json_column: str = "langchain_metadata", distance_strategy: DistanceStrategy = DEFAULT_DISTANCE_STRATEGY, @@ -414,16 +414,16 @@ async def afrom_documents( # type: ignore[override] """Create an AlloyDBVectorStore instance from documents. Args: - documents (List[Document]): Documents to add to the vector store. + documents (list[Document]): Documents to add to the vector store. embedding (Embeddings): Text embedding model to use. engine (AlloyDBEngine): Connection pool engine for managing connections to AlloyDB database. table_name (str): Name of an existing table. schema_name (str, optional): Name of the database schema. Defaults to "public". - ids: (Optional[List]): List of IDs to add to table records. Defaults to None. + ids: (Optional[list]): List of IDs to add to table records. Defaults to None. content_column (str, optional): Column that represent a Document’s page_content. Defaults to "content". embedding_column (str, optional): Column for embedding vectors. The embedding is generated from the document value. Defaults to "embedding". - metadata_columns (List[str], optional): Column(s) that represent a document's metadata. Defaults to an empty list. - ignore_metadata_columns (Optional[List[str]], optional): Column(s) to ignore in pre-existing tables for a document's metadata. Can not be used with metadata_columns. Defaults to None. + metadata_columns (list[str], optional): Column(s) that represent a document's metadata. Defaults to an empty list. + ignore_metadata_columns (Optional[list[str]], optional): Column(s) to ignore in pre-existing tables for a document's metadata. Can not be used with metadata_columns. Defaults to None. id_column (str, optional): Column that represents the Document's id. Defaults to "langchain_id". metadata_json_column (str, optional): Column to store metadata as JSON. Defaults to "langchain_metadata". distance_strategy (DistanceStrategy): Distance strategy to use for vector similarity search. Defaults to COSINE_DISTANCE. @@ -461,18 +461,18 @@ async def afrom_documents( # type: ignore[override] @classmethod def from_texts( # type: ignore[override] - cls: Type[AlloyDBVectorStore], - texts: List[str], + cls: type[AlloyDBVectorStore], + texts: list[str], embedding: Embeddings, engine: AlloyDBEngine, table_name: str, schema_name: str = "public", - metadatas: Optional[List[dict]] = None, - ids: Optional[List] = None, + metadatas: Optional[list[dict]] = None, + ids: Optional[list] = None, content_column: str = "content", embedding_column: str = "embedding", - metadata_columns: List[str] = [], - ignore_metadata_columns: Optional[List[str]] = None, + metadata_columns: list[str] = [], + ignore_metadata_columns: Optional[list[str]] = None, id_column: str = "langchain_id", metadata_json_column: str = "langchain_metadata", distance_strategy: DistanceStrategy = DEFAULT_DISTANCE_STRATEGY, @@ -485,17 +485,17 @@ def from_texts( # type: ignore[override] """Create an AlloyDBVectorStore instance from texts. Args: - texts (List[str]): Texts to add to the vector store. + texts (list[str]): Texts to add to the vector store. embedding (Embeddings): Text embedding model to use. engine (AlloyDBEngine): Connection pool engine for managing connections to AlloyDB database. table_name (str): Name of an existing table. schema_name (str, optional): Name of the database schema. Defaults to "public". - metadatas (Optional[List[dict]], optional): List of metadatas to add to table records. Defaults to None. - ids: (Optional[List]): List of IDs to add to table records. Defaults to None. + metadatas (Optional[list[dict]], optional): List of metadatas to add to table records. Defaults to None. + ids: (Optional[list]): List of IDs to add to table records. Defaults to None. content_column (str, optional): Column that represent a Document’s page_content. Defaults to "content". embedding_column (str, optional): Column for embedding vectors. The embedding is generated from the document value. Defaults to "embedding". - metadata_columns (List[str], optional): Column(s) that represent a document's metadata. Defaults to empty list. - ignore_metadata_columns (Optional[List[str]], optional): Column(s) to ignore in pre-existing tables for a document's metadata. Can not be used with metadata_columns. Defaults to None. + metadata_columns (list[str], optional): Column(s) that represent a document's metadata. Defaults to empty list. + ignore_metadata_columns (Optional[list[str]], optional): Column(s) to ignore in pre-existing tables for a document's metadata. Can not be used with metadata_columns. Defaults to None. id_column (str, optional): Column that represents the Document's id. Defaults to "langchain_id". metadata_json_column (str, optional): Column to store metadata as JSON. Defaults to "langchain_metadata". distance_strategy (DistanceStrategy): Distance strategy to use for vector similarity search. Defaults to COSINE_DISTANCE. @@ -533,17 +533,17 @@ def from_texts( # type: ignore[override] @classmethod def from_documents( # type: ignore[override] - cls: Type[AlloyDBVectorStore], - documents: List[Document], + cls: type[AlloyDBVectorStore], + documents: list[Document], embedding: Embeddings, engine: AlloyDBEngine, table_name: str, schema_name: str = "public", - ids: Optional[List] = None, + ids: Optional[list] = None, content_column: str = "content", embedding_column: str = "embedding", - metadata_columns: List[str] = [], - ignore_metadata_columns: Optional[List[str]] = None, + metadata_columns: list[str] = [], + ignore_metadata_columns: Optional[list[str]] = None, id_column: str = "langchain_id", metadata_json_column: str = "langchain_metadata", distance_strategy: DistanceStrategy = DEFAULT_DISTANCE_STRATEGY, @@ -556,16 +556,16 @@ def from_documents( # type: ignore[override] """Create an AlloyDBVectorStore instance from documents. Args: - documents (List[Document]): Documents to add to the vector store. + documents (list[Document]): Documents to add to the vector store. embedding (Embeddings): Text embedding model to use. engine (AlloyDBEngine): Connection pool engine for managing connections to AlloyDB database. table_name (str): Name of an existing table. schema_name (str, optional): Name of the database schema. Defaults to "public". - ids: (Optional[List]): List of IDs to add to table records. Defaults to None. + ids: (Optional[list]): List of IDs to add to table records. Defaults to None. content_column (str, optional): Column that represent a Document’s page_content. Defaults to "content". embedding_column (str, optional): Column for embedding vectors. The embedding is generated from the document value. Defaults to "embedding". - metadata_columns (List[str], optional): Column(s) that represent a document's metadata. Defaults to an empty list. - ignore_metadata_columns (Optional[List[str]], optional): Column(s) to ignore in pre-existing tables for a document's metadata. Can not be used with metadata_columns. Defaults to None. + metadata_columns (list[str], optional): Column(s) that represent a document's metadata. Defaults to an empty list. + ignore_metadata_columns (Optional[list[str]], optional): Column(s) to ignore in pre-existing tables for a document's metadata. Can not be used with metadata_columns. Defaults to None. id_column (str, optional): Column that represents the Document's id. Defaults to "langchain_id". metadata_json_column (str, optional): Column to store metadata as JSON. Defaults to "langchain_metadata". distance_strategy (DistanceStrategy): Distance strategy to use for vector similarity search. Defaults to COSINE_DISTANCE. @@ -607,7 +607,7 @@ def similarity_search( k: Optional[int] = None, filter: Optional[str] = None, **kwargs: Any, - ) -> List[Document]: + ) -> list[Document]: """Return docs selected by similarity search on query.""" return self._engine._run_as_sync( self.__vs.asimilarity_search(query, k, filter, **kwargs) @@ -619,7 +619,7 @@ def similarity_search_image( k: Optional[int] = None, filter: Optional[str] = None, **kwargs: Any, - ) -> List[Document]: + ) -> list[Document]: """Return docs selected by similarity search on image.""" return self._engine._run_as_sync( self.__vs.asimilarity_search_image(image_uri, k, filter, **kwargs) @@ -631,7 +631,7 @@ async def asimilarity_search( k: Optional[int] = None, filter: Optional[str] = None, **kwargs: Any, - ) -> List[Document]: + ) -> list[Document]: """Return docs selected by similarity search on query.""" return await self._engine._run_as_async( self.__vs.asimilarity_search(query, k, filter, **kwargs) @@ -643,7 +643,7 @@ async def asimilarity_search_image( k: Optional[int] = None, filter: Optional[str] = None, **kwargs: Any, - ) -> List[Document]: + ) -> list[Document]: """Return docs selected by similarity search on query.""" return await self._engine._run_as_async( self.__vs.asimilarity_search(image_uri, k, filter, **kwargs) @@ -666,7 +666,7 @@ async def asimilarity_search_with_score( k: Optional[int] = None, filter: Optional[str] = None, **kwargs: Any, - ) -> List[Tuple[Document, float]]: + ) -> list[tuple[Document, float]]: """Return docs and distance scores selected by similarity search on query.""" return await self._engine._run_as_async( self.__vs.asimilarity_search_with_score(query, k, filter, **kwargs) @@ -674,11 +674,11 @@ async def asimilarity_search_with_score( async def asimilarity_search_by_vector( self, - embedding: List[float], + embedding: list[float], k: Optional[int] = None, filter: Optional[str] = None, **kwargs: Any, - ) -> List[Document]: + ) -> list[Document]: """Return docs selected by vector similarity search.""" return await self._engine._run_as_async( self.__vs.asimilarity_search_by_vector(embedding, k, filter, **kwargs) @@ -686,11 +686,11 @@ async def asimilarity_search_by_vector( async def asimilarity_search_with_score_by_vector( self, - embedding: List[float], + embedding: list[float], k: Optional[int] = None, filter: Optional[str] = None, **kwargs: Any, - ) -> List[Tuple[Document, float]]: + ) -> list[tuple[Document, float]]: """Return docs and distance scores selected by vector similarity search.""" return await self._engine._run_as_async( self.__vs.asimilarity_search_with_score_by_vector( @@ -706,7 +706,7 @@ async def amax_marginal_relevance_search( lambda_mult: Optional[float] = None, filter: Optional[str] = None, **kwargs: Any, - ) -> List[Document]: + ) -> list[Document]: """Return docs selected using the maximal marginal relevance.""" return await self._engine._run_as_async( self.__vs.amax_marginal_relevance_search( @@ -716,13 +716,13 @@ async def amax_marginal_relevance_search( async def amax_marginal_relevance_search_by_vector( self, - embedding: List[float], + embedding: list[float], k: Optional[int] = None, fetch_k: Optional[int] = None, lambda_mult: Optional[float] = None, filter: Optional[str] = None, **kwargs: Any, - ) -> List[Document]: + ) -> list[Document]: """Return docs selected using the maximal marginal relevance.""" return await self._engine._run_as_async( self.__vs.amax_marginal_relevance_search_by_vector( @@ -732,13 +732,13 @@ async def amax_marginal_relevance_search_by_vector( async def amax_marginal_relevance_search_with_score_by_vector( self, - embedding: List[float], + embedding: list[float], k: Optional[int] = None, fetch_k: Optional[int] = None, lambda_mult: Optional[float] = None, filter: Optional[str] = None, **kwargs: Any, - ) -> List[Tuple[Document, float]]: + ) -> list[tuple[Document, float]]: """Return docs and distance scores selected using the maximal marginal relevance.""" return await self._engine._run_as_async( self.__vs.amax_marginal_relevance_search_with_score_by_vector( @@ -752,7 +752,7 @@ def similarity_search_with_score( k: Optional[int] = None, filter: Optional[str] = None, **kwargs: Any, - ) -> List[Tuple[Document, float]]: + ) -> list[tuple[Document, float]]: """Return docs and distance scores selected by similarity search on query.""" return self._engine._run_as_sync( self.__vs.asimilarity_search_with_score(query, k, filter, **kwargs) @@ -760,11 +760,11 @@ def similarity_search_with_score( def similarity_search_by_vector( self, - embedding: List[float], + embedding: list[float], k: Optional[int] = None, filter: Optional[str] = None, **kwargs: Any, - ) -> List[Document]: + ) -> list[Document]: """Return docs selected by vector similarity search.""" return self._engine._run_as_sync( self.__vs.asimilarity_search_by_vector(embedding, k, filter, **kwargs) @@ -772,11 +772,11 @@ def similarity_search_by_vector( def similarity_search_with_score_by_vector( self, - embedding: List[float], + embedding: list[float], k: Optional[int] = None, filter: Optional[str] = None, **kwargs: Any, - ) -> List[Tuple[Document, float]]: + ) -> list[tuple[Document, float]]: """Return docs and distance scores selected by similarity search on vector.""" return self._engine._run_as_sync( self.__vs.asimilarity_search_with_score_by_vector( @@ -792,7 +792,7 @@ def max_marginal_relevance_search( lambda_mult: Optional[float] = None, filter: Optional[str] = None, **kwargs: Any, - ) -> List[Document]: + ) -> list[Document]: """Return docs selected using the maximal marginal relevance.""" return self._engine._run_as_sync( self.__vs.amax_marginal_relevance_search( @@ -802,13 +802,13 @@ def max_marginal_relevance_search( def max_marginal_relevance_search_by_vector( self, - embedding: List[float], + embedding: list[float], k: Optional[int] = None, fetch_k: Optional[int] = None, lambda_mult: Optional[float] = None, filter: Optional[str] = None, **kwargs: Any, - ) -> List[Document]: + ) -> list[Document]: """Return docs selected using the maximal marginal relevance.""" return self._engine._run_as_sync( self.__vs.amax_marginal_relevance_search_by_vector( @@ -818,13 +818,13 @@ def max_marginal_relevance_search_by_vector( def max_marginal_relevance_search_with_score_by_vector( self, - embedding: List[float], + embedding: list[float], k: Optional[int] = None, fetch_k: Optional[int] = None, lambda_mult: Optional[float] = None, filter: Optional[str] = None, **kwargs: Any, - ) -> List[Tuple[Document, float]]: + ) -> list[tuple[Document, float]]: """Return docs and distance scores selected using the maximal marginal relevance.""" return self._engine._run_as_sync( self.__vs.amax_marginal_relevance_search_with_score_by_vector( diff --git a/tests/test_async_vectorstore.py b/tests/test_async_vectorstore.py index 1ea5ec25..8c95c18c 100644 --- a/tests/test_async_vectorstore.py +++ b/tests/test_async_vectorstore.py @@ -14,7 +14,7 @@ import os import uuid -from typing import List, Sequence +from typing import Sequence import pytest import pytest_asyncio @@ -46,7 +46,7 @@ class FakeImageEmbedding(DeterministicFakeEmbedding): - def embed_image(self, image_paths: List[str]) -> List[List[float]]: + def embed_image(self, image_paths: list[str]) -> list[list[float]]: return [self.embed_query(path) for path in image_paths] diff --git a/tests/test_async_vectorstore_search.py b/tests/test_async_vectorstore_search.py index 4684ac3e..2530fa58 100644 --- a/tests/test_async_vectorstore_search.py +++ b/tests/test_async_vectorstore_search.py @@ -14,7 +14,6 @@ import os import uuid -from typing import List import pytest import pytest_asyncio @@ -50,7 +49,7 @@ class FakeImageEmbedding(DeterministicFakeEmbedding): - def embed_image(self, image_paths: List[str]) -> List[List[float]]: + def embed_image(self, image_paths: list[str]) -> list[list[float]]: return [self.embed_query(path) for path in image_paths] diff --git a/tests/test_chatmessagehistory.py b/tests/test_chatmessagehistory.py index bcc0fbd6..7693567b 100644 --- a/tests/test_chatmessagehistory.py +++ b/tests/test_chatmessagehistory.py @@ -13,7 +13,7 @@ # limitations under the License. import os import uuid -from typing import Any, Generator +from typing import Any import pytest import pytest_asyncio diff --git a/tests/test_vectorstore.py b/tests/test_vectorstore.py index 2974812c..ee35ebb6 100644 --- a/tests/test_vectorstore.py +++ b/tests/test_vectorstore.py @@ -16,7 +16,7 @@ import os import uuid from threading import Thread -from typing import List, Sequence +from typing import Sequence import pytest import pytest_asyncio @@ -51,7 +51,7 @@ class FakeImageEmbedding(DeterministicFakeEmbedding): - def embed_image(self, image_paths: List[str]) -> List[List[float]]: + def embed_image(self, image_paths: list[str]) -> list[list[float]]: return [self.embed_query(path) for path in image_paths] diff --git a/tests/test_vectorstore_search.py b/tests/test_vectorstore_search.py index e17860dd..9cc0e056 100644 --- a/tests/test_vectorstore_search.py +++ b/tests/test_vectorstore_search.py @@ -14,7 +14,6 @@ import os import uuid -from typing import List import pytest import pytest_asyncio @@ -51,7 +50,7 @@ class FakeImageEmbedding(DeterministicFakeEmbedding): - def embed_image(self, image_paths: List[str]) -> List[List[float]]: + def embed_image(self, image_paths: list[str]) -> list[list[float]]: return [self.embed_query(path) for path in image_paths] diff --git a/tests/util_tests/test_pgvector_migrator.py b/tests/util_tests/test_pgvector_migrator.py index 49c0ef69..dc840022 100644 --- a/tests/util_tests/test_pgvector_migrator.py +++ b/tests/util_tests/test_pgvector_migrator.py @@ -16,7 +16,7 @@ import json import os import uuid -from typing import List, Optional, Sequence +from typing import Optional, Sequence from unittest import mock import pytest @@ -151,7 +151,7 @@ async def engine( await engine.close() @pytest.fixture(scope="module") - def sample_embeddings(self) -> List[float]: + def sample_embeddings(self) -> list[float]: return [0.1] * (VECTOR_SIZE - 1) + [0.2] def _create_metadata_for_collection( @@ -168,7 +168,7 @@ async def _create_collection( self, engine: AlloyDBEngine, collection_name: str, - sample_embeddings: List[float], + sample_embeddings: list[float], num_rows: int = 2, num_cols: int = 3, ) -> None: @@ -198,7 +198,7 @@ async def _create_collection( async def _create_pgvector_tables( self, engine: AlloyDBEngine, - sample_embeddings: List[float], + sample_embeddings: list[float], num_rows: int = 2, num_collections: int = 1, num_cols: int = 3,