Skip to content

Commit 6641529

Browse files
cbornetcbornet
committed
Add langchain standard integration tests
1 parent 667cdfb commit 6641529

File tree

5 files changed

+1069
-165
lines changed

5 files changed

+1069
-165
lines changed

‎langchain/pyproject.toml‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,7 @@ dev = [
3434
"mypy>=1.13.0",
3535
"pytest>=8.3.3",
3636
"ruff>=0.9.0,<0.10",
37+
"langchain-tests>=0.3.20",
3738
]
3839

3940
[tool.ruff.lint]

‎langchain/tests/conftest.py‎

Lines changed: 155 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,155 @@
1+
import json
2+
import logging
3+
import os
4+
from collections.abc import Iterator
5+
from pathlib import Path
6+
from typing import Literal
7+
8+
import pytest
9+
import urllib3
10+
import vectorize_client as v
11+
from vectorize_client import ApiClient
12+
13+
14+
@pytest.fixture(scope="session")
15+
def api_token() -> str:
16+
token = os.getenv("VECTORIZE_TOKEN")
17+
if not token:
18+
msg = "Please set the VECTORIZE_TOKEN environment variable"
19+
raise ValueError(msg)
20+
return token
21+
22+
23+
@pytest.fixture(scope="session")
24+
def org_id() -> str:
25+
org = os.getenv("VECTORIZE_ORG")
26+
if not org:
27+
msg = "Please set the VECTORIZE_ORG environment variable"
28+
raise ValueError(msg)
29+
return org
30+
31+
32+
@pytest.fixture(scope="session")
33+
def environment() -> Literal["prod", "dev", "local", "staging"]:
34+
env = os.getenv("VECTORIZE_ENV", "prod")
35+
if env not in ["prod", "dev", "local", "staging"]:
36+
msg = "Invalid VECTORIZE_ENV environment variable."
37+
raise ValueError(msg)
38+
return env
39+
40+
41+
@pytest.fixture(scope="session")
42+
def api_client(api_token: str, environment: str) -> Iterator[ApiClient]:
43+
header_name = None
44+
header_value = None
45+
if environment == "prod":
46+
host = "https://api.vectorize.io/v1"
47+
elif environment == "dev":
48+
host = "https://api-dev.vectorize.io/v1"
49+
elif environment == "local":
50+
host = "http://localhost:3000/api"
51+
header_name = "x-lambda-api-key"
52+
header_value = api_token
53+
else:
54+
host = "https://api-staging.vectorize.io/v1"
55+
56+
with v.ApiClient(
57+
v.Configuration(host=host, access_token=api_token, debug=True),
58+
header_name,
59+
header_value,
60+
) as api:
61+
yield api
62+
63+
64+
@pytest.fixture(scope="session")
65+
def pipeline_id(api_client: v.ApiClient, org_id: str) -> Iterator[str]:
66+
pipelines = v.PipelinesApi(api_client)
67+
68+
connectors_api = v.ConnectorsApi(api_client)
69+
response = connectors_api.create_source_connector(
70+
org_id,
71+
[
72+
v.CreateSourceConnector(
73+
name="from api", type=v.SourceConnectorType.FILE_UPLOAD
74+
)
75+
],
76+
)
77+
source_connector_id = response.connectors[0].id
78+
logging.info("Created source connector %s", source_connector_id)
79+
80+
uploads_api = v.UploadsApi(api_client)
81+
upload_response = uploads_api.start_file_upload_to_connector(
82+
org_id,
83+
source_connector_id,
84+
v.StartFileUploadToConnectorRequest(
85+
name="research.pdf",
86+
content_type="application/pdf",
87+
metadata=json.dumps({"created-from-api": True}),
88+
),
89+
)
90+
91+
http = urllib3.PoolManager()
92+
this_dir = Path(__file__).parent
93+
file_path = this_dir / "research.pdf"
94+
95+
with file_path.open("rb") as f:
96+
http_response = http.request(
97+
"PUT",
98+
upload_response.upload_url,
99+
body=f,
100+
headers={
101+
"Content-Type": "application/pdf",
102+
"Content-Length": str(file_path.stat().st_size),
103+
},
104+
)
105+
if http_response.status != 200:
106+
msg = "Upload failed:"
107+
raise ValueError(msg)
108+
else:
109+
logging.info("Upload successful")
110+
111+
ai_platforms = connectors_api.get_ai_platform_connectors(org_id)
112+
builtin_ai_platform = next(
113+
c.id for c in ai_platforms.ai_platform_connectors if c.type == "VECTORIZE"
114+
)
115+
logging.info("Using AI platform %s", builtin_ai_platform)
116+
117+
vector_databases = connectors_api.get_destination_connectors(org_id)
118+
builtin_vector_db = next(
119+
c.id for c in vector_databases.destination_connectors if c.type == "VECTORIZE"
120+
)
121+
logging.info("Using destination connector %s", builtin_vector_db)
122+
123+
pipeline_response = pipelines.create_pipeline(
124+
org_id,
125+
v.PipelineConfigurationSchema(
126+
source_connectors=[
127+
v.SourceConnectorSchema(
128+
id=source_connector_id,
129+
type=v.SourceConnectorType.FILE_UPLOAD,
130+
config={},
131+
)
132+
],
133+
destination_connector=v.DestinationConnectorSchema(
134+
id=builtin_vector_db,
135+
type=v.DestinationConnectorType.VECTORIZE,
136+
config={},
137+
),
138+
ai_platform=v.AIPlatformSchema(
139+
id=builtin_ai_platform,
140+
type=v.AIPlatformType.VECTORIZE,
141+
config=v.AIPlatformConfigSchema(),
142+
),
143+
pipeline_name="Test pipeline",
144+
schedule=v.ScheduleSchema(type=v.ScheduleSchemaType.MANUAL),
145+
),
146+
)
147+
pipeline_id = pipeline_response.data.id
148+
logging.info("Created pipeline %s", pipeline_id)
149+
150+
yield pipeline_id
151+
152+
try:
153+
pipelines.delete_pipeline(org_id, pipeline_id)
154+
except Exception:
155+
logging.exception("Failed to delete pipeline %s", pipeline_id)

‎langchain/tests/test_langchain_integration_tests.py‎

Lines changed: 39 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,39 @@
1+
from typing import Literal
2+
3+
import pytest
4+
from langchain_tests.integration_tests import RetrieversIntegrationTests
5+
6+
from langchain_vectorize import VectorizeRetriever
7+
8+
9+
class TestParrotRetrieverIntegration(RetrieversIntegrationTests):
10+
@pytest.fixture(autouse=True, scope="class")
11+
@classmethod
12+
def setup(
13+
cls,
14+
environment: Literal["prod", "dev", "local", "staging"],
15+
api_token: str,
16+
org_id: str,
17+
pipeline_id: str,
18+
) -> None:
19+
cls.environment = environment
20+
cls.api_token = api_token
21+
cls.org_id = org_id
22+
cls.pipeline_id = pipeline_id
23+
24+
@property
25+
def retriever_constructor(self) -> type[VectorizeRetriever]:
26+
return VectorizeRetriever
27+
28+
@property
29+
def retriever_constructor_params(self) -> dict:
30+
return {
31+
"environment": self.environment,
32+
"api_token": self.api_token,
33+
"organization": self.org_id,
34+
"pipeline_id": self.pipeline_id,
35+
}
36+
37+
@property
38+
def retriever_query_example(self) -> str:
39+
return "What are you?"

‎langchain/tests/test_retrievers.py‎

Lines changed: 0 additions & 154 deletions
Original file line numberDiff line numberDiff line change
@@ -1,163 +1,9 @@
1-
import json
2-
import logging
3-
import os
41
import time
5-
from collections.abc import Iterator
6-
from pathlib import Path
72
from typing import Literal
83

9-
import pytest
10-
import urllib3
11-
import vectorize_client as v
12-
from vectorize_client import ApiClient
13-
144
from langchain_vectorize.retrievers import VectorizeRetriever
155

166

17-
@pytest.fixture(scope="session")
18-
def api_token() -> str:
19-
token = os.getenv("VECTORIZE_TOKEN")
20-
if not token:
21-
msg = "Please set the VECTORIZE_TOKEN environment variable"
22-
raise ValueError(msg)
23-
return token
24-
25-
26-
@pytest.fixture(scope="session")
27-
def org_id() -> str:
28-
org = os.getenv("VECTORIZE_ORG")
29-
if not org:
30-
msg = "Please set the VECTORIZE_ORG environment variable"
31-
raise ValueError(msg)
32-
return org
33-
34-
35-
@pytest.fixture(scope="session")
36-
def environment() -> Literal["prod", "dev", "local", "staging"]:
37-
env = os.getenv("VECTORIZE_ENV", "prod")
38-
if env not in ["prod", "dev", "local", "staging"]:
39-
msg = "Invalid VECTORIZE_ENV environment variable."
40-
raise ValueError(msg)
41-
return env
42-
43-
44-
@pytest.fixture(scope="session")
45-
def api_client(api_token: str, environment: str) -> Iterator[ApiClient]:
46-
header_name = None
47-
header_value = None
48-
if environment == "prod":
49-
host = "https://api.vectorize.io/v1"
50-
elif environment == "dev":
51-
host = "https://api-dev.vectorize.io/v1"
52-
elif environment == "local":
53-
host = "http://localhost:3000/api"
54-
header_name = "x-lambda-api-key"
55-
header_value = api_token
56-
else:
57-
host = "https://api-staging.vectorize.io/v1"
58-
59-
with v.ApiClient(
60-
v.Configuration(host=host, access_token=api_token, debug=True),
61-
header_name,
62-
header_value,
63-
) as api:
64-
yield api
65-
66-
67-
@pytest.fixture(scope="session")
68-
def pipeline_id(api_client: v.ApiClient, org_id: str) -> Iterator[str]:
69-
pipelines = v.PipelinesApi(api_client)
70-
71-
connectors_api = v.ConnectorsApi(api_client)
72-
response = connectors_api.create_source_connector(
73-
org_id,
74-
[
75-
v.CreateSourceConnector(
76-
name="from api", type=v.SourceConnectorType.FILE_UPLOAD
77-
)
78-
],
79-
)
80-
source_connector_id = response.connectors[0].id
81-
logging.info("Created source connector %s", source_connector_id)
82-
83-
uploads_api = v.UploadsApi(api_client)
84-
upload_response = uploads_api.start_file_upload_to_connector(
85-
org_id,
86-
source_connector_id,
87-
v.StartFileUploadToConnectorRequest(
88-
name="research.pdf",
89-
content_type="application/pdf",
90-
metadata=json.dumps({"created-from-api": True}),
91-
),
92-
)
93-
94-
http = urllib3.PoolManager()
95-
this_dir = Path(__file__).parent
96-
file_path = this_dir / "research.pdf"
97-
98-
with file_path.open("rb") as f:
99-
http_response = http.request(
100-
"PUT",
101-
upload_response.upload_url,
102-
body=f,
103-
headers={
104-
"Content-Type": "application/pdf",
105-
"Content-Length": str(file_path.stat().st_size),
106-
},
107-
)
108-
if http_response.status != 200:
109-
msg = "Upload failed:"
110-
raise ValueError(msg)
111-
else:
112-
logging.info("Upload successful")
113-
114-
ai_platforms = connectors_api.get_ai_platform_connectors(org_id)
115-
builtin_ai_platform = next(
116-
c.id for c in ai_platforms.ai_platform_connectors if c.type == "VECTORIZE"
117-
)
118-
logging.info("Using AI platform %s", builtin_ai_platform)
119-
120-
vector_databases = connectors_api.get_destination_connectors(org_id)
121-
builtin_vector_db = next(
122-
c.id for c in vector_databases.destination_connectors if c.type == "VECTORIZE"
123-
)
124-
logging.info("Using destination connector %s", builtin_vector_db)
125-
126-
pipeline_response = pipelines.create_pipeline(
127-
org_id,
128-
v.PipelineConfigurationSchema(
129-
source_connectors=[
130-
v.SourceConnectorSchema(
131-
id=source_connector_id,
132-
type=v.SourceConnectorType.FILE_UPLOAD,
133-
config={},
134-
)
135-
],
136-
destination_connector=v.DestinationConnectorSchema(
137-
id=builtin_vector_db,
138-
type=v.DestinationConnectorType.VECTORIZE,
139-
config={},
140-
),
141-
ai_platform=v.AIPlatformSchema(
142-
id=builtin_ai_platform,
143-
type=v.AIPlatformType.VECTORIZE,
144-
config=v.AIPlatformConfigSchema(),
145-
),
146-
pipeline_name="Test pipeline",
147-
schedule=v.ScheduleSchema(type=v.ScheduleSchemaType.MANUAL),
148-
),
149-
)
150-
pipeline_id = pipeline_response.data.id
151-
logging.info("Created pipeline %s", pipeline_id)
152-
153-
yield pipeline_id
154-
155-
try:
156-
pipelines.delete_pipeline(org_id, pipeline_id)
157-
except Exception:
158-
logging.exception("Failed to delete pipeline %s", pipeline_id)
159-
160-
1617
def test_retrieve_init_args(
1628
environment: Literal["prod", "dev", "local", "staging"],
1639
api_token: str,

0 commit comments

Comments
 (0)