dm-aioaiagent 0.4.9__tar.gz → 0.5.1__tar.gz

{dm_aioaiagent-0.4.9 → dm_aioaiagent-0.5.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: dm-aioaiagent
-Version: 0.4.9
+Version: 0.5.1
 Summary: This is my custom aioaiagent client
 Home-page: https://pypi.org/project/dm-aioaiagent
 Author: dimka4621
@@ -58,6 +58,17 @@ 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.
@@ -179,24 +190,3 @@ def main():
 if __name__ == "__main__":
     main()
 ```
-
-### Set custom logger
-
-```python
-from dm_aioaiagent import DMAIAgent
-from dm_logger import FormatterConfig
-
-
-# set up custom logger for all clients
-DMAIAgent.set_logger_params(
-   {
-      "name": "my_name",
-      "formatter_config": FormatterConfig(
-         show_datetime=False,
-      )
-   }
-)
-```
-
-See more about DMLogger [here](https://github.com/MykhLibs/dm-logger)
-

{dm_aioaiagent-0.4.9 → dm_aioaiagent-0.5.1}/README.md

@@ -21,6 +21,17 @@ 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.
@@ -142,24 +153,3 @@ def main():
 if __name__ == "__main__":
     main()
 ```
-
-### Set custom logger
-
-```python
-from dm_aioaiagent import DMAIAgent
-from dm_logger import FormatterConfig
-
-
-# set up custom logger for all clients
-DMAIAgent.set_logger_params(
-   {
-      "name": "my_name",
-      "formatter_config": FormatterConfig(
-         show_datetime=False,
-      )
-   }
-)
-```
-
-See more about DMLogger [here](https://github.com/MykhLibs/dm-logger)
-

{dm_aioaiagent-0.4.9 → dm_aioaiagent-0.5.1}/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.9 → dm_aioaiagent-0.5.1}/dm_aioaiagent/ai_agent.py

@@ -12,58 +12,73 @@ from dm_logger import DMLogger
 
 from .types import *
 
-__all__ = ["DMAIAgent"]
-
 
 class DMAIAgent:
     MAX_MEMORY_MESSAGES = 20  # Only INT greater than 0
     _ALLOWED_ROLES = ("user", "ai")
-    _logger_params = None
 
     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._set_logger(agent_name)
-        self._input_output_logging = bool(input_output_logging)
+        self._logger = DMLogger(agent_name)
 
+        # 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,
@@ -86,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
@@ -127,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
@@ -140,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"]
@@ -171,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:]
@@ -187,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:
@@ -221,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
@@ -255,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()
 
@@ -265,14 +306,3 @@ class DMAIAgent:
                 f.write(image)
         except Exception as e:
             self._logger.error(e)
-
-    def _set_logger(self, agent_name: str) -> None:
-        params = {"name": agent_name}
-        if isinstance(self._logger_params, dict):
-            params.update(self._logger_params)
-        self._logger = DMLogger(**params)
-
-    @classmethod
-    def set_logger_params(cls, extra_params = None) -> None:
-        if isinstance(extra_params, dict) or extra_params is None:
-            cls._logger_params = extra_params

{dm_aioaiagent-0.4.9 → dm_aioaiagent-0.5.1}/dm_aioaiagent/async_ai_agent.py

@@ -6,24 +6,18 @@ from langchain_core.messages import AIMessage, ToolMessage
 from .ai_agent import DMAIAgent
 from .types import *
 
-__all__ = ["DMAioAIAgent"]
-
 
 class DMAioAIAgent(DMAIAgent):
-    _logger_params = None
-
-    def __init__(self, *args, agent_name: str = "AioAIAgent", **kwargs):
+    def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
 
-        if not isinstance(self._logger_params, dict):
-            self._logger_params = {}
-        if "name" not in self._logger_params:
-            self._logger_params["name"] = agent_name
-        self._set_logger(agent_name)
+    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)
 
-    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
+        last_message = new_messages[-1]
+        if isinstance(last_message, AIMessage):
+            return last_message.content
+        return last_message
 
     async def run_messages(
         self,
@@ -46,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:
@@ -57,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
@@ -70,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.9 → dm_aioaiagent-0.5.1}/dm_aioaiagent/types.py

@@ -1,7 +1,9 @@
-from typing import Literal, Union
-from typing_extensions import TypedDict
+from typing import Literal, Union, Type, 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 +27,7 @@ class TextMessage(TypedDict):
 
 InputMessage = Union[TextMessage, ImageMessage, BaseMessage]
 
+
 class State(TypedDict):
     messages: list[InputMessage]
     new_messages: list[BaseMessage]

{dm_aioaiagent-0.4.9 → dm_aioaiagent-0.5.1}/dm_aioaiagent.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: dm-aioaiagent
-Version: 0.4.9
+Version: 0.5.1
 Summary: This is my custom aioaiagent client
 Home-page: https://pypi.org/project/dm-aioaiagent
 Author: dimka4621
@@ -58,6 +58,17 @@ 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.
@@ -179,24 +190,3 @@ def main():
 if __name__ == "__main__":
     main()
 ```
-
-### Set custom logger
-
-```python
-from dm_aioaiagent import DMAIAgent
-from dm_logger import FormatterConfig
-
-
-# set up custom logger for all clients
-DMAIAgent.set_logger_params(
-   {
-      "name": "my_name",
-      "formatter_config": FormatterConfig(
-         show_datetime=False,
-      )
-   }
-)
-```
-
-See more about DMLogger [here](https://github.com/MykhLibs/dm-logger)
-

{dm_aioaiagent-0.4.9 → dm_aioaiagent-0.5.1}/setup.py

@@ -8,7 +8,7 @@ def readme():
 
 setup(
     name='dm-aioaiagent',
-    version='v0.4.9',
+    version='v0.5.1',
     author='dimka4621',
     author_email='mismartconfig@gmail.com',
     description='This is my custom aioaiagent client',