aient 1.1.90__py3-none-any.whl → 1.1.92__py3-none-any.whl

aient/architext/architext/core.py

@@ -1,165 +1,436 @@
+import pickle
+import base64
 import asyncio
+import logging
+import hashlib
+import mimetypes
+from dataclasses import dataclass
 from abc import ABC, abstractmethod
 from typing import List, Dict, Any, Optional, Union
 
 # 1. 核心数据结构: ContentBlock
+@dataclass
 class ContentBlock:
-    def __init__(self, name: str, content: str, provider: Optional['ContextProvider'] = None):
-        self.name = name; self.content = content; self.provider = provider
-    def __repr__(self): return f"Block(name='{self.name}')"
+    name: str
+    content: str
 
 # 2. 上下文提供者 (带缓存)
 class ContextProvider(ABC):
     def __init__(self, name: str):
         self.name = name; self._cached_content: Optional[str] = None; self._is_stale: bool = True
     def mark_stale(self): self._is_stale = True
-    async def _refresh(self):
+    async def refresh(self):
         if self._is_stale:
-            # 注意：我们将在这个方法上使用 mock，所以实际的 print 不再重要
-            # print(f"信息: 正在为上下文提供者 '{self.name}' 刷新内容...")
-            self._cached_content = await self._fetch_content()
+            self._cached_content = await self.render()
             self._is_stale = False
-        # else:
-            # print(f"调试: 上下文提供者 '{self.name}' 正在使用缓存内容。")
     @abstractmethod
-    async def _fetch_content(self) -> Optional[str]: raise NotImplementedError
-    async def render(self) -> Optional[ContentBlock]:
-        await self._refresh()
-        if self._cached_content is not None: return ContentBlock(self.name, self._cached_content, self)
+    async def render(self) -> Optional[str]: raise NotImplementedError
+    @abstractmethod
+    def update(self, *args, **kwargs): raise NotImplementedError
+    def get_content_block(self) -> Optional[ContentBlock]:
+        if self._cached_content is not None: return ContentBlock(self.name, self._cached_content)
         return None
 
 class Texts(ContextProvider):
-    def __init__(self, name: str, text: str): super().__init__(name); self._text = text
-    async def _fetch_content(self) -> str: return self._text
+    def __init__(self, text: str, name: Optional[str] = None):
+        self._text = text
+        if name is None:
+            h = hashlib.sha1(self._text.encode()).hexdigest()
+            _name = f"text_{h[:8]}"
+        else:
+            _name = name
+        super().__init__(_name)
+
+    def update(self, text: str):
+        self._text = text
+        self.mark_stale()
+
+    async def render(self) -> str: return self._text
 
 class Tools(ContextProvider):
     def __init__(self, tools_json: List[Dict]): super().__init__("tools"); self._tools_json = tools_json
-    async def _fetch_content(self) -> str: return f"<tools>{str(self._tools_json)}</tools>"
+    def update(self, tools_json: List[Dict]):
+        self._tools_json = tools_json
+        self.mark_stale()
+    async def render(self) -> str: return f"<tools>{str(self._tools_json)}</tools>"
 
 class Files(ContextProvider):
-    def __init__(self): super().__init__("files"); self._files: Dict[str, str] = {}
-    def update(self, path: str, content: str): self._files[path] = content; self.mark_stale()
-    async def _fetch_content(self) -> str:
-        if not self._files: return None
-        return "<files>\n" + "\n".join([f"<file path='{p}'>{c[:50]}...</file>" for p, c in self._files.items()]) + "\n</files>"
+    def __init__(self, *paths: Union[str, List[str]]):
+        super().__init__("files")
+        self._files: Dict[str, str] = {}
+
+        file_paths: List[str] = []
+        if paths:
+            # Handle the case where the first argument is a list of paths, e.g., Files(['a', 'b'])
+            if len(paths) == 1 and isinstance(paths[0], list):
+                file_paths.extend(paths[0])
+            # Handle the case where arguments are individual string paths, e.g., Files('a', 'b')
+            else:
+                file_paths.extend(paths)
+
+        if file_paths:
+            for path in file_paths:
+                try:
+                    with open(path, 'r', encoding='utf-8') as f:
+                        self._files[path] = f.read()
+                except FileNotFoundError:
+                    logging.warning(f"File not found during initialization: {path}. Skipping.")
+                except Exception as e:
+                    logging.error(f"Error reading file {path} during initialization: {e}")
 
-Item = Union[ContentBlock, ContextProvider]
+    async def refresh(self):
+        """
+        Overrides the default refresh behavior. It synchronizes the content of
+        all tracked files with the file system. If a file is not found, its
+        content is updated to reflect the error.
+        """
+        is_changed = False
+        for path in list(self._files.keys()):
+            try:
+                with open(path, 'r', encoding='utf-8') as f:
+                    new_content = f.read()
+                if self._files.get(path) != new_content:
+                    self._files[path] = new_content
+                    is_changed = True
+            except FileNotFoundError:
+                error_msg = f"[Error: File not found at path '{path}']"
+                if self._files.get(path) != error_msg:
+                    self._files[path] = error_msg
+                    is_changed = True
+            except Exception as e:
+                error_msg = f"[Error: Could not read file at path '{path}': {e}]"
+                if self._files.get(path) != error_msg:
+                    self._files[path] = error_msg
+                    is_changed = True
 
-# 3. 消息内容类与消息类
-class MessageContent:
-    def __init__(self, items: List[Item]): self._items: List[Item] = items
+        if is_changed:
+            self.mark_stale()
+
+        await super().refresh()
+
+    def update(self, path: str, content: Optional[str] = None):
+        """
+        Updates a single file. If content is provided, it updates the file in
+        memory. If content is None, it reads the file from disk.
+        """
+        if content is not None:
+            self._files[path] = content
+        else:
+            try:
+                with open(path, 'r', encoding='utf-8') as f:
+                    self._files[path] = f.read()
+            except FileNotFoundError:
+                logging.error(f"File not found for update: {path}.")
+                self._files[path] = f"[Error: File not found at path '{path}']"
+            except Exception as e:
+                logging.error(f"Error reading file for update {path}: {e}.")
+                self._files[path] = f"[Error: Could not read file at path '{path}': {e}]"
+        self.mark_stale()
     async def render(self) -> str:
-        tasks = []
-        for item in self._items:
-            if isinstance(item, ContextProvider):
-                tasks.append(item.render())
-            elif isinstance(item, ContentBlock):
-                # 为了统一处理，将 ContentBlock 也包装成一个已完成的 Future
-                future = asyncio.Future()
-                future.set_result(item)
-                tasks.append(future)
-
-        blocks = await asyncio.gather(*tasks)
+        if not self._files: return None
+        return "<latest_file_content>" + "\n".join([f"<file><file_path>{p}</file_path><file_content>{c}</file_content></file>" for p, c in self._files.items()]) + "\n</latest_file_content>"
+
+class Images(ContextProvider):
+    def __init__(self, url: str, name: Optional[str] = None):
+        super().__init__(name or url)
+        self.url = url
+    def update(self, url: str):
+        self.url = url
+        self.mark_stale()
+    async def render(self) -> Optional[str]:
+        if self.url.startswith("data:"):
+            return self.url
+        try:
+            with open(self.url, "rb") as image_file:
+                encoded_string = base64.b64encode(image_file.read()).decode('utf-8')
+                mime_type, _ = mimetypes.guess_type(self.url)
+                if not mime_type: mime_type = "application/octet-stream" # Fallback
+                return f"data:{mime_type};base64,{encoded_string}"
+        except FileNotFoundError:
+            logging.warning(f"Image file not found: {self.url}. Skipping.")
+            return None # Or handle error appropriately
+
+# 3. 消息类 (已合并 MessageContent)
+class Message(ABC):
+    def __init__(self, role: str, *initial_items: Union[ContextProvider, str, list]):
+        self.role = role
+        processed_items = []
+        for item in initial_items:
+            if isinstance(item, str):
+                processed_items.append(Texts(text=item))
+            elif isinstance(item, Message):
+                processed_items.extend(item.providers())
+            elif isinstance(item, ContextProvider):
+                processed_items.append(item)
+            elif isinstance(item, list):
+                for sub_item in item:
+                    if not isinstance(sub_item, dict) or 'type' not in sub_item:
+                        raise ValueError("List items must be dicts with a 'type' key.")
+
+                    item_type = sub_item['type']
+                    if item_type == 'text':
+                        processed_items.append(Texts(text=sub_item.get('text', '')))
+                    elif item_type == 'image_url':
+                        image_url = sub_item.get('image_url', {}).get('url')
+                        if image_url:
+                            processed_items.append(Images(url=image_url))
+                    else:
+                        raise ValueError(f"Unsupported item type in list: {item_type}")
+            else:
+                raise TypeError(f"Unsupported item type: {type(item)}. Must be str, ContextProvider, or list.")
+        self._items: List[ContextProvider] = processed_items
+        self._parent_messages: Optional['Messages'] = None
+
+    def _render_content(self) -> str:
+        blocks = [item.get_content_block() for item in self._items]
         return "\n\n".join(b.content for b in blocks if b and b.content)
-    def pop(self, name: str) -> Optional[Item]:
+
+    def pop(self, name: str) -> Optional[ContextProvider]:
+        popped_item = None
         for i, item in enumerate(self._items):
-            if hasattr(item, 'name') and item.name == name: return self._items.pop(i)
-        return None
-    def insert(self, index: int, item: Item): self._items.insert(index, item)
-    def append(self, item: Item): self._items.append(item)
-    def providers(self) -> List[ContextProvider]: return [item for item in self._items if isinstance(item, ContextProvider)]
-    def __repr__(self): return f"Content(items={[item.name for item in self._items if hasattr(item, 'name')]})"
+            if hasattr(item, 'name') and item.name == name:
+                popped_item = self._items.pop(i)
+                break
+        if popped_item and self._parent_messages:
+            self._parent_messages._notify_provider_removed(popped_item)
+        return popped_item
 
-class Message(ABC):
-    def __init__(self, role: str, *initial_items: Item): self.role = role; self.content = MessageContent(list(initial_items))
-    async def to_dict(self) -> Optional[Dict[str, Any]]:
-        rendered_content = await self.content.render()
-        if not rendered_content: return None
-        return {"role": self.role, "content": rendered_content}
+    def insert(self, index: int, item: ContextProvider):
+        self._items.insert(index, item)
+        if self._parent_messages:
+            self._parent_messages._notify_provider_added(item, self)
+
+    def append(self, item: ContextProvider):
+        self._items.append(item)
+        if self._parent_messages:
+            self._parent_messages._notify_provider_added(item, self)
+
+    def providers(self) -> List[ContextProvider]: return self._items
+
+    def __add__(self, other):
+        if isinstance(other, str):
+            new_items = self._items + [Texts(text=other)]
+            return type(self)(*new_items)
+        if isinstance(other, Message):
+            new_items = self._items + other.providers()
+            return type(self)(*new_items)
+        return NotImplemented
+
+    def __radd__(self, other):
+        if isinstance(other, str):
+            new_items = [Texts(text=other)] + self._items
+            return type(self)(*new_items)
+        if isinstance(other, Message):
+            new_items = other.providers() + self._items
+            return type(self)(*new_items)
+        return NotImplemented
+
+    def __getitem__(self, key: str) -> Any:
+        """
+        使得 Message 对象支持字典风格的访问 (e.g., message['content'])。
+        """
+        if key == 'role':
+            return self.role
+        elif key == 'content':
+            # 直接调用 to_dict 并提取 'content'，确保逻辑一致
+            rendered_dict = self.to_dict()
+            return rendered_dict.get('content') if rendered_dict else None
+        # 对于 tool_calls 等特殊属性，也通过 to_dict 获取
+        elif hasattr(self, key):
+            rendered_dict = self.to_dict()
+            if rendered_dict and key in rendered_dict:
+                return rendered_dict[key]
+
+        # 如果在对象本身或其 to_dict() 中都找不到，则引发 KeyError
+        if hasattr(self, key):
+             return getattr(self, key)
+        raise KeyError(f"'{key}'")
+
+    def __repr__(self): return f"Message(role='{self.role}', items={[i.name for i in self._items]})"
+    def __bool__(self) -> bool:
+        return bool(self._items)
+    def get(self, key: str, default: Any = None) -> Any:
+        """提供类似字典的 .get() 方法来访问属性。"""
+        return getattr(self, key, default)
+    def to_dict(self) -> Optional[Dict[str, Any]]:
+        is_multimodal = any(isinstance(p, Images) for p in self._items)
+
+        if not is_multimodal:
+            rendered_content = self._render_content()
+            if not rendered_content: return None
+            return {"role": self.role, "content": rendered_content}
+        else:
+            content_list = []
+            for item in self._items:
+                block = item.get_content_block()
+                if not block or not block.content: continue
+                if isinstance(item, Images):
+                    content_list.append({"type": "image_url", "image_url": {"url": block.content}})
+                else:
+                    content_list.append({"type": "text", "text": block.content})
+            if not content_list: return None
+            return {"role": self.role, "content": content_list}
 
 class SystemMessage(Message):
     def __init__(self, *items): super().__init__("system", *items)
 class UserMessage(Message):
     def __init__(self, *items): super().__init__("user", *items)
+class AssistantMessage(Message):
+    def __init__(self, *items): super().__init__("assistant", *items)
+
+class RoleMessage:
+    """A factory class that creates a specific message type based on the role."""
+    def __new__(cls, role: str, *items):
+        if role == 'system':
+            return SystemMessage(*items)
+        elif role == 'user':
+            return UserMessage(*items)
+        elif role == 'assistant':
+            return AssistantMessage(*items)
+        else:
+            raise ValueError(f"Invalid role: {role}. Must be 'system', 'user', or 'assistant'.")
+
+class ToolCalls(Message):
+    """Represents an assistant message that requests tool calls."""
+    def __init__(self, tool_calls: List[Any]):
+        super().__init__("assistant")
+        self.tool_calls = tool_calls
+
+    def to_dict(self) -> Dict[str, Any]:
+        # Duck-typing serialization for OpenAI's tool_call objects
+        serialized_calls = []
+        for tc in self.tool_calls:
+            try:
+                # Attempt to serialize based on openai-python > 1.0 tool_call structure
+                func = tc.function
+                serialized_calls.append({
+                    "id": tc.id,
+                    "type": tc.type,
+                    "function": { "name": func.name, "arguments": func.arguments }
+                })
+            except AttributeError:
+                if isinstance(tc, dict):
+                    serialized_calls.append(tc) # Assume it's already a serializable dict
+                else:
+                    raise TypeError(f"Unsupported tool_call type: {type(tc)}. It should be an OpenAI tool_call object or a dict.")
+
+        return {
+            "role": self.role,
+            "tool_calls": serialized_calls,
+            "content": None
+        }
 
-# 4. 上下文构建器 (内部使用)
-class ContextBuilder:
-    def __init__(self, providers: List[ContextProvider]): self.providers = {p.name: p for p in providers}
-    def get_provider(self, name: str) -> Optional[ContextProvider]: return self.providers.get(name)
+class ToolResults(Message):
+    """Represents a tool message with the result of a single tool call."""
+    def __init__(self, tool_call_id: str, content: str):
+        super().__init__("tool")
+        self.tool_call_id = tool_call_id
+        self.content = content
 
-# 5. 顶层容器: Messages
+    def to_dict(self) -> Dict[str, Any]:
+        return {
+            "role": self.role,
+            "tool_call_id": self.tool_call_id,
+            "content": self.content
+        }
+
+# 4. 顶层容器: Messages
 class Messages:
     def __init__(self, *initial_messages: Message):
-        self._messages: List[Message] = list(initial_messages)
-        all_providers = []
-        for msg in self._messages: all_providers.extend(msg.content.providers())
-        self._context_builder = ContextBuilder(all_providers)
-    def provider(self, name: str) -> Optional[ContextProvider]: return self._context_builder.get_provider(name)
-    def pop(self, name: str) -> Optional[Item]:
-        for message in self._messages:
-            popped_item = message.content.pop(name)
-            if popped_item: return popped_item
+        from typing import Tuple
+        self._messages: List[Message] = []
+        self._providers_index: Dict[str, Tuple[ContextProvider, Message]] = {}
+        if initial_messages:
+            for msg in initial_messages:
+                self.append(msg)
+
+    def _notify_provider_added(self, provider: ContextProvider, message: Message):
+        if provider.name not in self._providers_index:
+            self._providers_index[provider.name] = (provider, message)
+
+    def _notify_provider_removed(self, provider: ContextProvider):
+        if provider.name in self._providers_index:
+            del self._providers_index[provider.name]
+
+    def provider(self, name: str) -> Optional[ContextProvider]:
+        indexed = self._providers_index.get(name)
+        return indexed[0] if indexed else None
+
+    def pop(self, key: Optional[Union[str, int]] = None) -> Union[Optional[ContextProvider], Optional[Message]]:
+        # If no key is provided, pop the last message.
+        if key is None:
+            key = len(self._messages) - 1
+
+        if isinstance(key, str):
+            indexed = self._providers_index.get(key)
+            if not indexed:
+                return None
+            _provider, parent_message = indexed
+            return parent_message.pop(key)
+        elif isinstance(key, int):
+            try:
+                if key < 0: # Handle negative indices like -1
+                    key += len(self._messages)
+                if not (0 <= key < len(self._messages)):
+                    return None
+                popped_message = self._messages.pop(key)
+                popped_message._parent_messages = None
+                for provider in popped_message.providers():
+                    self._notify_provider_removed(provider)
+                return popped_message
+            except IndexError:
+                return None
+
         return None
-    async def render(self) -> List[Dict[str, Any]]:
-        tasks = [msg.to_dict() for msg in self._messages]
-        results = await asyncio.gather(*tasks)
+
+    async def refresh(self):
+        tasks = [provider.refresh() for provider, _ in self._providers_index.values()]
+        await asyncio.gather(*tasks)
+
+    def render(self) -> List[Dict[str, Any]]:
+        results = [msg.to_dict() for msg in self._messages]
         return [res for res in results if res]
+
+    async def render_latest(self) -> List[Dict[str, Any]]:
+        await self.refresh()
+        return self.render()
+
     def append(self, message: Message):
-        self._messages.append(message)
-        for p in message.content.providers():
-            if p.name not in self._context_builder.providers: self._context_builder.providers[p.name] = p
+        if self._messages and self._messages[-1].role == message.role:
+            last_message = self._messages[-1]
+            for provider in message.providers():
+                last_message.append(provider)
+        else:
+            message._parent_messages = self
+            self._messages.append(message)
+            for p in message.providers():
+                self._notify_provider_added(p, message)
+
+    def save(self, file_path: str):
+        """
+        Saves the entire Messages object to a file using pickle.
+        Warning: Deserializing data with pickle from an untrusted source is insecure.
+        """
+        with open(file_path, 'wb') as f:
+            pickle.dump(self, f)
+
+    @classmethod
+    def load(cls, file_path: str) -> Optional['Messages']:
+        """
+        Loads a Messages object from a file using pickle.
+        Returns the loaded object, or None if the file is not found or an error occurs.
+        Warning: Only load files from a trusted source.
+        """
+        try:
+            with open(file_path, 'rb') as f:
+                return pickle.load(f)
+        except FileNotFoundError:
+            logging.warning(f"File not found at {file_path}, returning empty Messages.")
+            return cls()
+        except (pickle.UnpicklingError, EOFError) as e:
+            logging.error(f"Could not deserialize file {file_path}: {e}")
+            return cls()
+
     def __getitem__(self, index: int) -> Message: return self._messages[index]
     def __len__(self) -> int: return len(self._messages)
     def __iter__(self): return iter(self._messages)
-
-
-# ==============================================================================
-# 6. 演示
-# ==============================================================================
-async def run_demo():
-    # --- 1. 初始化提供者 ---
-    system_prompt_provider = Texts("system_prompt", "你是一个AI助手。")
-    tools_provider = Tools(tools_json=[{"name": "read_file"}])
-    files_provider = Files()
-
-    # --- 2. 演示新功能：优雅地构建 Messages ---
-    print("\n>>> 场景 A: 使用新的、优雅的构造函数直接初始化 Messages")
-    messages = Messages(
-        SystemMessage(system_prompt_provider, tools_provider),
-        UserMessage(files_provider, Texts("user_input", "这是我的初始问题。"))
-    )
-
-    print("\n--- 渲染后的初始 Messages (首次渲染，全部刷新) ---")
-    for msg_dict in await messages.render(): print(msg_dict)
-    print("-" * 40)
-
-    # --- 3. 演示穿透更新 ---
-    print("\n>>> 场景 B: 穿透更新 File Provider，渲染时自动刷新")
-
-    # 直接通过 messages 对象穿透访问并更新 files provider
-    files_provider_instance = messages.provider("files")
-    if isinstance(files_provider_instance, Files):
-        files_provider_instance.update("file1.py", "这是新的文件内容！")
-
-    print("\n--- 再次渲染 Messages (只有文件提供者会刷新) ---")
-    for msg_dict in await messages.render(): print(msg_dict)
-    print("-" * 40)
-
-    # --- 4. 演示全局 Pop 和通过索引 Insert ---
-    print("\n>>> 场景 C: 全局 Pop 工具提供者，并 Insert 到 UserMessage 中")
-
-    # a. 全局弹出 'tools' Provider
-    popped_tools_provider = messages.pop("tools")
-
-    # b. 将弹出的 Provider 插入到第一个 UserMessage (索引为1) 的开头
-    if popped_tools_provider:
-        # 通过索引精确定位
-        messages[1].content.insert(0, popped_tools_provider)
-        print(f"\n已成功将 '{popped_tools_provider.name}' 提供者移动到用户消息。")
-
-    print("\n--- Pop 和 Insert 后渲染的 Messages (验证移动效果) ---")
-    for msg_dict in await messages.render(): print(msg_dict)
-    print("-" * 40)
-
-if __name__ == "__main__":
-    asyncio.run(run_demo())