chore: import upstream snapshot with attribution
This commit is contained in:
@@ -0,0 +1,66 @@
|
||||
# semantic_kernel.connectors.ai.nvidia
|
||||
|
||||
This connector enables integration with NVIDIA NIM API for text embeddings and chat completion. It allows you to use NVIDIA's models within the Semantic Kernel framework.
|
||||
|
||||
## Quick start
|
||||
|
||||
### Initialize the kernel
|
||||
```python
|
||||
import semantic_kernel as sk
|
||||
kernel = sk.Kernel()
|
||||
```
|
||||
|
||||
### Add NVIDIA text embedding service
|
||||
You can provide your API key directly or through environment variables
|
||||
```python
|
||||
from semantic_kernel.connectors.ai.nvidia import NvidiaTextEmbedding
|
||||
|
||||
embedding_service = NvidiaTextEmbedding(
|
||||
ai_model_id="nvidia/nv-embedqa-e5-v5", # Default model if not specified
|
||||
api_key="your-nvidia-api-key", # Can also use NVIDIA_API_KEY env variable
|
||||
service_id="nvidia-embeddings" # Optional service identifier
|
||||
)
|
||||
```
|
||||
|
||||
### Add the embedding service to the kernel
|
||||
```python
|
||||
kernel.add_service(embedding_service)
|
||||
```
|
||||
|
||||
### Generate embeddings for text
|
||||
```python
|
||||
texts = ["Hello, world!", "Semantic Kernel is awesome"]
|
||||
embeddings = await kernel.get_service("nvidia-embeddings").generate_embeddings(texts)
|
||||
```
|
||||
|
||||
### Add NVIDIA chat completion service
|
||||
```python
|
||||
from semantic_kernel.connectors.ai.nvidia import NvidiaChatCompletion
|
||||
|
||||
chat_service = NvidiaChatCompletion(
|
||||
ai_model_id="meta/llama-3.1-8b-instruct", # Default model if not specified
|
||||
api_key="your-nvidia-api-key", # Can also use NVIDIA_API_KEY env variable
|
||||
service_id="nvidia-chat" # Optional service identifier
|
||||
)
|
||||
kernel.add_service(chat_service)
|
||||
```
|
||||
|
||||
### Basic chat completion
|
||||
```python
|
||||
response = await kernel.invoke_prompt("Hello, how are you?")
|
||||
```
|
||||
|
||||
### Using with Chat Completion Agent
|
||||
```python
|
||||
from semantic_kernel.agents import ChatCompletionAgent
|
||||
from semantic_kernel.connectors.ai.nvidia import NvidiaChatCompletion
|
||||
|
||||
agent = ChatCompletionAgent(
|
||||
service=NvidiaChatCompletion(),
|
||||
name="SK-Assistant",
|
||||
instructions="You are a helpful assistant.",
|
||||
)
|
||||
response = await agent.get_response(messages="Write a haiku about Semantic Kernel.")
|
||||
print(response.content)
|
||||
```
|
||||
|
||||
@@ -0,0 +1,19 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
from semantic_kernel.connectors.ai.nvidia.prompt_execution_settings.nvidia_prompt_execution_settings import (
|
||||
NvidiaChatPromptExecutionSettings,
|
||||
NvidiaEmbeddingPromptExecutionSettings,
|
||||
NvidiaPromptExecutionSettings,
|
||||
)
|
||||
from semantic_kernel.connectors.ai.nvidia.services.nvidia_chat_completion import NvidiaChatCompletion
|
||||
from semantic_kernel.connectors.ai.nvidia.services.nvidia_text_embedding import NvidiaTextEmbedding
|
||||
from semantic_kernel.connectors.ai.nvidia.settings.nvidia_settings import NvidiaSettings
|
||||
|
||||
__all__ = [
|
||||
"NvidiaChatCompletion",
|
||||
"NvidiaChatPromptExecutionSettings",
|
||||
"NvidiaEmbeddingPromptExecutionSettings",
|
||||
"NvidiaPromptExecutionSettings",
|
||||
"NvidiaSettings",
|
||||
"NvidiaTextEmbedding",
|
||||
]
|
||||
@@ -0,0 +1 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
+73
@@ -0,0 +1,73 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
from typing import Annotated, Any, Literal
|
||||
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
from semantic_kernel.connectors.ai.prompt_execution_settings import PromptExecutionSettings
|
||||
|
||||
|
||||
class NvidiaPromptExecutionSettings(PromptExecutionSettings):
|
||||
"""Settings for NVIDIA prompt execution."""
|
||||
|
||||
format: Literal["json"] | None = None
|
||||
options: dict[str, Any] | None = None
|
||||
|
||||
|
||||
class NvidiaEmbeddingPromptExecutionSettings(NvidiaPromptExecutionSettings):
|
||||
"""Settings for NVIDIA embedding prompt execution."""
|
||||
|
||||
input: str | list[str] | None = None
|
||||
ai_model_id: Annotated[str | None, Field(serialization_alias="model")] = None
|
||||
encoding_format: Literal["float", "base64"] = "float"
|
||||
truncate: Literal["NONE", "START", "END"] = "NONE"
|
||||
input_type: Literal["passage", "query"] = "query" # required param with default value query
|
||||
user: str | None = None
|
||||
extra_headers: dict | None = None
|
||||
extra_body: dict | None = None
|
||||
timeout: float | None = None
|
||||
dimensions: Annotated[int | None, Field(gt=0)] = None
|
||||
|
||||
def prepare_settings_dict(self, **kwargs) -> dict[str, Any]:
|
||||
"""Override only for embeddings to exclude input_type and truncate."""
|
||||
return self.model_dump(
|
||||
exclude={"service_id", "extension_data", "structured_json_response", "input_type", "truncate"},
|
||||
exclude_none=True,
|
||||
by_alias=True,
|
||||
)
|
||||
|
||||
|
||||
class NvidiaChatPromptExecutionSettings(NvidiaPromptExecutionSettings):
|
||||
"""Settings for NVIDIA chat prompt execution."""
|
||||
|
||||
messages: list[dict[str, str]] | None = None
|
||||
ai_model_id: Annotated[str | None, Field(serialization_alias="model")] = None
|
||||
temperature: float | None = None
|
||||
top_p: float | None = None
|
||||
n: int | None = None
|
||||
stream: bool = False
|
||||
stop: str | list[str] | None = None
|
||||
max_tokens: int | None = None
|
||||
presence_penalty: float | None = None
|
||||
frequency_penalty: float | None = None
|
||||
logit_bias: dict[str, float] | None = None
|
||||
user: str | None = None
|
||||
tools: list[dict[str, Any]] | None = None
|
||||
tool_choice: str | dict[str, Any] | None = None
|
||||
response_format: (
|
||||
dict[Literal["type"], Literal["text", "json_object"]] | dict[str, Any] | type[BaseModel] | type | None
|
||||
) = None
|
||||
seed: int | None = None
|
||||
extra_headers: dict | None = None
|
||||
extra_body: dict | None = None
|
||||
timeout: float | None = None
|
||||
# NVIDIA-specific structured output support
|
||||
nvext: dict[str, Any] | None = None
|
||||
|
||||
def prepare_settings_dict(self, **kwargs) -> dict[str, Any]:
|
||||
"""Override only for embeddings to exclude input_type and truncate."""
|
||||
return self.model_dump(
|
||||
exclude={"service_id", "extension_data", "structured_json_response", "response_format"},
|
||||
exclude_none=True,
|
||||
by_alias=True,
|
||||
)
|
||||
@@ -0,0 +1 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
@@ -0,0 +1,313 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
import logging
|
||||
import sys
|
||||
from collections.abc import AsyncGenerator
|
||||
from typing import Any, Literal
|
||||
|
||||
from openai import AsyncOpenAI
|
||||
from openai.types.chat.chat_completion import ChatCompletion, Choice
|
||||
from openai.types.chat.chat_completion_chunk import ChatCompletionChunk
|
||||
from openai.types.chat.chat_completion_chunk import Choice as ChunkChoice
|
||||
from pydantic import ValidationError
|
||||
|
||||
from semantic_kernel.connectors.ai.chat_completion_client_base import ChatCompletionClientBase
|
||||
from semantic_kernel.connectors.ai.completion_usage import CompletionUsage
|
||||
from semantic_kernel.connectors.ai.nvidia.prompt_execution_settings.nvidia_prompt_execution_settings import (
|
||||
NvidiaChatPromptExecutionSettings,
|
||||
)
|
||||
from semantic_kernel.connectors.ai.nvidia.services.nvidia_handler import NvidiaHandler
|
||||
from semantic_kernel.connectors.ai.nvidia.services.nvidia_model_types import NvidiaModelTypes
|
||||
from semantic_kernel.connectors.ai.nvidia.settings.nvidia_settings import NvidiaSettings
|
||||
from semantic_kernel.connectors.ai.prompt_execution_settings import PromptExecutionSettings
|
||||
from semantic_kernel.contents import (
|
||||
AuthorRole,
|
||||
ChatMessageContent,
|
||||
FinishReason,
|
||||
FunctionCallContent,
|
||||
StreamingChatMessageContent,
|
||||
StreamingTextContent,
|
||||
TextContent,
|
||||
)
|
||||
from semantic_kernel.contents.chat_history import ChatHistory
|
||||
from semantic_kernel.exceptions.service_exceptions import ServiceInitializationError
|
||||
from semantic_kernel.utils.feature_stage_decorator import experimental
|
||||
from semantic_kernel.utils.telemetry.model_diagnostics.decorators import (
|
||||
trace_chat_completion,
|
||||
trace_streaming_chat_completion,
|
||||
)
|
||||
|
||||
if sys.version_info >= (3, 12):
|
||||
from typing import override # pragma: no cover
|
||||
else:
|
||||
from typing_extensions import override # pragma: no cover
|
||||
|
||||
logger: logging.Logger = logging.getLogger(__name__)
|
||||
|
||||
# Default NVIDIA chat model when none is specified
|
||||
DEFAULT_NVIDIA_CHAT_MODEL = "meta/llama-3.1-8b-instruct"
|
||||
|
||||
|
||||
@experimental
|
||||
class NvidiaChatCompletion(NvidiaHandler, ChatCompletionClientBase):
|
||||
"""NVIDIA Chat completion class.
|
||||
|
||||
This class does not support function calling. The SUPPORTS_FUNCTION_CALLING attribute
|
||||
is set to False (inherited from the base class).
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
ai_model_id: str | None = None,
|
||||
api_key: str | None = None,
|
||||
base_url: str | None = None,
|
||||
service_id: str | None = None,
|
||||
client: AsyncOpenAI | None = None,
|
||||
env_file_path: str | None = None,
|
||||
env_file_encoding: str | None = None,
|
||||
instruction_role: Literal["system", "user", "assistant", "developer"] | None = None,
|
||||
) -> None:
|
||||
"""Initialize an NvidiaChatCompletion service.
|
||||
|
||||
Args:
|
||||
ai_model_id (str): NVIDIA model name, see
|
||||
https://docs.api.nvidia.com/nim/reference/
|
||||
If not provided, defaults to DEFAULT_NVIDIA_CHAT_MODEL.
|
||||
service_id (str | None): Service ID tied to the execution settings.
|
||||
api_key (str | None): The optional API key to use. If provided will override,
|
||||
the env vars or .env file value.
|
||||
base_url (str | None): Custom API endpoint. (Optional)
|
||||
client (Optional[AsyncOpenAI]): An existing client to use. (Optional)
|
||||
env_file_path (str | None): Use the environment settings file as a fallback
|
||||
to environment variables. (Optional)
|
||||
env_file_encoding (str | None): The encoding of the environment settings file. (Optional)
|
||||
instruction_role (Literal["system", "user", "assistant", "developer"] | None): The role to use for
|
||||
'instruction' messages. Defaults to "system". (Optional)
|
||||
"""
|
||||
try:
|
||||
nvidia_settings = NvidiaSettings(
|
||||
api_key=api_key,
|
||||
base_url=base_url,
|
||||
chat_model_id=ai_model_id,
|
||||
env_file_path=env_file_path,
|
||||
env_file_encoding=env_file_encoding,
|
||||
)
|
||||
except ValidationError as ex:
|
||||
raise ServiceInitializationError("Failed to create NVIDIA settings.", ex) from ex
|
||||
|
||||
if not client and not nvidia_settings.api_key:
|
||||
raise ServiceInitializationError("The NVIDIA API key is required.")
|
||||
if not nvidia_settings.chat_model_id:
|
||||
# Default fallback model
|
||||
nvidia_settings.chat_model_id = DEFAULT_NVIDIA_CHAT_MODEL
|
||||
logger.warning(f"Default chat model set as: {nvidia_settings.chat_model_id}")
|
||||
|
||||
# Create client if not provided
|
||||
if not client:
|
||||
client = AsyncOpenAI(
|
||||
api_key=nvidia_settings.api_key.get_secret_value() if nvidia_settings.api_key else None,
|
||||
base_url=nvidia_settings.base_url,
|
||||
)
|
||||
|
||||
super().__init__(
|
||||
ai_model_id=nvidia_settings.chat_model_id,
|
||||
api_key=nvidia_settings.api_key.get_secret_value() if nvidia_settings.api_key else None,
|
||||
base_url=nvidia_settings.base_url,
|
||||
service_id=service_id or "",
|
||||
ai_model_type=NvidiaModelTypes.CHAT,
|
||||
client=client,
|
||||
instruction_role=instruction_role or "system",
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls: type["NvidiaChatCompletion"], settings: dict[str, Any]) -> "NvidiaChatCompletion":
|
||||
"""Initialize an NVIDIA service from a dictionary of settings.
|
||||
|
||||
Args:
|
||||
settings: A dictionary of settings for the service.
|
||||
"""
|
||||
return cls(
|
||||
ai_model_id=settings.get("ai_model_id"),
|
||||
api_key=settings.get("api_key"),
|
||||
base_url=settings.get("base_url"),
|
||||
service_id=settings.get("service_id"),
|
||||
env_file_path=settings.get("env_file_path"),
|
||||
)
|
||||
|
||||
@override
|
||||
def get_prompt_execution_settings_class(self) -> type["PromptExecutionSettings"]:
|
||||
return NvidiaChatPromptExecutionSettings
|
||||
|
||||
@override
|
||||
@trace_chat_completion("nvidia")
|
||||
async def _inner_get_chat_message_contents(
|
||||
self,
|
||||
chat_history: "ChatHistory",
|
||||
settings: "PromptExecutionSettings",
|
||||
) -> list["ChatMessageContent"]:
|
||||
if not isinstance(settings, NvidiaChatPromptExecutionSettings):
|
||||
settings = self.get_prompt_execution_settings_from_settings(settings)
|
||||
assert isinstance(settings, NvidiaChatPromptExecutionSettings) # nosec
|
||||
|
||||
settings.stream = False
|
||||
settings.messages = self._prepare_chat_history_for_request(chat_history)
|
||||
settings.ai_model_id = settings.ai_model_id or self.ai_model_id
|
||||
|
||||
# Handle structured output
|
||||
self._handle_structured_output(settings)
|
||||
|
||||
response = await self._send_request(settings)
|
||||
assert isinstance(response, ChatCompletion) # nosec
|
||||
response_metadata = self._get_metadata_from_chat_response(response)
|
||||
return [self._create_chat_message_content(response, choice, response_metadata) for choice in response.choices]
|
||||
|
||||
@override
|
||||
@trace_streaming_chat_completion("nvidia")
|
||||
async def _inner_get_streaming_chat_message_contents(
|
||||
self,
|
||||
chat_history: "ChatHistory",
|
||||
settings: "PromptExecutionSettings",
|
||||
function_invoke_attempt: int = 0,
|
||||
) -> AsyncGenerator[list["StreamingChatMessageContent"], Any]:
|
||||
if not isinstance(settings, NvidiaChatPromptExecutionSettings):
|
||||
settings = self.get_prompt_execution_settings_from_settings(settings)
|
||||
assert isinstance(settings, NvidiaChatPromptExecutionSettings) # nosec
|
||||
|
||||
settings.stream = True
|
||||
settings.messages = self._prepare_chat_history_for_request(chat_history)
|
||||
settings.ai_model_id = settings.ai_model_id or self.ai_model_id
|
||||
|
||||
# Handle structured output
|
||||
self._handle_structured_output(settings)
|
||||
|
||||
response = await self._send_request(settings)
|
||||
assert isinstance(response, AsyncGenerator) # nosec
|
||||
|
||||
async for chunk in response:
|
||||
if len(chunk.choices) == 0:
|
||||
continue
|
||||
chunk_metadata = self._get_metadata_from_chat_response(chunk)
|
||||
yield [
|
||||
self._create_streaming_chat_message_content(chunk, choice, chunk_metadata, function_invoke_attempt)
|
||||
for choice in chunk.choices
|
||||
]
|
||||
|
||||
def _create_chat_message_content(
|
||||
self, response: ChatCompletion, choice: Choice, response_metadata: dict[str, Any]
|
||||
) -> "ChatMessageContent":
|
||||
"""Create a chat message content object from a choice."""
|
||||
metadata = self._get_metadata_from_chat_choice(choice)
|
||||
metadata.update(response_metadata)
|
||||
|
||||
items: list[Any] = self._get_tool_calls_from_chat_choice(choice)
|
||||
items.extend(self._get_function_call_from_chat_choice(choice))
|
||||
if choice.message.content:
|
||||
items.append(TextContent(text=choice.message.content))
|
||||
|
||||
return ChatMessageContent(
|
||||
inner_content=response,
|
||||
ai_model_id=self.ai_model_id,
|
||||
metadata=metadata,
|
||||
role=AuthorRole(choice.message.role),
|
||||
items=items,
|
||||
finish_reason=(FinishReason(choice.finish_reason) if choice.finish_reason else None),
|
||||
)
|
||||
|
||||
def _create_streaming_chat_message_content(
|
||||
self,
|
||||
chunk: ChatCompletionChunk,
|
||||
choice: ChunkChoice,
|
||||
chunk_metadata: dict[str, Any],
|
||||
function_invoke_attempt: int,
|
||||
) -> StreamingChatMessageContent:
|
||||
"""Create a streaming chat message content object from a choice."""
|
||||
metadata = self._get_metadata_from_chat_choice(choice)
|
||||
metadata.update(chunk_metadata)
|
||||
|
||||
items: list[Any] = self._get_tool_calls_from_chat_choice(choice)
|
||||
items.extend(self._get_function_call_from_chat_choice(choice))
|
||||
if choice.delta and choice.delta.content is not None:
|
||||
items.append(StreamingTextContent(choice_index=choice.index, text=choice.delta.content))
|
||||
return StreamingChatMessageContent(
|
||||
choice_index=choice.index,
|
||||
inner_content=chunk,
|
||||
ai_model_id=self.ai_model_id,
|
||||
metadata=metadata,
|
||||
role=(AuthorRole(choice.delta.role) if choice.delta and choice.delta.role else AuthorRole.ASSISTANT),
|
||||
finish_reason=(FinishReason(choice.finish_reason) if choice.finish_reason else None),
|
||||
items=items,
|
||||
function_invoke_attempt=function_invoke_attempt,
|
||||
)
|
||||
|
||||
def _get_metadata_from_chat_response(self, response: ChatCompletion | ChatCompletionChunk) -> dict[str, Any]:
|
||||
"""Get metadata from a chat response."""
|
||||
return {
|
||||
"id": response.id,
|
||||
"created": response.created,
|
||||
"system_fingerprint": getattr(response, "system_fingerprint", None),
|
||||
"usage": CompletionUsage.from_openai(response.usage) if response.usage is not None else None,
|
||||
}
|
||||
|
||||
def _get_metadata_from_chat_choice(self, choice: Choice | ChunkChoice) -> dict[str, Any]:
|
||||
"""Get metadata from a chat choice."""
|
||||
return {
|
||||
"logprobs": getattr(choice, "logprobs", None),
|
||||
}
|
||||
|
||||
def _get_tool_calls_from_chat_choice(self, choice: Choice | ChunkChoice) -> list[FunctionCallContent]:
|
||||
"""Get tool calls from a chat choice."""
|
||||
content = choice.message if isinstance(choice, Choice) else choice.delta
|
||||
if content and (tool_calls := getattr(content, "tool_calls", None)) is not None:
|
||||
return [
|
||||
FunctionCallContent(
|
||||
id=tool.id,
|
||||
index=getattr(tool, "index", None),
|
||||
name=tool.function.name,
|
||||
arguments=tool.function.arguments,
|
||||
)
|
||||
for tool in tool_calls
|
||||
]
|
||||
return []
|
||||
|
||||
def _get_function_call_from_chat_choice(self, choice: Choice | ChunkChoice) -> list[FunctionCallContent]:
|
||||
"""Get function calls from a chat choice."""
|
||||
content = choice.message if isinstance(choice, Choice) else choice.delta
|
||||
if content and (function_call := getattr(content, "function_call", None)) is not None:
|
||||
return [
|
||||
FunctionCallContent(
|
||||
id="",
|
||||
name=function_call.name,
|
||||
arguments=function_call.arguments,
|
||||
)
|
||||
]
|
||||
return []
|
||||
|
||||
def _handle_structured_output(self, request_settings: NvidiaChatPromptExecutionSettings) -> None:
|
||||
"""Handle structured output for NVIDIA models using nvext parameter."""
|
||||
response_format = getattr(request_settings, "response_format", None)
|
||||
if response_format:
|
||||
# Convert Pydantic model to JSON schema for NVIDIA's guided_json
|
||||
if hasattr(response_format, "model_json_schema"):
|
||||
# It's a Pydantic model
|
||||
schema = response_format.model_json_schema()
|
||||
if not request_settings.extra_body:
|
||||
request_settings.extra_body = {}
|
||||
request_settings.extra_body["nvext"] = {"guided_json": schema}
|
||||
elif isinstance(response_format, dict):
|
||||
# It's already a dict, use it directly
|
||||
if not request_settings.extra_body:
|
||||
request_settings.extra_body = {}
|
||||
request_settings.extra_body["nvext"] = {"guided_json": response_format}
|
||||
|
||||
def _prepare_chat_history_for_request(
|
||||
self,
|
||||
chat_history: ChatHistory,
|
||||
role_key: str = "role",
|
||||
content_key: str = "content",
|
||||
) -> list[dict[str, str]]:
|
||||
"""Prepare chat history for request."""
|
||||
messages = []
|
||||
for message in chat_history.messages:
|
||||
message_dict = {role_key: message.role.value, content_key: message.content}
|
||||
messages.append(message_dict)
|
||||
return messages
|
||||
@@ -0,0 +1,123 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
import logging
|
||||
from abc import ABC
|
||||
from typing import Any, ClassVar, Union
|
||||
|
||||
from openai import AsyncOpenAI, AsyncStream
|
||||
from openai.types.chat.chat_completion import ChatCompletion
|
||||
from openai.types.chat.chat_completion_chunk import ChatCompletionChunk
|
||||
from openai.types.completion import Completion
|
||||
from openai.types.create_embedding_response import CreateEmbeddingResponse
|
||||
|
||||
from semantic_kernel.connectors.ai.nvidia.prompt_execution_settings.nvidia_prompt_execution_settings import (
|
||||
NvidiaChatPromptExecutionSettings,
|
||||
NvidiaEmbeddingPromptExecutionSettings,
|
||||
)
|
||||
from semantic_kernel.connectors.ai.nvidia.services.nvidia_model_types import NvidiaModelTypes
|
||||
from semantic_kernel.connectors.ai.prompt_execution_settings import PromptExecutionSettings
|
||||
from semantic_kernel.const import USER_AGENT
|
||||
from semantic_kernel.exceptions import ServiceResponseException
|
||||
from semantic_kernel.kernel_pydantic import KernelBaseModel
|
||||
|
||||
logger: logging.Logger = logging.getLogger(__name__)
|
||||
|
||||
RESPONSE_TYPE = Union[list[Any], ChatCompletion, Completion, AsyncStream[Any]]
|
||||
|
||||
|
||||
class NvidiaHandler(KernelBaseModel, ABC):
|
||||
"""Internal class for calls to Nvidia API's."""
|
||||
|
||||
MODEL_PROVIDER_NAME: ClassVar[str] = "nvidia"
|
||||
client: AsyncOpenAI
|
||||
ai_model_type: NvidiaModelTypes = NvidiaModelTypes.CHAT
|
||||
completion_tokens: int = 0
|
||||
total_tokens: int = 0
|
||||
prompt_tokens: int = 0
|
||||
|
||||
async def _send_request(self, settings: PromptExecutionSettings) -> RESPONSE_TYPE:
|
||||
"""Send a request to the Nvidia API."""
|
||||
if self.ai_model_type == NvidiaModelTypes.EMBEDDING:
|
||||
assert isinstance(settings, NvidiaEmbeddingPromptExecutionSettings) # nosec
|
||||
return await self._send_embedding_request(settings)
|
||||
if self.ai_model_type == NvidiaModelTypes.CHAT:
|
||||
assert isinstance(settings, NvidiaChatPromptExecutionSettings) # nosec
|
||||
return await self._send_chat_completion_request(settings)
|
||||
|
||||
raise NotImplementedError(f"Model type {self.ai_model_type} is not supported")
|
||||
|
||||
async def _send_embedding_request(self, settings: NvidiaEmbeddingPromptExecutionSettings) -> list[Any]:
|
||||
"""Send a request to the OpenAI embeddings endpoint."""
|
||||
try:
|
||||
# unsupported parameters are internally excluded from main dict and added to extra_body
|
||||
response = await self.client.embeddings.create(**settings.prepare_settings_dict())
|
||||
|
||||
self.store_usage(response)
|
||||
return [x.embedding for x in response.data]
|
||||
except Exception as ex:
|
||||
raise ServiceResponseException(
|
||||
f"{type(self)} service failed to generate embeddings",
|
||||
ex,
|
||||
) from ex
|
||||
|
||||
async def _send_chat_completion_request(
|
||||
self, settings: NvidiaChatPromptExecutionSettings
|
||||
) -> ChatCompletion | AsyncStream[Any]:
|
||||
"""Send a request to the NVIDIA chat completion endpoint."""
|
||||
try:
|
||||
settings_dict = settings.prepare_settings_dict()
|
||||
|
||||
# Handle structured output if nvext is present in extra_body
|
||||
if settings.extra_body and "nvext" in settings.extra_body:
|
||||
if "extra_body" not in settings_dict:
|
||||
settings_dict["extra_body"] = {}
|
||||
settings_dict["extra_body"]["nvext"] = settings.extra_body["nvext"]
|
||||
|
||||
response = await self.client.chat.completions.create(**settings_dict)
|
||||
self.store_usage(response)
|
||||
return response
|
||||
except Exception as ex:
|
||||
raise ServiceResponseException(
|
||||
f"{type(self)} service failed to complete the chat",
|
||||
ex,
|
||||
) from ex
|
||||
|
||||
def store_usage(
|
||||
self,
|
||||
response: ChatCompletion
|
||||
| Completion
|
||||
| AsyncStream[ChatCompletionChunk]
|
||||
| AsyncStream[Completion]
|
||||
| CreateEmbeddingResponse,
|
||||
):
|
||||
"""Store the usage information from the response."""
|
||||
if not isinstance(response, AsyncStream) and response.usage:
|
||||
logger.info(f"OpenAI usage: {response.usage}")
|
||||
self.prompt_tokens += response.usage.prompt_tokens
|
||||
self.total_tokens += response.usage.total_tokens
|
||||
if hasattr(response.usage, "completion_tokens"):
|
||||
self.completion_tokens += response.usage.completion_tokens
|
||||
|
||||
def to_dict(self) -> dict[str, str]:
|
||||
"""Create a dict of the service settings."""
|
||||
client_settings = {
|
||||
"api_key": self.client.api_key,
|
||||
"default_headers": {k: v for k, v in self.client.default_headers.items() if k != USER_AGENT},
|
||||
}
|
||||
if self.client.organization:
|
||||
client_settings["org_id"] = self.client.organization
|
||||
base = self.model_dump(
|
||||
exclude={
|
||||
"prompt_tokens",
|
||||
"completion_tokens",
|
||||
"total_tokens",
|
||||
"api_type",
|
||||
"ai_model_type",
|
||||
"service_id",
|
||||
"client",
|
||||
},
|
||||
by_alias=True,
|
||||
exclude_none=True,
|
||||
)
|
||||
base.update(client_settings)
|
||||
return base
|
||||
@@ -0,0 +1,10 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
from enum import Enum
|
||||
|
||||
|
||||
class NvidiaModelTypes(Enum):
|
||||
"""Nvidia model types, can be text, chat, or embedding."""
|
||||
|
||||
EMBEDDING = "embedding"
|
||||
CHAT = "chat"
|
||||
@@ -0,0 +1,167 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
import asyncio
|
||||
import copy
|
||||
import logging
|
||||
import sys
|
||||
from typing import Any
|
||||
|
||||
if sys.version_info >= (3, 12):
|
||||
from typing import override # pragma: no cover
|
||||
else:
|
||||
from typing_extensions import override # pragma: no cover
|
||||
|
||||
from numpy import array, ndarray
|
||||
from openai import AsyncOpenAI
|
||||
from pydantic import ValidationError
|
||||
|
||||
from semantic_kernel.connectors.ai.embedding_generator_base import EmbeddingGeneratorBase
|
||||
from semantic_kernel.connectors.ai.nvidia.prompt_execution_settings.nvidia_prompt_execution_settings import (
|
||||
NvidiaEmbeddingPromptExecutionSettings,
|
||||
)
|
||||
from semantic_kernel.connectors.ai.nvidia.services.nvidia_handler import NvidiaHandler
|
||||
from semantic_kernel.connectors.ai.nvidia.services.nvidia_model_types import NvidiaModelTypes
|
||||
from semantic_kernel.connectors.ai.nvidia.settings.nvidia_settings import NvidiaSettings
|
||||
from semantic_kernel.connectors.ai.prompt_execution_settings import PromptExecutionSettings
|
||||
from semantic_kernel.exceptions.service_exceptions import ServiceInitializationError
|
||||
from semantic_kernel.utils.feature_stage_decorator import experimental
|
||||
|
||||
logger: logging.Logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@experimental
|
||||
class NvidiaTextEmbedding(NvidiaHandler, EmbeddingGeneratorBase):
|
||||
"""Nvidia text embedding service."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
ai_model_id: str | None = None,
|
||||
api_key: str | None = None,
|
||||
base_url: str | None = None,
|
||||
client: AsyncOpenAI | None = None,
|
||||
env_file_path: str | None = None,
|
||||
service_id: str | None = None,
|
||||
) -> None:
|
||||
"""Initializes a new instance of the NvidiaTextEmbedding class.
|
||||
|
||||
Args:
|
||||
ai_model_id (str): NVIDIA model card string, see
|
||||
https://Nvidia.co/sentence-transformers
|
||||
api_key: NVIDIA API key, see https://console.NVIDIA.com/settings/keys
|
||||
(Env var NVIDIA_API_KEY)
|
||||
base_url: HttpsUrl | None - base_url: The url of the NVIDIA endpoint. The base_url consists of the endpoint,
|
||||
and more information refer https://docs.api.nvidia.com/nim/reference/
|
||||
use endpoint if you only want to supply the endpoint.
|
||||
(Env var NVIDIA_BASE_URL)
|
||||
client (Optional[AsyncOpenAI]): An existing client to use. (Optional)
|
||||
env_file_path (str | None): Use the environment settings file as
|
||||
a fallback to environment variables. (Optional)
|
||||
service_id (str): Service ID for the model. (optional)
|
||||
"""
|
||||
try:
|
||||
nvidia_settings = NvidiaSettings(
|
||||
api_key=api_key,
|
||||
base_url=base_url,
|
||||
embedding_model_id=ai_model_id,
|
||||
env_file_path=env_file_path,
|
||||
)
|
||||
except ValidationError as ex:
|
||||
raise ServiceInitializationError("Failed to create NVIDIA settings.", ex) from ex
|
||||
if not nvidia_settings.embedding_model_id:
|
||||
nvidia_settings.embedding_model_id = "nvidia/nv-embedqa-e5-v5"
|
||||
logger.warning(f"Default embedding model set as: {nvidia_settings.embedding_model_id}")
|
||||
if not nvidia_settings.api_key:
|
||||
logger.warning("API_KEY is missing, inference may fail.")
|
||||
if not client:
|
||||
client = AsyncOpenAI(
|
||||
api_key=nvidia_settings.api_key.get_secret_value() if nvidia_settings.api_key else None,
|
||||
base_url=nvidia_settings.base_url,
|
||||
)
|
||||
super().__init__(
|
||||
ai_model_id=nvidia_settings.embedding_model_id,
|
||||
api_key=nvidia_settings.api_key.get_secret_value() if nvidia_settings.api_key else None,
|
||||
ai_model_type=NvidiaModelTypes.EMBEDDING,
|
||||
service_id=service_id or nvidia_settings.embedding_model_id,
|
||||
env_file_path=env_file_path,
|
||||
client=client,
|
||||
)
|
||||
|
||||
@override
|
||||
async def generate_embeddings(
|
||||
self,
|
||||
texts: list[str],
|
||||
settings: "PromptExecutionSettings | None" = None,
|
||||
batch_size: int | None = None,
|
||||
**kwargs: Any,
|
||||
) -> ndarray:
|
||||
raw_embeddings = await self.generate_raw_embeddings(texts, settings, batch_size, **kwargs)
|
||||
return array(raw_embeddings)
|
||||
|
||||
@override
|
||||
async def generate_raw_embeddings(
|
||||
self,
|
||||
texts: list[str],
|
||||
settings: "PromptExecutionSettings | None" = None,
|
||||
batch_size: int | None = None,
|
||||
**kwargs: Any,
|
||||
) -> Any:
|
||||
"""Returns embeddings for the given texts in the unedited format.
|
||||
|
||||
Args:
|
||||
texts (List[str]): The texts to generate embeddings for.
|
||||
settings (NvidiaEmbeddingPromptExecutionSettings): The settings to use for the request.
|
||||
batch_size (int): The batch size to use for the request.
|
||||
kwargs (Dict[str, Any]): Additional arguments to pass to the request.
|
||||
"""
|
||||
if not settings:
|
||||
settings = NvidiaEmbeddingPromptExecutionSettings(ai_model_id=self.ai_model_id)
|
||||
else:
|
||||
if not isinstance(settings, NvidiaEmbeddingPromptExecutionSettings):
|
||||
settings = self.get_prompt_execution_settings_from_settings(settings)
|
||||
assert isinstance(settings, NvidiaEmbeddingPromptExecutionSettings) # nosec
|
||||
if settings.ai_model_id is None:
|
||||
settings.ai_model_id = self.ai_model_id
|
||||
for key, value in kwargs.items():
|
||||
setattr(settings, key, value)
|
||||
|
||||
# move input_type and truncate to extra-body
|
||||
if not settings.extra_body:
|
||||
settings.extra_body = {}
|
||||
settings.extra_body.setdefault("input_type", settings.input_type)
|
||||
if settings.truncate is not None:
|
||||
settings.extra_body.setdefault("truncate", settings.truncate)
|
||||
|
||||
raw_embeddings = []
|
||||
tasks = []
|
||||
|
||||
batch_size = batch_size or len(texts)
|
||||
for i in range(0, len(texts), batch_size):
|
||||
batch = texts[i : i + batch_size]
|
||||
batch_settings = copy.deepcopy(settings)
|
||||
batch_settings.input = batch
|
||||
tasks.append(self._send_request(settings=batch_settings))
|
||||
|
||||
results = await asyncio.gather(*tasks)
|
||||
for raw_embedding in results:
|
||||
assert isinstance(raw_embedding, list) # nosec
|
||||
raw_embeddings.extend(raw_embedding)
|
||||
|
||||
return raw_embeddings
|
||||
|
||||
def get_prompt_execution_settings_class(self) -> type["PromptExecutionSettings"]:
|
||||
"""Get the request settings class."""
|
||||
return NvidiaEmbeddingPromptExecutionSettings
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls: type["NvidiaTextEmbedding"], settings: dict[str, Any]) -> "NvidiaTextEmbedding":
|
||||
"""Initialize an Open AI service from a dictionary of settings.
|
||||
|
||||
Args:
|
||||
settings: A dictionary of settings for the service.
|
||||
"""
|
||||
return cls(
|
||||
ai_model_id=settings.get("ai_model_id"),
|
||||
api_key=settings.get("api_key"),
|
||||
env_file_path=settings.get("env_file_path"),
|
||||
service_id=settings.get("service_id"),
|
||||
)
|
||||
@@ -0,0 +1 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
@@ -0,0 +1,37 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
from typing import ClassVar
|
||||
|
||||
from pydantic import SecretStr
|
||||
|
||||
from semantic_kernel.kernel_pydantic import KernelBaseSettings
|
||||
|
||||
|
||||
class NvidiaSettings(KernelBaseSettings):
|
||||
"""Nvidia model settings.
|
||||
|
||||
The settings are first loaded from environment variables with the prefix 'NVIDIA_'. If the
|
||||
environment variables are not found, the settings can be loaded from a .env file with the
|
||||
encoding 'utf-8'. If the settings are not found in the .env file, the settings are ignored;
|
||||
however, validation will fail alerting that the settings are missing.
|
||||
|
||||
Optional settings for prefix 'NVIDIA_' are:
|
||||
- api_key: NVIDIA API key, see https://console.NVIDIA.com/settings/keys
|
||||
(Env var NVIDIA_API_KEY)
|
||||
- base_url: HttpsUrl | None - base_url: The url of the NVIDIA endpoint. The base_url consists of the endpoint,
|
||||
and more information refer https://docs.api.nvidia.com/nim/reference/
|
||||
use endpoint if you only want to supply the endpoint.
|
||||
(Env var NVIDIA_BASE_URL)
|
||||
- embedding_model_id: str | None - The NVIDIA embedding model ID to use, for example, nvidia/nv-embed-v1.
|
||||
(Env var NVIDIA_EMBEDDING_MODEL_ID)
|
||||
- chat_model_id: str | None - The NVIDIA chat model ID to use.
|
||||
(Env var NVIDIA_CHAT_MODEL_ID)
|
||||
- env_file_path: if provided, the .env settings are read from this file path location
|
||||
"""
|
||||
|
||||
env_prefix: ClassVar[str] = "NVIDIA_"
|
||||
|
||||
api_key: SecretStr | None = None
|
||||
base_url: str = "https://integrate.api.nvidia.com/v1"
|
||||
embedding_model_id: str | None = None
|
||||
chat_model_id: str | None = None
|
||||
Reference in New Issue
Block a user