dm-aioaiagent 0.1.0__py3-none-any.whl

dm_aioaiagent/__init__.py

@@ -0,0 +1,4 @@
+from dotenv import load_dotenv
+load_dotenv()
+from .ai_agent import DMAIAgent
+from .async_ai_agent import DMAioAIAgent

dm_aioaiagent/ai_agent.py

@@ -0,0 +1,176 @@
+import json
+import os
+from typing import Optional, Literal
+from typing_extensions import TypedDict
+from pydantic import BaseModel, Field
+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 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())
+
+
+class OutputState(TypedDict):
+    answer: str
+    context: list[Message]
+
+
+class DMAIAgent:
+    agent_name = "AIAgent"
+    _allowed_roles = ("user", "ai")
+
+    def __init__(
+        self,
+        system_message: str = "You are a helpful assistant.",
+        tools: list[BaseTool] = None,
+        *,
+        model: str = "gpt-4o-mini",
+        temperature: int = 1,
+        agent_name: str = None,
+        input_output_logging: bool = True
+    ):
+        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 = input_output_logging
+        self._is_tools_exists = bool(tools)
+
+        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
+
+        workflow = StateGraph(input=InputState, output=OutputState)
+        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()
+
+    def run(self, messages: list[Message]) -> OutputState:
+        return self._graph.invoke({"messages": messages})
+
+    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))
+        return state
+
+    def _invoke_llm_node(self, state: InputState) -> InputState:
+        self._logger.debug("Run node: Invoke LLM")
+        ai_response = self._agent.invoke({"messages": state.inner_state.messages})
+        state.inner_state.messages.append(ai_response)
+        return state
+
+    def _execute_tool_node(self, state: InputState) -> InputState:
+        self._logger.debug("Run node: Execute tool")
+        threads = []
+        for tool_call in state.inner_state.messages[-1].tool_calls:
+            tool_id = tool_call["id"]
+            tool_name = tool_call["name"]
+            tool_args = tool_call["args"]
+
+            def tool_callback(tool_id=tool_id, tool_name=tool_name, tool_args=tool_args) -> None:
+                self._logger.debug("Invoke tool", tool_id=tool_id, tool_name=tool_name, tool_args=tool_args)
+                if tool_name in self._tool_map:
+                    try:
+                        tool_response = self._tool_map[tool_name].run(tool_args)
+                    except Exception as e:
+                        self._logger.error(e, tool_id=tool_id)
+                        tool_response = "Tool executed with an error!"
+                else:
+                    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)
+
+            threads.append(Thread(target=tool_callback, daemon=True))
+
+        for t in threads:
+            t.start()
+        for t in threads:
+            t.join()
+
+        return state
+
+    def _exit_node(self, state: InputState) -> OutputState:
+        answer = state.inner_state.messages[-1].content if state.inner_state.messages else ""
+        if self._input_output_logging:
+            self._logger.debug(f"Answer:\n{answer}")
+        return OutputState(answer=answer, context=state.inner_state.context)
+
+    def _messages_router(self, state: InputState) -> str:
+        if self._is_tools_exists and state.inner_state.messages[-1].tool_calls:
+            route = "execute_tool"
+        else:
+            route = "exit"
+        return route
+
+    def print_graph(self) -> None:
+        self._graph.get_graph().print_ascii()
+
+    def save_graph_image(self, path: str) -> None:
+        try:
+            image = self._graph.get_graph().draw_mermaid_png()
+            with open(str(path), "wb") as f:
+                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/async_ai_agent.py

@@ -0,0 +1,55 @@
+import json
+import sys
+import asyncio
+from langchain_core.messages import ToolMessage
+
+from .ai_agent import DMAIAgent, InputState, OutputState
+
+__all__ = ["DMAioAIAgent"]
+
+if sys.platform == "win32":
+    asyncio.set_event_loop_policy(asyncio.WindowsSelectorEventLoopPolicy())
+
+
+class DMAioAIAgent(DMAIAgent):
+    agent_name = "AsyncAIAgent"
+
+    async def run(self, messages: list[dict[str, str]]) -> OutputState:
+        return await self._graph.ainvoke({"messages": messages})
+
+    async def _invoke_llm_node(self, state: InputState) -> InputState:
+        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)
+        return state
+
+    async def _execute_tool_node(self, state: InputState) -> InputState:
+        self._logger.debug("Run node: Execute tool")
+        tasks = []
+        for tool_call in state.inner_state.messages[-1].tool_calls:
+            tool_id = tool_call["id"]
+            tool_name = tool_call["name"]
+            tool_args = tool_call["args"]
+
+            async def tool_callback(tool_id=tool_id, tool_name=tool_name, tool_args=tool_args) -> None:
+                self._logger.debug("Invoke tool", tool_id=tool_id, tool_name=tool_name, tool_args=tool_args)
+                if tool_name in self._tool_map:
+                    try:
+                        tool_response = await self._tool_map[tool_name].arun(tool_args)
+                    except Exception as e:
+                        self._logger.error(e, tool_id=tool_id)
+                        tool_response = "Tool executed with an error!"
+                else:
+                    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)
+
+            tasks.append(asyncio.create_task(tool_callback()))
+
+        await asyncio.gather(*tasks)
+        return state

dm_aioaiagent-0.1.0.dist-info/METADATA

@@ -0,0 +1,106 @@
+Metadata-Version: 2.1
+Name: dm-aioaiagent
+Version: 0.1.0
+Summary: This is my custom aioaiagent client
+Home-page: https://pypi.org/project/dm-aioaiagent
+Author: dimka4621
+Author-email: mismartconfig@gmail.com
+Project-URL: GitHub, https://github.com/MykhLibs/dm-aioaiagent
+Keywords: dm aioaiagent
+Classifier: Programming Language :: Python :: 3.8
+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.3
+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
+
+# DM-aioaiagent
+
+## Urls
+
+* [PyPI](https://pypi.org/project/dm-aioaiagent)
+* [GitHub](https://github.com/MykhLibs/dm-aioaiagent)
+
+### * Package contains both `asynchronous` and `synchronous` clients
+
+## Usage
+
+### Example of using DMAioAIAgent without tools
+
+Analogue to `DMAioAIAgent` is the synchronous client `DMAIAgent`.
+
+```python
+import asyncio
+from dm_aioaiagent import DMAioAIAgent
+
+
+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
+    tools = [...]
+
+    # define a openai model, default is "gpt-4o-mini"
+    model_name = "gpt-4o"
+
+    # create an agent
+    ai_agent = DMAioAIAgent(system_message, tools, model=model_name)
+
+    # 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"},
+    ]
+
+    # start the agent
+    state = await ai_agent.run(messages)
+
+    # print answer
+    print(state["answer"])
+
+    # if you define tools, you can see the context of the tools
+    print(state["answer"])
+
+
+if __name__ == "__main__":
+    asyncio.run(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 agent
+ai_agent = DMAioAIAgent()
+
+# set up custom logger for this agent
+ai_agent.set_logger(MyLogger())
+```

dm_aioaiagent-0.1.0.dist-info/RECORD

@@ -0,0 +1,7 @@
+dm_aioaiagent/__init__.py,sha256=Hhgz18uMAnR0ezLOukoWCe1sFS4Q4LC26xbRKdu9FFw,118
+dm_aioaiagent/ai_agent.py,sha256=R08RvKc80yyxIgX7VRioqdwhQoLAdeTFtr2AJ5mPsvc,6956
+dm_aioaiagent/async_ai_agent.py,sha256=eGRg5IWR-lz8py_n6z4yB2cdKHwZNbcngIUxuOGN-hc,2396
+dm_aioaiagent-0.1.0.dist-info/METADATA,sha256=2oH3lfL4sigpzhY8EgK3Tzoe1pXOoGNj0lcW1L542Mc,2519
+dm_aioaiagent-0.1.0.dist-info/WHEEL,sha256=GV9aMThwP_4oNCtvEC2ec3qUYutgWeAzklro_0m4WJQ,91
+dm_aioaiagent-0.1.0.dist-info/top_level.txt,sha256=CbasLH0KI7zA77XwT6JDCnmRascxKNGvUVV9MgYjHAU,14
+dm_aioaiagent-0.1.0.dist-info/RECORD,,

dm_aioaiagent-0.1.0.dist-info/WHEEL

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

dm_aioaiagent-0.1.0.dist-info/top_level.txt

@@ -0,0 +1 @@
+dm_aioaiagent