From 7ac553f5be2de89a445947dcd7ddfa0cb76afcbf Mon Sep 17 00:00:00 2001 From: romanlutz Date: Mon, 18 May 2026 05:18:07 -0700 Subject: [PATCH] Delete unused _TextEmbedding ABC The _TextEmbedding abstract base class in pyrit/embedding/_text_embedding.py had no subclasses (OpenAITextEmbedding extends EmbeddingSupport directly and re-implements its own retry logic) and no importers anywhere in the repository. Remove the dead file. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- pyrit/embedding/_text_embedding.py | 58 ------------------------------ 1 file changed, 58 deletions(-) delete mode 100644 pyrit/embedding/_text_embedding.py diff --git a/pyrit/embedding/_text_embedding.py b/pyrit/embedding/_text_embedding.py deleted file mode 100644 index 71d06abc34..0000000000 --- a/pyrit/embedding/_text_embedding.py +++ /dev/null @@ -1,58 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# Licensed under the MIT license. - -import abc -from typing import Any, Union - -import tenacity -from openai import AzureOpenAI, OpenAI - -from pyrit.models import ( - EmbeddingData, - EmbeddingResponse, - EmbeddingSupport, - EmbeddingUsageInformation, -) - - -class _TextEmbedding(EmbeddingSupport, abc.ABC): - """Text embedding base class.""" - - _client: Union[OpenAI, AzureOpenAI] - _model: str - - def __init__(self) -> None: - super().__init__() - if not (hasattr(self, "_client") and hasattr(self, "_model")): - raise NotImplementedError( - "Text embedding client and model need to be provided by the implementing child class." - ) - - @tenacity.retry(wait=tenacity.wait_fixed(0.1), stop=tenacity.stop_after_delay(3)) - def generate_text_embedding(self, text: str, **kwargs: Any) -> EmbeddingResponse: - """ - Generate text embedding. - - Args: - text: The text to generate the embedding for - **kwargs: Additional arguments to pass to the LLM client API - - Returns: - The embedding response - """ - embedding_obj = self._client.embeddings.create(input=text, model=self._model, **kwargs) - return EmbeddingResponse( - model=embedding_obj.model, - object=embedding_obj.object, - data=[ - EmbeddingData( - embedding=embedding_obj.data[0].embedding, - index=embedding_obj.data[0].index, - object=embedding_obj.data[0].object, - ) - ], - usage=EmbeddingUsageInformation( - prompt_tokens=embedding_obj.usage.prompt_tokens, - total_tokens=embedding_obj.usage.total_tokens, - ), - )