dm-aioaiagent 0.3.6__tar.gz → 0.4.0__tar.gz

{dm_aioaiagent-0.3.6 → dm_aioaiagent-0.4.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: dm-aioaiagent
-Version: 0.3.6
+Version: 0.4.0
 Summary: This is my custom aioaiagent client
 Home-page: https://pypi.org/project/dm-aioaiagent
 Author: dimka4621
@@ -35,7 +35,7 @@ Requires-Dist: langchain-openai~=0.2.0
 
 Analogue to `DMAioAIAgent` is the synchronous client `DMAIAgent`.
 
-### Use agent *with* inner memory
+### Use agent *with* inner memory and run *single* message
 
 By default, agent use inner memory to store the conversation history.
 
@@ -61,35 +61,24 @@ async def main():
     # if you don't want to see the input and output messages from agent
     # you can set `input_output_logging=False` init argument
 
-    # 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"}
-    ]
+    answer = await ai_agent.run("Hello!")
 
     # call an agent
-    answer = await ai_agent.run(input_messages)
+    answer = await ai_agent.run("I want to know the weather in Kyiv")
 
     # get full conversation history
-    conversation_history = ai_agent.get_memory_messages()
+    conversation_history = ai_agent.memory_messages
 
     # clear conversation history
-    ai_agent.clear_memory()
+    ai_agent.clear_memory_messages()
 
 
 if __name__ == "__main__":
     asyncio.run(main())
 ```
 
-### Use agent *without* inner memory
+### Use agent *without* inner memory and run *multiple* messages
 
 If you want to control the memory of the agent, you can disable it by setting `is_memory_enabled=False`
 
@@ -114,13 +103,13 @@ async def main():
     # if you don't want to see the input and output messages from agent
     # you can set input_output_logging=False
 
-    # define the conversation message
+    # define the conversation message(s)
     messages = [
         {"role": "user", "content": "Hello!"}
     ]
 
     # call an agent
-    new_messages = await ai_agent.run(messages)
+    new_messages = await ai_agent.run_messages(messages)
 
     # add new_messages to messages
     messages.extend(new_messages)
@@ -131,7 +120,7 @@ async def main():
     )
 
     # call an agent
-    new_messages = await ai_agent.run(messages)
+    new_messages = await ai_agent.run_messages(messages)
 
 
 if __name__ == "__main__":
@@ -153,14 +142,15 @@ def main():
     img_content = OpenAIImageMessageContent(image_url="https://your.domain/image",
                                             text="Hello, what is shown in the photo?")
 
-    # define the conversation message
+    # define the conversation messages
     messages = [
         {"role": "user", "content": "Hello!"},
         {"role": "user", "content": img_content},
     ]
 
     # call an agent
-    answer = ai_agent.run(messages)
+    new_messages = ai_agent.run_messages(messages)
+    answer = new_messages[-1].content
 
 
 if __name__ == "__main__":

{dm_aioaiagent-0.3.6 → dm_aioaiagent-0.4.0}/README.md

@@ -11,7 +11,7 @@
 
 Analogue to `DMAioAIAgent` is the synchronous client `DMAIAgent`.
 
-### Use agent *with* inner memory
+### Use agent *with* inner memory and run *single* message
 
 By default, agent use inner memory to store the conversation history.
 
@@ -37,35 +37,24 @@ async def main():
     # if you don't want to see the input and output messages from agent
     # you can set `input_output_logging=False` init argument
 
-    # 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"}
-    ]
+    answer = await ai_agent.run("Hello!")
 
     # call an agent
-    answer = await ai_agent.run(input_messages)
+    answer = await ai_agent.run("I want to know the weather in Kyiv")
 
     # get full conversation history
-    conversation_history = ai_agent.get_memory_messages()
+    conversation_history = ai_agent.memory_messages
 
     # clear conversation history
-    ai_agent.clear_memory()
+    ai_agent.clear_memory_messages()
 
 
 if __name__ == "__main__":
     asyncio.run(main())
 ```
 
-### Use agent *without* inner memory
+### Use agent *without* inner memory and run *multiple* messages
 
 If you want to control the memory of the agent, you can disable it by setting `is_memory_enabled=False`
 
@@ -90,13 +79,13 @@ async def main():
     # if you don't want to see the input and output messages from agent
     # you can set input_output_logging=False
 
-    # define the conversation message
+    # define the conversation message(s)
     messages = [
         {"role": "user", "content": "Hello!"}
     ]
 
     # call an agent
-    new_messages = await ai_agent.run(messages)
+    new_messages = await ai_agent.run_messages(messages)
 
     # add new_messages to messages
     messages.extend(new_messages)
@@ -107,7 +96,7 @@ async def main():
     )
 
     # call an agent
-    new_messages = await ai_agent.run(messages)
+    new_messages = await ai_agent.run_messages(messages)
 
 
 if __name__ == "__main__":
@@ -129,14 +118,15 @@ def main():
     img_content = OpenAIImageMessageContent(image_url="https://your.domain/image",
                                             text="Hello, what is shown in the photo?")
 
-    # define the conversation message
+    # define the conversation messages
     messages = [
         {"role": "user", "content": "Hello!"},
         {"role": "user", "content": img_content},
     ]
 
     # call an agent
-    answer = ai_agent.run(messages)
+    new_messages = ai_agent.run_messages(messages)
+    answer = new_messages[-1].content
 
 
 if __name__ == "__main__":

{dm_aioaiagent-0.3.6 → dm_aioaiagent-0.4.0}/dm_aioaiagent/__init__.py

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

{dm_aioaiagent-0.3.6 → dm_aioaiagent-0.4.0}/dm_aioaiagent/ai_agent.py

@@ -1,4 +1,6 @@
 import os
+import uuid
+from typing import Any
 from pydantic import SecretStr
 from itertools import dropwhile
 from threading import Thread
@@ -14,10 +16,7 @@ __all__ = ["DMAIAgent"]
 
 
 class DMAIAgent:
-    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__(
@@ -28,16 +27,16 @@ class DMAIAgent:
         model: str = "gpt-4o-mini",
         temperature: int = 1,
         parallel_tool_calls: bool = True,
-        agent_name: str = None,
+        agent_name: str = "AIAgent",
         input_output_logging: bool = True,
         is_memory_enabled: bool = True,
+        max_memory_messages: int = MAX_MEMORY_MESSAGES,
         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
+        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."
     ):
-        self._logger = DMLogger(agent_name or self.AGENT_NAME)
+        self._logger = DMLogger(agent_name)
         self._input_output_logging = bool(input_output_logging)
 
         self._system_message = str(system_message)
@@ -48,39 +47,57 @@ class DMAIAgent:
         self._parallel_tool_calls = bool(parallel_tool_calls)
         self._llm_provider_api_key = str(llm_provider_api_key)
 
+        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)
-        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)
+        self._response_if_invalid_image = str(response_if_invalid_image)
 
+        self._check_langsmith_envs()
         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})
-        return state["response"]
+    def run(self, query: str, **kwargs) -> str:
+        new_messages = self.run_messages(messages=[{"role": "user", "content": query}], **kwargs)
+        return new_messages[-1].content
 
-    def get_memory_messages(
+    def run_messages(
         self,
-        memory_id: str = None,
+        messages: InputMessagesType,
         *,
-        without_tool_m: bool = False,
-        return_str_m: bool = False
-    ) -> Union[list[BaseMessage], list[str]]:
-        messages = self._memory.get(self._validate_memory_id(memory_id), [])
-        if without_tool_m:
-            messages = [m for m in messages if not (m.type == "tool" or (m.type == "ai" and m.tool_calls))]
-        if return_str_m:
-            messages = [m.content for m in messages]
-        return messages
-
-    def clear_memory(self, memory_id: str = None) -> None:
-        self._memory[self._validate_memory_id(memory_id)] = []
+        ls_metadata: dict[str, Any] = None,
+        ls_tags: list[str] = None,
+        ls_run_id: uuid.UUID = None,
+        ls_thread_id: uuid.UUID = None
+    ) -> list[BaseMessage]:
+        if ls_metadata is None:
+            ls_metadata = {}
+        if isinstance(ls_run_id, uuid.UUID):
+            ls_run_id = ls_run_id
+        if isinstance(ls_thread_id, uuid.UUID):
+            ls_metadata["thread_id"] = ls_thread_id
+
+        config_data = {
+            "metadata": ls_metadata,
+            "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})
+        return state["new_messages"]
+
+    @property
+    def memory_messages(self) -> list[BaseMessage]:
+        return self._memory_messages
+
+    def clear_memory_messages(self) -> None:
+        self._memory_messages.clear()
 
     def _prepare_messages_node(self, state: State) -> State:
-        state.input_messages = state.input_messages or [{"role": "user", "content": ""}]
-        for item in state.input_messages:
+        messages = state["messages"] or [{"role": "user", "content": ""}]
+        state["messages"] = []
+        for item in messages:
             if isinstance(item, dict):
                 role = item.get("role")
                 content = item.get("content")
@@ -90,35 +107,35 @@ class DMAIAgent:
                     MessageClass = AIMessage
                 else:
                     MessageClass = HumanMessage
-                state.messages.append(MessageClass(content))
+                state["messages"].append(MessageClass(content))
             elif isinstance(item, BaseMessage):
-                state.messages.append(item)
+                state["messages"].append(item)
 
         if self._input_output_logging:
-            log_kwargs = {} if state.memory_id is None else {"memory_id": state.memory_id}
-            self._logger.debug(f"Query:\n{state.messages[-1].content}", **log_kwargs)
+            self._logger.debug(f'Query:\n{state["messages"][-1].content}')
         if self._is_memory_enabled:
-            state.messages = self.get_memory_messages(state.memory_id) + state.messages
+            state["messages"] = self._memory_messages + state["messages"]
         return state
 
     def _invoke_llm_node(self, state: State, second_attempt: bool = False) -> State:
         self._logger.debug("Run node: Invoke LLM")
         try:
-            ai_response = self._agent.invoke({"messages": state.messages})
+            ai_response = self._agent.invoke({"messages": state["messages"]})
         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
-                state.messages.append(AIMessage(content=response))
+                state["messages"].append(AIMessage(content=response))
                 return state
             return self._invoke_llm_node(state, second_attempt=True)
-        state.messages.append(ai_response)
+        state["messages"].append(ai_response)
+        state["new_messages"].append(ai_response)
         return state
 
     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:
+        for tool_call in state["messages"][-1].tool_calls:
             tool_id = tool_call["id"]
             tool_name = tool_call["name"]
             tool_args = tool_call["args"]
@@ -136,7 +153,8 @@ class DMAIAgent:
                 self._logger.debug(f"Tool response:\n{tool_response}", tool_id=tool_id)
 
                 tool_message = ToolMessage(content=str(tool_response), name=tool_name, tool_call_id=tool_id)
-                state.messages.append(tool_message)
+                state["messages"].append(tool_message)
+                state["new_messages"].append(tool_message)
 
             threads.append(Thread(target=tool_callback, daemon=True))
 
@@ -144,34 +162,27 @@ class DMAIAgent:
             t.start()
         for t in threads:
             t.join()
-
         return state
 
     def _exit_node(self, state: State) -> State:
-        answer = state.messages[-1].content
         if self._input_output_logging:
-            log_kwargs = {} if state.memory_id is None else {"memory_id": state.memory_id}
-            self._logger.debug(f"Answer:\n{answer}", **log_kwargs)
+            self._logger.debug(f'Answer:\n{state["messages"][-1].content}')
 
         if self._is_memory_enabled:
-            memory_id = self._validate_memory_id(state.memory_id)
-            messages_to_memory = state.messages[-self._max_memory_messages:]
+            messages_to_memory = state["messages"][-self._max_memory_messages:]
             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))
+                self._memory_messages = list(dropwhile(lambda x: isinstance(x, ToolMessage), messages_to_memory))
             else:
-                self._memory[memory_id] = []
+                self._memory_messages.clear()
                 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):]
+                    self._memory_messages.append(mes)
         return state
 
     def _messages_router(self, state: State) -> str:
-        if self._is_tools_exists and state.messages[-1].tool_calls:
+        if self._is_tools_exists and state["messages"][-1].tool_calls:
             route = "execute_tool"
         else:
             route = "exit"
@@ -202,7 +213,6 @@ class DMAIAgent:
                                                    MessagesPlaceholder(variable_name="messages")])
 
         self._agent = prompt | llm
-        self._memory = {}
 
     def _init_graph(self) -> None:
         workflow = StateGraph(State)
@@ -221,8 +231,12 @@ class DMAIAgent:
         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
+    def _check_langsmith_envs() -> None:
+        if os.getenv("LANGCHAIN_API_KEY"):
+            if not os.getenv("LANGCHAIN_TRACING_V2"):
+                os.environ["LANGCHAIN_TRACING_V2"] = "true"
+            if not os.getenv("LANGCHAIN_ENDPOINT"):
+                os.environ["LANGCHAIN_ENDPOINT"] = "https://api.smith.langchain.com"
 
     @classmethod
     def _validate_max_memory_messages(cls, max_messages_in_memory: int) -> int:

{dm_aioaiagent-0.3.6 → dm_aioaiagent-0.4.0}/dm_aioaiagent/async_ai_agent.py

@@ -1,5 +1,7 @@
 import sys
+import uuid
 import asyncio
+from typing import Any
 from langchain_core.messages import AIMessage, ToolMessage
 
 from .ai_agent import DMAIAgent
@@ -14,28 +16,54 @@ __all__ = ["DMAioAIAgent"]
 class DMAioAIAgent(DMAIAgent):
     agent_name = "AsyncAIAgent"
 
-    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 run(self, query: str, **kwargs) -> str:
+        new_messages = await self.run_messages(messages=[{"role": "user", "content": query}], **kwargs)
+        return new_messages[-1].content
+
+    async def run_messages(
+        self,
+        messages: InputMessagesType,
+        *,
+        ls_metadata: dict[str, Any] = None,
+        ls_tags: list[str] = None,
+        ls_run_id: uuid.UUID = None,
+        ls_thread_id: uuid.UUID = None
+    ) -> list[BaseMessage]:
+        if ls_metadata is None:
+            ls_metadata = {}
+        if isinstance(ls_run_id, uuid.UUID):
+            ls_run_id = ls_run_id
+        if isinstance(ls_thread_id, uuid.UUID):
+            ls_metadata["thread_id"] = ls_thread_id
+
+        config_data = {
+            "metadata": ls_metadata,
+            "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})
+        return state["new_messages"]
 
     async def _invoke_llm_node(self, state: State, second_attempt: bool = False) -> State:
         self._logger.debug("Run node: Invoke LLM")
         try:
-            ai_response = await self._agent.ainvoke({"messages": state.messages})
+            ai_response = await self._agent.ainvoke({"messages": state["messages"]})
         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
-                state.messages.append(AIMessage(content=response))
+                state["messages"].append(AIMessage(content=response))
                 return state
             return await self._invoke_llm_node(state, second_attempt=True)
-        state.messages.append(ai_response)
+        state["messages"].append(ai_response)
+        state["new_messages"].append(ai_response)
         return state
 
     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:
+        for tool_call in state["messages"][-1].tool_calls:
             tool_id = tool_call["id"]
             tool_name = tool_call["name"]
             tool_args = tool_call["args"]
@@ -53,7 +81,8 @@ class DMAioAIAgent(DMAIAgent):
                 self._logger.debug(f"Tool response:\n{tool_response}", tool_id=tool_id)
 
                 tool_message = ToolMessage(content=str(tool_response), name=tool_name, tool_call_id=tool_id)
-                state.messages.append(tool_message)
+                state["messages"].append(tool_message)
+                state["new_messages"].append(tool_message)
 
             tasks.append(asyncio.create_task(tool_callback()))
 

dm_aioaiagent-0.4.0/dm_aioaiagent/types.py

@@ -0,0 +1,30 @@
+from typing import Literal, Union
+from typing_extensions import TypedDict
+from langchain_core.messages import BaseMessage
+
+
+class ImageMessageTextMessage(TypedDict):
+    type: Literal['text']
+    text: str
+
+
+class ImageMessageImageItem(TypedDict):
+    type: Literal['image_url']
+    image_url: dict
+
+
+class ImageMessage(TypedDict):
+    role: Literal["user"]
+    content: list[Union[ImageMessageTextMessage, ImageMessageImageItem]]
+
+
+class TextMessage(TypedDict):
+    role: Literal["user", "ai"]
+    content: str
+
+
+InputMessagesType = list[Union[TextMessage, ImageMessage, BaseMessage]]
+
+class State(TypedDict):
+    messages: InputMessagesType
+    new_messages: list[BaseMessage]

{dm_aioaiagent-0.3.6 → dm_aioaiagent-0.4.0}/dm_aioaiagent.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: dm-aioaiagent
-Version: 0.3.6
+Version: 0.4.0
 Summary: This is my custom aioaiagent client
 Home-page: https://pypi.org/project/dm-aioaiagent
 Author: dimka4621
@@ -35,7 +35,7 @@ Requires-Dist: langchain-openai~=0.2.0
 
 Analogue to `DMAioAIAgent` is the synchronous client `DMAIAgent`.
 
-### Use agent *with* inner memory
+### Use agent *with* inner memory and run *single* message
 
 By default, agent use inner memory to store the conversation history.
 
@@ -61,35 +61,24 @@ async def main():
     # if you don't want to see the input and output messages from agent
     # you can set `input_output_logging=False` init argument
 
-    # 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"}
-    ]
+    answer = await ai_agent.run("Hello!")
 
     # call an agent
-    answer = await ai_agent.run(input_messages)
+    answer = await ai_agent.run("I want to know the weather in Kyiv")
 
     # get full conversation history
-    conversation_history = ai_agent.get_memory_messages()
+    conversation_history = ai_agent.memory_messages
 
     # clear conversation history
-    ai_agent.clear_memory()
+    ai_agent.clear_memory_messages()
 
 
 if __name__ == "__main__":
     asyncio.run(main())
 ```
 
-### Use agent *without* inner memory
+### Use agent *without* inner memory and run *multiple* messages
 
 If you want to control the memory of the agent, you can disable it by setting `is_memory_enabled=False`
 
@@ -114,13 +103,13 @@ async def main():
     # if you don't want to see the input and output messages from agent
     # you can set input_output_logging=False
 
-    # define the conversation message
+    # define the conversation message(s)
     messages = [
         {"role": "user", "content": "Hello!"}
     ]
 
     # call an agent
-    new_messages = await ai_agent.run(messages)
+    new_messages = await ai_agent.run_messages(messages)
 
     # add new_messages to messages
     messages.extend(new_messages)
@@ -131,7 +120,7 @@ async def main():
     )
 
     # call an agent
-    new_messages = await ai_agent.run(messages)
+    new_messages = await ai_agent.run_messages(messages)
 
 
 if __name__ == "__main__":
@@ -153,14 +142,15 @@ def main():
     img_content = OpenAIImageMessageContent(image_url="https://your.domain/image",
                                             text="Hello, what is shown in the photo?")
 
-    # define the conversation message
+    # define the conversation messages
     messages = [
         {"role": "user", "content": "Hello!"},
         {"role": "user", "content": img_content},
     ]
 
     # call an agent
-    answer = ai_agent.run(messages)
+    new_messages = ai_agent.run_messages(messages)
+    answer = new_messages[-1].content
 
 
 if __name__ == "__main__":

{dm_aioaiagent-0.3.6 → dm_aioaiagent-0.4.0}/setup.py

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

dm_aioaiagent-0.3.6/dm_aioaiagent/types.py

@@ -1,41 +0,0 @@
-from typing import Optional, Literal, Union
-from typing_extensions import TypedDict
-from pydantic import BaseModel, Field
-from langchain_core.messages import BaseMessage
-
-
-class ImageMessageTextMessage(TypedDict):
-    type: Literal['text']
-    text: str
-
-
-class ImageMessageImageItem(TypedDict):
-    type: Literal['image_url']
-    image_url: dict
-
-
-ImageMessageContent = list[Union[ImageMessageTextMessage, ImageMessageImageItem]]
-
-
-class ImageMessage(TypedDict):
-    role: Literal["user"]
-    content: ImageMessageContent
-
-
-class TextMessage(TypedDict):
-    role: Literal["user", "ai"]
-    content: str
-
-
-Message = Union[TextMessage, ImageMessage]
-
-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="")