Create agents/wiki_agent.py

agents/wiki_agent.py

@@ -0,0 +1,111 @@
+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 END, START, StateGraph
+from langgraph.graph.message import add_messages
+from langgraph.graph.state import CompiledStateGraph
+from langgraph.prebuilt import ToolNode, tools_condition
+from pydantic import BaseModel, Field
+
+from config import settings
+from tools.tool_collection_wiki import ToolsCollection as WikiTool
+
+
+class AgentState(TypedDict):
+    messages: Annotated[list[AnyMessage], add_messages]
+
+
+class WikiAgent:
+    def __init__(self):
+        chat = ChatOpenAI(model="gpt-4o", verbose=True)
+        self.tools: list[StructuredTool] = WikiTool.get_tools(
+            [
+                "wikipedia_opensearch",
+                "get_page_title_excerpt_sections",
+                "get_page_section_content",
+            ]
+        )
+        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()
+
+    async def assistant(self, state: AgentState):
+        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("assistant", self.assistant)
+        builder.add_node("tools", ToolNode(self.tools))
+
+        # Define edges: these determine how the control flow moves
+        builder.add_edge(START, "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="""
+                        你是一個專門搜尋wikipedia的AI Agent,
+                        步驟一：使用 wikipedia_opensearch 工具找出與問題相關的頁面
+                        步驟二：使用 get_page_title_excerpt_sections 工具找出頁面的 excerpt 和 sections
+                        步驟三：根據步驟二的 excerpt 和 sections 結合用戶問題，判斷哪些 section 會有需要的答案，呼叫 get_page_section_content 工具取得這些 section 的所有內容。
+                        步驟四：總和前述步驟找出答案。
+                        """
+                    ),
+                    HumanMessage(content=message),
+                ]
+            },
+            config={"callbacks": [settings.LANGFUSE_HANDLER]},
+        )
+
+        # print("🎩 Agent's Response:")
+        # print(response["messages"][-1].content)
+        return response["messages"][-1].content
+
+
+class WikipediaEnToolAgentInput(BaseModel):
+    question: str = Field(description="The user question in natural language.")
+
+
+def wikipedia_en_tool_agent(question: str) -> str:
+    """
+    Invokes the WikiAgent asynchronously to answer a user-provided question using Wikipedia.
+
+    Args:
+        question (str): The user question in natural language.
+
+    Returns:
+        str: The answer or result generated by the WikiAgent.
+    """
+
+    import asyncio
+
+    return asyncio.run(WikiAgent().ainvoke(question))