add few shot example (#148)

tests/integration_tests/vectorstores/__init__.py

@@ -0,0 +1 @@
+"""Test vectorstores."""

tests/integration_tests/vectorstores/test_elasticsearch.py

@@ -0,0 +1,42 @@
+"""Test ElasticSearch functionality."""
+from typing import List
+
+from langchain.docstore.document import Document
+from langchain.embeddings.base import Embeddings
+from langchain.vectorstores.elastic_vector_search import ElasticVectorSearch
+
+
+class FakeEmbeddings(Embeddings):
+    """Fake embeddings functionality for testing."""
+
+    def embed_documents(self, texts: List[str]) -> List[List[float]]:
+        """Return simple embeddings."""
+        return [[1.0] * 9 + [i] for i in range(len(texts))]
+
+    def embed_query(self, text: str) -> List[float]:
+        """Return simple embeddings."""
+        return [1.0] * 9 + [0.0]
+
+
+def test_elasticsearch() -> None:
+    """Test end to end construction and search."""
+    texts = ["foo", "bar", "baz"]
+    docsearch = ElasticVectorSearch.from_texts(
+        texts, FakeEmbeddings(), elasticsearch_url="http://localhost:9200"
+    )
+    output = docsearch.similarity_search("foo", k=1)
+    assert output == [Document(page_content="foo")]
+
+
+def test_elasticsearch_with_metadatas() -> None:
+    """Test end to end construction and search."""
+    texts = ["foo", "bar", "baz"]
+    metadatas = [{"page": i} for i in range(len(texts))]
+    docsearch = ElasticVectorSearch.from_texts(
+        texts,
+        FakeEmbeddings(),
+        metadatas=metadatas,
+        elasticsearch_url="http://localhost:9200",
+    )
+    output = docsearch.similarity_search("foo", k=1)
+    assert output == [Document(page_content="foo", metadata={"page": 0})]

tests/integration_tests/vectorstores/test_faiss.py

@@ -37,6 +37,23 @@ def test_faiss() -> None:
     assert output == [Document(page_content="foo")]
 
 
+def test_faiss_with_metadatas() -> None:
+    """Test end to end construction and search."""
+    texts = ["foo", "bar", "baz"]
+    metadatas = [{"page": i} for i in range(len(texts))]
+    docsearch = FAISS.from_texts(texts, FakeEmbeddings(), metadatas=metadatas)
+    expected_docstore = InMemoryDocstore(
+        {
+            "0": Document(page_content="foo", metadata={"page": 0}),
+            "1": Document(page_content="bar", metadata={"page": 1}),
+            "2": Document(page_content="baz", metadata={"page": 2}),
+        }
+    )
+    assert docsearch.docstore.__dict__ == expected_docstore.__dict__
+    output = docsearch.similarity_search("foo", k=1)
+    assert output == [Document(page_content="foo", metadata={"page": 0})]
+
+
 def test_faiss_search_not_found() -> None:
     """Test what happens when document is not found."""
     texts = ["foo", "bar", "baz"]