symbolicai 0.20.2__py3-none-any.whl → 1.0.0__py3-none-any.whl

symai/backend/engines/neurosymbolic/engine_anthropic_claudeX_chat.py

@@ -2,20 +2,22 @@ import json
 import logging
 import re
 from copy import copy, deepcopy
-from typing import List, Optional
 
 import anthropic
 from anthropic._types import NOT_GIVEN
-from anthropic.types import (InputJSONDelta, Message,
-                             RawContentBlockDeltaEvent,
-                             RawContentBlockStartEvent,
-                             RawContentBlockStopEvent, TextBlock, TextDelta,
-                             ToolUseBlock)
+from anthropic.types import (
+    InputJSONDelta,
+    Message,
+    RawContentBlockDeltaEvent,
+    RawContentBlockStartEvent,
+    RawContentBlockStopEvent,
+    TextBlock,
+    TextDelta,
+    ToolUseBlock,
+)
 
 from ....components import SelfPrompt
-from ....misc.console import ConsoleStyle
-from ....symbol import Symbol
-from ....utils import CustomUserWarning, encode_media_frames
+from ....utils import UserMessage, encode_media_frames
 from ...base import Engine
 from ...mixin.anthropic import AnthropicMixin
 from ...settings import SYMAI_CONFIG
@@ -34,7 +36,7 @@ class TokenizerWrapper:
         return self.compute_tokens_func([{"role": "user", "content": text}])
 
 class ClaudeXChatEngine(Engine, AnthropicMixin):
-    def __init__(self, api_key: Optional[str] = None, model: Optional[str] = None):
+    def __init__(self, api_key: str | None = None, model: str | None = None):
         super().__init__()
         self.config = deepcopy(SYMAI_CONFIG)
         # In case we use EngineRepository.register to inject the api_key and model => dynamically change the engine at runtime
@@ -69,54 +71,76 @@ class ClaudeXChatEngine(Engine, AnthropicMixin):
             self.model = kwargs['NEUROSYMBOLIC_ENGINE_MODEL']
 
     def compute_required_tokens(self, messages) -> int:
+        claude_messages, system_content = self._build_claude_messages(messages)
+
+        if not claude_messages:
+            return 0
+
+        return self._count_claude_tokens(claude_messages, system_content)
+
+    def _build_claude_messages(self, messages):
         claude_messages = []
         system_content = None
 
+        for role, content_str in self._message_parts(messages):
+            if role == 'system':
+                system_content = content_str
+                continue
+
+            if role in ['user', 'assistant']:
+                message_content = self._build_message_content(content_str)
+                if message_content:
+                    claude_messages.append(self._create_claude_message(role, message_content))
+
+        return claude_messages, system_content
+
+    def _message_parts(self, messages):
         for msg in messages:
-            if not isinstance(msg, list):
-                msg = [msg]
-            for part in msg:
-                if isinstance(part, str):
-                    role = 'user'
-                    content_str = part
-                elif isinstance(part, dict):
-                    role = part.get('role')
-                    content_str = str(part.get('content', ''))
-                else:
-                    CustomUserWarning(f"Unsupported message part type: {type(part)}", raise_with=ValueError)
-
-                if role == 'system':
-                    system_content = content_str
-                    continue
-
-                if role in ['user', 'assistant']:
-                    message_content = []
-
-                    image_content = self._handle_image_content(content_str)
-                    message_content.extend(image_content)
-
-                    text_content = self._remove_vision_pattern(content_str)
-                    if text_content:
-                        message_content.append({
-                            "type": "text",
-                            "text": text_content
-                        })
-
-                    if message_content:
-                        if len(message_content) == 1 and message_content[0].get('type') == 'text':
-                            claude_messages.append({
-                                'role': role,
-                                'content': message_content[0]['text']
-                            })
-                        else:
-                            claude_messages.append({
-                                'role': role,
-                                'content': message_content
-                            })
+            msg_parts = msg if isinstance(msg, list) else [msg]
+            for part in msg_parts:
+                yield self._extract_message_details(part)
 
-        if not claude_messages:
-            return 0
+    def _extract_message_details(self, part):
+        if isinstance(part, str):
+            return 'user', part
+
+        if isinstance(part, dict):
+            role = part.get('role')
+            content_str = str(part.get('content', ''))
+            return role, content_str
+
+        msg = f"Unsupported message part type: {type(part)}"
+        UserMessage(msg, raise_with=ValueError)
+        raise ValueError(msg)
+
+    def _build_message_content(self, content_str: str) -> list:
+        message_content = []
 
+        image_content = self._handle_image_content(content_str)
+        message_content.extend(image_content)
+
+        text_content = self._remove_vision_pattern(content_str)
+        if text_content:
+            message_content.append({
+                "type": "text",
+                "text": text_content
+            })
+
+        return message_content
+
+    def _create_claude_message(self, role: str, message_content: list) -> dict:
+        if len(message_content) == 1 and message_content[0].get('type') == 'text':
+            return {
+                'role': role,
+                'content': message_content[0]['text']
+            }
+
+        return {
+            'role': role,
+            'content': message_content
+        }
+
+    def _count_claude_tokens(self, claude_messages: list, system_content: str | None) -> int:
         try:
             count_params = {
                 'model': self.model,
@@ -127,11 +151,11 @@ class ClaudeXChatEngine(Engine, AnthropicMixin):
             count_response = self.client.messages.count_tokens(**count_params)
             return count_response.input_tokens
         except Exception as e:
-            logging.error(f"Claude count_tokens failed: {e}")
-            CustomUserWarning(f"Error counting tokens for Claude: {str(e)}", raise_with=RuntimeError)
+            UserMessage(f"Claude count_tokens failed: {e}")
+            UserMessage(f"Error counting tokens for Claude: {e!s}", raise_with=RuntimeError)
 
-    def compute_remaining_tokens(self, prompts: list) -> int:
-        raise NotImplementedError('Method not implemented.')
+    def compute_remaining_tokens(self, _prompts: list) -> int:
+        UserMessage('Method not implemented.', raise_with=NotImplementedError)
 
     def _handle_image_content(self, content: str) -> list:
         """Handle image content by processing vision patterns and returning image file data."""
@@ -157,7 +181,7 @@ class ClaudeXChatEngine(Engine, AnthropicMixin):
                 elif len(buffer) == 1:
                     image_files.append({'data': buffer[0], 'media_type': f'image/{ext}', 'type': 'base64'})
                 else:
-                    CustomUserWarning(f'No frames found for image!')
+                    UserMessage('No frames found for image!')
         return image_files
 
     def _remove_vision_pattern(self, text: str) -> str:
@@ -180,21 +204,21 @@ class ClaudeXChatEngine(Engine, AnthropicMixin):
         except Exception as e:
             if anthropic.api_key is None or anthropic.api_key == '':
                 msg = 'Anthropic API key is not set. Please set it in the config file or pass it as an argument to the command method.'
-                logging.error(msg)
+                UserMessage(msg)
                 if self.config['NEUROSYMBOLIC_ENGINE_API_KEY'] is None or self.config['NEUROSYMBOLIC_ENGINE_API_KEY'] == '':
-                    CustomUserWarning(msg, raise_with=ValueError)
+                    UserMessage(msg, raise_with=ValueError)
                 anthropic.api_key = self.config['NEUROSYMBOLIC_ENGINE_API_KEY']
 
             callback = self.client.messages.create
-            kwargs['model'] = kwargs['model'] if 'model' in kwargs else self.model
+            kwargs['model'] = kwargs.get('model', self.model)
 
             if except_remedy is not None:
                 res = except_remedy(self, e, callback, argument)
             else:
-                CustomUserWarning(f'Error during generation. Caused by: {e}', raise_with=ValueError)
+                UserMessage(f'Error during generation. Caused by: {e}', raise_with=ValueError)
 
         if payload['stream']:
-            res = [_ for _ in res] # Unpack the iterator to a list
+            res = list(res) # Unpack the iterator to a list
         metadata = {'raw_output': res}
         response_data = self._collect_response(res)
 
@@ -210,11 +234,13 @@ class ClaudeXChatEngine(Engine, AnthropicMixin):
 
     def _prepare_raw_input(self, argument):
         if not argument.prop.processed_input:
-            raise ValueError('Need to provide a prompt instruction to the engine if `raw_input` is enabled!')
+            msg = 'Need to provide a prompt instruction to the engine if `raw_input` is enabled!'
+            UserMessage(msg)
+            raise ValueError(msg)
         system = NOT_GIVEN
         prompt = copy(argument.prop.processed_input)
-        if type(prompt) != list:
-            if type(prompt) != dict:
+        if not isinstance(prompt, list):
+            if not isinstance(prompt, dict):
                 prompt = {'role': 'user', 'content': str(prompt)}
             prompt = [prompt]
         if len(prompt) > 1:
@@ -233,85 +259,113 @@ class ClaudeXChatEngine(Engine, AnthropicMixin):
             argument.prop.prepared_input = self._prepare_raw_input(argument)
             return
 
-        _non_verbose_output = """<META_INSTRUCTION/>\nYou do not output anything else, like verbose preambles or post explanation, such as "Sure, let me...", "Hope that was helpful...", "Yes, I can help you with that...", etc. Consider well formatted output, e.g. for sentences use punctuation, spaces etc. or for code use indentation, etc. Never add meta instructions information to your output!\n\n"""
-        user:   str = ""
-        system: str = ""
+        non_verbose_output = """<META_INSTRUCTION/>\nYou do not output anything else, like verbose preambles or post explanation, such as "Sure, let me...", "Hope that was helpful...", "Yes, I can help you with that...", etc. Consider well formatted output, e.g. for sentences use punctuation, spaces etc. or for code use indentation, etc. Never add meta instructions information to your output!\n\n"""
+        image_files = self._handle_image_content(str(argument.prop.processed_input))
+        has_image = len(image_files) > 0
+
+        system = self._build_system_prompt(argument, has_image, non_verbose_output)
+        user_text, user_prompt, image_blocks = self._build_user_prompt(argument, image_files)
+        system, user_prompt = self._apply_self_prompt_if_needed(argument, system, user_text, image_blocks, user_prompt)
+
+        argument.prop.prepared_input = (system, [user_prompt])
 
+    def _build_system_prompt(self, argument, has_image: bool, non_verbose_output: str) -> str:
+        system = self._build_system_prefix(argument, non_verbose_output)
+        system = self._append_context_sections(system, argument)
+        system = self._append_instruction_section(system, argument, has_image)
+        return self._append_template_suffix(system, argument)
+
+    def _build_system_prefix(self, argument, non_verbose_output: str) -> str:
+        system = ""
         if argument.prop.suppress_verbose_output:
-            system += _non_verbose_output
+            system += non_verbose_output
+
         system = f'{system}\n' if system and len(system) > 0 else ''
 
         if argument.prop.response_format:
-            _rsp_fmt = argument.prop.response_format
-            assert _rsp_fmt.get('type') is not None, 'Response format type is required! Expected format `{"type": str}`! The str value will be passed to the engine. Refer to the Anthropic documentation for more information: https://docs.anthropic.com/en/docs/test-and-evaluate/strengthen-guardrails/increase-consistency#example-standardizing-customer-feedback'
-            system += _non_verbose_output
-            system += f'<RESPONSE_FORMAT/>\n{_rsp_fmt["type"]}\n\n'
+            response_format = argument.prop.response_format
+            assert response_format.get('type') is not None, 'Response format type is required! Expected format `{"type": str}`! The str value will be passed to the engine. Refer to the Anthropic documentation for more information: https://docs.anthropic.com/en/docs/test-and-evaluate/strengthen-guardrails/increase-consistency#example-standardizing-customer-feedback'
+            system += non_verbose_output
+            system += f'<RESPONSE_FORMAT/>\n{response_format["type"]}\n\n'
 
+        return system
+
+    def _append_context_sections(self, system: str, argument) -> str:
         ref = argument.prop.instance
-        static_ctxt, dyn_ctxt = ref.global_context
-        if len(static_ctxt) > 0:
-            system += f"<STATIC_CONTEXT/>\n{static_ctxt}\n\n"
+        static_context, dynamic_context = ref.global_context
+
+        if len(static_context) > 0:
+            system += f"<STATIC_CONTEXT/>\n{static_context}\n\n"
 
-        if len(dyn_ctxt) > 0:
-            system += f"<DYNAMIC_CONTEXT/>\n{dyn_ctxt}\n\n"
+        if len(dynamic_context) > 0:
+            system += f"<DYNAMIC_CONTEXT/>\n{dynamic_context}\n\n"
 
         payload = argument.prop.payload
         if argument.prop.payload:
-            system += f"<ADDITIONAL_CONTEXT/>\n{str(payload)}\n\n"
+            system += f"<ADDITIONAL_CONTEXT/>\n{payload!s}\n\n"
 
-        examples: List[str] = argument.prop.examples
+        examples: list[str] = argument.prop.examples
         if examples and len(examples) > 0:
-            system += f"<EXAMPLES/>\n{str(examples)}\n\n"
+            system += f"<EXAMPLES/>\n{examples!s}\n\n"
 
-        image_files = self._handle_image_content(str(argument.prop.processed_input))
+        return system
 
+    def _append_instruction_section(self, system: str, argument, has_image: bool) -> str:
         if argument.prop.prompt is not None and len(argument.prop.prompt) > 0:
-            val = str(argument.prop.prompt)
-            if len(image_files) > 0:
-                val = self._remove_vision_pattern(val)
-            system += f"<INSTRUCTION/>\n{val}\n\n"
+            instruction_value = str(argument.prop.prompt)
+            if has_image:
+                instruction_value = self._remove_vision_pattern(instruction_value)
+            system += f"<INSTRUCTION/>\n{instruction_value}\n\n"
+
+        return system
+
+    def _append_template_suffix(self, system: str, argument) -> str:
+        if argument.prop.template_suffix:
+            system += f' You will only generate content for the placeholder `{argument.prop.template_suffix!s}` following the instructions and the provided context information.\n\n'
 
-        suffix: str = str(argument.prop.processed_input)
+        return system
+
+    def _build_user_prompt(self, argument, image_files):
+        suffix = str(argument.prop.processed_input)
         if len(image_files) > 0:
             suffix = self._remove_vision_pattern(suffix)
 
-        user += f"{suffix}"
+        user_text = f"{suffix}"
+        if not user_text:
+            user_text = "N/A"
 
-        if not len(user):
-            # Anthropic doesn't allow empty user prompts; force it
-            user = "N/A"
+        image_blocks = [{'type': 'image', 'source': image_file} for image_file in image_files]
+        user_prompt = self._wrap_user_prompt_content(user_text, image_blocks)
+        return user_text, user_prompt, image_blocks
 
-        if argument.prop.template_suffix:
-            system += f' You will only generate content for the placeholder `{str(argument.prop.template_suffix)}` following the instructions and the provided context information.\n\n'
+    def _wrap_user_prompt_content(self, user_text: str, image_blocks: list[dict]) -> dict:
+        if len(image_blocks) > 0:
+            return {
+                "role": "user",
+                "content": [
+                    *image_blocks,
+                    {'type': 'text', 'text': user_text}
+                ]
+            }
 
-        if len(image_files) > 0:
-            images = [{ 'type': 'image', "source": im } for im in image_files]
-            user_prompt = { "role": "user", "content": [
-                *images,
-                { 'type': 'text', 'text': user }
-            ]}
-        else:
-            user_prompt = { "role": "user", "content": user }
-
-        # First check if the `Symbol` instance has the flag set, otherwise check if it was passed as an argument to a method
-        if argument.prop.instance._kwargs.get('self_prompt', False) or argument.prop.self_prompt:
-            self_prompter = SelfPrompt()
-
-            res = self_prompter({'user': user, 'system': system})
-            if res is None:
-                raise ValueError("Self-prompting failed!")
-
-            if len(image_files) > 0:
-                user_prompt = { "role": "user", "content": [
-                    *images,
-                    { 'type': 'text', 'text': res['user'] }
-                ]}
-            else:
-                user_prompt = { "role": "user", "content": res['user'] }
+        return {
+            "role": "user",
+            "content": user_text
+        }
 
-            system = res['system']
+    def _apply_self_prompt_if_needed(self, argument, system: str, user_text: str, image_blocks: list[dict], user_prompt: dict):
+        if not (argument.prop.instance._kwargs.get('self_prompt', False) or argument.prop.self_prompt):
+            return system, user_prompt
 
-        argument.prop.prepared_input = (system, [user_prompt])
+        self_prompter = SelfPrompt()
+        res = self_prompter({'user': user_text, 'system': system})
+        if res is None:
+            msg = "Self-prompting failed!"
+            UserMessage(msg)
+            raise ValueError(msg)
+
+        updated_user_prompt = self._wrap_user_prompt_content(res['user'], image_blocks)
+        return res['system'], updated_user_prompt
 
     def _prepare_request_payload(self, argument):
         kwargs = argument.kwargs
@@ -326,7 +380,7 @@ class ClaudeXChatEngine(Engine, AnthropicMixin):
         tool_choice = kwargs.get('tool_choice', NOT_GIVEN)
         metadata_anthropic = kwargs.get('metadata', NOT_GIVEN)
 
-        if stop != NOT_GIVEN and type(stop) != list:
+        if stop != NOT_GIVEN and not isinstance(stop, list):
             stop = [stop]
 
         #@NOTE: Anthropic fails if stop is not raw string, so cast it to r'…'
@@ -349,69 +403,95 @@ class ClaudeXChatEngine(Engine, AnthropicMixin):
 
     def _collect_response(self, res):
         if isinstance(res, list):
-            text_content = ''
-            tool_calls_raw = []
-            active_tool_calls = {}
-
-            for chunk in res:
-                if isinstance(chunk, RawContentBlockStartEvent):
-                    if isinstance(chunk.content_block, ToolUseBlock):
-                        active_tool_calls[chunk.index] = {
-                            'id': chunk.content_block.id,
-                            'name': chunk.content_block.name,
-                            'input_json_str': ""
-                        }
-                elif isinstance(chunk, RawContentBlockDeltaEvent):
-                    if isinstance(chunk.delta, TextDelta):
-                        text_content += chunk.delta.text
-                    elif isinstance(chunk.delta, InputJSONDelta):
-                        if chunk.index in active_tool_calls:
-                            active_tool_calls[chunk.index]['input_json_str'] += chunk.delta.partial_json
-                elif isinstance(chunk, RawContentBlockStopEvent):
-                    if chunk.index in active_tool_calls:
-                        tool_call_info = active_tool_calls.pop(chunk.index)
-                        try:
-                            tool_call_info['input'] = json.loads(tool_call_info['input_json_str'])
-                        except json.JSONDecodeError as e:
-                            logging.error(f"Failed to parse JSON for tool call {tool_call_info['name']}: {e}. Raw JSON: '{tool_call_info['input_json_str']}'")
-                            tool_call_info['input'] = {}
-                        tool_calls_raw.append(tool_call_info)
-
-            function_call_data = None
-            if tool_calls_raw:
-                if len(tool_calls_raw) > 1:
-                    CustomUserWarning("Multiple tool calls detected in the stream but only the first one will be processed.")
-                function_call_data = {
-                    'name': tool_calls_raw[0]['name'],
-                    'arguments': tool_calls_raw[0]['input']
-                }
+            return self._collect_streaming_response(res)
 
-            return {
-                "text": text_content,
-                "function_call": function_call_data
-            }
-
-        # Non-streamed response (res is a Message object)
         if isinstance(res, Message):
-            text_content = ''
-            function_call_data = None
-            hit_tool_use = False
-
-            for content_block in res.content:
-                if isinstance(content_block, TextBlock):
-                    text_content += content_block.text
-                elif isinstance(content_block, ToolUseBlock):
-                    if hit_tool_use:
-                        CustomUserWarning("Multiple tool use blocks detected in the response but only the first one will be processed.")
-                    else:
-                        function_call_data = {
-                            'name': content_block.name,
-                            'arguments': content_block.input
-                        }
-                        hit_tool_use = True
-            return {
-                "text": text_content,
-                "function_call": function_call_data
+            return self._collect_message_response(res)
+
+        UserMessage(f"Unexpected response type from Anthropic API: {type(res)}", raise_with=ValueError)
+        return {}
+
+    def _collect_streaming_response(self, res):
+        text_parts = []
+        tool_calls_raw = []
+        active_tool_calls = {}
+
+        for chunk in res:
+            if isinstance(chunk, RawContentBlockStartEvent):
+                self._start_tool_call(chunk, active_tool_calls)
+            elif isinstance(chunk, RawContentBlockDeltaEvent):
+                self._update_stream_chunk(chunk, text_parts, active_tool_calls)
+            elif isinstance(chunk, RawContentBlockStopEvent):
+                tool_call = self._finish_tool_call(chunk, active_tool_calls)
+                if tool_call is not None:
+                    tool_calls_raw.append(tool_call)
+
+        text_content = ''.join(text_parts)
+        function_call_data = self._build_function_call_data(tool_calls_raw)
+
+        return {
+            "text": text_content,
+            "function_call": function_call_data
+        }
+
+    def _start_tool_call(self, chunk, active_tool_calls: dict):
+        if isinstance(chunk.content_block, ToolUseBlock):
+            active_tool_calls[chunk.index] = {
+                'id': chunk.content_block.id,
+                'name': chunk.content_block.name,
+                'input_json_str': ""
             }
 
-        CustomUserWarning(f"Unexpected response type from Anthropic API: {type(res)}", raise_with=ValueError)
+    def _update_stream_chunk(self, chunk, text_parts: list, active_tool_calls: dict):
+        if isinstance(chunk.delta, TextDelta):
+            text_parts.append(chunk.delta.text)
+        elif isinstance(chunk.delta, InputJSONDelta) and chunk.index in active_tool_calls:
+            active_tool_calls[chunk.index]['input_json_str'] += chunk.delta.partial_json
+
+    def _finish_tool_call(self, chunk, active_tool_calls: dict):
+        if chunk.index not in active_tool_calls:
+            return None
+
+        tool_call_info = active_tool_calls.pop(chunk.index)
+        try:
+            tool_call_info['input'] = json.loads(tool_call_info['input_json_str'])
+        except json.JSONDecodeError as e:
+            UserMessage(f"Failed to parse JSON for tool call {tool_call_info['name']}: {e}. Raw JSON: '{tool_call_info['input_json_str']}'")
+            tool_call_info['input'] = {}
+        return tool_call_info
+
+    def _build_function_call_data(self, tool_calls_raw: list | None) -> dict | None:
+        if not tool_calls_raw:
+            return None
+
+        if len(tool_calls_raw) > 1:
+            UserMessage("Multiple tool calls detected in the stream but only the first one will be processed.")
+
+        tool_call = tool_calls_raw[0]
+        return {
+            'name': tool_call['name'],
+            'arguments': tool_call['input']
+        }
+
+    def _collect_message_response(self, res: Message):
+        text_parts = []
+        function_call_data = None
+        hit_tool_use = False
+
+        for content_block in res.content:
+            if isinstance(content_block, TextBlock):
+                text_parts.append(content_block.text)
+            elif isinstance(content_block, ToolUseBlock):
+                if hit_tool_use:
+                    UserMessage("Multiple tool use blocks detected in the response but only the first one will be processed.")
+                else:
+                    function_call_data = {
+                        'name': content_block.name,
+                        'arguments': content_block.input
+                    }
+                    hit_tool_use = True
+
+        return {
+            "text": ''.join(text_parts),
+            "function_call": function_call_data
+        }