Skip to content

Commit

Permalink
fix: Run Bedrock calls in executor for async
Browse files Browse the repository at this point in the history
  • Loading branch information
lou-k committed Jul 10, 2024
1 parent 0f8cfaf commit 7efc3e8
Showing 1 changed file with 5 additions and 3 deletions.
8 changes: 5 additions & 3 deletions packages/phoenix-evals/src/phoenix/evals/models/bedrock.py
Original file line number Diff line number Diff line change
@@ -1,3 +1,4 @@
import asyncio
import json
import logging
from dataclasses import dataclass, field
Expand All @@ -21,8 +22,8 @@ class BedrockModel(BaseModel):
AWS API are dynamically throttled when encountering rate limit errors. Requires the `boto3`
package to be installed.
Supports Async:
`boto3` does not support async calls
Supports Async: 🟡
`boto3` does not support async calls, so it's wrapped in an executor.
Args:
model_id (str): The model name to use.
Expand Down Expand Up @@ -109,7 +110,8 @@ def _generate(self, prompt: str, **kwargs: Dict[str, Any]) -> str:
return self._parse_output(response) or ""

async def _async_generate(self, prompt: str, **kwargs: Dict[str, Any]) -> str:
return self._generate(prompt, **kwargs)
loop = asyncio.get_event_loop()
return await loop.run_in_executor(None, self._generate, prompt, **kwargs)

def _rate_limited_completion(self, **kwargs: Any) -> Any:
"""Use tenacity to retry the completion call."""
Expand Down

0 comments on commit 7efc3e8

Please sign in to comment.