Enhance error handling in GeminiEmbedder and GeminiClient
- Added checks to raise exceptions when no embeddings or response text are returned, improving robustness. - Included type ignore comments for mypy compatibility in embed_content calls.
This commit is contained in:
parent
e311dab617
commit
b6c0bb9c54
2 changed files with 13 additions and 4 deletions
|
|
@ -61,18 +61,24 @@ class GeminiEmbedder(EmbedderClient):
|
||||||
# Generate embeddings
|
# Generate embeddings
|
||||||
result = await self.client.aio.models.embed_content(
|
result = await self.client.aio.models.embed_content(
|
||||||
model=self.config.embedding_model or DEFAULT_EMBEDDING_MODEL,
|
model=self.config.embedding_model or DEFAULT_EMBEDDING_MODEL,
|
||||||
contents=[input_data],
|
contents=[input_data], # type: ignore[arg-type] # mypy fails on broad union type
|
||||||
config=types.EmbedContentConfig(output_dimensionality=self.config.embedding_dim),
|
config=types.EmbedContentConfig(output_dimensionality=self.config.embedding_dim),
|
||||||
)
|
)
|
||||||
|
|
||||||
|
if not result.embeddings or len(result.embeddings) == 0 or not result.embeddings[0].values:
|
||||||
|
raise Exception('No embeddings returned')
|
||||||
|
|
||||||
return result.embeddings[0].values
|
return result.embeddings[0].values
|
||||||
|
|
||||||
async def create_batch(self, input_data_list: list[str]) -> list[list[float]]:
|
async def create_batch(self, input_data_list: list[str]) -> list[list[float]]:
|
||||||
# Generate embeddings
|
# Generate embeddings
|
||||||
result = await self.client.aio.models.embed_content(
|
result = await self.client.aio.models.embed_content(
|
||||||
model=self.config.embedding_model or DEFAULT_EMBEDDING_MODEL,
|
model=self.config.embedding_model or DEFAULT_EMBEDDING_MODEL,
|
||||||
contents=input_data_list,
|
contents=input_data_list, # type: ignore[arg-type] # mypy fails on broad union type
|
||||||
config=types.EmbedContentConfig(output_dimensionality=self.config.embedding_dim),
|
config=types.EmbedContentConfig(output_dimensionality=self.config.embedding_dim),
|
||||||
)
|
)
|
||||||
|
|
||||||
return [embedding.values for embedding in result.embeddings]
|
if not result.embeddings or len(result.embeddings) == 0:
|
||||||
|
raise Exception('No embeddings returned')
|
||||||
|
|
||||||
|
return [embedding.values if embedding.values else [] for embedding in result.embeddings]
|
||||||
|
|
|
||||||
|
|
@ -139,13 +139,16 @@ class GeminiClient(LLMClient):
|
||||||
# Generate content using the simple string approach
|
# Generate content using the simple string approach
|
||||||
response = await self.client.aio.models.generate_content(
|
response = await self.client.aio.models.generate_content(
|
||||||
model=self.model or DEFAULT_MODEL,
|
model=self.model or DEFAULT_MODEL,
|
||||||
contents=gemini_messages,
|
contents=gemini_messages, # type: ignore[arg-type] # mypy fails on broad union type
|
||||||
config=generation_config,
|
config=generation_config,
|
||||||
)
|
)
|
||||||
|
|
||||||
# If this was a structured output request, parse the response into the Pydantic model
|
# If this was a structured output request, parse the response into the Pydantic model
|
||||||
if response_model is not None:
|
if response_model is not None:
|
||||||
try:
|
try:
|
||||||
|
if not response.text:
|
||||||
|
raise Exception('No response text')
|
||||||
|
|
||||||
validated_model = response_model.model_validate(json.loads(response.text))
|
validated_model = response_model.model_validate(json.loads(response.text))
|
||||||
|
|
||||||
# Return as a dictionary for API consistency
|
# Return as a dictionary for API consistency
|
||||||
|
|
|
||||||
Loading…
Add table
Reference in a new issue