Skip to content

nvidia_nim

__all__ = ['NvidiaNimEmbeddingDriver'] module-attribute

NvidiaNimEmbeddingDriver

Bases: OpenAiEmbeddingDriver

Nvidia Embedding Driver. The API is OpenAI compatible, but requires an extra parameter 'input_type'.

Source code in griptape/drivers/embedding/nvidia_nim_embedding_driver.py 
11
12
13
14
15
16
17
18
19
20
21
22
23
@define
class NvidiaNimEmbeddingDriver(OpenAiEmbeddingDriver):
    """Nvidia Embedding Driver. The API is OpenAI compatible, but requires an extra parameter 'input_type'."""

    def try_embed_chunk(self, chunk: str, *, vector_operation: VectorOperation | None = None, **kwargs) -> list[float]:
        if vector_operation not in get_args(VectorOperation):
            raise ValueError(f"invalid value for vector_operation, must be one of {get_args(VectorOperation)}")

        extra_body = {
            "input_type": "query" if vector_operation == "query" else "passage",
        }

        return self.client.embeddings.create(**self._params(chunk), extra_body=extra_body).data[0].embedding

try_embed_chunk(chunk, *, vector_operation=None, **kwargs)

Source code in griptape/drivers/embedding/nvidia_nim_embedding_driver.py 
15
16
17
18
19
20
21
22
23
def try_embed_chunk(self, chunk: str, *, vector_operation: VectorOperation | None = None, **kwargs) -> list[float]:
    if vector_operation not in get_args(VectorOperation):
        raise ValueError(f"invalid value for vector_operation, must be one of {get_args(VectorOperation)}")

    extra_body = {
        "input_type": "query" if vector_operation == "query" else "passage",
    }

    return self.client.embeddings.create(**self._params(chunk), extra_body=extra_body).data[0].embedding