LagGraph とMongoDB AtlasでAIエージェントを構築

from langchain.agents import toolfrom langchain_mongodb.retrievers.full_text_search import MongoDBAtlasFullTextSearchRetrieverfrom config import vector_store, collection
@tooldef plot_search(user_query: str) -> str:    """    Retrieve information on the movie's plot to answer a user query by using vector search.    """        retriever = vector_store.as_retriever(        search_type="similarity",        search_kwargs={"k": 5}  # Retrieve top 5 most similar documents    )
    results = retriever.invoke(user_query)       # Concatenate the results into a string    context = "\n\n".join([f"{doc.metadata['title']}: {doc.page_content}" for doc in results])    return context
@tooldef title_search(user_query: str) -> str:    """    Retrieve movie plot content based on the provided title by using full-text search.    """        # Initialize the retriever    retriever = MongoDBAtlasFullTextSearchRetriever(        collection=collection,            # MongoDB Collection        search_field="title",             # Name of the field to search        search_index_name="search_index", # Name of the MongoDB Search index        top_k=1,                          # Number of top results to return           )     results = retriever.invoke(user_query)       for doc in results:        if doc:            return doc.metadata["fullplot"]        else:            return "Movie not found"
# List of search toolsSEARCH_TOOLS = [ plot_search, title_search ]

from langchain.agents import toolfrom langgraph.store.mongodb import MongoDBStore, create_vector_index_configfrom config import embedding_model, MONGODB_URI
# Vector search index configuration for memory collectionindex_config = create_vector_index_config(    embed=embedding_model,    dims=2048,    relevance_score_fn="dotProduct",    fields=["content"])
@tooldef save_memory(content: str) -> str:    """Save important information to memory."""    with MongoDBStore.from_conn_string(        conn_string=MONGODB_URI,        db_name="sample_mflix",        collection_name="memories",        index_config=index_config,        auto_index_timeout=60 # Wait a minute for vector index creation    ) as store:        store.put(            namespace=("user", "memories"),            key=f"memory_{hash(content)}",            value={"content": content}        )    return f"Memory saved: {content}"
@tooldef retrieve_memories(query: str) -> str:    """Retrieve relevant memories based on a query."""    with MongoDBStore.from_conn_string(        conn_string=MONGODB_URI,        db_name="sample_mflix",        collection_name="memories",        index_config=index_config    ) as store:        results = store.search(("user", "memories"), query=query, limit=3)
    if results:        memories = [result.value["content"] for result in results]        return f"Retrieved memories:\n" + "\n".join(memories)    return "No relevant memories found."
MEMORY_TOOLS = [save_memory, retrieve_memories]

from typing import Annotated, Dict, Listfrom typing_extensions import TypedDictfrom langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholderfrom langchain_core.messages import ToolMessagefrom langgraph.graph import StateGraph, START, ENDfrom langgraph.graph.message import add_messagesfrom langgraph.checkpoint.mongodb import MongoDBSaverfrom config import llm, mongo_clientfrom search_tools import SEARCH_TOOLSfrom memory_tools import MEMORY_TOOLS
# Define the graph stateclass GraphState(TypedDict):    messages: Annotated[list, add_messages]
# Define the LangGraph agentclass LangGraphAgent:    def __init__(self):        # Combine search tools with memory tools        self.tools = SEARCH_TOOLS + MEMORY_TOOLS        self.tools_by_name = {tool.name: tool for tool in self.tools}                # Create prompt template        self.prompt = ChatPromptTemplate.from_messages([            (                "system",                "You are a helpful AI chatbot."                " You are provided with tools to answer questions about movies."                " Think step-by-step and use these tools to get the information required to answer the user query."                " Do not re-run tools unless absolutely necessary."                " If you are not able to get enough information using the tools, reply with I DON'T KNOW."                " You have access to the following tools: {tool_names}."            ),            MessagesPlaceholder(variable_name="messages"),        ])                # Provide the tool names to the prompt        self.prompt = self.prompt.partial(tool_names=", ".join([tool.name for tool in self.tools]))                # Prepare the LLM with tools        bind_tools = llm.bind_tools(self.tools)        self.llm_with_tools = self.prompt | bind_tools                # Build the graph        self.app = self._build_graph()        def _build_graph(self):        """Build and compile the LangGraph workflow."""        # Instantiate the graph        graph = StateGraph(GraphState)                # Add nodes        graph.add_node("agent", self._agent_node)        graph.add_node("tools", self._tools_node)                # Add edges        graph.add_edge(START, "agent")        graph.add_edge("tools", "agent")                # Add conditional edge        graph.add_conditional_edges(            "agent",            self._route_tools,            {"tools": "tools", END: END},        )
        # Use the MongoDB checkpointer for short-term memory        checkpointer = MongoDBSaver(mongo_client, db_name = "sample_mflix")        return graph.compile(checkpointer=checkpointer)        def _agent_node(self, state: GraphState) -> Dict[str, List]:        """Agent node that processes messages and decides on tool usage."""        messages = state["messages"]        result = self.llm_with_tools.invoke(messages)        return {"messages": [result]}        def _tools_node(self, state: GraphState) -> Dict[str, List]:        """Tools node that executes the requested tools."""        result = []        messages = state["messages"]
        if not messages:            return {"messages": result}
        last_message = messages[-1]
        if not hasattr(last_message, "tool_calls") or not last_message.tool_calls:            return {"messages": result}
        tool_calls = last_message.tool_calls
        # Show which tools the agent chose to use        tool_names = [tool_call["name"] for tool_call in tool_calls]        print(f"🔧 Agent chose to use tool(s): {', '.join(tool_names)}")
        for tool_call in tool_calls:            try:                tool_name = tool_call["name"]                tool_args = tool_call["args"]                tool_id = tool_call["id"]                print(f"   → Executing {tool_name}")                                if tool_name not in self.tools_by_name:                    result.append(ToolMessage(content=f"Tool '{tool_name}' not found", tool_call_id=tool_id))                    continue
                tool = self.tools_by_name[tool_name]                observation = tool.invoke(tool_args)                result.append(ToolMessage(content=str(observation), tool_call_id=tool_id))
            except Exception as e:                result.append(ToolMessage(content=f"Tool error: {str(e)}", tool_call_id=tool_id))
        return {"messages": result}        def _route_tools(self, state: GraphState):        """        Uses a conditional_edge to route to the tools node if the last message        has tool calls. Otherwise, route to the end.        """        messages = state.get("messages", [])        if len(messages) > 0:            ai_message = messages[-1]        else:            raise ValueError(f"No messages found in input state to tool_edge: {state}")                if hasattr(ai_message, "tool_calls") and len(ai_message.tool_calls) > 0:            return "tools"        return END        def execute(self, user_input: str, thread_id: str) -> str:        """Execute the graph with user input."""        input_data = {"messages": [("user", user_input)]}        config = {"configurable": {"thread_id": thread_id}}
        outputs = list(self.app.stream(input_data, config))
        # Get the final answer        if outputs:            final_output = outputs[-1]            for _, value in final_output.items():                if "messages" in value and value["messages"]:                    return value["messages"][-1].content
        return "No response generated."

from agent import LangGraphAgentfrom config import mongo_client
def main():    """LangGraph and MongoDB agent with tools and memory."""    # Initialize agent (indexes are created during config import)    agent = LangGraphAgent()    thread_id = input("Enter a session ID: ").strip()    print("Ask me about movies! Type 'quit' to exit.")
    try:        while True:            user_query = input("\nYour question: ").strip()
            if user_query.lower() == 'quit':                break
            # Get response from agent            answer = agent.execute(user_query, thread_id)            print(f"\nAnswer: {answer}")
    finally:        mongo_client.close()
if __name__ == "__main__":    main()