Creating LangGraph Agents

from typing import Annotated, List, Optional
from langchain_core.messages import AnyMessage
from langgraph.graph import add_messages
from pydantic import BaseModel, Field
from collections.abc import MutableMapping

class DictLikeModel(BaseModel, MutableMapping):
    """Base model that acts like a dictionary."""
    def __getitem__(self, key):
        return getattr(self, key)

    def __setitem__(self, key, value):
        setattr(self, key, value)

    def __delitem__(self, key):
        delattr(self, key)

    def __len__(self):
        return len(self.__dict__)

class State(DictLikeModel):
    """Agent state with message history and context."""
    query: str = ""
    messages: Annotated[List[AnyMessage], add_messages] = Field(
        default_factory=list
    )
    current_datetime: Optional[str] = None
    mem0_user_id: Optional[str] = None
    memories: List[str] = Field(default_factory=list)
    memories_stored: bool = False
    conversation_id: Optional[str] = None

import asyncio
from contextlib import asynccontextmanager
from typing import Optional

from app.agents.core.graph_builder.checkpointer_manager import (
    get_checkpointer_manager,
)
from app.agents.core.nodes import (
    manage_system_prompts_node,
    trim_messages_node,
)
from app.agents.llm.client import init_llm
from app.override.langgraph_bigtool.create_agent import create_agent
from langchain_core.language_models import LanguageModelLike
from langgraph.checkpoint.memory import InMemorySaver

@asynccontextmanager
async def build_custom_graph(
    chat_llm: Optional[LanguageModelLike] = None,
    in_memory_checkpointer: bool = False,
):
    """Build a custom agent graph."""
    if chat_llm is None:
        chat_llm = init_llm()

    # Define tools for your agent
    tool_registry = {
        "my_tool": my_custom_tool,
    }

    # Get tool store
    from app.agents.tools.core.store import get_tools_store
    store = await get_tools_store()

    # Create agent with tools and hooks
    builder = create_agent(
        llm=chat_llm,
        agent_name="custom_agent",
        tool_registry=tool_registry,
        disable_retrieve_tools=False,
        initial_tool_ids=["my_tool"],
        pre_model_hooks=[
            manage_system_prompts_node,
            trim_messages_node,
        ],
    )

    # Setup checkpointer
    checkpointer_manager = await get_checkpointer_manager()

    if in_memory_checkpointer or not checkpointer_manager:
        checkpointer = InMemorySaver()
    else:
        checkpointer = checkpointer_manager.get_checkpointer()

    # Compile and yield graph
    graph = builder.compile(checkpointer=checkpointer, store=store)
    yield graph

from app.core.lazy_loader import MissingKeyStrategy, lazy_provider
from app.config.loggers import app_logger as logger

@lazy_provider(
    name="custom_agent",
    required_keys=[],
    strategy=MissingKeyStrategy.WARN,
    auto_initialize=False,
)
async def build_custom_agent():
    """Build and return the custom agent."""
    logger.debug("Building custom agent with lazy providers")

    async with build_custom_graph() as graph:
        logger.info("Custom agent built successfully")
    return graph

from app.agents.core.graph_manager import GraphManager

# In your initialization code
GraphManager.set_graph(custom_agent_graph, "custom_agent")

async def call_agent(
    request: MessageRequestWithHistory,
    conversation_id: str,
    user: dict,
    user_time: datetime,
    user_model_config: Optional[ModelConfig] = None,
) -> AsyncGenerator[str, None]:
    """Execute agent with SSE streaming."""
    graph, initial_state, config = await _core_agent_logic(
        request, conversation_id, user, user_time, user_model_config
    )

    return execute_graph_streaming(graph, initial_state, config)

async def call_agent_silent(
    request: MessageRequestWithHistory,
    conversation_id: str,
    user: dict,
    user_time: datetime,
) -> tuple[str, dict]:
    """Execute agent without streaming."""
    graph, initial_state, config = await _core_agent_logic(
        request, conversation_id, user, user_time
    )

    return await execute_graph_silent(graph, initial_state, config)

pre_model_hooks=[
    filter_messages_node,      # Filter message history
    manage_system_prompts_node, # Inject system prompts
    trim_messages_node,        # Limit context window
]

end_graph_hooks=[
    follow_up_actions_node,  # Suggest next actions
]

from app.utils.memory_utils import store_user_message_memory
import asyncio

# Store memory in background
if user_id and request.message:
    task = asyncio.create_task(
        store_user_message_memory(
            user_id, request.message, conversation_id
        )
    )
    _background_tasks.add(task)
    task.add_done_callback(_background_tasks.discard)

from app.helpers.agent_helpers import build_agent_config

config = build_agent_config(
    conversation_id=conversation_id,
    user=user,
    user_time=user_time,
    user_model_config=user_model_config,
    agent_name="custom_agent",
    selected_tool=request.selectedTool,
    tool_category=request.toolCategory,
)