mirror of
https://github.com/hwchase17/langchain
synced 2024-11-06 03:20:49 +00:00
04c458a270
Improve the integration tests for Pinecone by adding an `.env.example` file for local testing. Additionally, add some dev dependencies specifically for integration tests. This change also helps me understand how Pinecone deals with certain things, see related issues https://github.com/hwchase17/langchain/issues/2484 https://github.com/hwchase17/langchain/issues/2816
141 lines
4.8 KiB
Python
141 lines
4.8 KiB
Python
"""Test ElasticSearch functionality."""
|
|
import logging
|
|
import os
|
|
import uuid
|
|
from typing import Generator, List, Union
|
|
|
|
import pytest
|
|
from elasticsearch import Elasticsearch
|
|
|
|
from langchain.docstore.document import Document
|
|
from langchain.embeddings import OpenAIEmbeddings
|
|
from langchain.vectorstores.elastic_vector_search import ElasticVectorSearch
|
|
from tests.integration_tests.vectorstores.fake_embeddings import FakeEmbeddings
|
|
|
|
logging.basicConfig(level=logging.DEBUG)
|
|
|
|
"""
|
|
cd tests/integration_tests/vectorstores/docker-compose
|
|
docker-compose -f elasticsearch.yml up
|
|
"""
|
|
|
|
|
|
class TestElasticsearch:
|
|
@classmethod
|
|
def setup_class(cls) -> None:
|
|
if not os.getenv("OPENAI_API_KEY"):
|
|
raise ValueError("OPENAI_API_KEY environment variable is not set")
|
|
|
|
@pytest.fixture(scope="class", autouse=True)
|
|
def elasticsearch_url(self) -> Union[str, Generator[str, None, None]]:
|
|
"""Return the elasticsearch url."""
|
|
url = "http://localhost:9200"
|
|
yield url
|
|
es = Elasticsearch(hosts=url)
|
|
|
|
# Clear all indexes
|
|
index_names = es.indices.get(index="_all").keys()
|
|
for index_name in index_names:
|
|
# print(index_name)
|
|
es.indices.delete(index=index_name)
|
|
|
|
def test_similarity_search_without_metadata(self, elasticsearch_url: str) -> None:
|
|
"""Test end to end construction and search without metadata."""
|
|
texts = ["foo", "bar", "baz"]
|
|
docsearch = ElasticVectorSearch.from_texts(
|
|
texts, FakeEmbeddings(), elasticsearch_url=elasticsearch_url
|
|
)
|
|
output = docsearch.similarity_search("foo", k=1)
|
|
assert output == [Document(page_content="foo")]
|
|
|
|
def test_similarity_search_with_metadata(self, elasticsearch_url: str) -> None:
|
|
"""Test end to end construction and search with metadata."""
|
|
texts = ["foo", "bar", "baz"]
|
|
metadatas = [{"page": i} for i in range(len(texts))]
|
|
docsearch = ElasticVectorSearch.from_texts(
|
|
texts,
|
|
FakeEmbeddings(),
|
|
metadatas=metadatas,
|
|
elasticsearch_url=elasticsearch_url,
|
|
)
|
|
output = docsearch.similarity_search("foo", k=1)
|
|
assert output == [Document(page_content="foo", metadata={"page": 0})]
|
|
|
|
@pytest.mark.vcr(ignore_localhost=True)
|
|
def test_default_index_from_documents(
|
|
self,
|
|
documents: List[Document],
|
|
embedding_openai: OpenAIEmbeddings,
|
|
elasticsearch_url: str,
|
|
) -> None:
|
|
"""This test checks the construction of a default
|
|
ElasticSearch index using the 'from_documents'."""
|
|
|
|
elastic_vector_search = ElasticVectorSearch.from_documents(
|
|
documents=documents,
|
|
embedding=embedding_openai,
|
|
elasticsearch_url=elasticsearch_url,
|
|
)
|
|
|
|
search_result = elastic_vector_search.similarity_search("sharks")
|
|
|
|
print(search_result)
|
|
assert len(search_result) != 0
|
|
|
|
@pytest.mark.vcr(ignore_localhost=True)
|
|
def test_custom_index_from_documents(
|
|
self,
|
|
documents: List[Document],
|
|
embedding_openai: OpenAIEmbeddings,
|
|
elasticsearch_url: str,
|
|
) -> None:
|
|
"""This test checks the construction of a custom
|
|
ElasticSearch index using the 'from_documents'."""
|
|
|
|
index_name = f"custom_index_{uuid.uuid4().hex}"
|
|
elastic_vector_search = ElasticVectorSearch.from_documents(
|
|
documents=documents,
|
|
embedding=embedding_openai,
|
|
elasticsearch_url=elasticsearch_url,
|
|
index_name=index_name,
|
|
)
|
|
es = Elasticsearch(hosts=elasticsearch_url)
|
|
index_names = es.indices.get(index="_all").keys()
|
|
assert index_name in index_names
|
|
|
|
search_result = elastic_vector_search.similarity_search("sharks")
|
|
print(search_result)
|
|
|
|
assert len(search_result) != 0
|
|
|
|
@pytest.mark.vcr(ignore_localhost=True)
|
|
def test_custom_index_add_documents(
|
|
self,
|
|
documents: List[Document],
|
|
embedding_openai: OpenAIEmbeddings,
|
|
elasticsearch_url: str,
|
|
) -> None:
|
|
"""This test checks the construction of a custom
|
|
ElasticSearch index using the 'add_documents'."""
|
|
|
|
index_name = f"custom_index_{uuid.uuid4().hex}"
|
|
elastic_vector_search = ElasticVectorSearch(
|
|
embedding=embedding_openai,
|
|
elasticsearch_url=elasticsearch_url,
|
|
index_name=index_name,
|
|
)
|
|
es = Elasticsearch(hosts=elasticsearch_url)
|
|
elastic_vector_search.add_documents(documents)
|
|
|
|
index_names = es.indices.get(index="_all").keys()
|
|
assert index_name in index_names
|
|
|
|
search_result = elastic_vector_search.similarity_search("sharks")
|
|
print(search_result)
|
|
|
|
assert len(search_result) != 0
|
|
|
|
def test_custom_index_add_documents_to_exists_store(self) -> None:
|
|
# TODO: implement it
|
|
pass
|