mirror of
https://github.com/hwchase17/langchain
synced 2024-11-06 03:20:49 +00:00
921894960b
- Add langchain.llms.GooglePalm for text completion, - Add langchain.chat_models.ChatGooglePalm for chat completion, - Add langchain.embeddings.GooglePalmEmbeddings for sentence embeddings, - Add example field to HumanMessage and AIMessage so that users can feed in examples into the PaLM Chat API, - Add system and unit tests. Note async completion for the Text API is not yet supported and will be included in a future PR. Happy for feedback on any aspect of this PR, especially our choice of adding an example field to Human and AI Message objects to enable passing example messages to the API.
35 lines
1.0 KiB
Python
35 lines
1.0 KiB
Python
"""Test Google PaLM embeddings.
|
|
|
|
Note: This test must be run with the GOOGLE_API_KEY environment variable set to a
|
|
valid API key.
|
|
"""
|
|
from langchain.embeddings.google_palm import GooglePalmEmbeddings
|
|
|
|
|
|
def test_google_palm_embedding_documents() -> None:
|
|
"""Test Google PaLM embeddings."""
|
|
documents = ["foo bar"]
|
|
embedding = GooglePalmEmbeddings()
|
|
output = embedding.embed_documents(documents)
|
|
assert len(output) == 1
|
|
assert len(output[0]) == 768
|
|
|
|
|
|
def test_google_palm_embedding_documents_multiple() -> None:
|
|
"""Test Google PaLM embeddings."""
|
|
documents = ["foo bar", "bar foo", "foo"]
|
|
embedding = GooglePalmEmbeddings()
|
|
output = embedding.embed_documents(documents)
|
|
assert len(output) == 3
|
|
assert len(output[0]) == 768
|
|
assert len(output[1]) == 768
|
|
assert len(output[2]) == 768
|
|
|
|
|
|
def test_google_palm_embedding_query() -> None:
|
|
"""Test Google PaLM embeddings."""
|
|
document = "foo bar"
|
|
embedding = GooglePalmEmbeddings()
|
|
output = embedding.embed_query(document)
|
|
assert len(output) == 768
|