Add tests in CI

Author: cbornet

langchain/langchain_vectorize/retrievers.py

@@ -2,7 +2,7 @@
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Any, Optional
+from typing import TYPE_CHECKING, Any, Literal, Optional
 
 import vectorize_client
 from langchain_core.documents import Document
@@ -40,6 +40,8 @@ class VectorizeRetriever(BaseRetriever):
 
     api_token: str
     """The Vectorize API token."""
+    environment: Literal["prod", "dev", "local", "staging"] = "prod"
+    """The Vectorize API environment."""
     organization: Optional[str] = None  # noqa: UP007
     """The Vectorize organization ID."""
     pipeline_id: Optional[str] = None  # noqa: UP007
@@ -55,7 +57,23 @@ class VectorizeRetriever(BaseRetriever):
 
     @override
     def model_post_init(self, /, context: Any) -> None:
-        api = ApiClient(Configuration(access_token=self.api_token))
+        header_name = None
+        header_value = None
+        if self.environment == "prod":
+            host = "https://api.vectorize.io/v1"
+        elif self.environment == "dev":
+            host = "https://api-dev.vectorize.io/v1"
+        elif self.environment == "local":
+            host = "http://localhost:3000/api"
+            header_name = "x-lambda-api-key"
+            header_value = self.api_token
+        else:
+            host = "https://api-staging.vectorize.io/v1"
+        api = ApiClient(
+            Configuration(host=host, access_token=self.api_token, debug=True),
+            header_name,
+            header_value,
+        )
         self._pipelines = PipelinesApi(api)
 
     @staticmethod

langchain/tests/test_retrievers.py

@@ -3,6 +3,7 @@
 import os
 from collections.abc import Iterator
 from pathlib import Path
+from typing import Literal
 
 import pytest
 import urllib3
@@ -31,15 +32,23 @@ def org_id() -> str:
 
 
 @pytest.fixture(scope="session")
-def api_client(api_token: str) -> Iterator[ApiClient]:
+def environment() -> Literal["prod", "dev", "local", "staging"]:
     env = os.getenv("VECTORIZE_ENV", "prod")
+    if env not in ["prod", "dev", "local", "staging"]:
+        msg = "Invalid VECTORIZE_ENV environment variable."
+        raise ValueError(msg)
+    return env
+
+
+@pytest.fixture(scope="session")
+def api_client(api_token: str, environment: str) -> Iterator[ApiClient]:
     header_name = None
     header_value = None
-    if env == "prod":
+    if environment == "prod":
         host = "https://api.vectorize.io/v1"
-    elif env == "dev":
+    elif environment == "dev":
         host = "https://api-dev.vectorize.io/v1"
-    elif env == "local":
+    elif environment == "local":
         host = "http://localhost:3000/api"
         header_name = "x-lambda-api-key"
         header_value = api_token
@@ -148,16 +157,30 @@ def pipeline_id(api_client: v.ApiClient, org_id: str) -> Iterator[str]:
         logging.exception("Failed to delete pipeline %s", pipeline_id)
 
 
-def test_retrieve_init_args(api_token: str, org_id: str, pipeline_id: str) -> None:
+def test_retrieve_init_args(
+    environment: Literal["prod", "dev", "local", "staging"],
+    api_token: str,
+    org_id: str,
+    pipeline_id: str,
+) -> None:
     retriever = VectorizeRetriever(
-        api_token=api_token, organization=org_id, pipeline_id=pipeline_id, num_results=2
+        environment=environment,
+        api_token=api_token,
+        organization=org_id,
+        pipeline_id=pipeline_id,
+        num_results=2,
     )
     docs = retriever.invoke(input="What are you?")
     assert len(docs) == 2
 
 
-def test_retrieve_invoke_args(api_token: str, org_id: str, pipeline_id: str) -> None:
-    retriever = VectorizeRetriever(api_token=api_token)
+def test_retrieve_invoke_args(
+    environment: Literal["prod", "dev", "local", "staging"],
+    api_token: str,
+    org_id: str,
+    pipeline_id: str,
+) -> None:
+    retriever = VectorizeRetriever(environment=environment, api_token=api_token)
     docs = retriever.invoke(
         input="What are you?",
         organization=org_id,