Python: Migrate to Ollama Python SDK (#7165)

### Motivation and Context

<!-- Thank you for your contribution to the semantic-kernel repo!
Please help reviewers and future users, providing the following
information:
  1. Why is this change required?
  2. What problem does it solve?
  3. What scenario does it contribute to?
  4. If it fixes an open issue, please link to the issue here.
-->

### Description

<!-- Describe your changes, the overall approach, the underlying design.
These notes will help understanding how your code works. Thanks! -->
Currently, the Ollama AI connector uses `aiohttp` sessions to interface
with the local Ollama LLM service. However, we have made the decision to
transition to the official Ollama Python SDK.

Related to: https://github.com/microsoft/semantic-kernel/issues/6841,
https://github.com/microsoft/semantic-kernel/issues/7134

Test coverage report:

![image](https://github.com/microsoft/semantic-kernel/assets/12570346/8d414f71-fc1b-4870-b038-0f7d2985b985)

MyPy warnings: all cleared
### Contribution Checklist

<!-- Before submitting this PR, please make sure: -->

- [x] The code builds clean without any errors or warnings
- [x] The PR follows the [SK Contribution
Guidelines](https://github.com/microsoft/semantic-kernel/blob/main/CONTRIBUTING.md)
and the [pre-submission formatting
script](https://github.com/microsoft/semantic-kernel/blob/main/CONTRIBUTING.md#development-scripts)
raises no violations
- [x] All unit tests pass, and I have added new tests where possible
- [x] I didn't break anyone 😄
This commit is contained in:
Tao Chen
2024-07-16 22:38:21 +00:00
committed by GitHub
parent 4bbb694818
commit 79207ffbbc
18 changed files with 1072 additions and 305 deletions
@@ -2,25 +2,28 @@
from typing import Any, Literal
from pydantic import Field
from semantic_kernel.connectors.ai.prompt_execution_settings import PromptExecutionSettings
class OllamaPromptExecutionSettings(PromptExecutionSettings):
ai_model_id: str = Field("", serialization_alias="model")
"""Settings for Ollama prompt execution."""
format: Literal["json"] | None = None
options: dict[str, Any] | None = None
stream: bool = False
class OllamaTextPromptExecutionSettings(OllamaPromptExecutionSettings):
prompt: str | None = None
context: str | None = None
"""Settings for Ollama text prompt execution."""
system: str | None = None
template: str | None = None
raw: bool = False
context: str | None = None
raw: bool | None = None
class OllamaChatPromptExecutionSettings(OllamaPromptExecutionSettings):
messages: list[dict[str, str]] | None = None
"""Settings for Ollama chat prompt execution."""
class OllamaEmbeddingPromptExecutionSettings(OllamaPromptExecutionSettings):
"""Settings for Ollama embedding prompt execution."""
@@ -0,0 +1,29 @@
# Copyright (c) Microsoft. All rights reserved.
from typing import ClassVar
from semantic_kernel.kernel_pydantic import KernelBaseSettings
class OllamaSettings(KernelBaseSettings):
"""Ollama settings.
The settings are first loaded from environment variables with
the prefix 'OLLAMA_'.
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.
Required settings for prefix 'OLLAMA' are:
- model: str - Model name. (Env var OLLAMA_MODEL)
Optional settings for prefix 'OLLAMA' are:
- host: HttpsUrl - The endpoint of the Ollama service. (Env var OLLAMA_HOST)
"""
env_prefix: ClassVar[str] = "OLLAMA_"
model: str
host: str | None = None
@@ -0,0 +1,17 @@
# Copyright (c) Microsoft. All rights reserved.
from abc import ABC
from ollama import AsyncClient
from semantic_kernel.kernel_pydantic import KernelBaseModel
class OllamaBase(KernelBaseModel, ABC):
"""Ollama service base.
Args:
client [AsyncClient]: An Ollama client to use for the service.
"""
client: AsyncClient
@@ -1,16 +1,22 @@
# Copyright (c) Microsoft. All rights reserved.
import json
import logging
from collections.abc import AsyncGenerator
from typing import Any
import sys
from collections.abc import AsyncGenerator, AsyncIterator, Mapping
from typing import TYPE_CHECKING, Any
import aiohttp
from pydantic import HttpUrl
if sys.version_info >= (3, 12):
from typing import override # pragma: no cover
else:
from typing_extensions import override # pragma: no cover
from ollama import AsyncClient, Message
from pydantic import ValidationError
from semantic_kernel.connectors.ai.chat_completion_client_base import ChatCompletionClientBase
from semantic_kernel.connectors.ai.ollama.ollama_prompt_execution_settings import OllamaChatPromptExecutionSettings
from semantic_kernel.connectors.ai.ollama.utils import AsyncSession
from semantic_kernel.connectors.ai.ollama.ollama_settings import OllamaSettings
from semantic_kernel.connectors.ai.ollama.services.ollama_base import OllamaBase
from semantic_kernel.connectors.ai.text_completion_client_base import TextCompletionClientBase
from semantic_kernel.contents import AuthorRole
from semantic_kernel.contents.chat_history import ChatHistory
@@ -18,28 +24,60 @@ from semantic_kernel.contents.chat_message_content import ChatMessageContent
from semantic_kernel.contents.streaming_chat_message_content import StreamingChatMessageContent
from semantic_kernel.contents.streaming_text_content import StreamingTextContent
from semantic_kernel.contents.text_content import TextContent
from semantic_kernel.exceptions.service_exceptions import ServiceInitializationError, ServiceInvalidResponseError
if TYPE_CHECKING:
from semantic_kernel.connectors.ai.prompt_execution_settings import PromptExecutionSettings
logger: logging.Logger = logging.getLogger(__name__)
class OllamaChatCompletion(TextCompletionClientBase, ChatCompletionClientBase):
class OllamaChatCompletion(OllamaBase, TextCompletionClientBase, ChatCompletionClientBase):
"""Initializes a new instance of the OllamaChatCompletion class.
Make sure to have the ollama service running either locally or remotely.
Args:
ai_model_id (str): Ollama model name, see https://ollama.ai/library
url (Optional[Union[str, HttpUrl]]): URL of the Ollama server, defaults to http://localhost:11434/api/chat
session (Optional[aiohttp.ClientSession]): Optional client session to use for requests.
"""
url: HttpUrl = "http://localhost:11434/api/chat"
session: aiohttp.ClientSession | None = None
def __init__(
self,
service_id: str | None = None,
ai_model_id: str | None = None,
host: str | None = None,
client: AsyncClient | None = None,
env_file_path: str | None = None,
env_file_encoding: str | None = None,
) -> None:
"""Initialize an OllamaChatCompletion service.
Args:
service_id (Optional[str]): Service ID tied to the execution settings. (Optional)
ai_model_id (Optional[str]): The model name. (Optional)
host (Optional[str]): URL of the Ollama server, defaults to None and
will use the default Ollama service address: http://127.0.0.1:11434. (Optional)
client (Optional[AsyncClient]): A custom Ollama client to use for the service. (Optional)
env_file_path (str | None): Use the environment settings file as a fallback to using env vars.
env_file_encoding (str | None): The encoding of the environment settings file, defaults to 'utf-8'.
"""
try:
ollama_settings = OllamaSettings.create(
model=ai_model_id,
host=host,
env_file_path=env_file_path,
env_file_encoding=env_file_encoding,
)
except ValidationError as ex:
raise ServiceInitializationError("Failed to create Ollama settings.", ex) from ex
super().__init__(
service_id=service_id or ollama_settings.model,
ai_model_id=ollama_settings.model,
client=client or AsyncClient(host=ollama_settings.host),
)
async def get_chat_message_contents(
self,
chat_history: ChatHistory,
settings: OllamaChatPromptExecutionSettings,
settings: "PromptExecutionSettings",
**kwargs: Any,
) -> list[ChatMessageContent]:
"""This is the method that is called from the kernel to get a response from a chat-optimized LLM.
@@ -53,29 +91,35 @@ class OllamaChatCompletion(TextCompletionClientBase, ChatCompletionClientBase):
Returns:
List[ChatMessageContent]: A list of ChatMessageContent objects representing the response(s) from the LLM.
"""
if not settings.ai_model_id:
settings.ai_model_id = self.ai_model_id
settings.messages = self._prepare_chat_history_for_request(chat_history)
settings.stream = False
async with (
AsyncSession(self.session) as session,
session.post(str(self.url), json=settings.prepare_settings_dict()) as response,
):
response.raise_for_status()
response_object = await response.json()
return [
ChatMessageContent(
inner_content=response_object,
ai_model_id=self.ai_model_id,
role=AuthorRole.ASSISTANT,
content=response_object.get("message", {"content": None}).get("content", None),
)
]
settings = self.get_prompt_execution_settings_from_settings(settings)
prepared_chat_history = self._prepare_chat_history_for_request(chat_history)
response_object = await self.client.chat(
model=self.ai_model_id,
messages=prepared_chat_history,
stream=False,
**settings.prepare_settings_dict(),
)
if not isinstance(response_object, Mapping):
raise ServiceInvalidResponseError(
"Invalid response type from Ollama chat completion. "
f"Expected Mapping but got {type(response_object)}."
)
return [
ChatMessageContent(
inner_content=response_object,
ai_model_id=self.ai_model_id,
role=AuthorRole.ASSISTANT,
content=response_object.get("message", {"content": None}).get("content", None),
)
]
async def get_streaming_chat_message_contents(
self,
chat_history: ChatHistory,
settings: OllamaChatPromptExecutionSettings,
settings: "PromptExecutionSettings",
**kwargs: Any,
) -> AsyncGenerator[list[StreamingChatMessageContent], Any]:
"""Streams a text completion using an Ollama model.
@@ -85,73 +129,81 @@ class OllamaChatCompletion(TextCompletionClientBase, ChatCompletionClientBase):
Args:
chat_history (ChatHistory): A chat history that contains a list of chat messages,
that can be rendered into a set of messages, from system, user, assistant and function.
settings (OllamaChatPromptExecutionSettings): Request settings.
settings (PromptExecutionSettings): Request settings.
kwargs (Dict[str, Any]): The optional arguments.
Yields:
List[StreamingChatMessageContent]: Stream of StreamingChatMessageContent objects.
"""
if not settings.ai_model_id:
settings.ai_model_id = self.ai_model_id
settings.messages = self._prepare_chat_history_for_request(chat_history)
settings.stream = True
async with (
AsyncSession(self.session) as session,
session.post(str(self.url), json=settings.prepare_settings_dict()) as response,
):
response.raise_for_status()
async for line in response.content:
body = json.loads(line)
if body.get("done") and body.get("message", {}).get("content") is None:
break
yield [
StreamingChatMessageContent(
role=AuthorRole.ASSISTANT,
choice_index=0,
inner_content=body,
ai_model_id=self.ai_model_id,
content=body.get("message", {"content": None}).get("content", None),
)
]
if body.get("done"):
break
settings = self.get_prompt_execution_settings_from_settings(settings)
prepared_chat_history = self._prepare_chat_history_for_request(chat_history)
response_object = await self.client.chat(
model=self.ai_model_id,
messages=prepared_chat_history,
stream=True,
**settings.prepare_settings_dict(),
)
if not isinstance(response_object, AsyncIterator):
raise ServiceInvalidResponseError(
"Invalid response type from Ollama chat completion. "
f"Expected AsyncIterator but got {type(response_object)}."
)
async for part in response_object:
yield [
StreamingChatMessageContent(
role=AuthorRole.ASSISTANT,
choice_index=0,
inner_content=part,
ai_model_id=self.ai_model_id,
content=part.get("message", {"content": None}).get("content", None),
)
]
async def get_text_contents(
self,
prompt: str,
settings: OllamaChatPromptExecutionSettings,
settings: "PromptExecutionSettings",
) -> list[TextContent]:
"""This is the method that is called from the kernel to get a response from a text-optimized LLM.
Args:
prompt (str): A prompt to complete
settings (OllamaChatPromptExecutionSettings): Settings for the request.
settings (PromptExecutionSettings): Settings for the request.
Returns:
List["TextContent"]: The completion result(s).
"""
if not settings.ai_model_id:
settings.ai_model_id = self.ai_model_id
settings.messages = [{"role": AuthorRole.USER, "content": prompt}]
settings.stream = False
async with (
AsyncSession(self.session) as session,
session.post(str(self.url), json=settings.prepare_settings_dict()) as response,
):
response.raise_for_status()
response_object = await response.json()
return [
TextContent(
inner_content=response_object,
ai_model_id=self.ai_model_id,
text=response_object.get("message", {"content": None}).get("content", None),
)
]
settings = self.get_prompt_execution_settings_from_settings(settings)
prepared_chat_history = [Message(role=AuthorRole.USER.value, content=prompt)]
response_object = await self.client.chat(
model=self.ai_model_id,
messages=prepared_chat_history,
stream=False,
**settings.prepare_settings_dict(),
)
if not isinstance(response_object, Mapping):
raise ServiceInvalidResponseError(
"Invalid response type from Ollama chat completion. "
f"Expected Mapping but got {type(response_object)}."
)
return [
TextContent(
inner_content=response_object,
ai_model_id=self.ai_model_id,
text=response_object.get("message", {"content": None}).get("content", None),
)
]
async def get_streaming_text_contents(
self,
prompt: str,
settings: OllamaChatPromptExecutionSettings,
settings: "PromptExecutionSettings",
) -> AsyncGenerator[list[StreamingTextContent], Any]:
"""Streams a text completion using an Ollama model.
@@ -159,35 +211,38 @@ class OllamaChatCompletion(TextCompletionClientBase, ChatCompletionClientBase):
Args:
prompt (str): A chat history that contains the prompt to complete.
settings (OllamaChatPromptExecutionSettings): Request settings.
settings (PromptExecutionSettings): Request settings.
Yields:
List["StreamingTextContent"]: The result stream made up of StreamingTextContent objects.
"""
if not settings.ai_model_id:
settings.ai_model_id = self.ai_model_id
settings.messages = [{"role": AuthorRole.USER, "content": prompt}]
settings.stream = True
async with (
AsyncSession(self.session) as session,
session.post(str(self.url), json=settings.prepare_settings_dict()) as response,
):
response.raise_for_status()
async for line in response.content:
body = json.loads(line)
if body.get("done") and body.get("message", {}).get("content") is None:
break
yield [
StreamingTextContent(
choice_index=0,
inner_content=body,
ai_model_id=self.ai_model_id,
text=body.get("message", {"content": None}).get("content", None),
)
]
if body.get("done"):
break
settings = self.get_prompt_execution_settings_from_settings(settings)
prepared_chat_history = [Message(role=AuthorRole.USER.value, content=prompt)]
def get_prompt_execution_settings_class(self) -> "OllamaChatPromptExecutionSettings":
response_object = await self.client.chat(
model=self.ai_model_id,
messages=prepared_chat_history,
stream=True,
**settings.prepare_settings_dict(),
)
if not isinstance(response_object, AsyncIterator):
raise ServiceInvalidResponseError(
"Invalid response type from Ollama chat completion. "
f"Expected AsyncIterator but got {type(response_object)}."
)
async for part in response_object:
yield [
StreamingTextContent(
choice_index=0,
inner_content=part,
ai_model_id=self.ai_model_id,
text=part.get("message", {"content": None}).get("content", None),
)
]
@override
def get_prompt_execution_settings_class(self) -> type["PromptExecutionSettings"]:
"""Get the request settings class."""
return OllamaChatPromptExecutionSettings
@@ -1,39 +1,78 @@
# Copyright (c) Microsoft. All rights reserved.
import json
import logging
from collections.abc import AsyncGenerator
from typing import Any
import sys
from collections.abc import AsyncGenerator, AsyncIterator, Mapping
from typing import TYPE_CHECKING, Any
import aiohttp
from pydantic import HttpUrl
if sys.version_info >= (3, 12):
from typing import override # pragma: no cover
else:
from typing_extensions import override # pragma: no cover
from ollama import AsyncClient
from pydantic import ValidationError
from semantic_kernel.connectors.ai.ollama.ollama_prompt_execution_settings import OllamaTextPromptExecutionSettings
from semantic_kernel.connectors.ai.ollama.utils import AsyncSession
from semantic_kernel.connectors.ai.ollama.ollama_settings import OllamaSettings
from semantic_kernel.connectors.ai.ollama.services.ollama_base import OllamaBase
from semantic_kernel.connectors.ai.text_completion_client_base import TextCompletionClientBase
from semantic_kernel.contents.streaming_text_content import StreamingTextContent
from semantic_kernel.contents.text_content import TextContent
from semantic_kernel.exceptions.service_exceptions import ServiceInitializationError, ServiceInvalidResponseError
if TYPE_CHECKING:
from semantic_kernel.connectors.ai.prompt_execution_settings import PromptExecutionSettings
logger: logging.Logger = logging.getLogger(__name__)
class OllamaTextCompletion(TextCompletionClientBase):
class OllamaTextCompletion(OllamaBase, TextCompletionClientBase):
"""Initializes a new instance of the OllamaTextCompletion class.
Make sure to have the ollama service running either locally or remotely.
Args:
ai_model_id (str): Ollama model name, see https://ollama.ai/library
url (Optional[Union[str, HttpUrl]]): URL of the Ollama server, defaults to http://localhost:11434/api/generate
"""
url: HttpUrl = "http://localhost:11434/api/generate"
session: aiohttp.ClientSession | None = None
def __init__(
self,
service_id: str | None = None,
ai_model_id: str | None = None,
host: str | None = None,
client: AsyncClient | None = None,
env_file_path: str | None = None,
env_file_encoding: str | None = None,
) -> None:
"""Initialize an OllamaChatCompletion service.
Args:
service_id (Optional[str]): Service ID tied to the execution settings. (Optional)
ai_model_id (Optional[str]): The model name. (Optional)
host (Optional[str]): URL of the Ollama server, defaults to None and
will use the default Ollama service address: http://127.0.0.1:11434. (Optional)
client (Optional[AsyncClient]): A custom Ollama client to use for the service. (Optional)
env_file_path (str | None): Use the environment settings file as a fallback to using env vars.
env_file_encoding (str | None): The encoding of the environment settings file, defaults to 'utf-8'.
"""
try:
ollama_settings = OllamaSettings.create(
model=ai_model_id,
host=host,
env_file_path=env_file_path,
env_file_encoding=env_file_encoding,
)
except ValidationError as ex:
raise ServiceInitializationError("Failed to create Ollama settings.", ex) from ex
super().__init__(
service_id=service_id or ollama_settings.model,
ai_model_id=ollama_settings.model,
client=client or AsyncClient(host=ollama_settings.host),
)
async def get_text_contents(
self,
prompt: str,
settings: OllamaTextPromptExecutionSettings,
settings: "PromptExecutionSettings",
) -> list[TextContent]:
"""This is the method that is called from the kernel to get a response from a text-optimized LLM.
@@ -44,23 +83,29 @@ class OllamaTextCompletion(TextCompletionClientBase):
Returns:
List[TextContent]: A list of TextContent objects representing the response(s) from the LLM.
"""
if not settings.ai_model_id:
settings.ai_model_id = self.ai_model_id
settings.prompt = prompt
settings.stream = False
async with (
AsyncSession(self.session) as session,
session.post(self.url, json=settings.prepare_settings_dict()) as response,
):
response.raise_for_status()
inner_content = await response.json()
text = inner_content["response"]
return [TextContent(inner_content=inner_content, ai_model_id=self.ai_model_id, text=text)]
settings = self.get_prompt_execution_settings_from_settings(settings)
response_object = await self.client.generate(
model=self.ai_model_id,
prompt=prompt,
stream=False,
**settings.prepare_settings_dict(),
)
if not isinstance(response_object, Mapping):
raise ServiceInvalidResponseError(
"Invalid response type from Ollama chat completion. "
f"Expected Mapping but got {type(response_object)}."
)
inner_content = response_object
text = inner_content["response"]
return [TextContent(inner_content=inner_content, ai_model_id=self.ai_model_id, text=text)]
async def get_streaming_text_contents(
self,
prompt: str,
settings: OllamaTextPromptExecutionSettings,
settings: "PromptExecutionSettings",
) -> AsyncGenerator[list[StreamingTextContent], Any]:
"""Streams a text completion using an Ollama model.
@@ -74,27 +119,29 @@ class OllamaTextCompletion(TextCompletionClientBase):
Yields:
List[StreamingTextContent]: Completion result.
"""
if not settings.ai_model_id:
settings.ai_model_id = self.ai_model_id
settings.prompt = prompt
settings.stream = True
async with (
AsyncSession(self.session) as session,
session.post(self.url, json=settings.prepare_settings_dict()) as response,
):
response.raise_for_status()
async for line in response.content:
body = json.loads(line)
if body.get("done") and body.get("response") is None:
break
yield [
StreamingTextContent(
choice_index=0, inner_content=body, ai_model_id=self.ai_model_id, text=body.get("response")
)
]
if body.get("done"):
break
settings = self.get_prompt_execution_settings_from_settings(settings)
def get_prompt_execution_settings_class(self) -> "OllamaTextPromptExecutionSettings":
response_object = await self.client.generate(
model=self.ai_model_id,
prompt=prompt,
stream=True,
**settings.prepare_settings_dict(),
)
if not isinstance(response_object, AsyncIterator):
raise ServiceInvalidResponseError(
"Invalid response type from Ollama chat completion. "
f"Expected AsyncIterator but got {type(response_object)}."
)
async for part in response_object:
yield [
StreamingTextContent(
choice_index=0, inner_content=part, ai_model_id=self.ai_model_id, text=part.get("response")
)
]
@override
def get_prompt_execution_settings_class(self) -> type["PromptExecutionSettings"]:
"""Get the request settings class."""
return OllamaTextPromptExecutionSettings
@@ -2,51 +2,108 @@
import logging
import sys
from typing import Any
from typing import TYPE_CHECKING, Any
from ollama import AsyncClient
from pydantic import ValidationError
from semantic_kernel.connectors.ai.ollama.ollama_prompt_execution_settings import OllamaEmbeddingPromptExecutionSettings
from semantic_kernel.connectors.ai.ollama.ollama_settings import OllamaSettings
from semantic_kernel.connectors.ai.ollama.services.ollama_base import OllamaBase
from semantic_kernel.exceptions.service_exceptions import ServiceInitializationError
if TYPE_CHECKING:
from semantic_kernel.connectors.ai.prompt_execution_settings import PromptExecutionSettings
if sys.version_info >= (3, 12):
from typing import override # pragma: no cover
else:
from typing_extensions import override # pragma: no cover
import aiohttp
from numpy import array, ndarray
from pydantic import HttpUrl
from semantic_kernel.connectors.ai.embeddings.embedding_generator_base import EmbeddingGeneratorBase
from semantic_kernel.connectors.ai.ollama.utils import AsyncSession
from semantic_kernel.utils.experimental_decorator import experimental_class
logger: logging.Logger = logging.getLogger(__name__)
@experimental_class
class OllamaTextEmbedding(EmbeddingGeneratorBase):
class OllamaTextEmbedding(OllamaBase, EmbeddingGeneratorBase):
"""Ollama embeddings client.
Make sure to have the ollama service running either locally or remotely.
Args:
ai_model_id (str): Ollama model name, see https://ollama.ai/library
url (Optional[Union[str, HttpUrl]]): URL of the Ollama server, defaults to http://localhost:11434/api/embeddings
session (Optional[aiohttp.ClientSession]): Optional client session to use for requests.
"""
url: HttpUrl = "http://localhost:11434/api/embeddings"
session: aiohttp.ClientSession | None = None
def __init__(
self,
service_id: str | None = None,
ai_model_id: str | None = None,
host: str | None = None,
client: AsyncClient | None = None,
env_file_path: str | None = None,
env_file_encoding: str | None = None,
) -> None:
"""Initialize an OllamaChatCompletion service.
Args:
service_id (Optional[str]): Service ID tied to the execution settings. (Optional)
ai_model_id (Optional[str]): The model name. (Optional)
host (Optional[str]): URL of the Ollama server, defaults to None and
will use the default Ollama service address: http://127.0.0.1:11434. (Optional)
client (Optional[AsyncClient]): A custom Ollama client to use for the service. (Optional)
env_file_path (str | None): Use the environment settings file as a fallback to using env vars.
env_file_encoding (str | None): The encoding of the environment settings file, defaults to 'utf-8'.
"""
try:
ollama_settings = OllamaSettings.create(
model=ai_model_id,
host=host,
env_file_path=env_file_path,
env_file_encoding=env_file_encoding,
)
except ValidationError as ex:
raise ServiceInitializationError("Failed to create Ollama settings.", ex) from ex
super().__init__(
service_id=service_id or ollama_settings.model,
ai_model_id=ollama_settings.model,
client=client or AsyncClient(host=ollama_settings.host),
)
@override
async def generate_embeddings(self, texts: list[str], **kwargs: Any) -> ndarray:
async def generate_embeddings(
self,
texts: list[str],
settings: "PromptExecutionSettings | None" = None,
**kwargs: Any,
) -> ndarray:
raw_embeddings = await self.generate_raw_embeddings(texts, settings, **kwargs)
return array(raw_embeddings)
@override
async def generate_raw_embeddings(
self,
texts: list[str],
settings: "PromptExecutionSettings | None" = None,
**kwargs: Any,
) -> Any:
if not settings:
settings = OllamaEmbeddingPromptExecutionSettings()
else:
settings = self.get_prompt_execution_settings_from_settings(settings)
result = []
for text in texts:
async with (
AsyncSession(self.session) as session,
session.post(
self.url,
json={"model": self.ai_model_id, "prompt": text, "options": kwargs},
) as response,
):
response.raise_for_status()
response = await response.json()
result.append(response["embedding"])
return array(result)
response_object = await self.client.embeddings(
model=self.ai_model_id,
prompt=text,
**settings.prepare_settings_dict(),
)
result.append(response_object["embedding"])
return result
@override
def get_prompt_execution_settings_class(self) -> type["PromptExecutionSettings"]:
return OllamaEmbeddingPromptExecutionSettings
@@ -1,17 +0,0 @@
# Copyright (c) Microsoft. All rights reserved.
import aiohttp
class AsyncSession:
def __init__(self, session: aiohttp.ClientSession = None):
"""Initialize the AsyncSession."""
self._session = session if session else aiohttp.ClientSession()
async def __aenter__(self):
"""Enter the context manager."""
return await self._session.__aenter__()
async def __aexit__(self, *args, **kwargs):
"""Exit the context manager."""
await self._session.close()
@@ -32,7 +32,7 @@ class OpenAITextEmbeddingBase(OpenAIHandler, EmbeddingGeneratorBase):
**kwargs: Any,
) -> ndarray:
raw_embeddings = await self.generate_raw_embeddings(texts, settings, batch_size, **kwargs)
return array([array(emb) for emb in raw_embeddings])
return array(raw_embeddings)
@override
async def generate_raw_embeddings(