mirror of
https://github.com/hwchase17/langchain
synced 2024-11-06 03:20:49 +00:00
7b96de3d5d
- **Description:** Going forward, we have a own API `pip install gradientai`. Therefore gradually removing the self-build packages in llamaindex, haystack and langchain. - **Issue:** None. - **Dependencies:** `pip install gradientai` - **Tag maintainer:** @michaelfeil
129 lines
3.7 KiB
Python
129 lines
3.7 KiB
Python
import sys
|
|
from typing import Any, Dict, List
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
import pytest
|
|
|
|
from langchain_community.embeddings import GradientEmbeddings
|
|
|
|
_MODEL_ID = "my_model_valid_id"
|
|
_GRADIENT_SECRET = "secret_valid_token_123456"
|
|
_GRADIENT_WORKSPACE_ID = "valid_workspace_12345"
|
|
_GRADIENT_BASE_URL = "https://api.gradient.ai/api"
|
|
_DOCUMENTS = [
|
|
"pizza",
|
|
"another long pizza",
|
|
"a document",
|
|
"another long pizza",
|
|
"super long document with many tokens",
|
|
]
|
|
|
|
|
|
class GradientEmbeddingsModel(MagicMock):
|
|
"""MockGradientModel."""
|
|
|
|
def embed(self, inputs: List[Dict[str, str]]) -> Any:
|
|
"""Just duplicate the query m times."""
|
|
output = MagicMock()
|
|
|
|
embeddings = []
|
|
for i, inp in enumerate(inputs):
|
|
# verify correct ordering
|
|
inp = inp["input"]
|
|
if "pizza" in inp:
|
|
v = [1.0, 0.0, 0.0]
|
|
elif "document" in inp:
|
|
v = [0.0, 0.9, 0.0]
|
|
else:
|
|
v = [0.0, 0.0, -1.0]
|
|
if len(inp) > 10:
|
|
v[2] += 0.1
|
|
output_inner = MagicMock()
|
|
output_inner.embedding = v
|
|
embeddings.append(output_inner)
|
|
|
|
output.embeddings = embeddings
|
|
return output
|
|
|
|
async def aembed(self, *args) -> Any:
|
|
return self.embed(*args)
|
|
|
|
|
|
class MockGradient(MagicMock):
|
|
"""Mock Gradient package."""
|
|
|
|
def __init__(self, access_token: str, workspace_id, host):
|
|
assert access_token == _GRADIENT_SECRET
|
|
assert workspace_id == _GRADIENT_WORKSPACE_ID
|
|
assert host == _GRADIENT_BASE_URL
|
|
|
|
def get_embeddings_model(self, slug: str) -> GradientEmbeddingsModel:
|
|
assert slug == _MODEL_ID
|
|
return GradientEmbeddingsModel()
|
|
|
|
def close(self) -> None:
|
|
"""Mock Gradient close."""
|
|
return
|
|
|
|
|
|
class MockGradientaiPackage(MagicMock):
|
|
"""Mock Gradientai package."""
|
|
|
|
Gradient = MockGradient
|
|
__version__ = "1.4.0"
|
|
|
|
|
|
def test_gradient_llm_sync() -> None:
|
|
with patch.dict(sys.modules, {"gradientai": MockGradientaiPackage()}):
|
|
embedder = GradientEmbeddings(
|
|
gradient_api_url=_GRADIENT_BASE_URL,
|
|
gradient_access_token=_GRADIENT_SECRET,
|
|
gradient_workspace_id=_GRADIENT_WORKSPACE_ID,
|
|
model=_MODEL_ID,
|
|
)
|
|
assert embedder.gradient_access_token == _GRADIENT_SECRET
|
|
assert embedder.gradient_api_url == _GRADIENT_BASE_URL
|
|
assert embedder.gradient_workspace_id == _GRADIENT_WORKSPACE_ID
|
|
assert embedder.model == _MODEL_ID
|
|
|
|
response = embedder.embed_documents(_DOCUMENTS)
|
|
want = [
|
|
[1.0, 0.0, 0.0], # pizza
|
|
[1.0, 0.0, 0.1], # pizza + long
|
|
[0.0, 0.9, 0.0], # doc
|
|
[1.0, 0.0, 0.1], # pizza + long
|
|
[0.0, 0.9, 0.1], # doc + long
|
|
]
|
|
|
|
assert response == want
|
|
|
|
|
|
def test_gradient_wrong_setup() -> None:
|
|
with pytest.raises(Exception):
|
|
GradientEmbeddings(
|
|
gradient_api_url=_GRADIENT_BASE_URL,
|
|
gradient_access_token="", # empty
|
|
gradient_workspace_id=_GRADIENT_WORKSPACE_ID,
|
|
model=_MODEL_ID,
|
|
)
|
|
|
|
|
|
def test_gradient_wrong_setup2() -> None:
|
|
with pytest.raises(Exception):
|
|
GradientEmbeddings(
|
|
gradient_api_url=_GRADIENT_BASE_URL,
|
|
gradient_access_token=_GRADIENT_SECRET,
|
|
gradient_workspace_id="", # empty
|
|
model=_MODEL_ID,
|
|
)
|
|
|
|
|
|
def test_gradient_wrong_setup3() -> None:
|
|
with pytest.raises(Exception):
|
|
GradientEmbeddings(
|
|
gradient_api_url="-", # empty
|
|
gradient_access_token=_GRADIENT_SECRET,
|
|
gradient_workspace_id=_GRADIENT_WORKSPACE_ID,
|
|
model=_MODEL_ID,
|
|
)
|