dm-aioaiagent 0.2.0__py3-none-any.whl → 0.3.1__py3-none-any.whl

dm_aioaiagent/__init__.py

@@ -2,3 +2,4 @@ from dotenv import load_dotenv
 load_dotenv()
 from .ai_agent import DMAIAgent
 from .async_ai_agent import DMAioAIAgent
+from .types import Message

dm_aioaiagent/ai_agent.py

@@ -1,42 +1,22 @@
-import json
 import os
-from typing import Optional, Literal, Union
-from typing_extensions import TypedDict
-from pydantic import BaseModel, Field
+from itertools import dropwhile
 from threading import Thread
 from langchain_openai import ChatOpenAI
 from langchain_core.tools import BaseTool
 from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
-from langchain_core.messages import BaseMessage, SystemMessage, HumanMessage, AIMessage, ToolMessage
+from langchain_core.messages import SystemMessage, HumanMessage, AIMessage, ToolMessage
 from langgraph.graph import StateGraph
 from dm_logger import DMLogger
 
-__all__ = ["DMAIAgent"]
-
-
-class Message(TypedDict):
-    role: Literal["user", "ai"]
-    content: str
-
-
-class InnerState(BaseModel):
-    messages: list[BaseMessage] = Field(default=[])
-    context: list[Message] = Field(default=[])
-
-
-class InputState(BaseModel):
-    messages: list[Message]
-    inner_state: Optional[InnerState] = Field(default=InnerState())
-
+from .types import *
 
-class OutputState(TypedDict):
-    answer: str
-    context: list[Message]
+__all__ = ["DMAIAgent"]
 
 
 class DMAIAgent:
     agent_name = "AIAgent"
     _allowed_roles = ("user", "ai")
+    MAX_MEMORY_MESSAGES = 20 # Only INT greater than 0
 
     def __init__(
         self,
@@ -47,15 +27,17 @@ class DMAIAgent:
         temperature: int = 1,
         agent_name: str = None,
         input_output_logging: bool = True,
-        return_context: bool = False
+        is_memory_enabled: bool = True,
+        max_memory_messages: int = None,
     ):
         if not os.getenv("OPENAI_API_KEY"):
-            raise EnvironmentError("OPENAI_API_KEY environment variable is not set!")
+            raise EnvironmentError("'OPENAI_API_KEY' environment variable is not set!")
 
         self._logger = DMLogger(agent_name or self.agent_name)
-        self._input_output_logging = input_output_logging
-        self._return_context = return_context
         self._is_tools_exists = bool(tools)
+        self._input_output_logging = bool(input_output_logging)
+        self._is_memory_enabled = bool(is_memory_enabled)
+        self._max_memory_messages = self._validate_max_memory_messages(max_memory_messages)
 
         prompt = ChatPromptTemplate.from_messages([SystemMessage(content=system_message),
                                                    MessagesPlaceholder(variable_name="messages")])
@@ -64,8 +46,9 @@ class DMAIAgent:
             self._tool_map = {t.name: t for t in tools}
             llm = llm.bind_tools(tools)
         self._agent = prompt | llm
+        self._memory = {}
 
-        workflow = StateGraph(input=InputState, output=OutputState)
+        workflow = StateGraph(State)
         workflow.add_node("Prepare messages", self._prepare_messages_node)
         workflow.add_node("Invoke LLM", self._invoke_llm_node)
         workflow.add_node("Execute tool", self._execute_tool_node)
@@ -80,40 +63,49 @@ class DMAIAgent:
         workflow.set_finish_point("Exit")
         self._graph = workflow.compile()
 
-    def run(self, messages: list[Message]) -> Union[str, OutputState]:
-        state = self._graph.invoke({"messages": messages})
-        if self._return_context:
-            return state
-        return state["answer"]
+    def run(self, input_messages: InputMessagesType, memory_id: str = None) -> ResponseType:
+        state = self._graph.invoke({"input_messages": input_messages, "memory_id": memory_id})
+        return state["response"]
+
+    def get_memory_messages(self, memory_id: str = None) -> list[BaseMessage]:
+        return self._memory.get(self._validate_memory_id(memory_id), [])
+
+    def clear_memory(self, memory_id: str = None) -> None:
+        self._memory[self._validate_memory_id(memory_id)] = []
+
+    def _prepare_messages_node(self, state: State) -> State:
+        state.memory_id = self._validate_memory_id(state.memory_id)
+        state.input_messages = state.input_messages or [{"role": "user", "content": ""}]
+        for item in state.input_messages:
+            if isinstance(item, dict):
+                role = item.get("role")
+                content = item.get("content")
+                if not role or role not in self._allowed_roles or not content:
+                    continue
+                if role == "ai":
+                    MessageClass = AIMessage
+                else:
+                    MessageClass = HumanMessage
+                state.messages.append(MessageClass(content))
+            elif isinstance(item, BaseMessage):
+                state.messages.append(item)
 
-    def _prepare_messages_node(self, state: InputState) -> InputState:
-        state.messages = state.messages or [{"role": "user", "content": "Привіт"}]
-        state.inner_state = InnerState()
         if self._input_output_logging:
-            self._logger.debug(input_messages=state.messages)
-
-        for item in state.messages:
-            role = item.get("role")
-            content = item.get("content")
-            if not role or role not in self._allowed_roles or not content:
-                continue
-            if role == "ai":
-                MessageClass = AIMessage
-            else:
-                MessageClass = HumanMessage
-            state.inner_state.messages.append(MessageClass(content))
+            self._logger.debug(f"Query:\n{state.messages[-1].content}", memory_id=state.memory_id)
+        if self._is_memory_enabled:
+            state.messages = self.get_memory_messages(state.memory_id) + state.messages
         return state
 
-    def _invoke_llm_node(self, state: InputState) -> InputState:
+    def _invoke_llm_node(self, state: State) -> State:
         self._logger.debug("Run node: Invoke LLM")
-        ai_response = self._agent.invoke({"messages": state.inner_state.messages})
-        state.inner_state.messages.append(ai_response)
+        ai_response = self._agent.invoke({"messages": state.messages})
+        state.messages.append(ai_response)
         return state
 
-    def _execute_tool_node(self, state: InputState) -> InputState:
+    def _execute_tool_node(self, state: State) -> State:
         self._logger.debug("Run node: Execute tool")
         threads = []
-        for tool_call in state.inner_state.messages[-1].tool_calls:
+        for tool_call in state.messages[-1].tool_calls:
             tool_id = tool_call["id"]
             tool_name = tool_call["name"]
             tool_args = tool_call["args"]
@@ -130,11 +122,8 @@ class DMAIAgent:
                     tool_response = f"Tool not found!"
                 self._logger.debug(f"Tool response:\n{tool_response}", tool_id=tool_id)
 
-                state.inner_state.context.append({"tool_name": tool_name,
-                                                  "tool_args": json.dumps(tool_args, ensure_ascii=False),
-                                                  "tool_response": tool_response})
                 tool_message = ToolMessage(content=str(tool_response), name=tool_name, tool_call_id=tool_id)
-                state.inner_state.messages.append(tool_message)
+                state.messages.append(tool_message)
 
             threads.append(Thread(target=tool_callback, daemon=True))
 
@@ -145,19 +134,37 @@ class DMAIAgent:
 
         return state
 
-    def _exit_node(self, state: InputState) -> OutputState:
-        answer = state.inner_state.messages[-1].content if state.inner_state.messages else ""
+    def _exit_node(self, state: State) -> State:
+        answer = state.messages[-1].content
         if self._input_output_logging:
-            self._logger.debug(f"Answer:\n{answer}")
-        return OutputState(answer=answer, context=state.inner_state.context)
+            self._logger.debug(f"Answer:\n{answer}", memory_id=state.memory_id)
 
-    def _messages_router(self, state: InputState) -> str:
-        if self._is_tools_exists and state.inner_state.messages[-1].tool_calls:
+        if self._is_memory_enabled:
+            messages_to_memory = state.messages[-self._max_memory_messages:]
+            # drop ToolsMessages from start of list
+            self._memory[state.memory_id] = list(dropwhile(lambda x: isinstance(x, ToolMessage), messages_to_memory))
+            state.response = answer
+        else:
+            state.response = state.messages[len(state.input_messages):]
+        return state
+
+    def _messages_router(self, state: State) -> str:
+        if self._is_tools_exists and state.messages[-1].tool_calls:
             route = "execute_tool"
         else:
             route = "exit"
         return route
 
+    @staticmethod
+    def _validate_memory_id(memory_id: Union[str, None]) -> Union[str, int]:
+        return str(memory_id) if memory_id else 0
+
+    @classmethod
+    def _validate_max_memory_messages(cls, max_messages_in_memory: int) -> int:
+        if isinstance(max_messages_in_memory, int) and max_messages_in_memory > 0:
+            return max_messages_in_memory
+        return cls.MAX_MEMORY_MESSAGES
+
     def print_graph(self) -> None:
         self._graph.get_graph().print_ascii()
 

dm_aioaiagent/async_ai_agent.py

@@ -1,10 +1,9 @@
-import json
 import sys
 import asyncio
-from typing import Union
 from langchain_core.messages import ToolMessage
 
-from .ai_agent import DMAIAgent, InputState, OutputState, Message
+from .ai_agent import DMAIAgent
+from .types import *
 
 if sys.platform == "win32":
     asyncio.set_event_loop_policy(asyncio.WindowsSelectorEventLoopPolicy())
@@ -15,22 +14,20 @@ __all__ = ["DMAioAIAgent"]
 class DMAioAIAgent(DMAIAgent):
     agent_name = "AsyncAIAgent"
 
-    async def run(self, messages: list[Message]) -> Union[str, OutputState]:
-        state = await self._graph.ainvoke({"messages": messages})
-        if self._return_context:
-            return state
-        return state["answer"]
+    async def run(self, input_messages: InputMessagesType, memory_id: str = None) -> ResponseType:
+        state = await self._graph.ainvoke({"input_messages": input_messages, "memory_id": memory_id})
+        return state["response"]
 
-    async def _invoke_llm_node(self, state: InputState) -> InputState:
+    async def _invoke_llm_node(self, state: State) -> State:
         self._logger.debug("Run node: Invoke LLM")
-        ai_response = await self._agent.ainvoke({"messages": state.inner_state.messages})
-        state.inner_state.messages.append(ai_response)
+        ai_response = await self._agent.ainvoke({"messages": state.messages})
+        state.messages.append(ai_response)
         return state
 
-    async def _execute_tool_node(self, state: InputState) -> InputState:
+    async def _execute_tool_node(self, state: State) -> State:
         self._logger.debug("Run node: Execute tool")
         tasks = []
-        for tool_call in state.inner_state.messages[-1].tool_calls:
+        for tool_call in state.messages[-1].tool_calls:
             tool_id = tool_call["id"]
             tool_name = tool_call["name"]
             tool_args = tool_call["args"]
@@ -47,11 +44,8 @@ class DMAioAIAgent(DMAIAgent):
                     tool_response = f"Tool '{tool_name}' not found!"
                 self._logger.debug(f"Tool response:\n{tool_response}", tool_id=tool_id)
 
-                state.inner_state.context.append({"tool_name": tool_name,
-                                                  "tool_args": json.dumps(tool_args, ensure_ascii=False),
-                                                  "tool_response": tool_response})
                 tool_message = ToolMessage(content=str(tool_response), name=tool_name, tool_call_id=tool_id)
-                state.inner_state.messages.append(tool_message)
+                state.messages.append(tool_message)
 
             tasks.append(asyncio.create_task(tool_callback()))
 

dm_aioaiagent/types.py

@@ -0,0 +1,20 @@
+from typing import Optional, Literal, Union
+from typing_extensions import TypedDict
+from pydantic import BaseModel, Field
+from langchain_core.messages import BaseMessage
+
+
+class Message(TypedDict):
+    role: Literal["user", "ai"]
+    content: str
+
+
+InputMessagesType = list[Union[Message, BaseMessage]]
+ResponseType = Union[str, list[BaseMessage]]
+
+
+class State(BaseModel):
+    input_messages: InputMessagesType
+    memory_id: Union[str, int, None] = Field(default=0)
+    messages: Optional[list[BaseMessage]] = Field(default_factory=list)
+    response: ResponseType = Field(default="")

{dm_aioaiagent-0.2.0.dist-info → dm_aioaiagent-0.3.1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: dm-aioaiagent
-Version: 0.2.0
+Version: 0.3.1
 Summary: This is my custom aioaiagent client
 Home-page: https://pypi.org/project/dm-aioaiagent
 Author: dimka4621
@@ -12,14 +12,14 @@ Classifier: License :: OSI Approved :: MIT License
 Classifier: Operating System :: OS Independent
 Requires-Python: >=3.9
 Description-Content-Type: text/markdown
-Requires-Dist: dm-logger==0.5.2
-Requires-Dist: python-dotenv==1.0.1
-Requires-Dist: pydantic==2.9.2
-Requires-Dist: langchain==0.3.0
-Requires-Dist: langchain-core==0.3.5
-Requires-Dist: langgraph==0.2.23
-Requires-Dist: langchain-community==0.3.0
-Requires-Dist: langchain-openai==0.2.0
+Requires-Dist: dm-logger~=0.5.2
+Requires-Dist: python-dotenv>=1.0.0
+Requires-Dist: pydantic<3.0.0,>=2.9.2
+Requires-Dist: langchain~=0.3.0
+Requires-Dist: langchain-core~=0.3.5
+Requires-Dist: langgraph~=0.2.23
+Requires-Dist: langchain-community~=0.3.0
+Requires-Dist: langchain-openai~=0.2.0
 
 # DM-aioaiagent
 
@@ -34,6 +34,12 @@ Requires-Dist: langchain-openai==0.2.0
 
 Analogue to `DMAioAIAgent` is the synchronous client `DMAIAgent`.
 
+### Use agent *with* inner memory
+
+By default, agent use inner memory to store the conversation history.
+
+(You can set *max count messages in memory* by `max_memory_messages` init argument)
+
 ```python
 import asyncio
 from dm_aioaiagent import DMAioAIAgent
@@ -52,24 +58,40 @@ async def main():
     # create an agent
     ai_agent = DMAioAIAgent(system_message, tools, model=model_name)
     # if you don't want to see the input and output messages from agent
-    # you can set input_output_logging=False
+    # you can set `input_output_logging=False` init argument
 
-    # define the conversation messages
-    messages = [
-       {"role": "user", "content": "Hello!"},
-       {"role": "ai", "content": "How can I help you?"},
-       {"role": "user", "content": "I want to know the weather in Kyiv"},
+    # define the conversation message
+    input_messages = [
+        {"role": "user", "content": "Hello!"},
+    ]
+
+    # call an agent
+    # specify `memory_id` argument to store the conversation history by your custom id
+    answer = await ai_agent.run(input_messages)
+
+    # define the next conversation message
+    input_messages = [
+        {"role": "user", "content": "I want to know the weather in Kyiv"}
     ]
 
     # call an agent
-    answer = await ai_agent.run(messages)
+    answer = await ai_agent.run(input_messages)
+
+    # get full conversation history
+    conversation_history = ai_agent.get_memory_messages()
+
+    # clear conversation history
+    ai_agent.clear_memory()
 
 
 if __name__ == "__main__":
     asyncio.run(main())
 ```
 
-### Return context of the tools with answer
+### Use agent *without* inner memory
+
+If you want to control the memory of the agent, you can disable it by setting `is_memory_enabled=False`
+
 ```python
 import asyncio
 from dm_aioaiagent import DMAioAIAgent
@@ -79,25 +101,36 @@ async def main():
     # define a system message
     system_message = "Your custom system message with role, backstory and goal"
 
-    # define a list of tools, if you want to use them
+    # (optional) define a list of tools, if you want to use them
     tools = [...]
 
+    # define a openai model, default is "gpt-4o-mini"
+    model_name = "gpt-4o"
+
     # create an agent
-    ai_agent = DMAioAIAgent(system_message, tools, return_context=True)
+    ai_agent = DMAioAIAgent(system_message, tools, model=model_name,
+                            is_memory_enabled=False)
+    # if you don't want to see the input and output messages from agent
+    # you can set input_output_logging=False
 
-    # define the conversation messages
+    # define the conversation message
     messages = [
-       {"role": "user", "content": "Hello!"},
-       {"role": "ai", "content": "How can I help you?"},
-       {"role": "user", "content": "I want to know the weather in Kyiv"},
+        {"role": "user", "content": "Hello!"}
     ]
 
-    # start the agent
-    state = await ai_agent.run(messages)
+    # call an agent
+    new_messages = await ai_agent.run(messages)
+
+    # add new_messages to messages
+    messages.extend(new_messages)
 
-    # if you define tools, you can see the context of the tools
-    answer = state["answer"]
-    context = state["context"]
+    # define the next conversation message
+    messages.append(
+        {"role": "user", "content": "I want to know the weather in Kyiv"}
+    )
+
+    # call an agent
+    new_messages = await ai_agent.run(messages)
 
 
 if __name__ == "__main__":

dm_aioaiagent-0.3.1.dist-info/RECORD

@@ -0,0 +1,8 @@
+dm_aioaiagent/__init__.py,sha256=8B0XQR-XE2VieU7-LOHCjqDYv0gSYqkgAkj_eq6XAhI,145
+dm_aioaiagent/ai_agent.py,sha256=YqM2KpLZVMqFVWBpmqEzgJSLQSycvB1g9ndxdg2aL4k,7846
+dm_aioaiagent/async_ai_agent.py,sha256=n7OfJUHRwjk92gae9NkpPKQ21Xost2-JsQ0JyzYTXT4,2146
+dm_aioaiagent/types.py,sha256=Xvx0x1GxLAvCot_3CZGgD8BFWwnTEc1lg7moLxkkJoc,587
+dm_aioaiagent-0.3.1.dist-info/METADATA,sha256=BeQ8eugmPWTmAmL7WJwWUc64URWW-5IB6Qg8Fovkq1Q,4299
+dm_aioaiagent-0.3.1.dist-info/WHEEL,sha256=OVMc5UfuAQiSplgO0_WdW7vXVGAt9Hdd6qtN4HotdyA,91
+dm_aioaiagent-0.3.1.dist-info/top_level.txt,sha256=CbasLH0KI7zA77XwT6JDCnmRascxKNGvUVV9MgYjHAU,14
+dm_aioaiagent-0.3.1.dist-info/RECORD,,

{dm_aioaiagent-0.2.0.dist-info → dm_aioaiagent-0.3.1.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (75.1.0)
+Generator: setuptools (75.2.0)
 Root-Is-Purelib: true
 Tag: py3-none-any
 

dm_aioaiagent-0.2.0.dist-info/RECORD

@@ -1,7 +0,0 @@
-dm_aioaiagent/__init__.py,sha256=Hhgz18uMAnR0ezLOukoWCe1sFS4Q4LC26xbRKdu9FFw,118
-dm_aioaiagent/ai_agent.py,sha256=SuCOLx8n2rLCOnpilZ5k7KM6KqD4lU-BGy3ZN43gvA8,7149
-dm_aioaiagent/async_ai_agent.py,sha256=vrQLNWdvAGQhSZDNMffe9VFT5BPd2WQQhlnx4f3u-UY,2525
-dm_aioaiagent-0.2.0.dist-info/METADATA,sha256=PxBpBaUMp2Hgig-wRkpIvXq-EGot7mjfQK7s38-mRIE,3354
-dm_aioaiagent-0.2.0.dist-info/WHEEL,sha256=GV9aMThwP_4oNCtvEC2ec3qUYutgWeAzklro_0m4WJQ,91
-dm_aioaiagent-0.2.0.dist-info/top_level.txt,sha256=CbasLH0KI7zA77XwT6JDCnmRascxKNGvUVV9MgYjHAU,14
-dm_aioaiagent-0.2.0.dist-info/RECORD,,