Merge pull request #23 from atomiechen/dev

support extra properties, tool calls and images

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2024 Atomie CHEN
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

README.md

@@ -19,9 +19,10 @@ A handy toolkit for using LLM, with both [***development support***](https://pyp
 - **Easy chain**: You can chain `hprompt` files to construct dynamic logic.
 
 <p float="left" align="center">
-  <img src="https://raw.githubusercontent.com/atomiechen/vscode-handyllm/main/demo/example.jpg" width="60%" />
+  <img src="https://raw.githubusercontent.com/atomiechen/vscode-handyllm/main/demo/example.png" width="70%" />
 </p>
 
+
 **Other features:**
 
 ☯️ Unified API design with both sync and async support
@@ -52,3 +53,9 @@ Please check [HandyLLM VSCode extension](https://marketplace.visualstudio.com/it
 
 Please check out our [wiki](https://github.com/atomiechen/HandyLLM/wiki) for comprehensive guides ^_^
 
+
+
+## License
+
+[HandyLLM](https://github.com/atomiechen/HandyLLM) © 2024 by [Atomie CHEN](https://github.com/atomiechen) is licensed under the MIT License - see the [LICENSE](https://github.com/atomiechen/HandyLLM/blob/main/LICENSE) file for details.
+

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "HandyLLM"
-version = "0.7.1"
+version = "0.7.2"
 authors = [
   { name="Atomie CHEN", email="atomic_cwh@163.com" },
 ]

src/handyllm/hprompt.py

@@ -36,8 +36,8 @@
 from .prompt_converter import PromptConverter
 from .openai_client import ClientMode, OpenAIClient
 from .utils import (
-    astream_chat_with_role, astream_completions, 
-    stream_chat_with_role, stream_completions, 
+    astream_chat_all, astream_completions, 
+    stream_chat_all, stream_completions, 
 )
 from ._str_enum import AutoStrEnum
 
@@ -282,7 +282,9 @@ class HandyPrompt(ABC):
     def __init__(
         self, data: Union[str, list], request: Optional[dict] = None, 
         meta: Optional[Union[dict, RunConfig]] = None, 
-        base_path: Optional[PathType] = None):
+        base_path: Optional[PathType] = None,
+        response: Optional[dict] = None,
+        ):
         self.data = data
         self.request = request or {}
         # parse meta to run_config
@@ -291,6 +293,7 @@ def __init__(
         else:
             self.run_config = RunConfig.from_dict(meta or {}, base_path=base_path)
         self.base_path = base_path
+        self.response = response
 
     def __str__(self) -> str:
         return str(self.data)
@@ -554,8 +557,12 @@ def _parse_var_map(self, run_config: RunConfig):
 
 class ChatPrompt(HandyPrompt):
 
-    def __init__(self, chat: list, request: dict, meta: Union[dict, RunConfig], base_path: Optional[PathType] = None):
-        super().__init__(chat, request, meta, base_path)
+    def __init__(
+        self, chat: list, request: dict, meta: Union[dict, RunConfig], 
+        base_path: Optional[PathType] = None,
+        response: Optional[dict] = None,
+        ):
+        super().__init__(chat, request, meta, base_path, response)
 
     @property
     def chat(self) -> list:
@@ -582,20 +589,6 @@ def _eval_data(self, run_config: RunConfig) -> list:
         else:
             return self.chat
 
-    def _stream_chat_proc(self, response, fd: Optional[io.IOBase] = None) -> tuple[str, str]:
-        # stream response to fd
-        role = ""
-        content = ""
-        for r, text in stream_chat_with_role(response):
-            if r != role:
-                role = r
-                if fd:
-                    fd.write(f"${role}$\n")
-            elif fd:
-                fd.write(text)
-            content += text
-        return role, content
-
     def _run_with_client(
         self, client: OpenAIClient, 
         run_config: RunConfig,
@@ -614,36 +607,24 @@ def _run_with_client(
                 with open(run_config.output_path, 'w', encoding='utf-8') as fout:
                     # dump frontmatter
                     fout.write(self._dumps_frontmatter(new_request, run_config, base_path))
-                    role, content = self._stream_chat_proc(response, fout)
+                    role, content, tool_calls = converter.stream_chat2raw(stream_chat_all(response), fout)
             elif run_config.output_fd:
                 # dump frontmatter, no base_path
                 run_config.output_fd.write(self._dumps_frontmatter(new_request, run_config))
                 # stream response to a file descriptor
-                role, content = self._stream_chat_proc(response, run_config.output_fd)
+                role, content, tool_calls = converter.stream_chat2raw(stream_chat_all(response), run_config.output_fd)
             else:
-                role, content = self._stream_chat_proc(response)
+                role, content, tool_calls = converter.stream_chat2raw(stream_chat_all(response))
         else:
             role = response['choices'][0]['message']['role']
-            content = response['choices'][0]['message']['content']
+            content = response['choices'][0]['message'].get('content')
+            tool_calls = response['choices'][0]['message'].get('tool_calls')
         return ChatPrompt(
-            [{"role": role, "content": content}],
-            new_request, run_config, base_path
+            [{"role": role, "content": content, "tool_calls": tool_calls}],
+            new_request, run_config, base_path,
+            response=response
         )
 
-    async def _astream_chat_proc(self, response, fd: Optional[io.IOBase] = None) -> tuple[str, str]:
-        # stream response to fd
-        role = ""
-        content = ""
-        async for r, text in astream_chat_with_role(response):
-            if r != role:
-                role = r
-                if fd:
-                    fd.write(f"${role}$\n")
-            elif fd:
-                fd.write(text)
-            content += text
-        return role, content
-
     async def _arun_with_client(
         self, client: OpenAIClient, 
         run_config: RunConfig,
@@ -661,19 +642,21 @@ async def _arun_with_client(
                 # stream response to a file
                 with open(run_config.output_path, 'w', encoding='utf-8') as fout:
                     fout.write(self._dumps_frontmatter(new_request, run_config, base_path))
-                    role, content = await self._astream_chat_proc(response, fout)
+                    role, content, tool_calls = await converter.astream_chat2raw(astream_chat_all(response), fout)
             elif run_config.output_fd:
                 # stream response to a file descriptor
                 run_config.output_fd.write(self._dumps_frontmatter(new_request, run_config))
-                role, content = await self._astream_chat_proc(response, run_config.output_fd)
+                role, content, tool_calls = await converter.astream_chat2raw(astream_chat_all(response), run_config.output_fd)
             else:
-                role, content = await self._astream_chat_proc(response)
+                role, content, tool_calls = await converter.astream_chat2raw(astream_chat_all(response))
         else:
             role = response['choices'][0]['message']['role']
-            content = response['choices'][0]['message']['content']
+            content = response['choices'][0]['message'].get('content')
+            tool_calls = response['choices'][0]['message'].get('tool_calls')
         return ChatPrompt(
-            [{"role": role, "content": content}],
-            new_request, run_config, base_path
+            [{"role": role, "content": content, "tool_calls": tool_calls}],
+            new_request, run_config, base_path,
+            response=response
         )
 
     def __add__(self, other: Union[str, list, ChatPrompt]):
@@ -719,8 +702,12 @@ def __iadd__(self, other: Union[str, list, ChatPrompt]):
 
 class CompletionsPrompt(HandyPrompt):
 
-    def __init__(self, prompt: str, request: dict, meta: Union[dict, RunConfig], base_path: PathType = None):
-        super().__init__(prompt, request, meta, base_path)
+    def __init__(
+        self, prompt: str, request: dict, meta: Union[dict, RunConfig], 
+        base_path: PathType = None,
+        response: Optional[dict] = None,
+        ):
+        super().__init__(prompt, request, meta, base_path, response)
 
     @property
     def prompt(self) -> str:
@@ -775,7 +762,9 @@ def _run_with_client(
                 content = self._stream_completions_proc(response)
         else:
             content = response['choices'][0]['text']
-        return CompletionsPrompt(content, new_request, run_config, base_path)
+        return CompletionsPrompt(
+            content, new_request, run_config, base_path, response=response
+            )
 
     async def _astream_completions_proc(self, response, fd: Optional[io.IOBase] = None) -> str:
         # stream response to fd
@@ -812,7 +801,9 @@ async def _arun_with_client(
                 content = await self._astream_completions_proc(response)
         else:
             content = response['choices'][0]['text']
-        return CompletionsPrompt(content, new_request, run_config, base_path)
+        return CompletionsPrompt(
+            content, new_request, run_config, base_path, response=response
+            )
 
     def __add__(self, other: Union[str, CompletionsPrompt]):
         # support concatenation with string or another CompletionsPrompt

src/handyllm/prompt_converter.py

@@ -1,16 +1,21 @@
+import io
 import re
+from typing import Optional
+import yaml
+
 
 class PromptConverter:
 
-    role_keys = ['system', 'user', 'assistant']
+    role_keys = ['system', 'user', 'assistant', 'tool']
 
     def __init__(self):
         self.substitute_map = {}
 
     @property
     def split_pattern(self):
         # build a regex pattern to split the prompt by role keys
-        return r'^\$(' + '|'.join(self.role_keys) + r')\$$'
+        # return r'^\$(' + '|'.join(self.role_keys) + r')\$$'
+        return r'^\$(' + '|'.join(self.role_keys) + r')\$[^\S\r\n]*({[^}]*?})?[^\S\r\n]*$'
 
     def detect(self, raw_prompt: str):
         # detect the role keys in the prompt
@@ -38,10 +43,32 @@ def raw2chat(self, raw_prompt: str):
         # convert plain text to chat format
         chat = []
         blocks = re.split(self.split_pattern, raw_prompt, flags=re.MULTILINE)
-        for idx in range(1, len(blocks), 2):
-            key = blocks[idx]
-            value = blocks[idx+1]
-            chat.append({"role": key, "content": value.strip()})
+        for idx in range(1, len(blocks), 3):
+            role = blocks[idx]
+            extra = blocks[idx+1]
+            content = blocks[idx+2]
+            if content:
+                content = content.strip()
+            msg = {"role": role, "content": content}
+            if extra:
+                # remove curly braces
+                key_values_pairs = re.findall(r'(\w+)\s*=\s*("[^"]*"|\'[^\']*\')', extra[1:-1])
+                # parse extra properties
+                extra_properties = {}
+                for key, value in key_values_pairs:
+                    # remove quotes of the value
+                    extra_properties[key] = value[1:-1]
+                if 'type' in extra_properties:
+                    type_of_msg = extra_properties.pop('type')
+                    if type_of_msg == 'tool_calls':
+                        msg['tool_calls'] = yaml.safe_load(content)
+                        msg['content'] = None
+                    elif type_of_msg == 'content_array':
+                        # parse content array
+                        msg['content'] = yaml.safe_load(content)
+                for key in extra_properties:
+                    msg[key] = extra_properties[key]
+            chat.append(msg)
 
         return chat
 
@@ -56,9 +83,85 @@ def chat2raw(chat):
         # convert chat format to plain text
         messages = []
         for message in chat:
-            messages.append(f"${message['role']}$\n{message['content']}")
+            role = message.get('role')
+            content = message.get('content')
+            tool_calls = message.get('tool_calls')
+            extra_properties = {key: message[key] for key in message if key not in ['role', 'content', 'tool_calls']}
+            if tool_calls:
+                extra_properties['type'] = 'tool_calls'
+                content = yaml.dump(tool_calls)
+            elif isinstance(content, list):
+                extra_properties['type'] = 'content_array'
+                content = yaml.dump(content)
+            if extra_properties:
+                extra = " {" + " ".join([f'{key}="{extra_properties[key]}"' for key in extra_properties]) + "}"
+            else:
+                extra = ""
+            messages.append(f"${role}${extra}\n{content}")
         raw_prompt = "\n\n".join(messages)
         return raw_prompt
+
+    @staticmethod
+    def stream_chat2raw(gen_sync, fd: Optional[io.IOBase] = None) -> tuple[str, str]:
+        # stream response to fd
+        role = ""
+        content = ""
+        tool_calls = []
+        role_completed = False
+        for r, text, tool_call in gen_sync:
+            if r != role:
+                role = r
+                if fd:
+                    fd.write(f"${role}$")  # do not add newline
+            if tool_call:
+                if not role_completed:
+                    if fd:
+                        fd.write(' {type="tool_calls"}\n')
+                    role_completed = True
+                tool_calls.append(tool_call)  # do not stream, wait for the end
+            elif text:
+                if not role_completed:
+                    if fd:
+                        fd.write('\n')
+                    role_completed = True
+                if fd:
+                    fd.write(text)
+                content += text
+        if tool_calls and fd:
+            # dump tool calls
+            fd.write(yaml.dump(tool_calls))
+        return role, content, tool_calls
+
+    @staticmethod
+    async def astream_chat2raw(gen_async, fd: Optional[io.IOBase] = None) -> tuple[str, str]:
+        # stream response to fd
+        role = ""
+        content = ""
+        tool_calls = []
+        role_completed = False
+        async for r, text, tool_call in gen_async:
+            if r != role:
+                role = r
+                if fd:
+                    fd.write(f"${role}$")  # do not add newline
+            if tool_call:
+                if not role_completed:
+                    if fd:
+                        fd.write(' {type="tool_calls"}\n')
+                    role_completed = True
+                tool_calls.append(tool_call)  # do not stream, wait for the end
+            elif text:
+                if not role_completed:
+                    if fd:
+                        fd.write('\n')
+                    role_completed = True
+                if fd:
+                    fd.write(text)
+                content += text
+        if tool_calls and fd:
+            # dump tool calls
+            fd.write(yaml.dump(tool_calls))
+        return role, content, tool_calls
 
     @classmethod
     def chat2rawfile(cls, chat, raw_prompt_path: str):
@@ -72,15 +175,15 @@ def chat_replace_variables(chat, variable_map: dict, inplace=False):
         if inplace:
             for message in chat:
                 for var, value in variable_map.items():
-                    if var in message['content']:
+                    if message.get('content') and var in message['content']:
                         message['content'] = message['content'].replace(var, value)
             return chat
         else:
             new_chat = []
             for message in chat:
-                new_message = {"role": message['role'], "content": message['content']}
+                new_message = message.copy()
                 for var, value in variable_map.items():
-                    if var in new_message['content']:
+                    if new_message.get('content') and var in new_message['content']:
                         new_message['content'] = new_message['content'].replace(var, value)
                 new_chat.append(new_message)
             return new_chat