dm-aioaiagent 0.3.3__py3-none-any.whl → 0.3.5__py3-none-any.whl

dm_aioaiagent/__init__.py

@@ -2,5 +2,5 @@ from dotenv import load_dotenv
 load_dotenv()
 from .ai_agent import DMAIAgent
 from .async_ai_agent import DMAioAIAgent
-from .image_message_content_builder import ImageMessageContentBuilder
+from .openai_image_message_content import OpenAIImageMessageContent
 from .types import Message

dm_aioaiagent/ai_agent.py

@@ -1,7 +1,7 @@
 import os
+from pydantic import SecretStr
 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 SystemMessage, HumanMessage, AIMessage, ToolMessage
@@ -14,11 +14,11 @@ __all__ = ["DMAIAgent"]
 
 
 class DMAIAgent:
-    agent_name = "AIAgent"
-    _allowed_roles = ("user", "ai")
-    _response_if_request_fail = "I can't provide a response right now. Please try again later."
-    _response_if_invalid_image = "The image is unavailable or the link is incorrect."
+    AGENT_NAME = "AIAgent"
     MAX_MEMORY_MESSAGES = 20  # Only INT greater than 0
+    RESPONSE_IF_REQUEST_FAIL = "I can't provide a response right now. Please try again later."
+    RESPONSE_IF_INVALID_IMAGE = "The image is unavailable or the link is incorrect."
+    _ALLOWED_ROLES = ("user", "ai")
 
     def __init__(
         self,
@@ -30,44 +30,30 @@ class DMAIAgent:
         agent_name: str = None,
         input_output_logging: bool = True,
         is_memory_enabled: bool = True,
+        save_tools_responses_in_memory: bool = True,
         max_memory_messages: int = None,
+        llm_provider_api_key: str = "",
         response_if_request_fail: str = None,
         response_if_invalid_image: str = None
     ):
-        if not os.getenv("OPENAI_API_KEY"):
-            raise EnvironmentError("'OPENAI_API_KEY' environment variable is not set!")
+        self._logger = DMLogger(agent_name or self.AGENT_NAME)
+        self._input_output_logging = bool(input_output_logging)
 
-        self._logger = DMLogger(agent_name or self.agent_name)
+        self._system_message = str(system_message)
+        self._tools = tools or []
         self._is_tools_exists = bool(tools)
-        self._input_output_logging = bool(input_output_logging)
+        self._model = str(model)
+        self._temperature = int(temperature)
+        self._llm_provider_api_key = str(llm_provider_api_key)
+
         self._is_memory_enabled = bool(is_memory_enabled)
+        self._save_tools_responses_in_memory = bool(save_tools_responses_in_memory)
         self._max_memory_messages = self._validate_max_memory_messages(max_memory_messages)
-        self._response_if_request_fail = str(response_if_request_fail or self._response_if_request_fail)
-        self._response_if_invalid_image = str(response_if_invalid_image or self._response_if_invalid_image)
+        self._response_if_request_fail = str(response_if_request_fail or self.RESPONSE_IF_REQUEST_FAIL)
+        self._response_if_invalid_image = str(response_if_invalid_image or self.RESPONSE_IF_INVALID_IMAGE)
 
-        prompt = ChatPromptTemplate.from_messages([SystemMessage(content=system_message),
-                                                   MessagesPlaceholder(variable_name="messages")])
-        llm = ChatOpenAI(model=str(model), temperature=int(temperature))
-        if self._is_tools_exists:
-            self._tool_map = {t.name: t for t in tools}
-            llm = llm.bind_tools(tools)
-        self._agent = prompt | llm
-        self._memory = {}
-
-        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)
-        workflow.add_node("Exit", self._exit_node)
-
-        workflow.add_edge("Prepare messages", "Invoke LLM")
-        workflow.add_conditional_edges(source="Invoke LLM",
-                                       path=self._messages_router,
-                                       path_map={"execute_tool": "Execute tool", "exit": "Exit"})
-        workflow.add_edge("Execute tool", "Invoke LLM")
-        workflow.set_entry_point("Prepare messages")
-        workflow.set_finish_point("Exit")
-        self._graph = workflow.compile()
+        self._init_agent()
+        self._init_graph()
 
     def run(self, input_messages: InputMessagesType, memory_id: str = None) -> ResponseType:
         state = self._graph.invoke({"input_messages": input_messages, "memory_id": memory_id})
@@ -96,7 +82,7 @@ class DMAIAgent:
             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:
+                if not role or role not in self._ALLOWED_ROLES or not content:
                     continue
                 if role == "ai":
                     MessageClass = AIMessage
@@ -168,8 +154,15 @@ class DMAIAgent:
         if self._is_memory_enabled:
             memory_id = self._validate_memory_id(state.memory_id)
             messages_to_memory = state.messages[-self._max_memory_messages:]
-            # drop ToolsMessages from start of list
-            self._memory[memory_id] = list(dropwhile(lambda x: isinstance(x, ToolMessage), messages_to_memory))
+            if self._save_tools_responses_in_memory:
+                # drop ToolsMessages from start of list
+                self._memory[memory_id] = list(dropwhile(lambda x: isinstance(x, ToolMessage), messages_to_memory))
+            else:
+                self._memory[memory_id] = []
+                for mes in messages_to_memory:
+                    if isinstance(mes, ToolMessage) or (isinstance(mes, AIMessage) and mes.tool_calls):
+                        continue
+                    self._memory[memory_id].append(mes)
             state.response = answer
         else:
             state.response = state.messages[len(state.input_messages):]
@@ -182,6 +175,49 @@ class DMAIAgent:
             route = "exit"
         return route
 
+    def _init_agent(self) -> None:
+        if self._llm_provider_api_key:
+            self._llm_provider_api_key = SecretStr(self._llm_provider_api_key)
+
+        if self._model.startswith("gpt"):
+            from langchain_openai import ChatOpenAI
+
+            api_key = SecretStr(self._llm_provider_api_key or os.getenv("OPENAI_API_KEY"))
+            llm = ChatOpenAI(model_name=self._model, temperature=self._temperature, openai_api_key=api_key)
+        elif self._model.startswith("claude"):
+            from langchain_anthropic import ChatAnthropic
+
+            api_key = SecretStr(self._llm_provider_api_key or os.getenv("ANTHROPIC_API_KEY"))
+            llm = ChatAnthropic(model=self._model, temperature=self._temperature, anthropic_api_key=api_key)
+        else:
+            raise ValueError(f"{self.__class__.__name__} not support this model: '{self._model}'")
+
+        if self._is_tools_exists:
+            self._tool_map = {t.name: t for t in self._tools}
+            llm = llm.bind_tools(self._tools)
+
+        prompt = ChatPromptTemplate.from_messages([SystemMessage(content=self._system_message),
+                                                   MessagesPlaceholder(variable_name="messages")])
+
+        self._agent = prompt | llm
+        self._memory = {}
+
+    def _init_graph(self) -> None:
+        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)
+        workflow.add_node("Exit", self._exit_node)
+
+        workflow.add_edge("Prepare messages", "Invoke LLM")
+        workflow.add_conditional_edges(source="Invoke LLM",
+                                       path=self._messages_router,
+                                       path_map={"execute_tool": "Execute tool", "exit": "Exit"})
+        workflow.add_edge("Execute tool", "Invoke LLM")
+        workflow.set_entry_point("Prepare messages")
+        workflow.set_finish_point("Exit")
+        self._graph = workflow.compile()
+
     @staticmethod
     def _validate_memory_id(memory_id: Union[str, None]) -> Union[str, int]:
         return str(memory_id) if memory_id else 0

dm_aioaiagent/{image_message_content_builder.py → openai_image_message_content.py}

@@ -1,4 +1,4 @@
-class ImageMessageContentBuilder(list):
+class OpenAIImageMessageContent(list):
     def __init__(self, image_url: str, text: str = None):
         content = []
         if isinstance(text, str):

{dm_aioaiagent-0.3.3.dist-info → dm_aioaiagent-0.3.5.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: dm-aioaiagent
-Version: 0.3.3
+Version: 0.3.5
 Summary: This is my custom aioaiagent client
 Home-page: https://pypi.org/project/dm-aioaiagent
 Author: dimka4621
@@ -141,7 +141,8 @@ if __name__ == "__main__":
 ### Image vision
 
 ```python
-from dm_aioaiagent import DMAIAgent, ImageMessageContentBuilder
+from dm_aioaiagent import DMAIAgent, OpenAIImageMessageContent
+
 
 def main():
     # create an agent
@@ -149,8 +150,8 @@ def main():
 
     # create an image message content
     # NOTE: text argument is optional
-    img_content = ImageMessageContentBuilder(image_url="https://your.domain/image",
-                                             text="Hello, what is shown in the photo?")
+    img_content = OpenAIImageMessageContent(image_url="https://your.domain/image",
+                                            text="Hello, what is shown in the photo?")
 
     # define the conversation message
     messages = [
@@ -163,7 +164,7 @@ def main():
 
 
 if __name__ == "__main__":
-   main()
+    main()
 ```
 
 ### Set custom logger

dm_aioaiagent-0.3.5.dist-info/RECORD

@@ -0,0 +1,9 @@
+dm_aioaiagent/__init__.py,sha256=ffLO0jGxiY40VYptRjJXCh7t0XKAccGtBXW5FafbFgA,213
+dm_aioaiagent/ai_agent.py,sha256=ZW85hzBpWVFgnnaOlMcpayFyxtH9mtIrc43ZVzdKBBY,10815
+dm_aioaiagent/async_ai_agent.py,sha256=qg7LZRxC7MSzjjburEX09T5EF15ZXAWAd_9RB8vA0oE,2599
+dm_aioaiagent/openai_image_message_content.py,sha256=EP_i0ERCz7c4KOM8UXp2-AI91ntGC7PGQBc5MMNspcs,434
+dm_aioaiagent/types.py,sha256=H2_iICmWr6u9d1-6BtSxdt6qV6Jm8v7Zh8das5kV6I4,989
+dm_aioaiagent-0.3.5.dist-info/METADATA,sha256=B9lKnewVEzH_XCegUY8LsbQqHBWSW7eu1icdstFbH7o,5031
+dm_aioaiagent-0.3.5.dist-info/WHEEL,sha256=PZUExdf71Ui_so67QXpySuHtCi3-J3wvF4ORK6k_S8U,91
+dm_aioaiagent-0.3.5.dist-info/top_level.txt,sha256=CbasLH0KI7zA77XwT6JDCnmRascxKNGvUVV9MgYjHAU,14
+dm_aioaiagent-0.3.5.dist-info/RECORD,,

{dm_aioaiagent-0.3.3.dist-info → dm_aioaiagent-0.3.5.dist-info}/WHEEL

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

dm_aioaiagent-0.3.3.dist-info/RECORD

@@ -1,9 +0,0 @@
-dm_aioaiagent/__init__.py,sha256=pq9gL6E1VN4Tkx9PD83AIW9e_-N5REqJeC0TE8bZua4,215
-dm_aioaiagent/ai_agent.py,sha256=TTMSL2B-fdv_q7ytO38kC5fd8uB1Vk71ev7tS5tJMwE,9284
-dm_aioaiagent/async_ai_agent.py,sha256=qg7LZRxC7MSzjjburEX09T5EF15ZXAWAd_9RB8vA0oE,2599
-dm_aioaiagent/image_message_content_builder.py,sha256=mAJnsWdnpYpaFAdcCsVwgWH480Q4WfVEpM-Y6KdjyEc,435
-dm_aioaiagent/types.py,sha256=H2_iICmWr6u9d1-6BtSxdt6qV6Jm8v7Zh8das5kV6I4,989
-dm_aioaiagent-0.3.3.dist-info/METADATA,sha256=i06u5vC73yDmx73r5lUCr0Okw1vtf1FwKJgjpuZiKY0,5032
-dm_aioaiagent-0.3.3.dist-info/WHEEL,sha256=P9jw-gEje8ByB7_hXoICnHtVCrEwMQh-630tKvQWehc,91
-dm_aioaiagent-0.3.3.dist-info/top_level.txt,sha256=CbasLH0KI7zA77XwT6JDCnmRascxKNGvUVV9MgYjHAU,14
-dm_aioaiagent-0.3.3.dist-info/RECORD,,