Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion integrations/google_ai/pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -81,7 +81,7 @@ typing = "mypy --install-types --non-interactive --explicit-package-bases {args:


[tool.ruff]
target-version = "py38"
target-version = "py39"
line-length = 120

[tool.ruff.lint]
Expand Down
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
import json
from typing import Any, Dict, List, Optional, Union
from typing import Any, Optional, Union

import google.generativeai as genai
from google.ai.generativelanguage import Content, Part
Expand Down Expand Up @@ -146,9 +146,9 @@ def __init__(
*,
api_key: Secret = Secret.from_env_var("GOOGLE_API_KEY"), # noqa: B008
model: str = "gemini-2.0-flash",
generation_config: Optional[Union[GenerationConfig, Dict[str, Any]]] = None,
safety_settings: Optional[Dict[HarmCategory, HarmBlockThreshold]] = None,
tools: Optional[List[Tool]] = None,
generation_config: Optional[Union[GenerationConfig, dict[str, Any]]] = None,
safety_settings: Optional[dict[HarmCategory, HarmBlockThreshold]] = None,
tools: Optional[list[Tool]] = None,
tool_config: Optional[content_types.ToolConfigDict] = None,
streaming_callback: Optional[StreamingCallbackT] = None,
):
Expand Down Expand Up @@ -192,7 +192,7 @@ def __init__(
self._model = GenerativeModel(self._model_name)
self._streaming_callback = streaming_callback

def to_dict(self) -> Dict[str, Any]:
def to_dict(self) -> dict[str, Any]:
"""
Serializes the component to a dictionary.

Expand All @@ -218,7 +218,7 @@ def to_dict(self) -> Dict[str, Any]:
return data

@classmethod
def from_dict(cls, data: Dict[str, Any]) -> "GoogleAIGeminiChatGenerator":
def from_dict(cls, data: dict[str, Any]) -> "GoogleAIGeminiChatGenerator":
"""
Deserializes the component from a dictionary.

Expand Down Expand Up @@ -255,13 +255,13 @@ def _convert_to_google_tool(tool: Tool) -> FunctionDeclaration:

return FunctionDeclaration(name=tool.name, description=tool.description, parameters=parameters)

@component.output_types(replies=List[ChatMessage])
@component.output_types(replies=list[ChatMessage])
def run(
self,
messages: List[ChatMessage],
messages: list[ChatMessage],
streaming_callback: Optional[StreamingCallbackT] = None,
*,
tools: Optional[List[Tool]] = None,
tools: Optional[list[Tool]] = None,
):
"""
Generates text based on the provided messages.
Expand Down Expand Up @@ -308,13 +308,13 @@ def run(

return {"replies": replies}

@component.output_types(replies=List[ChatMessage])
@component.output_types(replies=list[ChatMessage])
async def run_async(
self,
messages: List[ChatMessage],
messages: list[ChatMessage],
streaming_callback: Optional[StreamingCallbackT] = None,
*,
tools: Optional[List[Tool]] = None,
tools: Optional[list[Tool]] = None,
):
"""
Async version of the run method. Generates text based on the provided messages.
Expand Down Expand Up @@ -367,7 +367,7 @@ async def run_async(
@staticmethod
def _convert_response_to_messages(
response_body: Union[GenerateContentResponse, AsyncGenerateContentResponse],
) -> List[ChatMessage]:
) -> list[ChatMessage]:
"""
Converts the Google AI response to a list of `ChatMessage` instances.

Expand Down Expand Up @@ -408,7 +408,7 @@ def _convert_response_to_messages(
@staticmethod
def _stream_response_and_convert_to_messages(
stream: GenerateContentResponse, streaming_callback: StreamingCallbackT
) -> List[ChatMessage]:
) -> list[ChatMessage]:
"""
Streams the Google AI response and converts it to a list of `ChatMessage` instances.

Expand Down Expand Up @@ -461,7 +461,7 @@ def _stream_response_and_convert_to_messages(
@staticmethod
async def _stream_response_and_convert_to_messages_async(
stream: AsyncGenerateContentResponse, streaming_callback: AsyncStreamingCallbackT
) -> List[ChatMessage]:
) -> list[ChatMessage]:
"""
Streams the Google AI response and converts it to a list of `ChatMessage` instances.

Expand Down
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
from typing import Any, Callable, Dict, List, Optional, Union
from typing import Any, Callable, Optional, Union

import google.generativeai as genai
from google.ai.generativelanguage import Content, Part
Expand Down Expand Up @@ -77,8 +77,8 @@ def __init__(
*,
api_key: Secret = Secret.from_env_var("GOOGLE_API_KEY"), # noqa: B008
model: str = "gemini-2.0-flash",
generation_config: Optional[Union[GenerationConfig, Dict[str, Any]]] = None,
safety_settings: Optional[Dict[HarmCategory, HarmBlockThreshold]] = None,
generation_config: Optional[Union[GenerationConfig, dict[str, Any]]] = None,
safety_settings: Optional[dict[HarmCategory, HarmBlockThreshold]] = None,
streaming_callback: Optional[Callable[[StreamingChunk], None]] = None,
):
"""
Expand Down Expand Up @@ -107,7 +107,7 @@ def __init__(
self._model = GenerativeModel(self.model_name)
self.streaming_callback = streaming_callback

def _generation_config_to_dict(self, config: Union[GenerationConfig, Dict[str, Any]]) -> Dict[str, Any]:
def _generation_config_to_dict(self, config: Union[GenerationConfig, dict[str, Any]]) -> dict[str, Any]:
if isinstance(config, dict):
return config
return {
Expand All @@ -119,7 +119,7 @@ def _generation_config_to_dict(self, config: Union[GenerationConfig, Dict[str, A
"stop_sequences": config.stop_sequences,
}

def to_dict(self) -> Dict[str, Any]:
def to_dict(self) -> dict[str, Any]:
"""
Serializes the component to a dictionary.

Expand All @@ -142,7 +142,7 @@ def to_dict(self) -> Dict[str, Any]:
return data

@classmethod
def from_dict(cls, data: Dict[str, Any]) -> "GoogleAIGeminiGenerator":
def from_dict(cls, data: dict[str, Any]) -> "GoogleAIGeminiGenerator":
"""
Deserializes the component from a dictionary.

Expand Down Expand Up @@ -180,7 +180,7 @@ def _convert_part(self, part: Union[str, ByteStream, Part]) -> Part:
msg = f"Unsupported type {type(part)} for part {part}"
raise ValueError(msg)

@component.output_types(replies=List[str])
@component.output_types(replies=list[str])
def run(
self,
parts: Variadic[Union[str, ByteStream, Part]],
Expand Down Expand Up @@ -212,7 +212,7 @@ def run(

return {"replies": replies}

def _get_response(self, response_body: GenerateContentResponse) -> List[str]:
def _get_response(self, response_body: GenerateContentResponse) -> list[str]:
"""
Extracts the responses from the Google AI request.
:param response_body: The response body from the Google AI request.
Expand All @@ -227,7 +227,7 @@ def _get_response(self, response_body: GenerateContentResponse) -> List[str]:

def _get_stream_response(
self, stream: GenerateContentResponse, streaming_callback: Callable[[StreamingChunk], None]
) -> List[str]:
) -> list[str]:
"""
Extracts the responses from the Google AI streaming response.
:param stream: The streaming response from the Google AI request.
Expand Down
2 changes: 1 addition & 1 deletion integrations/google_vertex/pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -81,7 +81,7 @@ typing = "mypy --install-types --non-interactive --explicit-package-bases {args:


[tool.ruff]
target-version = "py38"
target-version = "py39"
line-length = 120

[tool.ruff.lint]
Expand Down
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
import math
import time
from typing import Any, Dict, List, Literal, Optional
from typing import Any, Literal, Optional

import vertexai
from haystack import component, default_from_dict, default_to_dict, logging
Expand Down Expand Up @@ -71,7 +71,7 @@ def __init__(
retries: int = 3,
progress_bar: bool = True,
truncate_dim: Optional[int] = None,
meta_fields_to_embed: Optional[List[str]] = None,
meta_fields_to_embed: Optional[list[str]] = None,
embedding_separator: str = "\n",
) -> None:
"""
Expand Down Expand Up @@ -132,7 +132,7 @@ def resolve_secret(secret: Optional[Secret]) -> Optional[str]:
self.embedder = TextEmbeddingModel.from_pretrained(self.model)
self.task_type = task_type

def _prepare_texts_to_embed(self, documents: List[Document]) -> List[str]:
def _prepare_texts_to_embed(self, documents: list[Document]) -> list[str]:
"""
Prepare the texts to embed by concatenating the Document text with the metadata fields to embed.
"""
Expand All @@ -145,7 +145,7 @@ def _prepare_texts_to_embed(self, documents: List[Document]) -> List[str]:
texts_to_embed.append(text_to_embed)
return texts_to_embed

def get_text_embedding_input(self, batch: List[Document]) -> List[TextEmbeddingInput]:
def get_text_embedding_input(self, batch: list[Document]) -> list[TextEmbeddingInput]:
"""
Converts a batch of Document objects into a list of TextEmbeddingInput objects.

Expand All @@ -158,7 +158,7 @@ def get_text_embedding_input(self, batch: List[Document]) -> List[TextEmbeddingI
texts_to_embed = self._prepare_texts_to_embed(documents=batch)
return [TextEmbeddingInput(text=content, task_type=self.task_type) for content in texts_to_embed]

def embed_batch_by_smaller_batches(self, batch: List[str], subbatch=1) -> List[List[float]]:
def embed_batch_by_smaller_batches(self, batch: list[str], subbatch=1) -> list[list[float]]:
"""
Embeds a batch of text strings by dividing them into smaller sub-batches.
Args:
Expand Down Expand Up @@ -190,7 +190,7 @@ def embed_batch_by_smaller_batches(self, batch: List[str], subbatch=1) -> List[L

return embeddings_batch

def embed_batch(self, batch: List[str]) -> List[List[float]]:
def embed_batch(self, batch: list[str]) -> list[list[float]]:
"""
Generate embeddings for a batch of text strings.

Expand All @@ -205,8 +205,8 @@ def embed_batch(self, batch: List[str]) -> List[List[float]]:

return embeddings

@component.output_types(documents=List[Document])
def run(self, documents: List[Document]):
@component.output_types(documents=list[Document])
def run(self, documents: list[Document]):
"""
Processes all documents in batches while adhering to the API's token limit per request.

Expand Down Expand Up @@ -276,7 +276,7 @@ def run(self, documents: List[Document]):

return {"documents": documents}

def to_dict(self) -> Dict[str, Any]:
def to_dict(self) -> dict[str, Any]:
"""
Serializes the component to a dictionary.

Expand All @@ -300,7 +300,7 @@ def to_dict(self) -> Dict[str, Any]:
)

@classmethod
def from_dict(cls, data: Dict[str, Any]) -> "VertexAIDocumentEmbedder":
def from_dict(cls, data: dict[str, Any]) -> "VertexAIDocumentEmbedder":
"""
Deserializes the component from a dictionary.

Expand Down
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
from typing import Any, Dict, List, Literal, Optional, Union
from typing import Any, Literal, Optional, Union

import vertexai
from haystack import Document, component, default_from_dict, default_to_dict, logging
Expand Down Expand Up @@ -84,8 +84,8 @@ def resolve_secret(secret: Optional[Secret]) -> Optional[str]:
self.embedder = TextEmbeddingModel.from_pretrained(self.model)
self.task_type = task_type

@component.output_types(embedding=List[float])
def run(self, text: Union[List[Document], List[str], str]):
@component.output_types(embedding=list[float])
def run(self, text: Union[list[Document], list[str], str]):
"""
Processes text in batches while adhering to the API's token limit per request.

Expand All @@ -106,7 +106,7 @@ def run(self, text: Union[List[Document], List[str], str]):
embeddings = self.embedder.get_embeddings(text_embed_input)[0].values
return {"embedding": embeddings}

def to_dict(self) -> Dict[str, Any]:
def to_dict(self) -> dict[str, Any]:
"""
Serializes the component to a dictionary.

Expand All @@ -121,7 +121,7 @@ def to_dict(self) -> Dict[str, Any]:
)

@classmethod
def from_dict(cls, data: Dict[str, Any]) -> "VertexAITextEmbedder":
def from_dict(cls, data: dict[str, Any]) -> "VertexAITextEmbedder":
"""
Deserializes the component from a dictionary.

Expand Down
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
from typing import Any, Dict, List, Optional
from typing import Any, Optional

import vertexai
from haystack import logging
Expand Down Expand Up @@ -68,7 +68,7 @@ def __init__(

self._model = ImageTextModel.from_pretrained(self._model_name)

def to_dict(self) -> Dict[str, Any]:
def to_dict(self) -> dict[str, Any]:
"""
Serializes the component to a dictionary.

Expand All @@ -80,7 +80,7 @@ def to_dict(self) -> Dict[str, Any]:
)

@classmethod
def from_dict(cls, data: Dict[str, Any]) -> "VertexAIImageCaptioner":
def from_dict(cls, data: dict[str, Any]) -> "VertexAIImageCaptioner":
"""
Deserializes the component from a dictionary.

Expand All @@ -91,7 +91,7 @@ def from_dict(cls, data: Dict[str, Any]) -> "VertexAIImageCaptioner":
"""
return default_from_dict(cls, data)

@component.output_types(captions=List[str])
@component.output_types(captions=list[str])
def run(self, image: ByteStream):
"""Prompts the model to generate captions for the given image.

Expand Down
Loading
Loading