firts agent concept

src/naliiabot/bot/agent/agent.py

@@ -0,0 +1,81 @@
+from langgraph.graph import StateGraph, START, END
+from langchain.messages import SystemMessage, ToolMessage
+from typing import Literal
+from .schemas import MessagesState
+
+
+class Agent:
+
+    def __init__(self, model):
+        self._model = model
+        self._tools = []
+        self._tools_by_name = []
+        self.model_with_tools = self._model.bind_tools(
+            self._tools)
+        self.agent = self._compiler_agent()
+
+    def invoke(self, messages):
+        self.agent.invoke(
+            messages
+        )
+
+    def _compiler_agent(self):
+        agent = self._build_agent()
+
+        return agent.compile()
+
+    def _build_agent(self):
+        agent_builder = StateGraph(MessagesState)
+        agent_builder.add_node("llm_call", self.llm_call)
+        agent_builder.add_node("tool_node", self.tool_node)
+        agent_builder.add_edge(START, "llm_call")
+
+        return agent_builder
+
+    def llm_call(self, state: dict):
+        """LLM decides whether to call a tool or not"""
+
+        return {
+            "messages": [
+                self.model_with_tools.invoke(
+                    [
+                        SystemMessage(
+                            content=(
+                                "You are a helpful assistant tasked with "
+                                "performing arithmetic on a set of inputs."
+                            )
+                        )
+                    ]
+                    + state["messages"]
+                )
+            ],
+            "llm_calls": state.get('llm_calls', 0) + 1
+        }
+
+    def tool_node(self, state: dict):
+        """Performs the tool call"""
+
+        result = []
+        for tool_call in state["messages"][-1].tool_calls:
+            tool = self._tools_by_name[tool_call["name"]]
+            observation = tool.invoke(tool_call["args"])
+            result.append(
+                ToolMessage(
+                    content=observation,
+                    tool_call_id=tool_call["id"]
+                ))
+        return {"messages": result}
+
+    def should_continue(state: MessagesState) -> Literal["tool_node", END]:
+        """
+        Decide if we should continue the loop or stop base
+        d upon whether the LLM made a tool call
+        """
+
+        messages = state["messages"]
+        last_message = messages[-1]
+
+        if last_message.tool_calls:
+            return "tool_node"
+
+        return END