Skip to content

Commit 5606b08

Browse files
jph00jph00
committed
fixes #9 
1 parent 5cc725b commit 5606b08

File tree

3 files changed

+444
-820
lines changed

3 files changed

+444
-820
lines changed

‎msglm/_modidx.py

Lines changed: 2 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -10,18 +10,15 @@
1010
'msglm.core.AnthropicMsg.img_msg': ('core.html#anthropicmsg.img_msg', 'msglm/core.py'),
1111
'msglm.core.AnthropicMsg.is_sdk_obj': ('core.html#anthropicmsg.is_sdk_obj', 'msglm/core.py'),
1212
'msglm.core.AnthropicMsg.pdf_msg': ('core.html#anthropicmsg.pdf_msg', 'msglm/core.py'),
13+
'msglm.core.AnthropicMsg.text_msg': ('core.html#anthropicmsg.text_msg', 'msglm/core.py'),
1314
'msglm.core.Msg': ('core.html#msg', 'msglm/core.py'),
1415
'msglm.core.Msg.__call__': ('core.html#msg.__call__', 'msglm/core.py'),
15-
'msglm.core.Msg.find_block': ('core.html#msg.find_block', 'msglm/core.py'),
16-
'msglm.core.Msg.img_msg': ('core.html#msg.img_msg', 'msglm/core.py'),
17-
'msglm.core.Msg.is_sdk_obj': ('core.html#msg.is_sdk_obj', 'msglm/core.py'),
1816
'msglm.core.Msg.mk_content': ('core.html#msg.mk_content', 'msglm/core.py'),
19-
'msglm.core.Msg.pdf_msg': ('core.html#msg.pdf_msg', 'msglm/core.py'),
20-
'msglm.core.Msg.text_msg': ('core.html#msg.text_msg', 'msglm/core.py'),
2117
'msglm.core.OpenAiMsg': ('core.html#openaimsg', 'msglm/core.py'),
2218
'msglm.core.OpenAiMsg.find_block': ('core.html#openaimsg.find_block', 'msglm/core.py'),
2319
'msglm.core.OpenAiMsg.img_msg': ('core.html#openaimsg.img_msg', 'msglm/core.py'),
2420
'msglm.core.OpenAiMsg.is_sdk_obj': ('core.html#openaimsg.is_sdk_obj', 'msglm/core.py'),
21+
'msglm.core.OpenAiMsg.text_msg': ('core.html#openaimsg.text_msg', 'msglm/core.py'),
2522
'msglm.core._add_cache_control': ('core.html#_add_cache_control', 'msglm/core.py'),
2623
'msglm.core._is_img': ('core.html#_is_img', 'msglm/core.py'),
2724
'msglm.core._is_pdf': ('core.html#_is_pdf', 'msglm/core.py'),

‎msglm/core.py

Lines changed: 92 additions & 82 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@
33
# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/00_core.ipynb.
44

55
# %% auto 0
6-
__all__ = ['mk_msg_openai', 'mk_msgs_openai', 'mk_msg', 'mk_msgs', 'Msg', 'AnthropicMsg', 'OpenAiMsg', 'mk_msg_anthropic',
6+
__all__ = ['mk_msg_openai', 'mk_msgs_openai', 'Msg', 'OpenAiMsg', 'AnthropicMsg', 'mk_msg', 'mk_msgs', 'mk_msg_anthropic',
77
'mk_msgs_anthropic', 'mk_ant_doc']
88

99
# %% ../nbs/00_core.ipynb
@@ -22,16 +22,59 @@ def _mk_img(data:bytes)->tuple:
2222
mtype = mimetypes.types_map["."+imghdr.what(None, h=data)]
2323
return img, mtype
2424

25+
# %% ../nbs/00_core.ipynb
26+
class Msg:
27+
"Helper class to create a message for the OpenAI and Anthropic APIs."
28+
pass
29+
30+
# %% ../nbs/00_core.ipynb
31+
class OpenAiMsg(Msg):
32+
"Helper class to create a message for the OpenAI API."
33+
pass
34+
35+
# %% ../nbs/00_core.ipynb
36+
class AnthropicMsg(Msg):
37+
"Helper class to create a message for the Anthropic API."
38+
pass
39+
2540
# %% ../nbs/00_core.ipynb
2641
def _is_img(data): return isinstance(data, bytes) and bool(imghdr.what(None, data))
2742

2843
# %% ../nbs/00_core.ipynb
2944
def _is_pdf(data): return isinstance(data, bytes) and data.startswith(b'%PDF-')
3045

3146
# %% ../nbs/00_core.ipynb
32-
def _mk_pdf(data:bytes)->str:
33-
"Convert pdf bytes to a base64 encoded pdf"
34-
return base64.standard_b64encode(data).decode("utf-8")
47+
@patch
48+
def mk_content(self:Msg, content, text_only=False)->dict:
49+
if isinstance(content, str): return self.text_msg(content, text_only=text_only)
50+
if _is_img(content): return self.img_msg(content)
51+
if _is_pdf(content): return self.pdf_msg(content)
52+
return content
53+
54+
# %% ../nbs/00_core.ipynb
55+
@patch
56+
def img_msg(self:OpenAiMsg, data:bytes)->dict:
57+
"Convert `data` to an image message"
58+
img, mtype = _mk_img(data)
59+
return {"type": "input_image", "image_url": f"data:{mtype};base64,{img}"}
60+
61+
@patch
62+
def text_msg(self:OpenAiMsg, s:str, text_only=False)->dict:
63+
"Convert `s` to a text message"
64+
return s if text_only else {"type": "input_text", "text":s}
65+
66+
# %% ../nbs/00_core.ipynb
67+
@patch
68+
def img_msg(self:AnthropicMsg, data:bytes)->dict:
69+
"Convert `data` to an image message"
70+
img, mtype = _mk_img(data)
71+
r = {"type": "base64", "media_type": mtype, "data":img}
72+
return {"type": "image", "source": r}
73+
74+
@patch
75+
def text_msg(self:AnthropicMsg, s:str, text_only=False)->dict:
76+
"Convert `s` to a text message"
77+
return s if text_only else {"type": "text", "text":s}
3578

3679
# %% ../nbs/00_core.ipynb
3780
def mk_msg(content:Union[list,str], role:str="user", *args, api:str="openai", **kw)->dict:
@@ -41,94 +84,62 @@ def mk_msg(content:Union[list,str], role:str="user", *args, api:str="openai", **
4184
msg = m()(role, content, text_only=text_only, **kw)
4285
return dict2obj(msg, list_func=list)
4386

87+
# %% ../nbs/00_core.ipynb
88+
def _mk_pdf(data:bytes)->str:
89+
"Convert pdf bytes to a base64 encoded pdf"
90+
return base64.standard_b64encode(data).decode("utf-8")
91+
92+
# %% ../nbs/00_core.ipynb
93+
@patch
94+
def pdf_msg(self:AnthropicMsg, data: bytes) -> dict:
95+
"Convert `data` to a pdf message"
96+
r = {"type": "base64", "media_type": "application/pdf", "data":_mk_pdf(data)}
97+
return {"type": "document", "source": r}
98+
4499
# %% ../nbs/00_core.ipynb
45100
def mk_msgs(msgs: list, *args, api:str="openai", **kw) -> list:
46101
"Create a list of messages compatible with OpenAI/Anthropic."
47102
if isinstance(msgs, str): msgs = [msgs]
48103
return [mk_msg(o, ('user', 'assistant')[i % 2], *args, api=api, **kw) for i, o in enumerate(msgs)]
49104

50105
# %% ../nbs/00_core.ipynb
51-
class Msg:
52-
"Helper class to create a message for the OpenAI and Anthropic APIs."
53-
sdk_obj_support=False # is an SDK object a valid message?
54-
def __call__(self, role:str, content:[list,str], text_only:bool=False, **kw)->dict:
55-
"Create an OpenAI/Anthropic compatible message with `role` and `content`."
56-
if self.sdk_obj_support and self.is_sdk_obj(content): return self.find_block(content)
57-
if hasattr(content, "content"): content, role = content.content, content.role
58-
content = self.find_block(content)
59-
if content is not None and not isinstance(content, list): content = [content]
60-
content = [self.mk_content(o, text_only=text_only) for o in content] if content else ''
61-
return dict(role=role, content=content[0] if text_only else content, **kw)
62-
63-
def is_sdk_obj(self, r)-> bool:
64-
"Check if `r` is an SDK object."
65-
raise NotImplemented
66-
67-
def find_block(self, r)->dict:
68-
"Find the message in `r`."
69-
raise NotImplemented
70-
71-
def text_msg(self, s:str, text_only:bool=False, **kw):
72-
"Convert `s` to a text message"
73-
return s if text_only else {"type":"text", "text":s}
74-
75-
def img_msg(self, *args, **kw)->dict:
76-
"Convert bytes to an image message"
77-
raise NotImplemented
78-
79-
def pdf_msg(self, *args, **kw)->dict:
80-
"Convert bytes to a pdf message"
81-
raise NotImplemented
82-
83-
def mk_content(self, content:[str, bytes], text_only:bool=False) -> dict:
84-
"Create the appropriate data structure based the content type."
85-
if isinstance(content, str): return self.text_msg(content, text_only=text_only)
86-
if _is_img(content): return self.img_msg(content)
87-
if _is_pdf(content): return self.pdf_msg(content)
88-
return content
106+
@patch
107+
def __call__(self:Msg, role:str, content:[list,str], text_only:bool=False, **kw)->dict:
108+
"Create an OpenAI/Anthropic compatible message with `role` and `content`."
109+
if self.sdk_obj_support and self.is_sdk_obj(content): return self.find_block(content)
110+
if hasattr(content, "content"): content, role = content.content, content.role
111+
content = self.find_block(content)
112+
if content is not None and not isinstance(content, list): content = [content]
113+
content = [self.mk_content(o, text_only=text_only) for o in content] if content else ''
114+
return dict(role=role, content=content[0] if text_only else content, **kw)
89115

90116
# %% ../nbs/00_core.ipynb
91-
class AnthropicMsg(Msg):
92-
sdk_obj_support=False
93-
def img_msg(self, data: bytes) -> dict:
94-
"Convert `data` to an image message"
95-
img, mtype = _mk_img(data)
96-
r = {"type": "base64", "media_type": mtype, "data":img}
97-
return {"type": "image", "source": r}
98-
99-
def pdf_msg(self, data: bytes) -> dict:
100-
"Convert `data` to a pdf message"
101-
r = {"type": "base64", "media_type": "application/pdf", "data":_mk_pdf(data)}
102-
return {"type": "document", "source": r}
103-
104-
def is_sdk_obj(self, r)-> bool:
105-
"Check if `r` is an SDK object."
106-
return isinstance(r, abc.Mapping)
107-
108-
def find_block(self, r):
109-
"Find the message in `r`."
110-
return r.get('content', r) if self.is_sdk_obj(r) else r
117+
AnthropicMsg.sdk_obj_support=False
118+
OpenAiMsg.sdk_obj_support=True
111119

112120
# %% ../nbs/00_core.ipynb
113-
class OpenAiMsg(Msg):
114-
sdk_obj_support=True
115-
def img_msg(self, data: bytes) -> dict:
116-
"Convert `data` to an image message"
117-
img, mtype = _mk_img(data)
118-
r = {"url": f"data:{mtype};base64,{img}"}
119-
return {"type": "image_url", "image_url": r}
120-
121-
def is_sdk_obj(self, r)-> bool:
122-
"Check if `r` is an SDK object."
123-
return type(r).__module__ != "builtins"
124-
125-
def find_block(self, r):
126-
"Find the message in `r`."
127-
if not self.is_sdk_obj(r): return r
128-
m = nested_idx(r, "choices", 0)
129-
if not m: return m
130-
if hasattr(m, "message"): return m.message
131-
return m.delta
121+
@patch
122+
def is_sdk_obj(self:AnthropicMsg, r)-> bool:
123+
"Check if `r` is an SDK object."
124+
return isinstance(r, abc.Mapping)
125+
126+
@patch
127+
def find_block(self:AnthropicMsg, r):
128+
"Find the message in `r`."
129+
return r.get('content', r) if self.is_sdk_obj(r) else r
130+
131+
# %% ../nbs/00_core.ipynb
132+
@patch
133+
def is_sdk_obj(self:OpenAiMsg, r)-> bool:
134+
"Check if `r` is an SDK object."
135+
return type(r).__module__ != "builtins"
136+
137+
@patch
138+
def find_block(self:OpenAiMsg, r):
139+
"Find the message in `r`."
140+
if not self.is_sdk_obj(r): return r
141+
if hasattr(r, "output"): return r.output[0]
142+
return r.delta
132143

133144
# %% ../nbs/00_core.ipynb
134145
mk_msg_openai = partial(mk_msg, api="openai")
@@ -150,7 +161,6 @@ def _remove_cache_ckpts(msg):
150161
else: delattr(msg["content"][-1], 'cache_control') if hasattr(msg["content"][-1], 'cache_control') else None
151162
return msg
152163

153-
154164
@delegates(mk_msg)
155165
def mk_msg_anthropic(*args, cache=False, **kwargs):
156166
"Create an Anthropic compatible message."

0 commit comments

Comments
 (0)