Create agents/common_agent.py

agents/common_agent.py

@@ -0,0 +1,114 @@
+from typing import Annotated, Any, Sequence, TypedDict
+
+from langchain.tools import StructuredTool
+from langchain_core.messages import AnyMessage, HumanMessage, SystemMessage
+from langchain_core.messages.base import BaseMessage
+from langchain_core.prompt_values import PromptValue
+from langchain_core.runnables.base import Runnable
+from langchain_openai import ChatOpenAI
+from langgraph.graph import START, StateGraph
+from langgraph.graph.message import add_messages
+from langgraph.graph.state import CompiledStateGraph
+from langgraph.prebuilt import ToolNode, tools_condition
+
+from config import settings
+from tools.encyclopedia import EncyclopediaRetriever
+from tools.tool_collection_common import ToolsCollection
+
+retriver = EncyclopediaRetriever(["gaia"], settings.PROJ_PATH)
+
+
+class AgentState(TypedDict):
+    messages: Annotated[list[AnyMessage], add_messages]
+
+
+class AgenticRAG:
+    def __init__(self):
+        chat = ChatOpenAI(model="gpt-4o", verbose=True)
+        self.tools: list[StructuredTool] = ToolsCollection.get_tools(
+            [
+                "search_tool",
+                "get_weather",
+                "hub_stats_tool",
+                "wikipedia_en_tool_agent",
+                "EncyclopediaRetriever",
+            ]
+        )
+        self.chat_with_tools: Runnable[
+            PromptValue
+            | str
+            | Sequence[
+                BaseMessage | list[str] | tuple[str, str] | str | dict[str, Any]
+            ],
+            BaseMessage,
+        ] = chat.bind_tools(self.tools)
+
+        self.agent = self.build_agent()
+
+    def retriver(self, state: AgentState):
+        question = state["messages"][-1].content
+        result = retriver.get_related_question(question)
+
+        return {
+            "messages": [
+                HumanMessage(
+                    content=f"以下有多組可能的問題及最終答案，請檢查是否有該問題`{question}`的答案，含有該問題答案時，挑出答案並回傳"
+                    + f"\n{result}"
+                )
+            ],
+        }
+
+    async def assistant(self, state: AgentState) -> dict[str, list[BaseMessage]]:
+        # print("\n=================", state["messages"], "=================\n")
+        result_message: BaseMessage = await self.chat_with_tools.ainvoke(
+            state["messages"]
+        )
+
+        return {
+            "messages": [result_message],
+        }
+
+    def build_agent(self) -> CompiledStateGraph:
+        builder = StateGraph(AgentState)
+
+        builder.add_node("retriver", self.retriver)
+        builder.add_node("assistant", self.assistant)
+        builder.add_node("tools", ToolNode(self.tools))
+
+        # Define edges: these determine how the control flow moves
+        builder.add_edge(START, "retriver")
+        builder.add_edge("retriver", "assistant")
+        builder.add_conditional_edges(
+            source="assistant",
+            path=tools_condition,
+        )
+        builder.add_edge("tools", "assistant")
+        agent: CompiledStateGraph = builder.compile()
+
+        return agent
+
+    async def ainvoke(self, message: str) -> dict[list[BaseMessage], str, Any]:
+        response = await self.agent.ainvoke(
+            {
+                "messages": [
+                    SystemMessage(
+                        content="""
+                        你是一個AI助理，專門回答問題，當根據現有資訊無法得出答案時，優先使用外部工具嘗試找出答案。
+                        
+                        當你使用外部工具時，以外部工具提供給你的答案為最準。
+                        雖然你有出眾的語言能力，回答時請精簡，不要解釋為什麼是這個答案，也不需要提供參考資訊，直接告訴我結果就好。
+                        以這個問題為例
+                        提問：How many studio albums were published by xxx between 1990 and 2009 (included)?
+                        原始回覆：Between 1990 and 2009 (included), xxx published n studio albums.
+                        希望的回覆：n
+                        """
+                    ),
+                    HumanMessage(content=message),
+                ]
+            },
+            config={"callbacks": [settings.LANGFUSE_HANDLER]},
+        )
+
+        # print("🎩 Agent's Response:")
+        # print(response["messages"][-1].content)
+        return response["messages"][-1].content