8000 models - anthropic - document - plain text (#141) · lgigit200/sdk-python@5d785a1 · GitHub
[go: up one dir, main page]
Skip to content

Commit 5d785a1

Browse files
pgrayypgrayy
authored
models - anthropic - document - plain text (strands-agents#141)
1 parent 58bfecb commit 5d785a1

File tree

2 files changed

+45
-25
lines changed

2 files changed

+45
-25
lines changed

‎src/strands/models/anthropic.py

Lines changed: 7 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -97,13 +97,16 @@ def _format_request_message_content(self, content: ContentBlock) -> dict[str, An
9797
Anthropic formatted content block.
9898
"""
9999
if "document" in content:
100+
mime_type = mimetypes.types_map.get(f".{content['document']['format']}", "application/octet-stream")
100101
return {
101102
"source": {
102-
"data": base64.b64encode(content["document"]["source"]["bytes"]).decode("utf-8"),
103-
"media_type": mimetypes.types_map.get(
104-
f".{content['document']['format']}", "application/octet-stream"
103+
"data": (
104+
content["document"]["source"]["bytes"].decode("utf-8")
105+
if mime_type == "text/plain"
106+
else base64.b64encode(content["document"]["source"]["bytes"]).decode("utf-8")
105107
),
106-
"type": "base64",
108+
"media_type": mime_type,
109+
"type": "text" if mime_type == "text/plain" else "base64",
107110
},
108111
"title": content["document"]["name"],
109112
"type": "document",

‎tests/strands/models/test_anthropic.py

Lines changed: 38 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -102,19 +102,46 @@ def test_format_request_with_system_prompt(model, messages, model_id, max_tokens
102102
assert tru_request == exp_request
103103

104104

105-
def test_format_request_with_document(model, model_id, max_tokens):
105+
@pytest.mark.parametrize(
106+
("content", "formatted_content"),
107+
[
108+
# PDF
109+
(
110+
{
111+
"document": {"format": "pdf", "name": "test doc", "source": {"bytes": b"pdf"}},
112+
},
113+
{
114+
"source": {
115+
"data": "cGRm",
116+
"media_type": "application/pdf",
117+
"type": "base64",
118+
},
119+
"title": "test doc",
120+
"type": "document",
121+
},
122+
),
123+
# Plain text
124+
(
125+
{
126+
"document": {"format": "txt", "name": "test doc", "source": {"bytes": b"txt"}},
127+
},
128+
{
129+
"source": {
130+
"data": "txt",
131+
"media_type": "text/plain",
132+
"type": "text",
133+
},
134+
"title": "test doc",
135+
"type": "document",
136+
},
137+
),
138+
],
139+
)
140+
def test_format_request_with_document(content, formatted_content, model, model_id, max_tokens):
106141
messages = [
107142
{
108143
"role": "user",
109-
"content": [
110-
{
111-
"document": {
112-
"format": "pdf",
113-
"name": "test-doc",
114-
"source": {"bytes": b"base64encodeddoc"},
115-
},
116-
},
117-
],
144+
"content": [content],
118145
},
119146
]
120147

@@ -124,17 +151,7 @@ def test_format_request_with_document(model, model_id, max_tokens):
124151
"messages": [
125152
{
126153
"role": "user",
127-
"content": [
128-
{
129-
"source": {
130-
"data": "YmFzZTY0ZW5jb2RlZGRvYw==",
131-
"media_type": "application/pdf",
132-
"type": "base64",
133-
},
134-
"title": "test-doc",
135-
"type": "document",
136-
},
137-
],
154+
"content": [formatted_content],
138155
},
139156
],
140157
"model": model_id,

0 commit comments

Comments
 (0)
0