diff --git a/README.md b/README.md index eae08b3..b8129d9 100644 --- a/README.md +++ b/README.md @@ -27,6 +27,10 @@ | **Parallel summarization** | `fast_start` method uses `ThreadPoolExecutor` to speed up web content processing | | **Prompt-driven summaries** | Summarization is driven by customizable system prompts and token-limited chunking for accurate context | | **Simple, Pythonic design** | Agents are lightweight Python classes with callable message interfaces—no metaclasses or hidden magic | +| **Memory, BaseMemory** | BaseMemory provides foundational memory management for agents, storing conversation history and message objects. It supports adding, retrieving, and clearing memory, offering a flexible design for session-based context, interaction history, or task-specific memory across multiple agent invocations. Ideal for scenarios where the agent needs to recall past interactions for continuity. | + + + --- diff --git a/example/AutoAgent.ipynb b/example/AutoAgent.ipynb index c0b9d5e..8422b6e 100644 --- a/example/AutoAgent.ipynb +++ b/example/AutoAgent.ipynb @@ -13,12 +13,12 @@ "base_url= \"\" # use your own base_url from api provider or local provider like ollama.\n", "api_key = \"\" # use your own api_key.\n", "provider = \"openai\" # openai for openai like provider (vLLM or openrouter) and ollama for local use.\n", - "model = \"gpt-4o-mini\"\n" + "model = \"gpt-4o-mini\"" ] }, { "cell_type": "code", - "execution_count": 7, + "execution_count": null, "id": "6057cd04", "metadata": {}, "outputs": [ diff --git a/example/BaseMemory.ipynb b/example/BaseMemory.ipynb new file mode 100644 index 0000000..4a7c3db --- /dev/null +++ b/example/BaseMemory.ipynb @@ -0,0 +1,118 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": null, + "id": "1da71a7f", + "metadata": {}, + "outputs": [], + "source": [ + "from iragent.agent import AgentFactory\n", + "from iragent.models import AutoAgentManager\n", + "\n", + "base_url= \"\" # use your own base_url from api provider or local provider like ollama.\n", + "api_key = \"\" # use your own api_key.\n", + "provider = \"openai\" # openai for openai like provider (vLLM or openrouter) and ollama for local use.\n", + "model = \"gpt-4o-mini\"\n" + ] + }, + { + "cell_type": "markdown", + "id": "d9777f2b", + "metadata": {}, + "source": [ + "### BaseMemory\n", + "\n", + "This class of memory is a class for keeping history of chat for each agent, so the agents know the last conversation." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "ea2ca3b6", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Routing from user -> time_reader \n", + " content: what time is it?\n", + "Routing from time_reader -> date_exctractor \n", + " content: The current time is 11:57 PM.\n", + "Routing from date_exctractor -> date_converter \n", + " content: 11:57 PM\n" + ] + }, + { + "data": { + "text/plain": [ + "'23:57 \\n[#finish#]'" + ] + }, + "execution_count": 2, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "from iragent.memory import BaseMemory\n", + "from iragent.tools import get_time_now, simple_termination\n", + "\n", + "factory = AgentFactory(base_url,api_key, model, provider)\n", + "\n", + "agent1 = factory.create_agent(name=\"time_reader\",\n", + " system_prompt=\"You are that one who can read time. there is a fucntion named get_time_now(), you can call it whether user ask about time or date.\",\n", + " fn=[get_time_now],\n", + " memory=BaseMemory\n", + " )\n", + "agent2 = factory.create_agent(name=\"date_exctractor\", \n", + " system_prompt= \"You are that one who extract time from date. only return time.\", memory=BaseMemory)\n", + "agent3 = factory.create_agent(name=\"date_converter\", \n", + " system_prompt= \"You are that one who write the time in Persian. when you wrote time, then in new line write [#finish#]\", memory=BaseMemory)\n", + "\n", + "manager = AutoAgentManager(\n", + " init_message=\"what time is it?\",\n", + " agents= [agent1,agent2,agent3],\n", + " first_agent=agent1,\n", + " max_round=5,\n", + " termination_fn=simple_termination,\n", + " termination_word=\"[#finish#]\"\n", + ")\n", + "\n", + "\n", + "res = manager.start()\n", + "res.content" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "ece0343b", + "metadata": {}, + "outputs": [], + "source": [] + } + ], + "metadata": { + "kernelspec": { + "display_name": "venv", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.12.9" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/iragent/__pycache__/agent.cpython-312.pyc b/iragent/__pycache__/agent.cpython-312.pyc index 6a7d101..3898102 100644 Binary files a/iragent/__pycache__/agent.cpython-312.pyc and b/iragent/__pycache__/agent.cpython-312.pyc differ diff --git a/iragent/__pycache__/memory.cpython-312.pyc b/iragent/__pycache__/memory.cpython-312.pyc new file mode 100644 index 0000000..9b8788c Binary files /dev/null and b/iragent/__pycache__/memory.cpython-312.pyc differ diff --git a/iragent/__pycache__/models.cpython-312.pyc b/iragent/__pycache__/models.cpython-312.pyc index e5fc977..e459526 100644 Binary files a/iragent/__pycache__/models.cpython-312.pyc and b/iragent/__pycache__/models.cpython-312.pyc differ diff --git a/iragent/agent.py b/iragent/agent.py index 6585bb1..19a0865 100644 --- a/iragent/agent.py +++ b/iragent/agent.py @@ -31,6 +31,7 @@ def __init__( fn: List[Callable] = [], provider: str = "openai", response_format: str = None, + memory = None ): ## The platform we use for loading the large lanuage models. you should peak ```ollama``` or ```openai``` as provider. self.provider = provider @@ -59,18 +60,37 @@ def __init__( # Support Structured output self.response_format = response_format + # Set Memory + self.memory = memory() if memory is not None else None + def call_message(self, message: Message, **kwargs) -> str: - msgs = [ - {"role": "system", "content": self.system_prompt}, - {"role": "user", "content": message.content}, - ] + + msgs = [{"role": "system", "content": self.system_prompt}] + + # If agent has a history + if self.memory: + history = self.memory.get_history() + if history: + msgs.extend(history) + + user_msg = {"role": "user", "content": message.content} + msgs.append(user_msg) + + # Add to memory if it is first time + if self.memory: + self.memory.add_history(user_msg) if self.provider == "openai": - return self._call_openai(msgs=msgs, message=message, **kwargs) - if self.provider == "ollama": - return self._call_ollama_v2(msgs=msgs, message=message) + res = self._call_openai(msgs=msgs, message=message, **kwargs) + elif self.provider == "ollama": + res = self._call_ollama_v2(msgs=msgs, message=message) else: raise ValueError(f"Unsupported provider: {self.provider}") + + # Add Assistant + self.memory.add_history({"role": "assistant", "content": res.content}) + self.memory.add_message(res) + return res def _call_ollama(self, msgs: List[Dict], message: Message) -> Message: """! diff --git a/iragent/memory.py b/iragent/memory.py new file mode 100644 index 0000000..17d63a5 --- /dev/null +++ b/iragent/memory.py @@ -0,0 +1,54 @@ +from .message import Message + + +class BaseMemory: + """ + BaseMemory is a foundational memory management class for conversational agents. + + It stores two types of information: + - `history`: a list of message dictionaries representing role-based dialogue turns + (e.g., user and assistant messages). + - `messages`: a list of raw Message objects, useful for storing additional metadata or original input/output. + + This class supports adding, retrieving, and clearing both types of memory and + is designed to be extended for more advanced memory strategies. + + Attributes: + history (list[dict]): List of role-content dictionaries used in LLM context (e.g., {"role": "user", "content": "Hi"}). + messages (list[Message]): List of Message objects (structured user inputs/outputs). + + Methods: + add_history(msg): Adds a single dict or list of dicts to the conversation history. + get_history(): Returns the stored conversation history as a list of dicts. + clear_history(): Clears the conversation history. + + add_message(msg): Adds a Message object to the internal message list. + get_messages(): Returns the stored messages as a list. + clear_messages(): Clears all stored messages. + """ + def __init__(self) -> None: + self.history = [] + self.messages = [] + + def add_history(self, msg: dict | list[dict]) -> None: + if isinstance(msg, dict): + self.history.append(msg) + elif isinstance(msg, list) and all(isinstance(m, dict) for m in msg): + self.history.extend(msg) + else: + raise TypeError("msg must be a dict or a list of dicts.") + + def get_history(self) -> list[dict]: + return self.history + + def clear_history(self) -> None: + self.history.clear() + + def add_message(self, msg: Message) -> None: + self.messages.append(msg) + + def get_messages(self) -> list[dict]: + return self.messages + + def clear_messages(self) -> None: + self.messages.clear() \ No newline at end of file diff --git a/iragent/models.py b/iragent/models.py index e30f8d8..d51303b 100644 --- a/iragent/models.py +++ b/iragent/models.py @@ -5,6 +5,7 @@ from tqdm import tqdm from .agent import Agent +from .memory import BaseMemory from .message import Message from .prompts import AUTO_AGENT_PROMPT, SUMMARIZER_PROMPT from .utility import chunker, fetch_url @@ -81,6 +82,7 @@ def __init__( api_key=first_agent.api_key, temprature=0.1, max_token=32, + memory=BaseMemory ) self.termination_fn = termination_fn self.max_round = max_round diff --git a/pyproject.toml b/pyproject.toml index 3b1743f..8c42166 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "iragent" -version = "0.1.2" +version = "0.1.3" description = "A simple multi-agent framework" authors = [{ name = "Parsa Bakhtiari", email = "spacenavard1@gmail.com" }] readme = "README.md"