dm-aioaiagent 0.4.8__tar.gz → 0.5.0__tar.gz

{dm_aioaiagent-0.4.8 → dm_aioaiagent-0.5.0}/PKG-INFO

@@ -1,6 +1,6 @@
-Metadata-Version: 2.2
+Metadata-Version: 2.4
 Name: dm-aioaiagent
-Version: 0.4.8
+Version: 0.5.0
 Summary: This is my custom aioaiagent client
 Home-page: https://pypi.org/project/dm-aioaiagent
 Author: dimka4621
@@ -12,7 +12,7 @@ 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: dm-logger~=0.6.2
 Requires-Dist: python-dotenv>=1.0.0
 Requires-Dist: pydantic<3.0.0,>=2.9.2
 Requires-Dist: langchain~=0.3.0
@@ -48,6 +48,27 @@ Dynamic: summary
 
 Analogue to `DMAioAIAgent` is the synchronous client `DMAIAgent`.
 
+### Windows Setup
+
+```python
+import asyncio
+import sys
+
+if sys.platform == "win32":
+    asyncio.set_event_loop_policy(asyncio.WindowsSelectorEventLoopPolicy())
+```
+
+### Api Key Setup
+
+You can set your OpenAI API key in the environment variable `OPENAI_API_KEY` or pass it as an argument to the agent.
+
+**Use load_dotenv to load the `.env` file.**
+
+```python
+from dotenv import load_dotenv
+load_dotenv()
+```
+
 ### Use agent *with* inner memory and run *single* message
 
 By default, agent use inner memory to store the conversation history.
@@ -169,33 +190,3 @@ def main():
 if __name__ == "__main__":
     main()
 ```
-
-### Set custom logger
-
-_If you want set up custom logger_
-
-```python
-from dm_aioaiagent import DMAioAIAgent
-
-
-# create custom logger
-class MyLogger:
-    def debug(self, message):
-        pass
-
-    def info(self, message):
-        pass
-
-    def warning(self, message):
-        print(message)
-
-    def error(self, message):
-        print(message)
-
-
-# create an agent
-ai_agent = DMAioAIAgent()
-
-# set up custom logger for this agent
-ai_agent.set_logger(MyLogger())
-```

{dm_aioaiagent-0.4.8 → dm_aioaiagent-0.5.0}/README.md

@@ -11,6 +11,27 @@
 
 Analogue to `DMAioAIAgent` is the synchronous client `DMAIAgent`.
 
+### Windows Setup
+
+```python
+import asyncio
+import sys
+
+if sys.platform == "win32":
+    asyncio.set_event_loop_policy(asyncio.WindowsSelectorEventLoopPolicy())
+```
+
+### Api Key Setup
+
+You can set your OpenAI API key in the environment variable `OPENAI_API_KEY` or pass it as an argument to the agent.
+
+**Use load_dotenv to load the `.env` file.**
+
+```python
+from dotenv import load_dotenv
+load_dotenv()
+```
+
 ### Use agent *with* inner memory and run *single* message
 
 By default, agent use inner memory to store the conversation history.
@@ -132,33 +153,3 @@ def main():
 if __name__ == "__main__":
     main()
 ```
-
-### Set custom logger
-
-_If you want set up custom logger_
-
-```python
-from dm_aioaiagent import DMAioAIAgent
-
-
-# create custom logger
-class MyLogger:
-    def debug(self, message):
-        pass
-
-    def info(self, message):
-        pass
-
-    def warning(self, message):
-        print(message)
-
-    def error(self, message):
-        print(message)
-
-
-# create an agent
-ai_agent = DMAioAIAgent()
-
-# set up custom logger for this agent
-ai_agent.set_logger(MyLogger())
-```

{dm_aioaiagent-0.4.8 → dm_aioaiagent-0.5.0}/dm_aioaiagent/__init__.py

@@ -1,5 +1,3 @@
-from dotenv import load_dotenv
-load_dotenv()
 from .ai_agent import DMAIAgent
 from .async_ai_agent import DMAioAIAgent
 from .openai_image_message_content import OpenAIImageMessageContent

{dm_aioaiagent-0.4.8 → dm_aioaiagent-0.5.0}/dm_aioaiagent/ai_agent.py

@@ -12,8 +12,6 @@ from dm_logger import DMLogger
 
 from .types import *
 
-__all__ = ["DMAIAgent"]
-
 
 class DMAIAgent:
     MAX_MEMORY_MESSAGES = 20  # Only INT greater than 0
@@ -21,48 +19,66 @@ class DMAIAgent:
 
     def __init__(
         self,
+        # general
         system_message: str = "You are a helpful assistant.",
-        tools: list[BaseTool] = None,
-        *,
         model: str = "gpt-4o-mini",
         temperature: float = None,
-        parallel_tool_calls: bool = None,
+        output_schema: OutputSchemaType = None,  # IF you set output_schema, tools and memory will be disabled
         agent_name: str = "AIAgent",
-        input_output_logging: bool = True,
+        # tools
+        tools: list[BaseTool] = None,
+        parallel_tool_calls: bool = None,
+        # memory
         is_memory_enabled: bool = True,
-        max_memory_messages: int = MAX_MEMORY_MESSAGES,
         save_tools_responses_in_memory: bool = True,
-        llm_provider_api_key: str = "",
-        llm_provider_base_url: str = "",
+        max_memory_messages: int = MAX_MEMORY_MESSAGES,
+        # other
+        input_output_logging: bool = True,
         response_if_request_fail: str = "I can't provide a response right now. Please try again later.",
-        response_if_invalid_image: str = "The image is unavailable or the link is incorrect."
+        response_if_invalid_image: str = "The image is unavailable or the link is incorrect.",
+        # llm provider
+        llm_provider_api_key: str = "",
+        llm_provider_base_url: str = ""
     ):
         self._logger = DMLogger(agent_name)
-        self._input_output_logging = bool(input_output_logging)
 
+        # general
         self._system_message = str(system_message)
-        self._tools = tools or []
-        self._is_tools_exists = bool(tools)
         self._model = str(model)
         self._temperature = temperature
+        self._output_schema = self._validate_output_schema(output_schema)
+        if self._output_schema:
+            tools = []
+            parallel_tool_calls = None
+            is_memory_enabled = False
+        # tools
+        self._tools = tools or []
+        self._is_tools_exists = bool(tools)
         self._parallel_tool_calls = parallel_tool_calls
-        self._llm_provider_api_key = str(llm_provider_api_key)
-        self._llm_provider_base_url = str(llm_provider_base_url)
-
+        # memory
         self._memory_messages = []
         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)
+        # other
+        self._input_output_logging = bool(input_output_logging)
         self._response_if_request_fail = str(response_if_request_fail)
         self._response_if_invalid_image = str(response_if_invalid_image)
+        # llm provider
+        self._llm_provider_api_key = str(llm_provider_api_key)
+        self._llm_provider_base_url = str(llm_provider_base_url)
 
         self._check_langsmith_envs()
         self._init_agent()
         self._init_graph()
 
-    def run(self, query: str, **kwargs) -> str:
-        new_messages = self.run_messages(messages=[{"role": "user", "content": query}], **kwargs)
-        return new_messages[-1].content
+    def run(self, query: str, **kwargs) -> Union[str, TypedDict, BaseModel]:
+        new_messages = self.run_messages(messages=[TextMessage(role="user", content=query)], **kwargs)
+
+        last_message = new_messages[-1]
+        if isinstance(last_message, AIMessage):
+            return last_message.content
+        return last_message
 
     def run_messages(
         self,
@@ -85,8 +101,10 @@ class DMAIAgent:
             "tags": ls_tags,
             "run_id": ls_run_id
         }
-        state = self._graph.invoke(input={"messages": messages, "new_messages": []},
-                                   config={k: v for k, v in config_data.items() if v})
+        state = self._graph.invoke(
+            input={"messages": messages, "new_messages": []},
+            config={k: v for k, v in config_data.items() if v}
+        )
         return state["new_messages"]
 
     @property
@@ -126,12 +144,18 @@ class DMAIAgent:
         except Exception as e:
             self._logger.error(e)
             if second_attempt:
-                response = self._response_if_invalid_image if "invalid_image_url" in str(e) else self._response_if_request_fail
+                if "invalid_image_url" in str(e):
+                    response = self._response_if_invalid_image
+                else:
+                    response = self._response_if_request_fail
+
                 ai_response = AIMessage(content=response)
                 state["messages"].append(ai_response)
                 state["new_messages"].append(ai_response)
                 return state
+
             return self._invoke_llm_node(state, second_attempt=True)
+
         state["messages"].append(ai_response)
         state["new_messages"].append(ai_response)
         return state
@@ -139,6 +163,7 @@ class DMAIAgent:
     def _execute_tool_node(self, state: State) -> State:
         self._logger.debug("Run node: Execute tool")
         threads = []
+
         for tool_call in state["messages"][-1].tool_calls:
             tool_id = tool_call["id"]
             tool_name = tool_call["name"]
@@ -170,7 +195,10 @@ class DMAIAgent:
 
     def _exit_node(self, state: State) -> State:
         if self._input_output_logging:
-            self._logger.debug(f'Answer:\n{state["messages"][-1].content}')
+            answer = state["messages"][-1]
+            if isinstance(answer, AIMessage):
+                answer = answer.content
+            self._logger.debug(f'Answer:\n{answer}')
 
         if self._is_memory_enabled:
             messages_to_memory = state["messages"][-self._max_memory_messages:]
@@ -186,16 +214,18 @@ class DMAIAgent:
         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
+        if self._output_schema:
+            return "exit"
+        elif self._is_tools_exists and state["messages"][-1].tool_calls:
+            return "execute_tool"
+        return "exit"
 
     def _init_agent(self) -> None:
         base_kwargs = {"model": self._model}
         if isinstance(self._temperature, float):
             base_kwargs["temperature"] = self._temperature
+        else:
+            ValueError("Temperature must be a float value.")
         if self._llm_provider_api_key:
             base_kwargs["api_key"] = SecretStr(self._llm_provider_api_key)
         if self._llm_provider_base_url:
@@ -220,6 +250,9 @@ class DMAIAgent:
             self._tool_map = {t.name: t for t in self._tools}
             llm = llm.bind_tools(self._tools, **bind_tool_kwargs)
 
+        if self._output_schema:
+            llm = llm.with_structured_output(self._output_schema)
+
         prompt = ChatPromptTemplate.from_messages([SystemMessage(content=self._system_message),
                                                    MessagesPlaceholder(variable_name="messages")])
         self._agent = prompt | llm
@@ -254,6 +287,15 @@ class DMAIAgent:
             return max_messages_in_memory
         return cls.MAX_MEMORY_MESSAGES
 
+    @staticmethod
+    def _validate_output_schema(schema: OutputSchemaType) -> OutputSchemaType:
+        if schema is None:
+            return None
+        if isinstance(schema, type) and \
+            (type(schema).__name__ == "_TypedDictMeta" or issubclass(schema, BaseModel)):
+            return schema
+        raise ValueError("Output schema must be a TypedDict or BaseModel type, or None.")
+
     def print_graph(self) -> None:
         self._graph.get_graph().print_ascii()
 
@@ -264,14 +306,3 @@ class DMAIAgent:
                 f.write(image)
         except Exception as e:
             self._logger.error(e)
-
-    def set_logger(self, logger) -> None:
-        if (
-            hasattr(logger, "debug") and callable(logger.debug) and
-            hasattr(logger, "info") and callable(logger.info) and
-            hasattr(logger, "warning") and callable(logger.warning) and
-            hasattr(logger, "error") and callable(logger.error)
-        ):
-            self._logger = logger
-        else:
-            print("Invalid logger")

{dm_aioaiagent-0.4.8 → dm_aioaiagent-0.5.0}/dm_aioaiagent/async_ai_agent.py

@@ -6,13 +6,18 @@ from langchain_core.messages import AIMessage, ToolMessage
 from .ai_agent import DMAIAgent
 from .types import *
 
-__all__ = ["DMAioAIAgent"]
-
 
 class DMAioAIAgent(DMAIAgent):
-    async def run(self, query: str, *args, **kwargs) -> str:
-        new_messages = await self.run_messages(messages=[{"role": "user", "content": query}], *args, **kwargs)
-        return new_messages[-1].content
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+    async def run(self, query: str, *args, **kwargs) -> Union[str, TypedDict, BaseModel]:
+        new_messages = await self.run_messages(messages=[TextMessage(role="user", content=query)], *args, **kwargs)
+
+        last_message = new_messages[-1]
+        if isinstance(last_message, AIMessage):
+            return last_message.content
+        return last_message
 
     async def run_messages(
         self,
@@ -35,8 +40,10 @@ class DMAioAIAgent(DMAIAgent):
             "tags": ls_tags,
             "run_id": ls_run_id
         }
-        state = await self._graph.ainvoke(input={"messages": messages, "new_messages": []},
-                                          config={k: v for k, v in config_data.items() if v})
+        state = await self._graph.ainvoke(
+            input={"messages": messages, "new_messages": []},
+            config={k: v for k, v in config_data.items() if v}
+        )
         return state["new_messages"]
 
     async def _invoke_llm_node(self, state: State, second_attempt: bool = False) -> State:
@@ -46,12 +53,18 @@ class DMAioAIAgent(DMAIAgent):
         except Exception as e:
             self._logger.error(e)
             if second_attempt:
-                response = self._response_if_invalid_image if "invalid_image_url" in str(e) else self._response_if_request_fail
+                if "invalid_image_url" in str(e):
+                    response = self._response_if_invalid_image
+                else:
+                    response = self._response_if_request_fail
+
                 ai_response = AIMessage(content=response)
                 state["messages"].append(ai_response)
                 state["new_messages"].append(ai_response)
                 return state
+
             return await self._invoke_llm_node(state, second_attempt=True)
+
         state["messages"].append(ai_response)
         state["new_messages"].append(ai_response)
         return state
@@ -59,6 +72,7 @@ class DMAioAIAgent(DMAIAgent):
     async def _execute_tool_node(self, state: State) -> State:
         self._logger.debug("Run node: Execute tool")
         tasks = []
+
         for tool_call in state["messages"][-1].tool_calls:
             tool_id = tool_call["id"]
             tool_name = tool_call["name"]

{dm_aioaiagent-0.4.8 → dm_aioaiagent-0.5.0}/dm_aioaiagent/types.py

@@ -1,7 +1,10 @@
-from typing import Literal, Union
+from typing import Literal, Union, Type
 from typing_extensions import TypedDict
+from pydantic import BaseModel
 from langchain_core.messages import BaseMessage
 
+OutputSchemaType = Union[Type[TypedDict], Type[BaseModel], None]
+
 
 class ImageMessageTextItem(TypedDict):
     type: Literal['text']
@@ -25,6 +28,7 @@ class TextMessage(TypedDict):
 
 InputMessage = Union[TextMessage, ImageMessage, BaseMessage]
 
+
 class State(TypedDict):
     messages: list[InputMessage]
     new_messages: list[BaseMessage]

{dm_aioaiagent-0.4.8 → dm_aioaiagent-0.5.0}/dm_aioaiagent.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
-Metadata-Version: 2.2
+Metadata-Version: 2.4
 Name: dm-aioaiagent
-Version: 0.4.8
+Version: 0.5.0
 Summary: This is my custom aioaiagent client
 Home-page: https://pypi.org/project/dm-aioaiagent
 Author: dimka4621
@@ -12,7 +12,7 @@ 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: dm-logger~=0.6.2
 Requires-Dist: python-dotenv>=1.0.0
 Requires-Dist: pydantic<3.0.0,>=2.9.2
 Requires-Dist: langchain~=0.3.0
@@ -48,6 +48,27 @@ Dynamic: summary
 
 Analogue to `DMAioAIAgent` is the synchronous client `DMAIAgent`.
 
+### Windows Setup
+
+```python
+import asyncio
+import sys
+
+if sys.platform == "win32":
+    asyncio.set_event_loop_policy(asyncio.WindowsSelectorEventLoopPolicy())
+```
+
+### Api Key Setup
+
+You can set your OpenAI API key in the environment variable `OPENAI_API_KEY` or pass it as an argument to the agent.
+
+**Use load_dotenv to load the `.env` file.**
+
+```python
+from dotenv import load_dotenv
+load_dotenv()
+```
+
 ### Use agent *with* inner memory and run *single* message
 
 By default, agent use inner memory to store the conversation history.
@@ -169,33 +190,3 @@ def main():
 if __name__ == "__main__":
     main()
 ```
-
-### Set custom logger
-
-_If you want set up custom logger_
-
-```python
-from dm_aioaiagent import DMAioAIAgent
-
-
-# create custom logger
-class MyLogger:
-    def debug(self, message):
-        pass
-
-    def info(self, message):
-        pass
-
-    def warning(self, message):
-        print(message)
-
-    def error(self, message):
-        print(message)
-
-
-# create an agent
-ai_agent = DMAioAIAgent()
-
-# set up custom logger for this agent
-ai_agent.set_logger(MyLogger())
-```

{dm_aioaiagent-0.4.8 → dm_aioaiagent-0.5.0}/dm_aioaiagent.egg-info/requires.txt

@@ -1,4 +1,4 @@
-dm-logger~=0.5.2
+dm-logger~=0.6.2
 python-dotenv>=1.0.0
 pydantic<3.0.0,>=2.9.2
 langchain~=0.3.0

{dm_aioaiagent-0.4.8 → dm_aioaiagent-0.5.0}/setup.py

@@ -8,7 +8,7 @@ def readme():
 
 setup(
     name='dm-aioaiagent',
-    version='v0.4.8',
+    version='v0.5.0',
     author='dimka4621',
     author_email='mismartconfig@gmail.com',
     description='This is my custom aioaiagent client',
@@ -17,7 +17,7 @@ setup(
     url='https://pypi.org/project/dm-aioaiagent',
     packages=find_packages(),
     install_requires=[
-        'dm-logger~=0.5.2',
+        'dm-logger~=0.6.2',
         'python-dotenv>=1.0.0',
         'pydantic>=2.9.2, < 3.0.0',
         'langchain~=0.3.0',