Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add text extraction to tool call processing in BedrockLLM #96

Open
wants to merge 2 commits into
base: main
Choose a base branch
from
Open

Add text extraction to tool call processing in BedrockLLM #96

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
389f130
extract_tool_calls now is extract_tool_calls_and_text. In this way yo…
Cerrix Jul 2, 2024
82a8e89
extract_tool_calls now is extract_tool_calls_and_text. In this way yo…
Cerrix Jul 2, 2024
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
11 changes: 7 additions & 4 deletions libs/aws/langchain_aws/llms/bedrock.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -193,15 +193,19 @@ def _get_invocation_metrics_chunk(chunk: Dict[str, Any]) -> GenerationChunk:
return GenerationChunk(text="", generation_info=generation_info)


def extract_tool_calls(content: List[dict]) -> List[ToolCall]:
def extract_tool_calls_and_text(content: List[dict]) -> Tuple[List[ToolCall], str]:
text = ""
tool_calls = []
for block in content:
if block["type"] == "text":
text = block["text"]
continue
if block["type"] != "tool_use":
continue
tool_calls.append(
ToolCall(name=block["name"], args=block["input"], id=block["id"])
)
return tool_calls
return tool_calls, text


class AnthropicTool(TypedDict):
Expand Down Expand Up @@ -266,7 +270,6 @@ def prepare_output(cls, provider: str, response: Any) -> dict:
text = ""
tool_calls = []
response_body = json.loads(response.get("body").read().decode())

if provider == "anthropic":
if "completion" in response_body:
text = response_body.get("completion")
Expand All @@ -275,7 +278,7 @@ def prepare_output(cls, provider: str, response: Any) -> dict:
if len(content) == 1 and content[0]["type"] == "text":
text = content[0]["text"]
elif any(block["type"] == "tool_use" for block in content):
tool_calls = extract_tool_calls(content)
tool_calls, text = extract_tool_calls_and_text(content)

else:
if provider == "ai21":
Expand Down
Loading