Merge pull request #3 from mpskex/master

fixed unittest & lint
myscale · Jun 14, 2023 · 8de5059 · 8de5059
2 parents 33f122f + ff1618e
commit 8de5059
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 12 deletions.
diff --git a/langchain/retrievers/self_query/myscale.py b/langchain/retrievers/self_query/myscale.py
@@ -14,15 +14,15 @@
 
 def DEFAULT_COMPOSER(op_name: str) -> Callable:
     def f(*args: Any) -> str:
-        args: Tuple[str] = map(str, args)
-        return f" {op_name} ".join(args)
+        args_: map[str] = map(str, args)
+        return f" {op_name} ".join(args_)
     return f
 
 
 def FUNCTION_COMPOSER(op_name: str) -> Callable:
     def f(*args: Any) -> str:
-        args: Tuple[str] = map(str, args)
-        return f"{op_name}({','.join(args)})"
+        args_: map[str] = map(str, args)
+        return f"{op_name}({','.join(args_)})"
     return f
 
 

diff --git a/langchain/vectorstores/mongodb_atlas.py b/langchain/vectorstores/mongodb_atlas.py
@@ -9,10 +9,8 @@
     Iterable,
     List,
     Optional,
-    Sequence,
     Mapping,
     Tuple,
-    TypeVar,
     Union,
 )
 
@@ -23,8 +21,6 @@
 if TYPE_CHECKING:
     from pymongo.collection import Collection
 
-MongoDBDocumentType = TypeVar("MongoDBDocumentType", bound=Dict[str, Any])
-
 logger = logging.getLogger(__name__)
 
 DEFAULT_INSERT_BATCH_SIZE = 100
@@ -53,7 +49,7 @@ class MongoDBAtlasVectorSearch(VectorStore):
 
     def __init__(
         self,
-        collection: Collection[MongoDBDocumentType],
+        collection: Collection[Dict[str, Any]],
         embedding: Embeddings,
         *,
         index_name: str = "default",
@@ -131,7 +127,7 @@ def _insert_texts(self, texts: List[str], metadatas: List[Dict[str, Any]]) -> Li
             return []
         # Embed and create the documents
         embeddings = self._embedding.embed_documents(texts)
-        to_insert: List[MongoDBDocumentType] = [
+        to_insert: List[Dict[str, Any]] = [
             {self._text_key: t, self._embedding_key: embedding, **m}
             for t, m, embedding in zip(texts, metadatas, embeddings)
         ]
@@ -174,7 +170,7 @@ def similarity_search_with_score(
         }
         if pre_filter:
             knn_beta["filter"] = pre_filter
-        pipeline: Sequence[Mapping[str, Any]] = [
+        pipeline: List[Mapping[str, Any]] = [
             {
                 "$search": {
                     "index": self._index_name,
@@ -235,7 +231,7 @@ def from_texts(
         texts: List[str],
         embedding: Embeddings,
         metadatas: Optional[List[dict]] = None,
-        collection: Optional[Collection[MongoDBDocumentType]] = None,
+        collection: Optional[Collection[Dict[str, Any]]] = None,
         **kwargs: Any,
     ) -> MongoDBAtlasVectorSearch:
         """Construct MongoDBAtlasVectorSearch wrapper from raw documents.