first commit
This commit is contained in:
DigiJ
@@ -0,0 +1,4 @@
|
||||
from ._text_canvas import TextCanvas
|
||||
from ._text_canvas_memory import TextCanvasMemory
|
||||
|
||||
__all__ = ["TextCanvas", "TextCanvasMemory"]
|
||||
50
agent_dhal/agentdhal_extensions/memory/canvas/_canvas.py
Normal file
50
agent_dhal/agentdhal_extensions/memory/canvas/_canvas.py
Normal file
@@ -0,0 +1,50 @@
|
||||
from abc import ABC, abstractmethod
|
||||
from typing import Any, Dict, Union
|
||||
|
||||
|
||||
class BaseCanvas(ABC):
|
||||
"""
|
||||
An abstract protocol for "canvas" objects that maintain
|
||||
revision history for file-like data. Concrete subclasses
|
||||
can handle text, images, structured data, etc.
|
||||
|
||||
.. warning::
|
||||
|
||||
This is an experimental API and may change in the future.
|
||||
|
||||
"""
|
||||
|
||||
@abstractmethod
|
||||
def list_files(self) -> Dict[str, int]:
|
||||
"""
|
||||
Returns a dict of filename -> latest revision number.
|
||||
"""
|
||||
raise NotImplementedError
|
||||
|
||||
@abstractmethod
|
||||
def get_latest_content(self, filename: str) -> Union[str, bytes, Any]:
|
||||
"""
|
||||
Returns the latest version of a file's content.
|
||||
"""
|
||||
raise NotImplementedError
|
||||
|
||||
@abstractmethod
|
||||
def add_or_update_file(self, filename: str, new_content: Union[str, bytes, Any]) -> None:
|
||||
"""
|
||||
Creates or updates the file content with a new revision.
|
||||
"""
|
||||
raise NotImplementedError
|
||||
|
||||
@abstractmethod
|
||||
def get_diff(self, filename: str, from_revision: int, to_revision: int) -> str:
|
||||
"""
|
||||
Returns a diff (in some format) between two revisions.
|
||||
"""
|
||||
raise NotImplementedError
|
||||
|
||||
@abstractmethod
|
||||
def apply_patch(self, filename: str, patch_data: Union[str, bytes, Any]) -> None:
|
||||
"""
|
||||
Applies a patch/diff to the latest revision and increments the revision.
|
||||
"""
|
||||
raise NotImplementedError
|
||||
@@ -0,0 +1,64 @@
|
||||
from agentdhal_core import CancellationToken
|
||||
from agentdhal_core.tools import BaseTool
|
||||
from pydantic import BaseModel
|
||||
|
||||
from ._text_canvas import TextCanvas
|
||||
|
||||
|
||||
class UpdateFileArgs(BaseModel):
|
||||
filename: str
|
||||
new_content: str
|
||||
|
||||
|
||||
class UpdateFileResult(BaseModel):
|
||||
status: str
|
||||
|
||||
|
||||
class UpdateFileTool(BaseTool[UpdateFileArgs, UpdateFileResult]):
|
||||
"""
|
||||
Overwrites or creates a file in the canvas.
|
||||
"""
|
||||
|
||||
def __init__(self, canvas: TextCanvas):
|
||||
super().__init__(
|
||||
args_type=UpdateFileArgs,
|
||||
return_type=UpdateFileResult,
|
||||
name="update_file",
|
||||
description="Create/update a file on the canvas with the provided content.",
|
||||
)
|
||||
self._canvas = canvas
|
||||
|
||||
async def run(self, args: UpdateFileArgs, cancellation_token: CancellationToken) -> UpdateFileResult:
|
||||
self._canvas.add_or_update_file(args.filename, args.new_content)
|
||||
return UpdateFileResult(status="OK")
|
||||
|
||||
|
||||
class ApplyPatchArgs(BaseModel):
|
||||
filename: str
|
||||
patch_text: str
|
||||
|
||||
|
||||
class ApplyPatchResult(BaseModel):
|
||||
status: str
|
||||
|
||||
|
||||
class ApplyPatchTool(BaseTool[ApplyPatchArgs, ApplyPatchResult]):
|
||||
"""
|
||||
Applies a unified diff patch to the given file on the canvas.
|
||||
"""
|
||||
|
||||
def __init__(self, canvas: TextCanvas):
|
||||
super().__init__(
|
||||
args_type=ApplyPatchArgs,
|
||||
return_type=ApplyPatchResult,
|
||||
name="apply_patch",
|
||||
description=(
|
||||
"Apply a unified diff patch to an existing file on the canvas. "
|
||||
"The patch must be in diff/patch format. The file must exist or be created first."
|
||||
),
|
||||
)
|
||||
self._canvas = canvas
|
||||
|
||||
async def run(self, args: ApplyPatchArgs, cancellation_token: CancellationToken) -> ApplyPatchResult:
|
||||
self._canvas.apply_patch(args.filename, args.patch_text)
|
||||
return ApplyPatchResult(status="PATCH APPLIED")
|
||||
192
agent_dhal/agentdhal_extensions/memory/canvas/_text_canvas.py
Normal file
192
agent_dhal/agentdhal_extensions/memory/canvas/_text_canvas.py
Normal file
@@ -0,0 +1,192 @@
|
||||
import difflib
|
||||
from typing import Any, Dict, List, Union
|
||||
|
||||
try: # pragma: no cover
|
||||
from unidiff import PatchSet
|
||||
except ModuleNotFoundError: # pragma: no cover
|
||||
PatchSet = None # type: ignore
|
||||
|
||||
from ._canvas import BaseCanvas
|
||||
|
||||
|
||||
class FileRevision:
|
||||
"""Tracks the history of one file's content."""
|
||||
|
||||
__slots__ = ("content", "revision")
|
||||
|
||||
def __init__(self, content: str, revision: int) -> None:
|
||||
self.content: str = content
|
||||
self.revision: int = revision # e.g. an integer, a timestamp, or git hash
|
||||
|
||||
|
||||
class TextCanvas(BaseCanvas):
|
||||
"""An in‑memory canvas that stores *text* files with full revision history.
|
||||
|
||||
.. warning::
|
||||
|
||||
This is an experimental API and may change in the future.
|
||||
|
||||
Besides the original CRUD‑like operations, this enhanced implementation adds:
|
||||
|
||||
* **apply_patch** – applies patches using the ``unidiff`` library for accurate
|
||||
hunk application and context line validation.
|
||||
* **get_revision_content** – random access to any historical revision.
|
||||
* **get_revision_diffs** – obtain the list of diffs applied between every
|
||||
consecutive pair of revisions so that a caller can replay or audit the
|
||||
full change history.
|
||||
"""
|
||||
|
||||
# ----------------------------------------------------------------------------------
|
||||
# Construction helpers
|
||||
# ----------------------------------------------------------------------------------
|
||||
|
||||
def __init__(self) -> None:
|
||||
# For each file we keep an *ordered* list of FileRevision where the last
|
||||
# element is the most recent. Using a list keeps the memory footprint
|
||||
# small and preserves order without any extra bookkeeping.
|
||||
self._files: Dict[str, List[FileRevision]] = {}
|
||||
|
||||
# ----------------------------------------------------------------------------------
|
||||
# Internal utilities
|
||||
# ----------------------------------------------------------------------------------
|
||||
|
||||
def _latest_idx(self, filename: str) -> int:
|
||||
"""Return the index (not revision number) of the newest revision."""
|
||||
return len(self._files.get(filename, [])) - 1
|
||||
|
||||
def _ensure_file(self, filename: str) -> None:
|
||||
if filename not in self._files:
|
||||
raise ValueError(f"File '{filename}' does not exist on the canvas; create it first.")
|
||||
|
||||
# ----------------------------------------------------------------------------------
|
||||
# Revision inspection helpers
|
||||
# ----------------------------------------------------------------------------------
|
||||
|
||||
def get_revision_content(self, filename: str, revision: int) -> str: # NEW 🚀
|
||||
"""Return the exact content stored in *revision*.
|
||||
|
||||
If the revision does not exist an empty string is returned so that
|
||||
downstream code can handle the "not found" case without exceptions.
|
||||
"""
|
||||
for rev in self._files.get(filename, []):
|
||||
if rev.revision == revision:
|
||||
return rev.content
|
||||
return ""
|
||||
|
||||
def get_revision_diffs(self, filename: str) -> List[str]: # NEW 🚀
|
||||
"""Return a *chronological* list of unified‑diffs for *filename*.
|
||||
|
||||
Each element in the returned list represents the diff that transformed
|
||||
revision *n* into revision *n+1* (starting at revision 1 → 2).
|
||||
"""
|
||||
revisions = self._files.get(filename, [])
|
||||
diffs: List[str] = []
|
||||
for i in range(1, len(revisions)):
|
||||
older, newer = revisions[i - 1], revisions[i]
|
||||
diff = difflib.unified_diff(
|
||||
older.content.splitlines(keepends=True),
|
||||
newer.content.splitlines(keepends=True),
|
||||
fromfile=f"{filename}@r{older.revision}",
|
||||
tofile=f"{filename}@r{newer.revision}",
|
||||
)
|
||||
diffs.append("".join(diff))
|
||||
return diffs
|
||||
|
||||
# ----------------------------------------------------------------------------------
|
||||
# BaseCanvas interface implementation
|
||||
# ----------------------------------------------------------------------------------
|
||||
|
||||
def list_files(self) -> Dict[str, int]:
|
||||
"""Return a mapping of *filename → latest revision number*."""
|
||||
return {fname: revs[-1].revision for fname, revs in self._files.items() if revs}
|
||||
|
||||
def get_latest_content(self, filename: str) -> str: # noqa: D401 – keep API identical
|
||||
"""Return the most recent content or an empty string if the file is new."""
|
||||
revs = self._files.get(filename, [])
|
||||
return revs[-1].content if revs else ""
|
||||
|
||||
def add_or_update_file(self, filename: str, new_content: Union[str, bytes, Any]) -> None:
|
||||
"""Create *filename* or append a new revision containing *new_content*."""
|
||||
if isinstance(new_content, bytes):
|
||||
new_content = new_content.decode("utf-8")
|
||||
if not isinstance(new_content, str):
|
||||
raise ValueError(f"Expected str or bytes, got {type(new_content)}")
|
||||
if filename not in self._files:
|
||||
self._files[filename] = [FileRevision(new_content, 1)]
|
||||
else:
|
||||
last_rev_num = self._files[filename][-1].revision
|
||||
self._files[filename].append(FileRevision(new_content, last_rev_num + 1))
|
||||
|
||||
def get_diff(self, filename: str, from_revision: int, to_revision: int) -> str:
|
||||
"""Return a unified diff between *from_revision* and *to_revision*."""
|
||||
revisions = self._files.get(filename, [])
|
||||
if not revisions:
|
||||
return ""
|
||||
# Fetch the contents for the requested revisions.
|
||||
from_content = self.get_revision_content(filename, from_revision)
|
||||
to_content = self.get_revision_content(filename, to_revision)
|
||||
if from_content == "" and to_content == "": # one (or both) revision ids not found
|
||||
return ""
|
||||
diff = difflib.unified_diff(
|
||||
from_content.splitlines(keepends=True),
|
||||
to_content.splitlines(keepends=True),
|
||||
fromfile=f"{filename}@r{from_revision}",
|
||||
tofile=f"{filename}@r{to_revision}",
|
||||
)
|
||||
return "".join(diff)
|
||||
|
||||
def apply_patch(self, filename: str, patch_data: Union[str, bytes, Any]) -> None:
|
||||
"""Apply *patch_text* (unified diff) to the latest revision and save a new revision.
|
||||
|
||||
Uses the *unidiff* library to accurately apply hunks and validate context lines.
|
||||
"""
|
||||
if isinstance(patch_data, bytes):
|
||||
patch_data = patch_data.decode("utf-8")
|
||||
if not isinstance(patch_data, str):
|
||||
raise ValueError(f"Expected str or bytes, got {type(patch_data)}")
|
||||
self._ensure_file(filename)
|
||||
original_content = self.get_latest_content(filename)
|
||||
|
||||
if PatchSet is None:
|
||||
raise ImportError(
|
||||
"The 'unidiff' package is required for patch application. Install with 'pip install unidiff'."
|
||||
)
|
||||
|
||||
patch = PatchSet(patch_data)
|
||||
# Our canvas stores exactly one file per patch operation so we
|
||||
# use the first (and only) patched_file object.
|
||||
if not patch:
|
||||
raise ValueError("Empty patch text provided.")
|
||||
patched_file = patch[0]
|
||||
working_lines = original_content.splitlines(keepends=True)
|
||||
line_offset = 0
|
||||
for hunk in patched_file:
|
||||
# Calculate the slice boundaries in the *current* working copy.
|
||||
start = hunk.source_start - 1 + line_offset
|
||||
end = start + hunk.source_length
|
||||
# Build the replacement block for this hunk.
|
||||
replacement: List[str] = []
|
||||
for line in hunk:
|
||||
if line.is_added or line.is_context:
|
||||
replacement.append(line.value)
|
||||
# removed lines (line.is_removed) are *not* added.
|
||||
# Replace the slice with the hunk‑result.
|
||||
working_lines[start:end] = replacement
|
||||
line_offset += len(replacement) - (end - start)
|
||||
new_content = "".join(working_lines)
|
||||
|
||||
# Finally commit the new revision.
|
||||
self.add_or_update_file(filename, new_content)
|
||||
|
||||
# ----------------------------------------------------------------------------------
|
||||
# Convenience helpers
|
||||
# ----------------------------------------------------------------------------------
|
||||
|
||||
def get_all_contents_for_context(self) -> str: # noqa: D401 – keep public API stable
|
||||
"""Return a summarised view of every file and its *latest* revision."""
|
||||
out: List[str] = ["=== CANVAS FILES ==="]
|
||||
for fname, revs in self._files.items():
|
||||
latest = revs[-1]
|
||||
out.append(f"File: {fname} (rev {latest.revision}):\n{latest.content}\n")
|
||||
out.append("=== END OF CANVAS ===")
|
||||
return "\n".join(out)
|
||||
@@ -0,0 +1,229 @@
|
||||
from typing import Any, Optional
|
||||
|
||||
from agentdhal_core import CancellationToken
|
||||
from agentdhal_core.memory import (
|
||||
Memory,
|
||||
MemoryContent,
|
||||
MemoryMimeType,
|
||||
MemoryQueryResult,
|
||||
UpdateContextResult,
|
||||
)
|
||||
from agentdhal_core.model_context import ChatCompletionContext
|
||||
from agentdhal_core.models import SystemMessage
|
||||
|
||||
from ._canvas_writer import ApplyPatchTool, UpdateFileTool
|
||||
from ._text_canvas import TextCanvas
|
||||
|
||||
|
||||
class TextCanvasMemory(Memory):
|
||||
"""
|
||||
A memory implementation that uses a Canvas for storing file-like content.
|
||||
Inserts the current state of the canvas into the ChatCompletionContext on each turn.
|
||||
|
||||
.. warning::
|
||||
|
||||
This is an experimental API and may change in the future.
|
||||
|
||||
The TextCanvasMemory provides a persistent, file-like storage mechanism that can be used
|
||||
by agents to read and write content. It automatically injects the current state of all files
|
||||
in the canvas into the model context before each inference.
|
||||
|
||||
This is particularly useful for:
|
||||
- Allowing agents to create and modify documents over multiple turns
|
||||
- Enabling collaborative document editing between multiple agents
|
||||
- Maintaining persistent state across conversation turns
|
||||
- Working with content too large to fit in a single message
|
||||
|
||||
The canvas provides tools for:
|
||||
- Creating or updating files with new content
|
||||
- Applying patches (unified diff format) to existing files
|
||||
|
||||
Examples:
|
||||
|
||||
**Example: Using TextCanvasMemory with an AssistantAgent**
|
||||
|
||||
The following example demonstrates how to create a TextCanvasMemory and use it with
|
||||
an AssistantAgent to write and update a story file.
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
import asyncio
|
||||
from agentdhal_core import CancellationToken
|
||||
from agentdhal_extensions.models.openai import OpenAIChatCompletionClient
|
||||
from agentdhal_agentchat.agents import AssistantAgent
|
||||
from agentdhal_agentchat.messages import TextMessage
|
||||
from agentdhal_extensions.memory.canvas import TextCanvasMemory
|
||||
|
||||
|
||||
async def main():
|
||||
# Create a model client
|
||||
model_client = OpenAIChatCompletionClient(
|
||||
model="gpt-4o",
|
||||
# api_key = "your_openai_api_key"
|
||||
)
|
||||
|
||||
# Create the canvas memory
|
||||
text_canvas_memory = TextCanvasMemory()
|
||||
|
||||
# Get tools for working with the canvas
|
||||
update_file_tool = text_canvas_memory.get_update_file_tool()
|
||||
apply_patch_tool = text_canvas_memory.get_apply_patch_tool()
|
||||
|
||||
# Create an agent with the canvas memory and tools
|
||||
writer_agent = AssistantAgent(
|
||||
name="Writer",
|
||||
model_client=model_client,
|
||||
description="A writer agent that creates and updates stories.",
|
||||
system_message='''
|
||||
You are a Writer Agent. Your focus is to generate a story based on the user's request.
|
||||
|
||||
Instructions for using the canvas:
|
||||
|
||||
- The story should be stored on the canvas in a file named "story.md".
|
||||
- If "story.md" does not exist, create it by calling the 'update_file' tool.
|
||||
- If "story.md" already exists, generate a unified diff (patch) from the current
|
||||
content to the new version, and call the 'apply_patch' tool to apply the changes.
|
||||
|
||||
IMPORTANT: Do not include the full story text in your chat messages.
|
||||
Only write the story content to the canvas using the tools.
|
||||
''',
|
||||
tools=[update_file_tool, apply_patch_tool],
|
||||
memory=[text_canvas_memory],
|
||||
)
|
||||
|
||||
# Send a message to the agent
|
||||
await writer_agent.on_messages(
|
||||
[TextMessage(content="Write a short story about a bunny and a sunflower.", source="user")],
|
||||
CancellationToken(),
|
||||
)
|
||||
|
||||
# Retrieve the content from the canvas
|
||||
story_content = text_canvas_memory.canvas.get_latest_content("story.md")
|
||||
print("Story content from canvas:")
|
||||
print(story_content)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
asyncio.run(main())
|
||||
|
||||
**Example: Using TextCanvasMemory with multiple agents**
|
||||
|
||||
The following example shows how to use TextCanvasMemory with multiple agents
|
||||
collaborating on the same document.
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
import asyncio
|
||||
from agentdhal_extensions.models.openai import OpenAIChatCompletionClient
|
||||
from agentdhal_agentchat.agents import AssistantAgent
|
||||
from agentdhal_agentchat.teams import RoundRobinGroupChat
|
||||
from agentdhal_agentchat.conditions import TextMentionTermination
|
||||
from agentdhal_extensions.memory.canvas import TextCanvasMemory
|
||||
|
||||
|
||||
async def main():
|
||||
# Create a model client
|
||||
model_client = OpenAIChatCompletionClient(
|
||||
model="gpt-4o",
|
||||
# api_key = "your_openai_api_key"
|
||||
)
|
||||
|
||||
# Create the shared canvas memory
|
||||
text_canvas_memory = TextCanvasMemory()
|
||||
update_file_tool = text_canvas_memory.get_update_file_tool()
|
||||
apply_patch_tool = text_canvas_memory.get_apply_patch_tool()
|
||||
|
||||
# Create a writer agent
|
||||
writer_agent = AssistantAgent(
|
||||
name="Writer",
|
||||
model_client=model_client,
|
||||
description="A writer agent that creates stories.",
|
||||
system_message="You write children's stories on the canvas in story.md.",
|
||||
tools=[update_file_tool, apply_patch_tool],
|
||||
memory=[text_canvas_memory],
|
||||
)
|
||||
|
||||
# Create a critique agent
|
||||
critique_agent = AssistantAgent(
|
||||
name="Critique",
|
||||
model_client=model_client,
|
||||
description="A critique agent that provides feedback on stories.",
|
||||
system_message="You review the story.md file and provide constructive feedback.",
|
||||
memory=[text_canvas_memory],
|
||||
)
|
||||
|
||||
# Create a team with both agents
|
||||
team = RoundRobinGroupChat(
|
||||
participants=[writer_agent, critique_agent],
|
||||
termination_condition=TextMentionTermination("TERMINATE"),
|
||||
max_turns=10,
|
||||
)
|
||||
|
||||
# Run the team on a task
|
||||
await team.run(task="Create a children's book about a bunny and a sunflower")
|
||||
|
||||
# Get the final story
|
||||
story = text_canvas_memory.canvas.get_latest_content("story.md")
|
||||
print(story)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
asyncio.run(main())
|
||||
"""
|
||||
|
||||
def __init__(self, canvas: Optional[TextCanvas] = None):
|
||||
super().__init__()
|
||||
self.canvas = canvas if canvas is not None else TextCanvas()
|
||||
|
||||
async def update_context(self, model_context: ChatCompletionContext) -> UpdateContextResult:
|
||||
"""
|
||||
Inject the entire canvas summary (or a selected subset) as reference data.
|
||||
Here, we just put it into a system message, but you could customize.
|
||||
"""
|
||||
snapshot = self.canvas.get_all_contents_for_context()
|
||||
if snapshot.strip():
|
||||
msg = SystemMessage(content=snapshot)
|
||||
await model_context.add_message(msg)
|
||||
|
||||
# Return it for debugging/logging
|
||||
memory_content = MemoryContent(content=snapshot, mime_type=MemoryMimeType.TEXT)
|
||||
return UpdateContextResult(memories=MemoryQueryResult(results=[memory_content]))
|
||||
|
||||
return UpdateContextResult(memories=MemoryQueryResult(results=[]))
|
||||
|
||||
async def query(
|
||||
self, query: str | MemoryContent, cancellation_token: Optional[CancellationToken] = None, **kwargs: Any
|
||||
) -> MemoryQueryResult:
|
||||
"""
|
||||
Potentially search for matching filenames or file content.
|
||||
This example returns empty.
|
||||
"""
|
||||
return MemoryQueryResult(results=[])
|
||||
|
||||
async def add(self, content: MemoryContent, cancellation_token: Optional[CancellationToken] = None) -> None:
|
||||
"""
|
||||
Example usage: Possibly interpret content as a patch or direct file update.
|
||||
Could also be done by a specialized "CanvasTool" instead.
|
||||
"""
|
||||
# NO-OP here, leaving actual changes to the CanvasTool
|
||||
pass
|
||||
|
||||
async def clear(self) -> None:
|
||||
"""Clear the entire canvas by replacing it with a new empty instance."""
|
||||
# Create a new TextCanvas instance instead of calling __init__ directly
|
||||
self.canvas = TextCanvas()
|
||||
|
||||
async def close(self) -> None:
|
||||
pass
|
||||
|
||||
def get_update_file_tool(self) -> UpdateFileTool:
|
||||
"""
|
||||
Returns an UpdateFileTool instance that works with this memory's canvas.
|
||||
"""
|
||||
return UpdateFileTool(self.canvas)
|
||||
|
||||
def get_apply_patch_tool(self) -> ApplyPatchTool:
|
||||
"""
|
||||
Returns an ApplyPatchTool instance that works with this memory's canvas.
|
||||
"""
|
||||
return ApplyPatchTool(self.canvas)
|
||||
Reference in New Issue
Block a user